diff --git a/.github/workflows/build-push.yml b/.github/workflows/build-push.yml index 56ca2860ad1..545495ca712 100644 --- a/.github/workflows/build-push.yml +++ b/.github/workflows/build-push.yml @@ -78,15 +78,15 @@ jobs: - service_name: "build-agent-local-sandbox-amd64" image_name_env: "DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME" artifact_context: "local-sandbox" - build_context: "{{defaultContext}}:dify-agent" - file: "docker/local-sandbox/Dockerfile" + build_context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" platform: linux/amd64 runs_on: depot-ubuntu-24.04-4 - service_name: "build-agent-local-sandbox-arm64" image_name_env: "DIFY_AGENT_LOCAL_SANDBOX_IMAGE_NAME" artifact_context: "local-sandbox" - build_context: "{{defaultContext}}:dify-agent" - file: "docker/local-sandbox/Dockerfile" + build_context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" platform: linux/arm64 runs_on: depot-ubuntu-24.04-4 @@ -120,6 +120,7 @@ jobs: file: ${{ matrix.file }} platforms: ${{ matrix.platform }} build-args: COMMIT_SHA=${{ fromJSON(steps.meta.outputs.json).labels['org.opencontainers.image.revision'] }} + build-contexts: ${{ matrix.build_contexts }} labels: ${{ steps.meta.outputs.labels }} outputs: type=image,name=${{ env[matrix.image_name_env] }},push-by-digest=true,name-canonical=true,push=true @@ -155,8 +156,8 @@ jobs: build_context: "{{defaultContext}}" file: "dify-agent/Dockerfile" - service_name: "validate-agent-local-sandbox-amd64" - build_context: "{{defaultContext}}:dify-agent" - file: "docker/local-sandbox/Dockerfile" + build_context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" steps: - name: Set up Docker Buildx uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 @@ -167,6 +168,7 @@ jobs: push: false context: ${{ matrix.build_context }} file: ${{ matrix.file }} + build-contexts: ${{ matrix.build_contexts }} platforms: linux/amd64 create-manifest: diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index 0646734af89..9fb71a42cfd 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -13,6 +13,11 @@ on: - dify-agent/README.md - dify-agent/src/** - web/Dockerfile + - dify-agent-runtime/docker/Dockerfile + - dify-agent-runtime/go.mod + - dify-agent-runtime/go.sum + - dify-agent-runtime/cmd/** + - dify-agent-runtime/internal/** concurrency: group: docker-build-${{ github.head_ref || github.run_id }} @@ -48,6 +53,16 @@ jobs: runs_on: depot-ubuntu-24.04-4 context: "{{defaultContext}}" file: "web/Dockerfile" + - service_name: "local-sandbox-amd64" + platform: linux/amd64 + runs_on: depot-ubuntu-24.04-4 + context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" + - service_name: "local-sandbox-arm64" + platform: linux/arm64 + runs_on: depot-ubuntu-24.04-4 + context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" steps: - name: Set up Depot CLI uses: depot/setup-action@15c09a5f77a0840ad4bce955686522a257853461 # v1.7.1 @@ -59,6 +74,7 @@ jobs: push: false context: ${{ matrix.context }} file: ${{ matrix.file }} + build-contexts: ${{ matrix.build_contexts }} platforms: ${{ matrix.platform }} build-docker-fork: @@ -75,6 +91,9 @@ jobs: - service_name: "web-amd64" context: "{{defaultContext}}" file: "web/Dockerfile" + - service_name: "local-sandbox-amd64" + context: "{{defaultContext}}:dify-agent-runtime" + file: "docker/Dockerfile" steps: - name: Set up Docker Buildx uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 @@ -85,4 +104,5 @@ jobs: push: false context: ${{ matrix.context }} file: ${{ matrix.file }} + build-contexts: ${{ matrix.build_contexts }} platforms: linux/amd64 diff --git a/.github/workflows/main-ci.yml b/.github/workflows/main-ci.yml index 924d825f0fd..59e8d856975 100644 --- a/.github/workflows/main-ci.yml +++ b/.github/workflows/main-ci.yml @@ -45,6 +45,7 @@ jobs: web-changed: ${{ steps.changes.outputs.web }} vdb-changed: ${{ steps.changes.outputs.vdb }} migration-changed: ${{ steps.changes.outputs.migration }} + sandbox-runtime-changed: ${{ steps.changes.outputs.sandbox-runtime }} steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2 @@ -130,6 +131,9 @@ jobs: - 'docker/volumes/**' - 'api/uv.lock' - 'api/pyproject.toml' + sandbox-runtime: + - 'dify-agent-runtime/**' + - '.github/workflows/sandbox-runtime-tests.yml' migration: - 'api/migrations/**' - 'api/.env.example' @@ -508,3 +512,63 @@ jobs: echo "DB migration tests were not required, but the skip job finished with result: $SKIP_RESULT" >&2 exit 1 + + sandbox-runtime-tests-run: + name: Run Sandbox Runtime Tests + needs: + - pre_job + - check-changes + if: needs.pre_job.outputs.should_skip != 'true' && needs.check-changes.outputs.sandbox-runtime-changed == 'true' + uses: ./.github/workflows/sandbox-runtime-tests.yml + secrets: inherit + + sandbox-runtime-tests-skip: + name: Skip Sandbox Runtime Tests + needs: + - pre_job + - check-changes + if: needs.pre_job.outputs.should_skip != 'true' && needs.check-changes.outputs.sandbox-runtime-changed != 'true' + runs-on: depot-ubuntu-24.04 + steps: + - name: Report skipped sandbox runtime tests + run: echo "No sandbox-runtime-related changes detected; skipping sandbox runtime tests." + + sandbox-runtime-tests: + name: Sandbox Runtime Tests + if: ${{ always() }} + needs: + - pre_job + - check-changes + - sandbox-runtime-tests-run + - sandbox-runtime-tests-skip + runs-on: depot-ubuntu-24.04 + steps: + - name: Finalize Sandbox Runtime Tests status + env: + SHOULD_SKIP_WORKFLOW: ${{ needs.pre_job.outputs.should_skip }} + TESTS_CHANGED: ${{ needs.check-changes.outputs.sandbox-runtime-changed }} + RUN_RESULT: ${{ needs.sandbox-runtime-tests-run.result }} + SKIP_RESULT: ${{ needs.sandbox-runtime-tests-skip.result }} + run: | + if [[ "$SHOULD_SKIP_WORKFLOW" == 'true' ]]; then + echo "Sandbox runtime tests were skipped because this workflow run duplicated a successful or newer run." + exit 0 + fi + + if [[ "$TESTS_CHANGED" == 'true' ]]; then + if [[ "$RUN_RESULT" == 'success' ]]; then + echo "Sandbox runtime tests ran successfully." + exit 0 + fi + + echo "Sandbox runtime tests were required but finished with result: $RUN_RESULT" >&2 + exit 1 + fi + + if [[ "$SKIP_RESULT" == 'success' ]]; then + echo "Sandbox runtime tests were skipped because no sandbox-runtime-related files changed." + exit 0 + fi + + echo "Sandbox runtime tests were not required, but the skip job finished with result: $SKIP_RESULT" >&2 + exit 1 diff --git a/.github/workflows/sandbox-runtime-tests.yml b/.github/workflows/sandbox-runtime-tests.yml new file mode 100644 index 00000000000..03c5eb1337e --- /dev/null +++ b/.github/workflows/sandbox-runtime-tests.yml @@ -0,0 +1,98 @@ +name: Sandbox Runtime Tests + +on: + workflow_call: + +permissions: + contents: read + +concurrency: + group: sandbox-runtime-tests-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +jobs: + sandbox-runtime-unit: + name: Sandbox Runtime Unit Tests + runs-on: depot-ubuntu-24.04 + defaults: + run: + shell: bash + working-directory: dify-agent-runtime + + steps: + - name: Checkout code + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Setup Go + uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 + with: + go-version-file: dify-agent-runtime/go.mod + cache-dependency-path: dify-agent-runtime/go.sum + + - name: Run unit tests + run: go test -race -count=1 ./... + + sandbox-runtime-lint: + name: Sandbox Runtime Lint + runs-on: depot-ubuntu-24.04 + defaults: + run: + shell: bash + working-directory: dify-agent-runtime + + steps: + - name: Checkout code + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Setup Go + uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 + with: + go-version-file: dify-agent-runtime/go.mod + cache-dependency-path: dify-agent-runtime/go.sum + + - name: Run golangci-lint + uses: golangci/golangci-lint-action@4afd733a84b1f43292c63897423277bb7f4313a9 # v6.5.0 + with: + working-directory: dify-agent-runtime + version: latest + + sandbox-runtime-integration: + name: Sandbox Runtime Integration Tests + runs-on: depot-ubuntu-24.04 + defaults: + run: + shell: bash + working-directory: dify-agent-runtime + + steps: + - name: Checkout code + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Setup Go + uses: actions/setup-go@d35c59abb061a4a6fb18e82ac0862c26744d6ab5 # v5.5.0 + with: + go-version-file: dify-agent-runtime/go.mod + cache-dependency-path: dify-agent-runtime/go.sum + + - name: Build and start runtime + run: make integration-up + + - name: Run integration tests + run: make integration-test + + - name: Dump container logs on failure + if: failure() + run: make integration-logs || true + + - name: Stop runtime container + if: always() + run: make integration-down diff --git a/.gitignore b/.gitignore index fcec9e4721c..54fdb704d8c 100644 --- a/.gitignore +++ b/.gitignore @@ -261,3 +261,7 @@ scripts/stress-test/reports/ .context/ # Vitest local reports web/.vitest-reports/ + +# dify-agent-runtime +dify-agent-runtime/bin/ +dify-agent-runtime/.integration-state diff --git a/Makefile b/Makefile index 801c59cc281..07199d4e16a 100644 --- a/Makefile +++ b/Makefile @@ -2,6 +2,7 @@ DOCKER_REGISTRY=langgenius WEB_IMAGE=$(DOCKER_REGISTRY)/dify-web API_IMAGE=$(DOCKER_REGISTRY)/dify-api +SANDBOX_RUNTIME_IMAGE=$(DOCKER_REGISTRY)/dify-agent-local-sandbox VERSION=latest DOCKER_DIR=docker DOCKER_MIDDLEWARE_ENV=$(DOCKER_DIR)/middleware.env @@ -160,6 +161,13 @@ build-api: docker build -t $(API_IMAGE):$(VERSION) -f api/Dockerfile . @echo "API Docker image built successfully: $(API_IMAGE):$(VERSION)" +build-sandbox-runtime: + @echo "Building sandbox runtime Docker image: $(SANDBOX_RUNTIME_IMAGE):$(VERSION)..." + docker build -t $(SANDBOX_RUNTIME_IMAGE):$(VERSION) \ + -f dify-agent-runtime/docker/Dockerfile \ + dify-agent-runtime + @echo "Sandbox runtime Docker image built successfully: $(SANDBOX_RUNTIME_IMAGE):$(VERSION)" + # Push Docker images push-web: @echo "Pushing web Docker image: $(WEB_IMAGE):$(VERSION)..." @@ -171,14 +179,20 @@ push-api: docker push $(API_IMAGE):$(VERSION) @echo "API Docker image pushed successfully: $(API_IMAGE):$(VERSION)" +push-sandbox-runtime: + @echo "Pushing sandbox runtime Docker image: $(SANDBOX_RUNTIME_IMAGE):$(VERSION)..." + docker push $(SANDBOX_RUNTIME_IMAGE):$(VERSION) + @echo "Sandbox runtime Docker image pushed successfully: $(SANDBOX_RUNTIME_IMAGE):$(VERSION)" + # Build all images -build-all: build-web build-api +build-all: build-web build-api build-sandbox-runtime # Push all images -push-all: push-web push-api +push-all: push-web push-api push-sandbox-runtime build-push-api: build-api push-api build-push-web: build-web push-web +build-push-sandbox-runtime: build-sandbox-runtime push-sandbox-runtime # Build and push all images build-push-all: build-all push-all @@ -206,9 +220,10 @@ help: @echo "Docker Build Targets:" @echo " make build-web - Build web Docker image" @echo " make build-api - Build API Docker image" + @echo " make build-sandbox-runtime - Build sandbox runtime Docker image" @echo " make build-all - Build all Docker images" @echo " make push-all - Push all Docker images" @echo " make build-push-all - Build and push all Docker images" # Phony targets -.PHONY: build-web build-api push-web push-api build-all push-all build-push-all dev-setup prepare-docker prepare-web prepare-api dev-clean help format check lint api-contract-lint type-check test test-all +.PHONY: build-web build-api build-sandbox-runtime push-web push-api push-sandbox-runtime build-all push-all build-push-all build-push-sandbox-runtime dev-setup prepare-docker prepare-web prepare-api dev-clean help format check lint api-contract-lint type-check test test-all diff --git a/api/uv.lock b/api/uv.lock index f2c1942ff86..358c1e97c1a 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1283,21 +1283,16 @@ name = "dify-agent" version = "1.16.0rc1" source = { editable = "../dify-agent" } dependencies = [ - { name = "anyio" }, { name = "httpx" }, { name = "httpx2" }, { name = "pydantic" }, { name = "pydantic-ai-slim" }, - { name = "typer" }, { name = "typing-extensions" }, ] [package.metadata] requires-dist = [ - { name = "aiosqlite", marker = "extra == 'shellctl-server'", specifier = ">=0.21.0,<1.0.0" }, - { name = "anyio", specifier = ">=4.12.1,<5.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "fastapi", marker = "extra == 'shellctl-server'", specifier = "==0.136.0" }, { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, { name = "grpclib", extras = ["protobuf"], marker = "extra == 'grpc'", specifier = ">=0.4.9,<0.5.0" }, { name = "httpx", specifier = "==0.28.1" }, @@ -1311,13 +1306,10 @@ requires-dist = [ { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" }, { name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" }, { name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" }, - { name = "sqlmodel", marker = "extra == 'shellctl-server'", specifier = ">=0.0.24,<0.1.0" }, - { name = "typer", specifier = ">=0.16.1,<0.17" }, { name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = "==0.46.0" }, - { name = "uvicorn", extras = ["standard"], marker = "extra == 'shellctl-server'", specifier = "==0.46.0" }, ] -provides-extras = ["grpc", "server", "shellctl-server"] +provides-extras = ["grpc", "server"] [package.metadata.requires-dev] dev = [ diff --git a/dify-agent-runtime/Makefile b/dify-agent-runtime/Makefile new file mode 100644 index 00000000000..64550ef16ce --- /dev/null +++ b/dify-agent-runtime/Makefile @@ -0,0 +1,108 @@ +.PHONY: build clean test lint proto gen-cli-help integration integration-up integration-test integration-down +BIN_DIR := bin +PROTO_SRC := ../dify-agent/proto/dify/agent/stub/v1/agent_stub.proto +PROTO_OUT := gen +AGENT_CLI_HELP_JSON := ../dify-agent/src/dify_agent/layers/_agent_cli_help.json + +build: $(BIN_DIR)/shellctl $(BIN_DIR)/shellctl-sanitize-pty $(BIN_DIR)/shellctl-runner-exit $(BIN_DIR)/dify-agent gen-cli-help + +$(BIN_DIR)/shellctl: $(shell find cmd/shellctl internal -name '*.go') + go build -o $@ ./cmd/shellctl + +$(BIN_DIR)/shellctl-sanitize-pty: $(shell find cmd/sanitize-pty internal/sanitize -name '*.go') + go build -o $@ ./cmd/sanitize-pty + +$(BIN_DIR)/shellctl-runner-exit: $(shell find cmd/runner-exit internal/runner_exit -name '*.go') + go build -o $@ ./cmd/runner-exit + +$(BIN_DIR)/dify-agent: $(shell find cmd/dify-agent-cli internal/agentcli internal/stubclient -name '*.go') $(shell find gen -name '*.go' 2>/dev/null) + go build -o $@ ./cmd/dify-agent-cli + +# Regenerate the JSON help snapshot from the Go command tree. Runs as part of +# `build` so the checked-in data never drifts from the CLI. +gen-cli-help: $(BIN_DIR)/dify-agent + $(BIN_DIR)/dify-agent __dump-cli-help > $(AGENT_CLI_HELP_JSON) + +test: lint + go test ./... + +lint: + golangci-lint run ./... + +proto: + @mkdir -p $(PROTO_OUT) + protoc \ + --proto_path=../dify-agent/proto \ + --go_out=$(PROTO_OUT) --go_opt=paths=source_relative \ + --go-grpc_out=$(PROTO_OUT) --go-grpc_opt=paths=source_relative \ + $(PROTO_SRC) + +clean: + rm -rf $(BIN_DIR) + +# --- 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. +# +# State is written to .integration-state so that separate make invocations +# (integration-up, integration-test, integration-logs, integration-down) +# can reference the same container. `make integration` runs the full +# lifecycle in a single invocation with automatic cleanup. + +IMAGE_NAME := dify-agent-runtime:test +STATE_FILE := .integration-state + +integration-up: + @echo "Building runtime image..." + docker build -t $(IMAGE_NAME) -f docker/Dockerfile . + $(eval TEST_ID := $(shell printf '%05d' $$((RANDOM % 100000)))) + $(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)..." + docker run -d --name $(CONTAINER_NAME) \ + -p $(HOST_PORT):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)"; \ + 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 + +integration-test: + @test -f $(STATE_FILE) || { echo "ERROR: run 'make integration-up' first" >&2; exit 1; } + @. ./$(STATE_FILE); \ + SHELLCTL_GO_URL=http://localhost:$$HOST_PORT \ + SHELLCTL_TEST_TOKEN=$$AUTH_TOKEN \ + go test -tags=integration -v -count=1 -timeout=300s ./tests/... + +integration-logs: + @test -f $(STATE_FILE) || { echo "ERROR: no state file" >&2; exit 1; } + @. ./$(STATE_FILE); docker logs $$CONTAINER_NAME + +integration-down: + @if [ -f $(STATE_FILE) ]; then \ + . ./$(STATE_FILE); \ + docker rm -f $$CONTAINER_NAME 2>/dev/null || true; \ + rm -f $(STATE_FILE); \ + fi + +integration: + @$(MAKE) integration-up + @trap '$(MAKE) integration-down' EXIT; \ + $(MAKE) integration-test diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md new file mode 100644 index 00000000000..6b96ca76435 --- /dev/null +++ b/dify-agent-runtime/README.md @@ -0,0 +1,79 @@ +# dify-agent-runtime + +Go implementation of the shellctl server and runtime utilities. + +This is a rewrite of the Python `shellctl` package (`dify-agent/src/shellctl/` and +`dify-agent/src/shellctl_runtime/`). The original Python code is kept as reference. + +## Architecture + +``` +cmd/ + shellctl/ — main server binary (shellctl serve) + sanitize-pty/ — tmux pipe-pane PTY sanitizer (stdin→stdout filter) + runner-exit/ — post-drain SQLite exit recorder + +internal/ + sanitize/ — PTY ANSI stripping + CR normalization + runner_exit/ — SQLite CAS update for job exit + server/ — HTTP API, job service, tmux controller, output reader +``` + +## Building + +```bash +make build +``` + +Produces three binaries in `bin/`: + +- `shellctl` — the main server (`shellctl serve --listen 0.0.0.0:5004`) +- `shellctl-sanitize-pty` — PTY sanitizer for tmux pipe-pane +- `shellctl-runner-exit` — exit state writer + +### Building docker image + +``` +docker build -f dify-agent-runtime/docker/Dockerfile \ + --build-context agent=./dify-agent \ + -t dify-agent-runtime:latest \ + dify-agent-runtime/ +``` + +### Runing docker container + +``` +docker run -d --name dify-agent-runtime \ + -p 15004:5004 \ + dify-agent-runtime:latest +``` + +## Help text generation + +Cli help can be generated and injected to dify-agent's system prompt. + +```sh +make gen-cli-help +``` + +## Testing + +```bash +make test +``` + +## Dependencies + +- Go 1.23+ +- `modernc.org/sqlite` (pure-Go SQLite driver, no CGO required) +- tmux (runtime dependency, not a build dependency) + +## Migration from Python + +The Go binaries are drop-in replacements for the Python console scripts: + +- `shellctl-sanitize-pty` replaces the Python `shellctl-sanitize-pty` entrypoint +- `shellctl-runner-exit` replaces the Python `shellctl-runner-exit` entrypoint +- `shellctl serve` replaces the Python `shellctl serve` (FastAPI/uvicorn) + +The HTTP API contract, SQLite schema, and filesystem artifact layout are identical. diff --git a/dify-agent-runtime/cmd/dify-agent-cli/dumphelp.go b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp.go new file mode 100644 index 00000000000..46827553e08 --- /dev/null +++ b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp.go @@ -0,0 +1,68 @@ +package main + +// dumphelp renders the visible cobra command tree into the JSON snapshot at +// dify-agent/src/dify_agent/layers/_agent_cli_help.json. The Go CLI is the +// single source of truth for both the command surface and its help text; the +// Python server loads this JSON and injects the strings into the agent prompt +// without executing a shell. Emitting JSON keeps the Go side free of any +// Python formatting/encoding concerns. +// +// It is invoked via the hidden `dify-agent __dump-cli-help` intercept (see +// main.go) and driven by `make gen-cli-help` / `go generate`. + +import ( + "bytes" + "encoding/json" + "io" + "strings" + + "github.com/spf13/cobra" +) + +// dumpCLIHelp writes the help snapshot as a JSON object mapping each +// space-joined command path (e.g. "config note push") to the verbatim +// `dify-agent --help` output. json.MarshalIndent sorts object keys, so +// the output is deterministic for clean diffs. +func dumpCLIHelp(w io.Writer) error { + paths := collectHelpPaths(newRootCommand(), nil) + help := make(map[string]string, len(paths)) + for _, path := range paths { + help[strings.Join(path, " ")] = captureHelp(path) + } + + data, err := json.MarshalIndent(help, "", " ") + if err != nil { + return err + } + _, err = w.Write(append(data, '\n')) + return err +} + +// collectHelpPaths walks the command tree and returns the path of every visible +// (non-hidden) command, excluding the root itself. Paths are relative to the +// root, e.g. ["config", "note", "push"]. +func collectHelpPaths(cmd *cobra.Command, prefix []string) [][]string { + var paths [][]string + for _, child := range cmd.Commands() { + if child.Hidden || child.Name() == "help" { + continue + } + childPath := append(append([]string{}, prefix...), child.Name()) + paths = append(paths, childPath) + paths = append(paths, collectHelpPaths(child, childPath)...) + } + return paths +} + +// captureHelp renders one command's `--help` output exactly as the CLI does. +// A fresh root per call avoids leaking parsed flag state between commands, +// mirroring the executeHelp test helper. +func captureHelp(path []string) string { + root := newRootCommand() + var buf bytes.Buffer + root.SetOut(&buf) + root.SetErr(&buf) + root.SetArgs(append(append([]string{}, path...), "--help")) + _ = root.Execute() + return strings.TrimRight(buf.String(), "\n") +} diff --git a/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go new file mode 100644 index 00000000000..a8a0c328f8d --- /dev/null +++ b/dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go @@ -0,0 +1,68 @@ +package main + +import ( + "bytes" + "encoding/json" + "os" + "testing" +) + +// checkedInHelpJSON is the JSON snapshot generated from this command tree, +// relative to the package directory (cmd/dify-agent-cli). +const checkedInHelpJSON = "../../../dify-agent/src/dify_agent/layers/_agent_cli_help.json" + +// TestDumpCLIHelpMatchesCheckedIn fails when the committed JSON help snapshot +// drifts from the current Go command tree. Regenerate with: +// +// make -C dify-agent-runtime gen-cli-help +func TestDumpCLIHelpMatchesCheckedIn(t *testing.T) { + want, err := os.ReadFile(checkedInHelpJSON) + if err != nil { + t.Skipf("checked-in help snapshot not found (%v); run `make gen-cli-help`", err) + } + + var got bytes.Buffer + if err := dumpCLIHelp(&got); err != nil { + t.Fatalf("dumpCLIHelp: %v", err) + } + + if got.String() != string(want) { + t.Errorf("%s is stale; regenerate with `make -C dify-agent-runtime gen-cli-help`", checkedInHelpJSON) + } +} + +// TestDumpCLIHelpCoversPromptCommands ensures the generated snapshot includes +// every command path the Python config layer injects into the prompt, and that +// the output is valid JSON keyed by space-joined command paths. +func TestDumpCLIHelpCoversPromptCommands(t *testing.T) { + var buf bytes.Buffer + if err := dumpCLIHelp(&buf); err != nil { + t.Fatalf("dumpCLIHelp: %v", err) + } + + var table map[string]string + if err := json.Unmarshal(buf.Bytes(), &table); err != nil { + t.Fatalf("generated snapshot is not valid JSON: %v", err) + } + + wantKeys := []string{ + "config", + "config manifest", + "config skills pull", + "config files pull", + "config note pull", + "config note push", + "config env push", + "config files push", + "config files delete", + "config skills push", + "config skills delete", + "file upload", + "file download", + } + for _, key := range wantKeys { + if _, ok := table[key]; !ok { + t.Errorf("generated snapshot missing prompt command key %q", key) + } + } +} diff --git a/dify-agent-runtime/cmd/dify-agent-cli/main.go b/dify-agent-runtime/cmd/dify-agent-cli/main.go new file mode 100644 index 00000000000..bcee7dc1fe4 --- /dev/null +++ b/dify-agent-runtime/cmd/dify-agent-cli/main.go @@ -0,0 +1,438 @@ +// dify-agent-cli is the Go replacement for the Python dify-agent CLI. +// It communicates with the Agent Stub server via HTTP or gRPC to provide +// connect, file, drive, and config operations inside the sandbox container. +package main + +import ( + "fmt" + "os" + "strings" + + "github.com/langgenius/dify/dify-agent-runtime/internal/agentcli" + "github.com/spf13/cobra" +) + +// knownRootCommands mirrors the Python CLI's _KNOWN_ROOT_COMMANDS. Anything else +// on the command line is treated as argv forwarded to an implicit connect. +var knownRootCommands = map[string]struct{}{ + "config": {}, + "connect": {}, + "drive": {}, + "file": {}, +} + +func main() { + if err := runCLI(os.Args[1:]); err != nil { + fmt.Fprintf(os.Stderr, "dify-agent: %v\n", err) + os.Exit(1) + } +} + +//go:generate sh -c "go run . __dump-cli-help > ../../../dify-agent/src/dify_agent/layers/_agent_cli_help.json" + +func runCLI(args []string) error { + // Hidden developer intercept: emit the JSON help snapshot from the cobra + // tree. Not registered as a cobra command, so the visible command surface + // stays identical to the Python CLI. See dumphelp.go. + if len(args) == 1 && args[0] == "__dump-cli-help" { + return dumpCLIHelp(os.Stdout) + } + + // `connect` forwards its argv verbatim, so bypass flag parsing unless the + // user explicitly asked for help (which should render connect's help). + if len(args) >= 1 && args[0] == "connect" && !isHelpRequest(args[1:]) { + jsonOutput, forwarded := parseConnectArgs(args[1:]) + return runConnect(forwarded, jsonOutput) + } + + jsonOutput, forwarded := extractRootJSONFlag(args) + if isUnknownBareCommand(forwarded) { + // Match Python: surface root help when the environment is missing, then + // still attempt connect so the missing-env error is reported. + if !agentcli.HasEnvironment() { + _ = newRootCommand().Help() + } + return runConnect(forwarded, jsonOutput) + } + + root := newRootCommand() + root.SetArgs(args) + return root.Execute() +} + +// --- command tree --- + +func newRootCommand() *cobra.Command { + root := &cobra.Command{ + Use: "dify-agent", + Short: "Forward shell-visible dify-agent commands to the Dify Agent Stub server.", + SilenceUsage: true, + SilenceErrors: true, + } + root.CompletionOptions.DisableDefaultCmd = true + // The Python CLI exposes no `help` command; hide cobra's auto-added one so + // the visible command surface stays identical. + root.SetHelpCommand(&cobra.Command{Hidden: true}) + root.AddCommand( + newConnectCommand(), + newFileCommand(), + newDriveCommand(), + newConfigCommand(), + ) + return root +} + +func newConnectCommand() *cobra.Command { + var jsonOutput bool + // connect is normally intercepted in runCLI; this definition exists so + // `dify-agent connect --help` renders framework help for the command. + cmd := &cobra.Command{ + Use: "connect [ARGV]...", + Short: "Establish one Agent Stub connection using the current environment.", + RunE: func(_ *cobra.Command, args []string) error { + return runConnect(args, jsonOutput) + }, + } + cmd.Flags().BoolVar(&jsonOutput, "json", false, "Emit the connection response as JSON.") + return cmd +} + +func newFileCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "file", + Short: "Upload or download workflow files through the Agent Stub.", + } + + upload := &cobra.Command{ + Use: "upload PATH", + Short: "Upload one sandbox-local file as a ToolFile output reference.", + Args: cobra.ExactArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunFileUpload(env, args[0]) + }), + } + + var downloadTo string + download := &cobra.Command{ + Use: "download TRANSFER_METHOD REFERENCE_OR_URL", + Short: "Download one workflow file mapping into the local sandbox directory.", + Args: cobra.ExactArgs(2), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunFileDownload(env, args[0], args[1], downloadTo) + }), + } + download.Flags().StringVar(&downloadTo, "to", "", "Local directory for the downloaded file.") + + cmd.AddCommand(upload, download) + return cmd +} + +func newDriveCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "drive", + Short: "List, pull, or push agent drive files through the Agent Stub.", + } + + var listJSON bool + list := &cobra.Command{ + Use: "list [REMOTE_PREFIX]", + Short: "List drive files visible to the current sandbox execution.", + Args: cobra.MaximumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + prefix := "" + if len(args) > 0 { + prefix = args[0] + } + return agentcli.RunDriveList(env, prefix, listJSON) + }), + } + list.Flags().BoolVar(&listJSON, "json", false, "Emit the drive manifest as JSON.") + + var pullTo string + var pullJSON bool + pull := &cobra.Command{ + Use: "pull [REMOTE]...", + Short: "Pull one or more drive keys/prefixes into one local directory tree.", + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + localBase := pullTo + if localBase == "" { + localBase = agentcli.ReadDriveBase() + } + return agentcli.RunDrivePull(env, args, localBase, pullJSON) + }), + } + pull.Flags().StringVar(&pullTo, "to", "", "Local base directory for pulled drive files.") + pull.Flags().BoolVar(&pullJSON, "json", false, "Emit the pull result as JSON.") + + var pushKind string + var pushJSON bool + push := &cobra.Command{ + Use: "push LOCAL_PATH REMOTE_PATH", + Short: "Upload one local file or directory into the agent drive.", + Args: cobra.ExactArgs(2), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunDrivePush(env, args[0], args[1], pushKind) + }), + } + push.Flags().StringVar(&pushKind, "kind", "", "Directory upload kind: skill or dir.") + push.Flags().BoolVar(&pushJSON, "json", false, "Accepted for consistency; drive push output is already emitted as JSON.") + + cmd.AddCommand(list, pull, push) + return cmd +} + +func newConfigCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "config", + Short: "Inspect or update Agent Soul-backed config assets through the Agent Stub.", + } + + manifest := &cobra.Command{ + Use: "manifest", + Short: "Show the current visible Agent config manifest as JSON.", + Args: cobra.NoArgs, + RunE: withEnv(func(env *agentcli.Environment, _ []string, _ *cobra.Command) error { + return agentcli.RunConfigManifest(env) + }), + } + + cmd.AddCommand( + manifest, + newConfigSkillsCommand(), + newConfigSkillPullAliasCommand(), + newConfigFilesCommand(), + newConfigFilePullAliasCommand(), + newConfigEnvCommand(), + newConfigNoteCommand(), + ) + return cmd +} + +func newConfigSkillsPullCommand() *cobra.Command { + var to string + var jsonOutput bool + cmd := &cobra.Command{ + Use: "pull [NAME]...", + Short: "Pull one or all visible config skills into ./.dify_conf/skills by default.", + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigSkillsPull(env, args, to, jsonOutput) + }), + } + cmd.Flags().StringVar(&to, "to", "", "Local directory for pulled config skills.") + cmd.Flags().BoolVar(&jsonOutput, "json", false, "Emit the pull result as JSON.") + return cmd +} + +func newConfigSkillsCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "skills", + Short: "Pull or update config skills through the Agent Stub.", + } + + push := &cobra.Command{ + Use: "push PATH...", + Short: "Upload one or more local skill directories into the current config manifest.", + Args: cobra.MinimumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigSkillsPush(env, args) + }), + } + + del := &cobra.Command{ + Use: "delete NAME...", + Short: "Delete one or more config skills by name without touching local directories.", + Args: cobra.MinimumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigSkillsDelete(env, args) + }), + } + + cmd.AddCommand(newConfigSkillsPullCommand(), push, del) + return cmd +} + +// newConfigSkillPullAliasCommand mirrors the hidden singular `config skill pull` +// alias, which is intentionally pull-only. +func newConfigSkillPullAliasCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "skill", + Short: "Pull config skills through the Agent Stub.", + Hidden: true, + } + cmd.AddCommand(newConfigSkillsPullCommand()) + return cmd +} + +func newConfigFilesPullCommand() *cobra.Command { + var to string + var jsonOutput bool + cmd := &cobra.Command{ + Use: "pull [NAME]...", + Short: "Pull one or all visible config files into ./.dify_conf/files by default.", + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigFilesPull(env, args, to, jsonOutput) + }), + } + cmd.Flags().StringVar(&to, "to", "", "Local directory for pulled config files.") + cmd.Flags().BoolVar(&jsonOutput, "json", false, "Emit the pull result as JSON.") + return cmd +} + +func newConfigFilesCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "files", + Short: "Pull or update config files through the Agent Stub.", + } + + push := &cobra.Command{ + Use: "push PATH...", + Short: "Upload one or more local files into the current config manifest.", + Args: cobra.MinimumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigFilesPush(env, args) + }), + } + + del := &cobra.Command{ + Use: "delete NAME...", + Short: "Delete one or more config files by name without touching local files.", + Args: cobra.MinimumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigFilesDelete(env, args) + }), + } + + cmd.AddCommand(newConfigFilesPullCommand(), push, del) + return cmd +} + +// newConfigFilePullAliasCommand mirrors the hidden singular `config file pull` +// alias, which is intentionally pull-only. +func newConfigFilePullAliasCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "file", + Short: "Pull config files through the Agent Stub.", + Hidden: true, + } + cmd.AddCommand(newConfigFilesPullCommand()) + return cmd +} + +func newConfigEnvCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "env", + Short: "Update config env variables visible to the current run.", + } + push := &cobra.Command{ + Use: "push PATH|-", + Short: "Set or delete config env entries from one local dotenv file or stdin.", + Args: cobra.ExactArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + return agentcli.RunConfigEnvPush(env, args[0]) + }), + } + cmd.AddCommand(push) + return cmd +} + +func newConfigNoteCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "note", + Short: "Pull or update the current config note.", + } + + var pullTo string + pull := &cobra.Command{ + Use: "pull", + Short: "Export the current config note into ./.dify_conf/note.md by default.", + Args: cobra.NoArgs, + RunE: withEnv(func(env *agentcli.Environment, _ []string, _ *cobra.Command) error { + return agentcli.RunConfigNotePull(env, pullTo) + }), + } + pull.Flags().StringVar(&pullTo, "to", "", "Local markdown file path.") + + push := &cobra.Command{ + Use: "push [PATH|-]", + Short: "Replace the current config note from one local text file or stdin.", + Args: cobra.MaximumNArgs(1), + RunE: withEnv(func(env *agentcli.Environment, args []string, _ *cobra.Command) error { + localPath := "" + if len(args) > 0 { + localPath = args[0] + } + return agentcli.RunConfigNotePush(env, localPath) + }), + } + + cmd.AddCommand(pull, push) + return cmd +} + +// --- connect dispatch helpers (mirror the Python CLI) --- + +func runConnect(argv []string, jsonOutput bool) error { + env, err := agentcli.ReadEnvironment() + if err != nil { + return err + } + return agentcli.RunConnect(env, argv, jsonOutput) +} + +// parseConnectArgs strips a leading `--json` flag and an optional `--` +// separator, matching Python's _parse_connect_args. +func parseConnectArgs(argv []string) (bool, []string) { + jsonOutput := false + remaining := append([]string{}, argv...) + if len(remaining) >= 1 && remaining[0] == "--json" { + jsonOutput = true + remaining = remaining[1:] + } + if len(remaining) >= 1 && remaining[0] == "--" { + remaining = remaining[1:] + } + return jsonOutput, remaining +} + +// extractRootJSONFlag consumes a leading root `--json` flag only when it +// precedes an unknown bare command, matching Python's _extract_root_json_flag. +func extractRootJSONFlag(argv []string) (bool, []string) { + if len(argv) >= 2 && argv[0] == "--json" && !isKnownRootCommand(argv[1]) { + return true, argv[1:] + } + return false, argv +} + +func isUnknownBareCommand(argv []string) bool { + if len(argv) == 0 { + return false + } + first := argv[0] + return !isKnownRootCommand(first) && !strings.HasPrefix(first, "-") +} + +func isKnownRootCommand(name string) bool { + _, ok := knownRootCommands[name] + return ok +} + +func isHelpRequest(argv []string) bool { + for _, v := range argv { + if v == "--help" || v == "-h" { + return true + } + } + return false +} + +// withEnv adapts a command handler that needs a validated Agent Stub +// environment into a cobra RunE, resolving the environment before dispatch. +func withEnv(fn func(env *agentcli.Environment, args []string, cmd *cobra.Command) error) func(*cobra.Command, []string) error { + return func(cmd *cobra.Command, args []string) error { + env, err := agentcli.ReadEnvironment() + if err != nil { + return err + } + return fn(env, args, cmd) + } +} diff --git a/dify-agent-runtime/cmd/dify-agent-cli/main_test.go b/dify-agent-runtime/cmd/dify-agent-cli/main_test.go new file mode 100644 index 00000000000..89b2677e96d --- /dev/null +++ b/dify-agent-runtime/cmd/dify-agent-cli/main_test.go @@ -0,0 +1,228 @@ +// Tests that `--help` renders correctly for every command and subcommand in +// the cobra tree, including the hidden singular `config skill|file` aliases. +// These lock the command surface so it stays identical to the original Python +// (Typer) CLI: each help output must show the command's usage path, its +// Short description, and its expected flags/subcommands. +package main + +import ( + "bytes" + "strings" + "testing" + + "github.com/spf13/cobra" +) + +// executeHelp builds a fresh root command, runs it with the given args, and +// returns the captured stdout/stderr. A fresh tree per call avoids leaking +// parsed flag state between cases. +func executeHelp(t *testing.T, args ...string) string { + t.Helper() + root := newRootCommand() + var buf bytes.Buffer + root.SetOut(&buf) + root.SetErr(&buf) + root.SetArgs(args) + if err := root.Execute(); err != nil { + t.Fatalf("execute %v: %v", args, err) + } + return buf.String() +} + +func TestCommandHelp(t *testing.T) { + cases := []struct { + name string + args []string + want []string + }{ + { + name: "root", + args: []string{"--help"}, + want: []string{"Usage:", "dify-agent", "config", "connect", "drive", "file"}, + }, + { + name: "connect", + args: []string{"connect", "--help"}, + want: []string{"dify-agent connect", "Establish one Agent Stub connection", "--json"}, + }, + { + name: "file", + args: []string{"file", "--help"}, + want: []string{"dify-agent file", "upload", "download"}, + }, + { + name: "file upload", + args: []string{"file", "upload", "--help"}, + want: []string{"dify-agent file upload", "Upload one sandbox-local file"}, + }, + { + name: "file download", + args: []string{"file", "download", "--help"}, + want: []string{"dify-agent file download", "Download one workflow file", "--to"}, + }, + { + name: "drive", + args: []string{"drive", "--help"}, + want: []string{"dify-agent drive", "list", "pull", "push"}, + }, + { + name: "drive list", + args: []string{"drive", "list", "--help"}, + want: []string{"dify-agent drive list", "List drive files", "--json"}, + }, + { + name: "drive pull", + args: []string{"drive", "pull", "--help"}, + want: []string{"dify-agent drive pull", "Pull one or more drive", "--to", "--json"}, + }, + { + name: "drive push", + args: []string{"drive", "push", "--help"}, + want: []string{"dify-agent drive push", "Upload one local file or directory", "--kind", "--json"}, + }, + { + name: "config", + args: []string{"config", "--help"}, + want: []string{"dify-agent config", "manifest", "skills", "files", "env", "note"}, + }, + { + name: "config manifest", + args: []string{"config", "manifest", "--help"}, + want: []string{"dify-agent config manifest", "Show the current visible Agent config manifest"}, + }, + { + name: "config skills", + args: []string{"config", "skills", "--help"}, + want: []string{"dify-agent config skills", "pull", "push", "delete"}, + }, + { + name: "config skills pull", + args: []string{"config", "skills", "pull", "--help"}, + want: []string{"dify-agent config skills pull", "Pull one or all visible config skills", "--to", "--json"}, + }, + { + name: "config skills push", + args: []string{"config", "skills", "push", "--help"}, + want: []string{"dify-agent config skills push", "Upload one or more local skill directories"}, + }, + { + name: "config skills delete", + args: []string{"config", "skills", "delete", "--help"}, + want: []string{"dify-agent config skills delete", "Delete one or more config skills"}, + }, + { + name: "config files", + args: []string{"config", "files", "--help"}, + want: []string{"dify-agent config files", "pull", "push", "delete"}, + }, + { + name: "config files pull", + args: []string{"config", "files", "pull", "--help"}, + want: []string{"dify-agent config files pull", "Pull one or all visible config files", "--to", "--json"}, + }, + { + name: "config files push", + args: []string{"config", "files", "push", "--help"}, + want: []string{"dify-agent config files push", "Upload one or more local files"}, + }, + { + name: "config files delete", + args: []string{"config", "files", "delete", "--help"}, + want: []string{"dify-agent config files delete", "Delete one or more config files"}, + }, + { + name: "config env", + args: []string{"config", "env", "--help"}, + want: []string{"dify-agent config env", "push"}, + }, + { + name: "config env push", + args: []string{"config", "env", "push", "--help"}, + want: []string{"dify-agent config env push", "Set or delete config env entries"}, + }, + { + name: "config note", + args: []string{"config", "note", "--help"}, + want: []string{"dify-agent config note", "pull", "push"}, + }, + { + name: "config note pull", + args: []string{"config", "note", "pull", "--help"}, + want: []string{"dify-agent config note pull", "Export the current config note", "--to"}, + }, + { + name: "config note push", + args: []string{"config", "note", "push", "--help"}, + want: []string{"dify-agent config note push", "Replace the current config note"}, + }, + { + name: "config skill pull (hidden alias)", + args: []string{"config", "skill", "pull", "--help"}, + want: []string{"dify-agent config skill pull", "Pull one or all visible config skills", "--to", "--json"}, + }, + { + name: "config file pull (hidden alias)", + args: []string{"config", "file", "pull", "--help"}, + want: []string{"dify-agent config file pull", "Pull one or all visible config files", "--to", "--json"}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out := executeHelp(t, tc.args...) + for _, want := range tc.want { + if !strings.Contains(out, want) { + t.Errorf("help for %v missing %q\n--- help output ---\n%s", tc.args, want, out) + } + } + }) + } +} + +// findCommand walks the cobra tree by command name (the first word of Use), +// returning nil when any path segment is missing. +func findCommand(root *cobra.Command, path ...string) *cobra.Command { + current := root + for _, name := range path { + var next *cobra.Command + for _, child := range current.Commands() { + if child.Name() == name { + next = child + break + } + } + if next == nil { + return nil + } + current = next + } + return current +} + +// TestHiddenAliasesRemainHidden verifies the singular `config skill|file` +// aliases exist (so their pull help still renders) but stay hidden from the +// parent's command listing, matching the Python CLI's hidden aliases. +func TestHiddenAliasesRemainHidden(t *testing.T) { + root := newRootCommand() + for _, alias := range [][]string{{"config", "skill"}, {"config", "file"}} { + cmd := findCommand(root, alias...) + if cmd == nil { + t.Errorf("hidden alias %v not registered", alias) + continue + } + if !cmd.Hidden { + t.Errorf("alias %v must be hidden", alias) + } + if findCommand(cmd, "pull") == nil { + t.Errorf("alias %v must expose a pull subcommand", alias) + } + } + + // The hidden singular aliases must not appear in `config --help`. + configHelp := executeHelp(t, "config", "--help") + for _, unwanted := range []string{" skill ", " file "} { + if strings.Contains(configHelp, unwanted) { + t.Errorf("config help unexpectedly lists hidden alias %q\n--- help output ---\n%s", unwanted, configHelp) + } + } +} diff --git a/dify-agent-runtime/cmd/runner-exit/main.go b/dify-agent-runtime/cmd/runner-exit/main.go new file mode 100644 index 00000000000..b1eb2344a8b --- /dev/null +++ b/dify-agent-runtime/cmd/runner-exit/main.go @@ -0,0 +1,30 @@ +// shellctl-runner-exit persists a drained runner exit into the shellctl SQLite DB. +// It runs after the tmux output pipe reaches EOF and output.log is fully flushed. +package main + +import ( + "flag" + "fmt" + "os" + + runnerexit "github.com/langgenius/dify/dify-agent-runtime/internal/runner_exit" +) + +func main() { + stateDir := flag.String("state-dir", "", "shellctl state directory") + jobID := flag.String("job-id", "", "job identifier") + exitCode := flag.Int("exit-code", 0, "runner process exit code") + endedAt := flag.String("ended-at", "", "ISO-8601 ended_at timestamp") + busyTimeoutMs := flag.Int("sqlite-busy-timeout-ms", 5000, "SQLite busy timeout in milliseconds") + flag.Parse() + + if *stateDir == "" || *jobID == "" || *endedAt == "" { + fmt.Fprintf(os.Stderr, "shellctl-runner-exit: --state-dir, --job-id, and --ended-at are required\n") + os.Exit(1) + } + + if err := runnerexit.RecordRunnerExit(*stateDir, *jobID, *exitCode, *endedAt, *busyTimeoutMs); err != nil { + fmt.Fprintf(os.Stderr, "shellctl-runner-exit: %v\n", err) + os.Exit(1) + } +} diff --git a/dify-agent-runtime/cmd/sanitize-pty/main.go b/dify-agent-runtime/cmd/sanitize-pty/main.go new file mode 100644 index 00000000000..298071f2f19 --- /dev/null +++ b/dify-agent-runtime/cmd/sanitize-pty/main.go @@ -0,0 +1,22 @@ +// shellctl-sanitize-pty is a stdin→stdout PTY sanitizer used by tmux pipe-pane. +// It strips ANSI escape sequences, normalizes carriage-return progress lines, +// and performs incremental UTF-8 decoding. +package main + +import ( + "flag" + "fmt" + "os" + + "github.com/langgenius/dify/dify-agent-runtime/internal/sanitize" +) + +func main() { + readyFile := flag.String("ready-file", "", "path to touch before reading stdin") + flag.Parse() + + if err := sanitize.Run(*readyFile, os.Stdin, os.Stdout); err != nil { + fmt.Fprintf(os.Stderr, "shellctl-sanitize-pty: %v\n", err) + os.Exit(1) + } +} diff --git a/dify-agent-runtime/cmd/shellctl/main.go b/dify-agent-runtime/cmd/shellctl/main.go new file mode 100644 index 00000000000..ca08d3a6c23 --- /dev/null +++ b/dify-agent-runtime/cmd/shellctl/main.go @@ -0,0 +1,89 @@ +// shellctl is the main server binary that exposes a REST API for managing +// tmux-backed shell jobs inside a sandbox container. +package main + +import ( + "context" + "flag" + "fmt" + "log" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/server" +) + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "shellctl: %v\n", err) + os.Exit(1) + } +} + +func run() error { + // Parse CLI flags + serveCmd := flag.NewFlagSet("serve", flag.ExitOnError) + listen := serveCmd.String("listen", "", "address to listen on (host:port)") + stateDir := serveCmd.String("state-dir", "", "state directory path") + token := serveCmd.String("token", "", "bearer auth token") + + if len(os.Args) < 2 || os.Args[1] != "serve" { + fmt.Fprintf(os.Stderr, "Usage: shellctl serve [flags]\n") + return nil + } + if err := serveCmd.Parse(os.Args[2:]); err != nil { + return err + } + + // Build config + config := server.DefaultConfig() + if *listen != "" { + config.Listen = *listen + } + if *stateDir != "" { + config.StateDir = *stateDir + } + if *token != "" { + config.AuthToken = *token + } + + // Initialize service + svc := server.NewService(config) + if err := svc.Initialize(); err != nil { + return fmt.Errorf("initialize: %w", err) + } + defer svc.Shutdown() + + svc.StartBackgroundGC() + svc.StartBackgroundPipeMonitor() + + // Create HTTP server + handler := server.Handler(svc, config) + srv := &http.Server{ + Addr: config.Listen, + Handler: handler, + ReadTimeout: 0, // Long-poll requires no read timeout + WriteTimeout: 0, + } + + // Graceful shutdown on signal + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + go func() { + <-ctx.Done() + log.Println("shutting down...") + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = srv.Shutdown(shutdownCtx) + }() + + log.Printf("shellctl serving on %s", config.Listen) + if err := srv.ListenAndServe(); err != http.ErrServerClosed { + return err + } + return nil +} diff --git a/dify-agent/docker/local-sandbox/Dockerfile b/dify-agent-runtime/docker/Dockerfile similarity index 53% rename from dify-agent/docker/local-sandbox/Dockerfile rename to dify-agent-runtime/docker/Dockerfile index 9996f1f53d7..97f8ee81c99 100644 --- a/dify-agent/docker/local-sandbox/Dockerfile +++ b/dify-agent-runtime/docker/Dockerfile @@ -1,18 +1,27 @@ -# Local sandbox image for shellctl-managed Dify Agent workspaces. +# Go shellctl sandbox image — provides a sandboxed runtime environment +# with common dev tools (Node.js, Python, pnpm, uv) for agent-managed jobs. # -# Build this from the dify-agent package root: -# docker build -f docker/local-sandbox/Dockerfile -t dify-agent-local-sandbox:local . -# -# This image merges the former shellctl-only image with the sandbox-visible -# Agent Stub client CLI. It runs `shellctl serve` by default, and shellctl-managed jobs -# can call `dify-agent ...` without installing extra packages at runtime. +# Build: +# docker build -f docker/Dockerfile -t dify-agent-local-sandbox:local . -FROM python:3.12-slim-bookworm AS base +# ── Go build stage ─────────────────────────────────────────────────────────── +FROM golang:1.26 AS go-builder + +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY . . +RUN CGO_ENABLED=0 go build -o /bin/shellctl ./cmd/shellctl && \ + CGO_ENABLED=0 go build -o /bin/shellctl-sanitize-pty ./cmd/sanitize-pty && \ + CGO_ENABLED=0 go build -o /bin/shellctl-runner-exit ./cmd/runner-exit && \ + CGO_ENABLED=0 go build -o /bin/dify-agent ./cmd/dify-agent-cli + +# ── Runtime stage ──────────────────────────────────────────────────────────── +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 -ARG DIFY_AGENT_TOOL_SPEC=.[grpc,shellctl-server] ENV PYTHONDONTWRITEBYTECODE=1 \ PYTHONUNBUFFERED=1 \ @@ -52,31 +61,11 @@ RUN apt-get update \ RUN python -m pip install --no-cache-dir "uv==${UV_VERSION}" -WORKDIR /opt/dify-agent - - -FROM base AS tools - -ARG DIFY_AGENT_TOOL_SPEC - -COPY pyproject.toml uv.lock README.md ./ -COPY src ./src - -RUN uv export --frozen --no-dev --all-extras --no-emit-project --no-hashes \ - > /tmp/dify-agent-constraints.txt \ - && UV_TOOL_DIR=/opt/dify-agent-tools/envs \ - UV_TOOL_BIN_DIR=/opt/dify-agent-tools/bin \ - uv tool install --force --python /usr/local/bin/python --no-python-downloads \ - --constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${DIFY_AGENT_TOOL_SPEC}" \ - && rm -f /tmp/dify-agent-constraints.txt - - -FROM base AS production - -ENV PATH="/opt/dify-agent-tools/bin:${PATH}" - -COPY --from=tools /opt/dify-agent-tools/envs /opt/dify-agent-tools/envs -COPY --from=tools /opt/dify-agent-tools/bin /opt/dify-agent-tools/bin +# Go binaries +COPY --from=go-builder /bin/shellctl /usr/local/bin/shellctl +COPY --from=go-builder /bin/shellctl-sanitize-pty /usr/local/bin/shellctl-sanitize-pty +COPY --from=go-builder /bin/shellctl-runner-exit /usr/local/bin/shellctl-runner-exit +COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent RUN useradd --create-home --shell /bin/sh dify \ && mkdir -p /mnt/drive \ diff --git a/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub.pb.go b/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub.pb.go new file mode 100644 index 00000000000..4cfc11903d8 --- /dev/null +++ b/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub.pb.go @@ -0,0 +1,515 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v5.28.3 +// source: dify/agent/stub/v1/agent_stub.proto + +package stubv1 + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type ConnectRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + ProtocolVersion int32 `protobuf:"varint,1,opt,name=protocol_version,json=protocolVersion,proto3" json:"protocol_version,omitempty"` + Argv []string `protobuf:"bytes,2,rep,name=argv,proto3" json:"argv,omitempty"` + MetadataJson string `protobuf:"bytes,3,opt,name=metadata_json,json=metadataJson,proto3" json:"metadata_json,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConnectRequest) Reset() { + *x = ConnectRequest{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConnectRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConnectRequest) ProtoMessage() {} + +func (x *ConnectRequest) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConnectRequest.ProtoReflect.Descriptor instead. +func (*ConnectRequest) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{0} +} + +func (x *ConnectRequest) GetProtocolVersion() int32 { + if x != nil { + return x.ProtocolVersion + } + return 0 +} + +func (x *ConnectRequest) GetArgv() []string { + if x != nil { + return x.Argv + } + return nil +} + +func (x *ConnectRequest) GetMetadataJson() string { + if x != nil { + return x.MetadataJson + } + return "" +} + +type ConnectResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + ConnectionId string `protobuf:"bytes,1,opt,name=connection_id,json=connectionId,proto3" json:"connection_id,omitempty"` + Status string `protobuf:"bytes,2,opt,name=status,proto3" json:"status,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConnectResponse) Reset() { + *x = ConnectResponse{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConnectResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConnectResponse) ProtoMessage() {} + +func (x *ConnectResponse) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConnectResponse.ProtoReflect.Descriptor instead. +func (*ConnectResponse) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{1} +} + +func (x *ConnectResponse) GetConnectionId() string { + if x != nil { + return x.ConnectionId + } + return "" +} + +func (x *ConnectResponse) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +type FileUploadRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Filename string `protobuf:"bytes,1,opt,name=filename,proto3" json:"filename,omitempty"` + Mimetype string `protobuf:"bytes,2,opt,name=mimetype,proto3" json:"mimetype,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FileUploadRequest) Reset() { + *x = FileUploadRequest{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FileUploadRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FileUploadRequest) ProtoMessage() {} + +func (x *FileUploadRequest) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FileUploadRequest.ProtoReflect.Descriptor instead. +func (*FileUploadRequest) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{2} +} + +func (x *FileUploadRequest) GetFilename() string { + if x != nil { + return x.Filename + } + return "" +} + +func (x *FileUploadRequest) GetMimetype() string { + if x != nil { + return x.Mimetype + } + return "" +} + +type FileUploadResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + UploadUrl string `protobuf:"bytes,1,opt,name=upload_url,json=uploadUrl,proto3" json:"upload_url,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FileUploadResponse) Reset() { + *x = FileUploadResponse{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FileUploadResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FileUploadResponse) ProtoMessage() {} + +func (x *FileUploadResponse) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FileUploadResponse.ProtoReflect.Descriptor instead. +func (*FileUploadResponse) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{3} +} + +func (x *FileUploadResponse) GetUploadUrl() string { + if x != nil { + return x.UploadUrl + } + return "" +} + +type FileMapping struct { + state protoimpl.MessageState `protogen:"open.v1"` + TransferMethod string `protobuf:"bytes,1,opt,name=transfer_method,json=transferMethod,proto3" json:"transfer_method,omitempty"` + Reference *string `protobuf:"bytes,2,opt,name=reference,proto3,oneof" json:"reference,omitempty"` + Url *string `protobuf:"bytes,3,opt,name=url,proto3,oneof" json:"url,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FileMapping) Reset() { + *x = FileMapping{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FileMapping) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FileMapping) ProtoMessage() {} + +func (x *FileMapping) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FileMapping.ProtoReflect.Descriptor instead. +func (*FileMapping) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{4} +} + +func (x *FileMapping) GetTransferMethod() string { + if x != nil { + return x.TransferMethod + } + return "" +} + +func (x *FileMapping) GetReference() string { + if x != nil && x.Reference != nil { + return *x.Reference + } + return "" +} + +func (x *FileMapping) GetUrl() string { + if x != nil && x.Url != nil { + return *x.Url + } + return "" +} + +type FileDownloadRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + File *FileMapping `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"` + ForExternal *bool `protobuf:"varint,2,opt,name=for_external,json=forExternal,proto3,oneof" json:"for_external,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FileDownloadRequest) Reset() { + *x = FileDownloadRequest{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FileDownloadRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FileDownloadRequest) ProtoMessage() {} + +func (x *FileDownloadRequest) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FileDownloadRequest.ProtoReflect.Descriptor instead. +func (*FileDownloadRequest) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{5} +} + +func (x *FileDownloadRequest) GetFile() *FileMapping { + if x != nil { + return x.File + } + return nil +} + +func (x *FileDownloadRequest) GetForExternal() bool { + if x != nil && x.ForExternal != nil { + return *x.ForExternal + } + return false +} + +type FileDownloadResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Filename string `protobuf:"bytes,1,opt,name=filename,proto3" json:"filename,omitempty"` + MimeType *string `protobuf:"bytes,2,opt,name=mime_type,json=mimeType,proto3,oneof" json:"mime_type,omitempty"` + Size int64 `protobuf:"varint,3,opt,name=size,proto3" json:"size,omitempty"` + DownloadUrl string `protobuf:"bytes,4,opt,name=download_url,json=downloadUrl,proto3" json:"download_url,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FileDownloadResponse) Reset() { + *x = FileDownloadResponse{} + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FileDownloadResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FileDownloadResponse) ProtoMessage() {} + +func (x *FileDownloadResponse) ProtoReflect() protoreflect.Message { + mi := &file_dify_agent_stub_v1_agent_stub_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FileDownloadResponse.ProtoReflect.Descriptor instead. +func (*FileDownloadResponse) Descriptor() ([]byte, []int) { + return file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP(), []int{6} +} + +func (x *FileDownloadResponse) GetFilename() string { + if x != nil { + return x.Filename + } + return "" +} + +func (x *FileDownloadResponse) GetMimeType() string { + if x != nil && x.MimeType != nil { + return *x.MimeType + } + return "" +} + +func (x *FileDownloadResponse) GetSize() int64 { + if x != nil { + return x.Size + } + return 0 +} + +func (x *FileDownloadResponse) GetDownloadUrl() string { + if x != nil { + return x.DownloadUrl + } + return "" +} + +var File_dify_agent_stub_v1_agent_stub_proto protoreflect.FileDescriptor + +const file_dify_agent_stub_v1_agent_stub_proto_rawDesc = "" + + "\n" + + "#dify/agent/stub/v1/agent_stub.proto\x12\x12dify.agent.stub.v1\"t\n" + + "\x0eConnectRequest\x12)\n" + + "\x10protocol_version\x18\x01 \x01(\x05R\x0fprotocolVersion\x12\x12\n" + + "\x04argv\x18\x02 \x03(\tR\x04argv\x12#\n" + + "\rmetadata_json\x18\x03 \x01(\tR\fmetadataJson\"N\n" + + "\x0fConnectResponse\x12#\n" + + "\rconnection_id\x18\x01 \x01(\tR\fconnectionId\x12\x16\n" + + "\x06status\x18\x02 \x01(\tR\x06status\"K\n" + + "\x11FileUploadRequest\x12\x1a\n" + + "\bfilename\x18\x01 \x01(\tR\bfilename\x12\x1a\n" + + "\bmimetype\x18\x02 \x01(\tR\bmimetype\"3\n" + + "\x12FileUploadResponse\x12\x1d\n" + + "\n" + + "upload_url\x18\x01 \x01(\tR\tuploadUrl\"\x86\x01\n" + + "\vFileMapping\x12'\n" + + "\x0ftransfer_method\x18\x01 \x01(\tR\x0etransferMethod\x12!\n" + + "\treference\x18\x02 \x01(\tH\x00R\treference\x88\x01\x01\x12\x15\n" + + "\x03url\x18\x03 \x01(\tH\x01R\x03url\x88\x01\x01B\f\n" + + "\n" + + "_referenceB\x06\n" + + "\x04_url\"\x83\x01\n" + + "\x13FileDownloadRequest\x123\n" + + "\x04file\x18\x01 \x01(\v2\x1f.dify.agent.stub.v1.FileMappingR\x04file\x12&\n" + + "\ffor_external\x18\x02 \x01(\bH\x00R\vforExternal\x88\x01\x01B\x0f\n" + + "\r_for_external\"\x99\x01\n" + + "\x14FileDownloadResponse\x12\x1a\n" + + "\bfilename\x18\x01 \x01(\tR\bfilename\x12 \n" + + "\tmime_type\x18\x02 \x01(\tH\x00R\bmimeType\x88\x01\x01\x12\x12\n" + + "\x04size\x18\x03 \x01(\x03R\x04size\x12!\n" + + "\fdownload_url\x18\x04 \x01(\tR\vdownloadUrlB\f\n" + + "\n" + + "_mime_type2\xc0\x02\n" + + "\x10AgentStubService\x12R\n" + + "\aConnect\x12\".dify.agent.stub.v1.ConnectRequest\x1a#.dify.agent.stub.v1.ConnectResponse\x12h\n" + + "\x17CreateFileUploadRequest\x12%.dify.agent.stub.v1.FileUploadRequest\x1a&.dify.agent.stub.v1.FileUploadResponse\x12n\n" + + "\x19CreateFileDownloadRequest\x12'.dify.agent.stub.v1.FileDownloadRequest\x1a(.dify.agent.stub.v1.FileDownloadResponseB\x1bZ\x19dify/agent/stub/v1;stubv1b\x06proto3" + +var ( + file_dify_agent_stub_v1_agent_stub_proto_rawDescOnce sync.Once + file_dify_agent_stub_v1_agent_stub_proto_rawDescData []byte +) + +func file_dify_agent_stub_v1_agent_stub_proto_rawDescGZIP() []byte { + file_dify_agent_stub_v1_agent_stub_proto_rawDescOnce.Do(func() { + file_dify_agent_stub_v1_agent_stub_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_dify_agent_stub_v1_agent_stub_proto_rawDesc), len(file_dify_agent_stub_v1_agent_stub_proto_rawDesc))) + }) + return file_dify_agent_stub_v1_agent_stub_proto_rawDescData +} + +var file_dify_agent_stub_v1_agent_stub_proto_msgTypes = make([]protoimpl.MessageInfo, 7) +var file_dify_agent_stub_v1_agent_stub_proto_goTypes = []any{ + (*ConnectRequest)(nil), // 0: dify.agent.stub.v1.ConnectRequest + (*ConnectResponse)(nil), // 1: dify.agent.stub.v1.ConnectResponse + (*FileUploadRequest)(nil), // 2: dify.agent.stub.v1.FileUploadRequest + (*FileUploadResponse)(nil), // 3: dify.agent.stub.v1.FileUploadResponse + (*FileMapping)(nil), // 4: dify.agent.stub.v1.FileMapping + (*FileDownloadRequest)(nil), // 5: dify.agent.stub.v1.FileDownloadRequest + (*FileDownloadResponse)(nil), // 6: dify.agent.stub.v1.FileDownloadResponse +} +var file_dify_agent_stub_v1_agent_stub_proto_depIdxs = []int32{ + 4, // 0: dify.agent.stub.v1.FileDownloadRequest.file:type_name -> dify.agent.stub.v1.FileMapping + 0, // 1: dify.agent.stub.v1.AgentStubService.Connect:input_type -> dify.agent.stub.v1.ConnectRequest + 2, // 2: dify.agent.stub.v1.AgentStubService.CreateFileUploadRequest:input_type -> dify.agent.stub.v1.FileUploadRequest + 5, // 3: dify.agent.stub.v1.AgentStubService.CreateFileDownloadRequest:input_type -> dify.agent.stub.v1.FileDownloadRequest + 1, // 4: dify.agent.stub.v1.AgentStubService.Connect:output_type -> dify.agent.stub.v1.ConnectResponse + 3, // 5: dify.agent.stub.v1.AgentStubService.CreateFileUploadRequest:output_type -> dify.agent.stub.v1.FileUploadResponse + 6, // 6: dify.agent.stub.v1.AgentStubService.CreateFileDownloadRequest:output_type -> dify.agent.stub.v1.FileDownloadResponse + 4, // [4:7] is the sub-list for method output_type + 1, // [1:4] is the sub-list for method input_type + 1, // [1:1] is the sub-list for extension type_name + 1, // [1:1] is the sub-list for extension extendee + 0, // [0:1] is the sub-list for field type_name +} + +func init() { file_dify_agent_stub_v1_agent_stub_proto_init() } +func file_dify_agent_stub_v1_agent_stub_proto_init() { + if File_dify_agent_stub_v1_agent_stub_proto != nil { + return + } + file_dify_agent_stub_v1_agent_stub_proto_msgTypes[4].OneofWrappers = []any{} + file_dify_agent_stub_v1_agent_stub_proto_msgTypes[5].OneofWrappers = []any{} + file_dify_agent_stub_v1_agent_stub_proto_msgTypes[6].OneofWrappers = []any{} + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_dify_agent_stub_v1_agent_stub_proto_rawDesc), len(file_dify_agent_stub_v1_agent_stub_proto_rawDesc)), + NumEnums: 0, + NumMessages: 7, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_dify_agent_stub_v1_agent_stub_proto_goTypes, + DependencyIndexes: file_dify_agent_stub_v1_agent_stub_proto_depIdxs, + MessageInfos: file_dify_agent_stub_v1_agent_stub_proto_msgTypes, + }.Build() + File_dify_agent_stub_v1_agent_stub_proto = out.File + file_dify_agent_stub_v1_agent_stub_proto_goTypes = nil + file_dify_agent_stub_v1_agent_stub_proto_depIdxs = nil +} diff --git a/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub_grpc.pb.go b/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub_grpc.pb.go new file mode 100644 index 00000000000..f7b649df503 --- /dev/null +++ b/dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub_grpc.pb.go @@ -0,0 +1,197 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.6.2 +// - protoc v5.28.3 +// source: dify/agent/stub/v1/agent_stub.proto + +package stubv1 + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.64.0 or later. +const _ = grpc.SupportPackageIsVersion9 + +const ( + AgentStubService_Connect_FullMethodName = "/dify.agent.stub.v1.AgentStubService/Connect" + AgentStubService_CreateFileUploadRequest_FullMethodName = "/dify.agent.stub.v1.AgentStubService/CreateFileUploadRequest" + AgentStubService_CreateFileDownloadRequest_FullMethodName = "/dify.agent.stub.v1.AgentStubService/CreateFileDownloadRequest" +) + +// AgentStubServiceClient is the client API for AgentStubService service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type AgentStubServiceClient interface { + Connect(ctx context.Context, in *ConnectRequest, opts ...grpc.CallOption) (*ConnectResponse, error) + CreateFileUploadRequest(ctx context.Context, in *FileUploadRequest, opts ...grpc.CallOption) (*FileUploadResponse, error) + CreateFileDownloadRequest(ctx context.Context, in *FileDownloadRequest, opts ...grpc.CallOption) (*FileDownloadResponse, error) +} + +type agentStubServiceClient struct { + cc grpc.ClientConnInterface +} + +func NewAgentStubServiceClient(cc grpc.ClientConnInterface) AgentStubServiceClient { + return &agentStubServiceClient{cc} +} + +func (c *agentStubServiceClient) Connect(ctx context.Context, in *ConnectRequest, opts ...grpc.CallOption) (*ConnectResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ConnectResponse) + err := c.cc.Invoke(ctx, AgentStubService_Connect_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *agentStubServiceClient) CreateFileUploadRequest(ctx context.Context, in *FileUploadRequest, opts ...grpc.CallOption) (*FileUploadResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(FileUploadResponse) + err := c.cc.Invoke(ctx, AgentStubService_CreateFileUploadRequest_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *agentStubServiceClient) CreateFileDownloadRequest(ctx context.Context, in *FileDownloadRequest, opts ...grpc.CallOption) (*FileDownloadResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(FileDownloadResponse) + err := c.cc.Invoke(ctx, AgentStubService_CreateFileDownloadRequest_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// AgentStubServiceServer is the server API for AgentStubService service. +// All implementations must embed UnimplementedAgentStubServiceServer +// for forward compatibility. +type AgentStubServiceServer interface { + Connect(context.Context, *ConnectRequest) (*ConnectResponse, error) + CreateFileUploadRequest(context.Context, *FileUploadRequest) (*FileUploadResponse, error) + CreateFileDownloadRequest(context.Context, *FileDownloadRequest) (*FileDownloadResponse, error) + mustEmbedUnimplementedAgentStubServiceServer() +} + +// UnimplementedAgentStubServiceServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedAgentStubServiceServer struct{} + +func (UnimplementedAgentStubServiceServer) Connect(context.Context, *ConnectRequest) (*ConnectResponse, error) { + return nil, status.Error(codes.Unimplemented, "method Connect not implemented") +} +func (UnimplementedAgentStubServiceServer) CreateFileUploadRequest(context.Context, *FileUploadRequest) (*FileUploadResponse, error) { + return nil, status.Error(codes.Unimplemented, "method CreateFileUploadRequest not implemented") +} +func (UnimplementedAgentStubServiceServer) CreateFileDownloadRequest(context.Context, *FileDownloadRequest) (*FileDownloadResponse, error) { + return nil, status.Error(codes.Unimplemented, "method CreateFileDownloadRequest not implemented") +} +func (UnimplementedAgentStubServiceServer) mustEmbedUnimplementedAgentStubServiceServer() {} +func (UnimplementedAgentStubServiceServer) testEmbeddedByValue() {} + +// UnsafeAgentStubServiceServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to AgentStubServiceServer will +// result in compilation errors. +type UnsafeAgentStubServiceServer interface { + mustEmbedUnimplementedAgentStubServiceServer() +} + +func RegisterAgentStubServiceServer(s grpc.ServiceRegistrar, srv AgentStubServiceServer) { + // If the following call panics, it indicates UnimplementedAgentStubServiceServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&AgentStubService_ServiceDesc, srv) +} + +func _AgentStubService_Connect_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ConnectRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AgentStubServiceServer).Connect(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AgentStubService_Connect_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AgentStubServiceServer).Connect(ctx, req.(*ConnectRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AgentStubService_CreateFileUploadRequest_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(FileUploadRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AgentStubServiceServer).CreateFileUploadRequest(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AgentStubService_CreateFileUploadRequest_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AgentStubServiceServer).CreateFileUploadRequest(ctx, req.(*FileUploadRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AgentStubService_CreateFileDownloadRequest_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(FileDownloadRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AgentStubServiceServer).CreateFileDownloadRequest(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AgentStubService_CreateFileDownloadRequest_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AgentStubServiceServer).CreateFileDownloadRequest(ctx, req.(*FileDownloadRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// AgentStubService_ServiceDesc is the grpc.ServiceDesc for AgentStubService service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var AgentStubService_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "dify.agent.stub.v1.AgentStubService", + HandlerType: (*AgentStubServiceServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "Connect", + Handler: _AgentStubService_Connect_Handler, + }, + { + MethodName: "CreateFileUploadRequest", + Handler: _AgentStubService_CreateFileUploadRequest_Handler, + }, + { + MethodName: "CreateFileDownloadRequest", + Handler: _AgentStubService_CreateFileDownloadRequest_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "dify/agent/stub/v1/agent_stub.proto", +} diff --git a/dify-agent-runtime/go.mod b/dify-agent-runtime/go.mod new file mode 100644 index 00000000000..80f71bfb98b --- /dev/null +++ b/dify-agent-runtime/go.mod @@ -0,0 +1,28 @@ +module github.com/langgenius/dify/dify-agent-runtime + +go 1.26 + +require ( + github.com/spf13/cobra v1.10.2 + google.golang.org/grpc v1.82.0 + google.golang.org/protobuf v1.36.11 + modernc.org/sqlite v1.37.1 +) + +require ( + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/ncruces/go-strftime v0.1.9 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/spf13/pflag v1.0.9 // indirect + golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 // indirect + golang.org/x/net v0.53.0 // indirect + golang.org/x/sys v0.43.0 // indirect + golang.org/x/text v0.36.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect + modernc.org/libc v1.65.7 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect +) diff --git a/dify-agent-runtime/go.sum b/dify-agent-runtime/go.sum new file mode 100644 index 00000000000..91aa0612cef --- /dev/null +++ b/dify-agent-runtime/go.sum @@ -0,0 +1,91 @@ +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= +github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= +github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= +github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= +go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0= +go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= +go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= +go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= +go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= +go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= +go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= +go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 h1:R84qjqJb5nVJMxqWYb3np9L5ZsaDtB+a39EqjV0JSUM= +golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0/go.mod h1:S9Xr4PYopiDyqSyp5NjCrhFrqg6A5zA2E/iPHPhqnS8= +golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI= +golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY= +golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA= +golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= +golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg= +golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164= +golang.org/x/tools v0.43.0 h1:12BdW9CeB3Z+J/I/wj34VMl8X+fEXBxVR90JeMX5E7s= +golang.org/x/tools v0.43.0/go.mod h1:uHkMso649BX2cZK6+RpuIPXS3ho2hZo4FVwfoy1vIk0= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU= +google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +modernc.org/cc/v4 v4.26.1 h1:+X5NtzVBn0KgsBCBe+xkDC7twLb/jNVj9FPgiwSQO3s= +modernc.org/cc/v4 v4.26.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= +modernc.org/ccgo/v4 v4.28.0 h1:rjznn6WWehKq7dG4JtLRKxb52Ecv8OUGah8+Z/SfpNU= +modernc.org/ccgo/v4 v4.28.0/go.mod h1:JygV3+9AV6SmPhDasu4JgquwU81XAKLd3OKTUDNOiKE= +modernc.org/fileutil v1.3.1 h1:8vq5fe7jdtEvoCf3Zf9Nm0Q05sH6kGx0Op2CPx1wTC8= +modernc.org/fileutil v1.3.1/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/libc v1.65.7 h1:Ia9Z4yzZtWNtUIuiPuQ7Qf7kxYrxP1/jeHZzG8bFu00= +modernc.org/libc v1.65.7/go.mod h1:011EQibzzio/VX3ygj1qGFt5kMjP0lHb0qCW5/D/pQU= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.37.1 h1:EgHJK/FPoqC+q2YBXg7fUmES37pCHFc97sI7zSayBEs= +modernc.org/sqlite v1.37.1/go.mod h1:XwdRtsE1MpiBcL54+MbKcaDvcuej+IYSMfLN6gSKV8g= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/dify-agent-runtime/internal/agentcli/archive.go b/dify-agent-runtime/internal/agentcli/archive.go new file mode 100644 index 00000000000..3a458009d71 --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/archive.go @@ -0,0 +1,137 @@ +package agentcli + +import ( + "archive/zip" + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// createZipArchive creates a zip file at archivePath containing all files in dirPath, +// excluding common transient directories and files. +func createZipArchive(archivePath string, dirPath string) (retErr error) { + outFile, err := os.Create(archivePath) + if err != nil { + return fmt.Errorf("create archive file: %w", err) + } + defer func() { + if cErr := outFile.Close(); cErr != nil && retErr == nil { + retErr = cErr + } + }() + + writer := zip.NewWriter(outFile) + defer func() { + if cErr := writer.Close(); cErr != nil && retErr == nil { + retErr = cErr + } + }() + + skipFiles := map[string]bool{".DS_Store": true, ".DIFY-SKILL-FULL.zip": true} + + return filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + + if info.IsDir() { + if shouldSkipDir(info.Name()) { + return filepath.SkipDir + } + return nil + } + + if skipFiles[info.Name()] { + return nil + } + + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("archive does not support symlinked files: %s", path) + } + + relPath, err := filepath.Rel(dirPath, path) + if err != nil { + return err + } + + header, err := zip.FileInfoHeader(info) + if err != nil { + return err + } + header.Name = filepath.ToSlash(relPath) + header.Method = zip.Deflate + + w, err := writer.CreateHeader(header) + if err != nil { + return err + } + + file, err := os.Open(path) + if err != nil { + return err + } + defer func() { _ = file.Close() }() + + _, err = io.Copy(w, file) + return err + }) +} + +// extractZip extracts a zip archive into targetDir with path safety checks. +func extractZip(archivePath string, targetDir string) error { + r, err := zip.OpenReader(archivePath) + if err != nil { + return fmt.Errorf("open zip: %w", err) + } + defer func() { _ = r.Close() }() + + absTarget, err := filepath.Abs(targetDir) + if err != nil { + return err + } + + for _, f := range r.File { + destPath := filepath.Join(absTarget, f.Name) + + // Safety: prevent path traversal + if !strings.HasPrefix(filepath.Clean(destPath)+string(os.PathSeparator), filepath.Clean(absTarget)+string(os.PathSeparator)) && + filepath.Clean(destPath) != filepath.Clean(absTarget) { + return fmt.Errorf("archive entry escapes target directory: %s", f.Name) + } + + if f.FileInfo().IsDir() { + if err := os.MkdirAll(destPath, 0o755); err != nil { + return err + } + continue + } + + if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil { + return err + } + + rc, err := f.Open() + if err != nil { + return err + } + + outFile, err := os.Create(destPath) + if err != nil { + _ = rc.Close() + return err + } + + _, err = io.Copy(outFile, rc) + cErr := outFile.Close() + _ = rc.Close() + if err != nil { + return err + } + if cErr != nil { + return cErr + } + } + return nil +} diff --git a/dify-agent-runtime/internal/agentcli/client.go b/dify-agent-runtime/internal/agentcli/client.go new file mode 100644 index 00000000000..328386cdac1 --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/client.go @@ -0,0 +1,45 @@ +package agentcli + +import "context" + +// StubClient abstracts Agent Stub control-plane and data-plane operations. +// Business logic depends only on this interface, never on HTTP/gRPC details. +type StubClient interface { + // Control-plane: available via gRPC or HTTP + Connect(ctx context.Context, argv []string, metadataJSON string) (*ConnectResponse, error) + CreateFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) + CreateFileDownloadURL(ctx context.Context, transferMethod string, reference, url *string, forExternal bool) (*FileDownloadResponse, error) + + // Drive operations (HTTP-only control-plane) + GetDriveManifest(ctx context.Context, prefix string, includeDownloadURL bool) (*DriveManifestResponse, error) + CommitDrive(ctx context.Context, items []DriveCommitItem) ([]byte, error) + + // Config operations (HTTP-only control-plane) + GetConfigManifest(ctx context.Context) ([]byte, error) + PullConfigSkill(ctx context.Context, name string) ([]byte, error) + PullConfigFile(ctx context.Context, name string) ([]byte, error) + PushConfig(ctx context.Context, payload any) ([]byte, error) + PatchConfigEnv(ctx context.Context, envText string) ([]byte, error) + PutConfigNote(ctx context.Context, note string) ([]byte, error) + + // Data-plane (always HTTP, signed URLs) + UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error) + DownloadFromURL(downloadURL string) ([]byte, error) + + Close() error +} + +// NewStubClient creates the appropriate client based on the endpoint scheme. +func NewStubClient(env *Environment) (StubClient, error) { + endpoint, err := ParseEndpoint(env.URL) + if err != nil { + return nil, err + } + + httpClient := newHTTPStubClient(env) + + if endpoint.IsGRPC { + return newGRPCStubClient(endpoint, httpClient) + } + return httpClient, nil +} diff --git a/dify-agent-runtime/internal/agentcli/client_grpc.go b/dify-agent-runtime/internal/agentcli/client_grpc.go new file mode 100644 index 00000000000..6a5f14b715d --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/client_grpc.go @@ -0,0 +1,69 @@ +package agentcli + +import ( + "context" + "fmt" + + "github.com/langgenius/dify/dify-agent-runtime/internal/stubclient" +) + +// grpcStubClient implements StubClient with gRPC for Connect/FileUpload/FileDownload +// and delegates all other operations to the embedded HTTP client. +type grpcStubClient struct { + *httpStubClient + grpc *stubclient.Client +} + +func newGRPCStubClient(endpoint *Endpoint, httpClient *httpStubClient) (*grpcStubClient, error) { + target := endpoint.Host + ":" + endpoint.Port + client, err := stubclient.Dial(target) + if err != nil { + return nil, fmt.Errorf("gRPC dial %s: %w", target, err) + } + return &grpcStubClient{ + httpStubClient: httpClient, + grpc: client, + }, nil +} + +func (c *grpcStubClient) Close() error { + return c.grpc.Close() +} + +func (c *grpcStubClient) Connect(ctx context.Context, argv []string, metadataJSON string) (*ConnectResponse, error) { + if metadataJSON == "" { + metadataJSON = "{}" + } + result, err := c.grpc.Connect(ctx, agentStubProtocolVersion, argv, metadataJSON) + if err != nil { + return nil, err + } + return &ConnectResponse{ + ConnectionID: result.ConnectionID, + Status: result.Status, + }, nil +} + +func (c *grpcStubClient) CreateFileUploadURL(ctx context.Context, filename, mimetype string) (string, error) { + result, err := c.grpc.CreateFileUpload(ctx, filename, mimetype) + if err != nil { + return "", err + } + if result.UploadURL == "" { + return "", fmt.Errorf("signed file upload response is missing upload_url") + } + return result.UploadURL, nil +} + +func (c *grpcStubClient) CreateFileDownloadURL(ctx context.Context, transferMethod string, reference, url *string, forExternal bool) (*FileDownloadResponse, error) { + result, err := c.grpc.CreateFileDownload(ctx, transferMethod, reference, url, forExternal) + if err != nil { + return nil, err + } + return &FileDownloadResponse{ + Filename: result.Filename, + MimeType: result.MimeType, + Size: result.Size, + DownloadURL: result.DownloadURL, + }, nil +} diff --git a/dify-agent-runtime/internal/agentcli/client_http.go b/dify-agent-runtime/internal/agentcli/client_http.go new file mode 100644 index 00000000000..2a6031a54fd --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/client_http.go @@ -0,0 +1,218 @@ +package agentcli + +import ( + "context" + "encoding/json" + "fmt" +) + +// httpStubClient implements StubClient using pure HTTP transport. +type httpStubClient struct { + http *HTTPClient +} + +func newHTTPStubClient(env *Environment) *httpStubClient { + return &httpStubClient{http: NewHTTPClient(env)} +} + +func (c *httpStubClient) Close() error { return nil } + +func (c *httpStubClient) Connect(_ context.Context, argv []string, metadataJSON string) (*ConnectResponse, error) { + var metadata any + if metadataJSON != "" { + if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil { + metadata = map[string]any{} + } + } else { + metadata = map[string]any{} + } + + payload := map[string]any{ + "argv": argv, + "metadata": metadata, + } + + body, statusCode, err := c.http.postJSON("/connections", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "connect"); err != nil { + return nil, err + } + + var resp ConnectResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("parse connect response: %w", err) + } + return &resp, nil +} + +func (c *httpStubClient) CreateFileUploadURL(_ context.Context, filename, mimetype string) (string, error) { + payload := map[string]string{ + "filename": filename, + "mimetype": mimetype, + } + body, statusCode, err := c.http.postJSON("/files/upload-request", payload) + if err != nil { + return "", err + } + if err := checkHTTPError(body, statusCode, "file upload request"); err != nil { + return "", err + } + + var resp struct { + UploadURL string `json:"upload_url"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return "", fmt.Errorf("parse upload response: %w", err) + } + if resp.UploadURL == "" { + return "", fmt.Errorf("signed file upload response is missing upload_url") + } + return resp.UploadURL, nil +} + +func (c *httpStubClient) CreateFileDownloadURL(_ context.Context, transferMethod string, reference, url *string, forExternal bool) (*FileDownloadResponse, error) { + fileMapping := map[string]any{ + "transfer_method": transferMethod, + } + if reference != nil { + fileMapping["reference"] = *reference + } + if url != nil { + fileMapping["url"] = *url + } + + payload := map[string]any{ + "file": fileMapping, + "for_external": forExternal, + } + body, statusCode, err := c.http.postJSON("/files/download-request", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "file download request"); err != nil { + return nil, err + } + + var resp FileDownloadResponse + if err := json.Unmarshal(body, &resp); err != nil { + return nil, fmt.Errorf("parse download response: %w", err) + } + return &resp, nil +} + +func (c *httpStubClient) GetDriveManifest(_ context.Context, prefix string, includeDownloadURL bool) (*DriveManifestResponse, error) { + params := map[string]string{ + "prefix": prefix, + } + if includeDownloadURL { + params["include_download_url"] = "true" + } else { + params["include_download_url"] = "false" + } + + body, statusCode, err := c.http.getJSON("/drive/manifest", params) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "drive manifest"); err != nil { + return nil, err + } + + var manifest DriveManifestResponse + if err := json.Unmarshal(body, &manifest); err != nil { + return nil, fmt.Errorf("parse drive manifest: %w", err) + } + return &manifest, nil +} + +func (c *httpStubClient) CommitDrive(_ context.Context, items []DriveCommitItem) ([]byte, error) { + payload := map[string]any{ + "items": items, + } + body, statusCode, err := c.http.postJSON("/drive/commit", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "drive commit"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) GetConfigManifest(_ context.Context) ([]byte, error) { + body, statusCode, err := c.http.getJSON("/config/manifest", nil) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config manifest"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) PullConfigSkill(_ context.Context, name string) ([]byte, error) { + body, statusCode, err := c.http.getRaw(fmt.Sprintf("/config/skills/%s/pull", name), nil) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config skill pull"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) PullConfigFile(_ context.Context, name string) ([]byte, error) { + body, statusCode, err := c.http.getRaw(fmt.Sprintf("/config/files/%s/pull", name), nil) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config file pull"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) PushConfig(_ context.Context, payload any) ([]byte, error) { + body, statusCode, err := c.http.postJSON("/config/push", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config push"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) PatchConfigEnv(_ context.Context, envText string) ([]byte, error) { + payload := map[string]string{"env_text": envText} + body, statusCode, err := c.http.patchJSON("/config/env", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config env update"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) PutConfigNote(_ context.Context, note string) ([]byte, error) { + payload := map[string]string{"note": note} + body, statusCode, err := c.http.putJSON("/config/note", payload) + if err != nil { + return nil, err + } + if err := checkHTTPError(body, statusCode, "config note update"); err != nil { + return nil, err + } + return body, nil +} + +func (c *httpStubClient) UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error) { + return c.http.uploadFile(uploadURL, filePath, filename, mimetype) +} + +func (c *httpStubClient) DownloadFromURL(downloadURL string) ([]byte, error) { + return c.http.downloadFromURL(downloadURL) +} diff --git a/dify-agent-runtime/internal/agentcli/config.go b/dify-agent-runtime/internal/agentcli/config.go new file mode 100644 index 00000000000..6b00d0147b3 --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/config.go @@ -0,0 +1,498 @@ +package agentcli + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" +) + +const defaultConfigBase = ".dify_conf" + +// RunConfigManifest executes the `config manifest` command. +func RunConfigManifest(env *Environment) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + body, err := client.GetConfigManifest(context.Background()) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigSkillsPull executes the `config skills pull` command. +func RunConfigSkillsPull(env *Environment, names []string, localDir string, jsonOutput bool) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + ctx := context.Background() + + // Get manifest to find available skills if no names specified + if len(names) == 0 { + body, err := client.GetConfigManifest(ctx) + if err != nil { + return err + } + var manifest struct { + Skills struct { + Items []struct { + Name string `json:"name"` + } `json:"items"` + } `json:"skills"` + } + if err := json.Unmarshal(body, &manifest); err != nil { + return fmt.Errorf("parse manifest: %w", err) + } + for _, item := range manifest.Skills.Items { + names = append(names, item.Name) + } + } + + targetDir := localDir + if targetDir == "" { + targetDir = filepath.Join(defaultConfigBase, "skills") + } + if err := os.MkdirAll(targetDir, 0o755); err != nil { + return fmt.Errorf("create target directory: %w", err) + } + + type pullItem struct { + Name string `json:"name"` + ArchivePath string `json:"archive_path"` + DirectoryPath string `json:"directory_path"` + SkillMD string `json:"skill_md"` + } + var items []pullItem + + for _, name := range names { + archiveBytes, err := client.PullConfigSkill(ctx, name) + if err != nil { + return err + } + + archivePath := filepath.Join(targetDir, name+".zip") + skillDir := filepath.Join(targetDir, name) + if err := os.MkdirAll(skillDir, 0o755); err != nil { + return fmt.Errorf("create skill directory: %w", err) + } + if err := os.WriteFile(archivePath, archiveBytes, 0o644); err != nil { + return fmt.Errorf("write archive: %w", err) + } + + // Extract archive + if err := extractZipArchive(archivePath, skillDir); err != nil { + return fmt.Errorf("extract skill archive: %w", err) + } + + skillMDPath := filepath.Join(skillDir, "SKILL.md") + skillMD := "" + if data, err := os.ReadFile(skillMDPath); err == nil { + skillMD = string(data) + } + + items = append(items, pullItem{ + Name: name, + ArchivePath: archivePath, + DirectoryPath: skillDir, + SkillMD: skillMD, + }) + } + + if jsonOutput { + out, _ := json.Marshal(map[string]any{"items": items}) + fmt.Println(string(out)) + return nil + } + + for i, item := range items { + if i > 0 { + fmt.Println() + } + fmt.Println(item.DirectoryPath) + fmt.Print(item.SkillMD) + } + return nil +} + +// RunConfigFilesPull executes the `config files pull` command. +func RunConfigFilesPull(env *Environment, names []string, localDir string, jsonOutput bool) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + ctx := context.Background() + + if len(names) == 0 { + body, err := client.GetConfigManifest(ctx) + if err != nil { + return err + } + var manifest struct { + Files struct { + Items []struct { + Name string `json:"name"` + } `json:"items"` + } `json:"files"` + } + if err := json.Unmarshal(body, &manifest); err != nil { + return fmt.Errorf("parse manifest: %w", err) + } + for _, item := range manifest.Files.Items { + names = append(names, item.Name) + } + } + + targetDir := localDir + if targetDir == "" { + targetDir = filepath.Join(defaultConfigBase, "files") + } + if err := os.MkdirAll(targetDir, 0o755); err != nil { + return fmt.Errorf("create target directory: %w", err) + } + + type fileItem struct { + Name string `json:"name"` + Path string `json:"path"` + } + var items []fileItem + + for _, name := range names { + payload, err := client.PullConfigFile(ctx, name) + if err != nil { + return err + } + + targetPath := filepath.Join(targetDir, name) + if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil { + return fmt.Errorf("create parent directory: %w", err) + } + if err := os.WriteFile(targetPath, payload, 0o644); err != nil { + return fmt.Errorf("write file: %w", err) + } + items = append(items, fileItem{Name: name, Path: targetPath}) + } + + if jsonOutput { + out, _ := json.Marshal(map[string]any{"items": items}) + fmt.Println(string(out)) + return nil + } + + for _, item := range items { + fmt.Println(item.Path) + } + return nil +} + +// RunConfigSkillsPush executes the `config skills push` command. +func RunConfigSkillsPush(env *Environment, paths []string) error { + if len(paths) == 0 { + return fmt.Errorf("at least one skill directory is required") + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + type skillPushItem struct { + Name string `json:"name"` + FileRef *DriveFileRef `json:"file_ref"` + } + var skills []skillPushItem + + for _, path := range paths { + absPath, err := filepath.Abs(path) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + info, err := os.Stat(absPath) + if err != nil || !info.IsDir() { + return fmt.Errorf("config skill path must be a directory: %s", absPath) + } + skillMDPath := filepath.Join(absPath, "SKILL.md") + if _, err := os.Stat(skillMDPath); os.IsNotExist(err) { + return fmt.Errorf("config skill directory must contain SKILL.md: %s", absPath) + } + + // Build archive and upload + archivePath, err := buildSkillArchive(absPath) + if err != nil { + return err + } + defer func() { _ = os.Remove(archivePath) }() + + name := filepath.Base(absPath) + commitItem, err := uploadAndPrepareCommitItem(client, archivePath, name) + if err != nil { + return err + } + + skills = append(skills, skillPushItem{ + Name: name, + FileRef: &DriveFileRef{Kind: commitItem.FileRef.Kind, ID: commitItem.FileRef.ID}, + }) + } + + payload := map[string]any{ + "skills": skills, + "files": []any{}, + } + + body, err := client.PushConfig(context.Background(), payload) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigFilesPush executes the `config files push` command. +func RunConfigFilesPush(env *Environment, paths []string) error { + if len(paths) == 0 { + return fmt.Errorf("at least one file path is required") + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + type filePushItem struct { + Name string `json:"name"` + FileRef *DriveFileRef `json:"file_ref"` + } + var files []filePushItem + + for _, path := range paths { + absPath, err := filepath.Abs(path) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + info, err := os.Stat(absPath) + if err != nil || info.IsDir() { + return fmt.Errorf("config file path must be a regular file: %s", absPath) + } + + name := filepath.Base(absPath) + commitItem, err := uploadAndPrepareCommitItem(client, absPath, name) + if err != nil { + return err + } + + files = append(files, filePushItem{ + Name: name, + FileRef: &DriveFileRef{Kind: commitItem.FileRef.Kind, ID: commitItem.FileRef.ID}, + }) + } + + payload := map[string]any{ + "files": files, + "skills": []any{}, + } + + body, err := client.PushConfig(context.Background(), payload) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigSkillsDelete executes the `config skills delete` command. +func RunConfigSkillsDelete(env *Environment, names []string) error { + if len(names) == 0 { + return fmt.Errorf("at least one skill name is required") + } + + type skillDeleteItem struct { + Name string `json:"name"` + FileRef *any `json:"file_ref"` + } + var skills []skillDeleteItem + for _, name := range names { + skills = append(skills, skillDeleteItem{Name: name, FileRef: nil}) + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + payload := map[string]any{ + "skills": skills, + "files": []any{}, + } + + body, err := client.PushConfig(context.Background(), payload) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigFilesDelete executes the `config files delete` command. +func RunConfigFilesDelete(env *Environment, names []string) error { + if len(names) == 0 { + return fmt.Errorf("at least one file name is required") + } + + type fileDeleteItem struct { + Name string `json:"name"` + FileRef *any `json:"file_ref"` + } + var files []fileDeleteItem + for _, name := range names { + files = append(files, fileDeleteItem{Name: name, FileRef: nil}) + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + payload := map[string]any{ + "files": files, + "skills": []any{}, + } + + body, err := client.PushConfig(context.Background(), payload) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigEnvPush executes the `config env push` command. +func RunConfigEnvPush(env *Environment, localPath string) error { + var envText string + if localPath == "-" { + data, err := os.ReadFile("/dev/stdin") + if err != nil { + return fmt.Errorf("read stdin: %w", err) + } + envText = string(data) + } else { + absPath, err := filepath.Abs(localPath) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + data, err := os.ReadFile(absPath) + if err != nil { + return fmt.Errorf("read file: %w", err) + } + envText = string(data) + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + body, err := client.PatchConfigEnv(context.Background(), envText) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// RunConfigNotePull executes the `config note pull` command. +func RunConfigNotePull(env *Environment, localPath string) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + body, err := client.GetConfigManifest(context.Background()) + if err != nil { + return err + } + + var manifest struct { + Note string `json:"note"` + } + if err := json.Unmarshal(body, &manifest); err != nil { + return fmt.Errorf("parse manifest: %w", err) + } + + targetPath := localPath + if targetPath == "" { + targetPath = filepath.Join(defaultConfigBase, "note.md") + } + absPath, err := filepath.Abs(targetPath) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + if err := os.WriteFile(absPath, []byte(manifest.Note), 0o644); err != nil { + return fmt.Errorf("write note: %w", err) + } + fmt.Println(absPath) + return nil +} + +// RunConfigNotePush executes the `config note push` command. +func RunConfigNotePush(env *Environment, localPath string) error { + var note string + if localPath == "-" { + data, err := os.ReadFile("/dev/stdin") + if err != nil { + return fmt.Errorf("read stdin: %w", err) + } + note = string(data) + } else { + targetPath := localPath + if targetPath == "" { + targetPath = filepath.Join(defaultConfigBase, "note.md") + } + absPath, err := filepath.Abs(targetPath) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + data, err := os.ReadFile(absPath) + if err != nil { + return fmt.Errorf("read file: %w", err) + } + note = string(data) + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + body, err := client.PutConfigNote(context.Background(), note) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +// extractZipArchive extracts a zip file into targetDir. +func extractZipArchive(archivePath string, targetDir string) error { + return extractZip(archivePath, targetDir) +} diff --git a/dify-agent-runtime/internal/agentcli/connect.go b/dify-agent-runtime/internal/agentcli/connect.go new file mode 100644 index 00000000000..6d17cc45e6f --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/connect.go @@ -0,0 +1,37 @@ +package agentcli + +import ( + "context" + "encoding/json" + "fmt" +) + +const agentStubProtocolVersion = 1 + +// ConnectResponse is the JSON output for `dify-agent connect --json`. +type ConnectResponse struct { + ConnectionID string `json:"connection_id"` + Status string `json:"status"` +} + +// RunConnect executes the connect command. +func RunConnect(env *Environment, argv []string, jsonOutput bool) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + resp, err := client.Connect(context.Background(), argv, "{}") + if err != nil { + return err + } + + if jsonOutput { + out, _ := json.Marshal(resp) + fmt.Println(string(out)) + } else { + fmt.Printf("connected %s\n", resp.ConnectionID) + } + return nil +} diff --git a/dify-agent-runtime/internal/agentcli/drive.go b/dify-agent-runtime/internal/agentcli/drive.go new file mode 100644 index 00000000000..8d5e19f265a --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/drive.go @@ -0,0 +1,353 @@ +package agentcli + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" +) + +// DriveItem represents one item in a drive manifest. +type DriveItem struct { + Key string `json:"key"` + Size *int64 `json:"size,omitempty"` + MimeType string `json:"mime_type,omitempty"` + Hash string `json:"hash,omitempty"` + DownloadURL *string `json:"download_url,omitempty"` +} + +// DriveManifestResponse is the drive manifest from the Agent Stub. +type DriveManifestResponse struct { + Items []DriveItem `json:"items"` +} + +// DrivePullResultItem represents one pulled drive file. +type DrivePullResultItem struct { + Key string `json:"key"` + LocalPath string `json:"local_path"` +} + +// DrivePullResult is the JSON output for `dify-agent drive pull --json`. +type DrivePullResult struct { + Items []DrivePullResultItem `json:"items"` +} + +// DriveCommitItem represents one file to commit into the drive. +type DriveCommitItem struct { + Key string `json:"key"` + FileRef DriveFileRef `json:"file_ref"` +} + +// DriveFileRef is the reference to an uploaded file. +type DriveFileRef struct { + Kind string `json:"kind"` + ID string `json:"id"` +} + +// DriveCommitResponse is the response from a drive commit. +type DriveCommitResponse struct { + Items []DriveItem `json:"items"` +} + +// RunDriveList executes the `drive list` command. +func RunDriveList(env *Environment, pathPrefix string, jsonOutput bool) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + manifest, err := client.GetDriveManifest(context.Background(), pathPrefix, false) + if err != nil { + return err + } + + if jsonOutput { + out, _ := json.Marshal(manifest) + fmt.Println(string(out)) + return nil + } + + for _, item := range manifest.Items { + size := "-" + if item.Size != nil { + size = fmt.Sprintf("%d", *item.Size) + } + mimeType := item.MimeType + if mimeType == "" { + mimeType = "-" + } + hash := item.Hash + if hash == "" { + hash = "-" + } + fmt.Printf("%s\t%s\t%s\t%s\n", size, mimeType, hash, item.Key) + } + return nil +} + +// RunDrivePull executes the `drive pull` command. +func RunDrivePull(env *Environment, targets []string, localBase string, jsonOutput bool) error { + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + if localBase == "" { + localBase = ReadDriveBase() + } + resolvedBase, err := filepath.Abs(localBase) + if err != nil { + return fmt.Errorf("resolve drive base: %w", err) + } + + if len(targets) == 0 { + targets = []string{""} + } + + ctx := context.Background() + resultItems := []DrivePullResultItem{} + + for _, target := range targets { + manifest, err := client.GetDriveManifest(ctx, target, true) + if err != nil { + return err + } + + if len(manifest.Items) == 0 { + continue + } + + localPath := resolveDriveDestination(resolvedBase, target) + resultItems = append(resultItems, DrivePullResultItem{Key: target, LocalPath: localPath}) + + for _, item := range manifest.Items { + if item.DownloadURL == nil || *item.DownloadURL == "" { + return fmt.Errorf("drive manifest item is missing download_url: %s", item.Key) + } + + destPath := resolveDriveDestination(resolvedBase, item.Key) + destDir := filepath.Dir(destPath) + if err := os.MkdirAll(destDir, 0o755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + + data, err := client.DownloadFromURL(*item.DownloadURL) + if err != nil { + return fmt.Errorf("download %s: %w", item.Key, err) + } + + if err := os.WriteFile(destPath, data, 0o644); err != nil { + return fmt.Errorf("write %s: %w", destPath, err) + } + } + } + + if jsonOutput { + out, _ := json.Marshal(DrivePullResult{Items: resultItems}) + fmt.Println(string(out)) + return nil + } + + for _, item := range resultItems { + fmt.Println(item.LocalPath) + } + return nil +} + +// RunDrivePush executes the `drive push` command. +func RunDrivePush(env *Environment, localPath string, drivePath string, kind string) error { + absPath, err := filepath.Abs(localPath) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + + info, err := os.Stat(absPath) + if err != nil { + return fmt.Errorf("local path not found: %s", absPath) + } + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + if info.IsDir() { + if kind == "" { + return fmt.Errorf("directory drive push requires --kind skill or --kind dir") + } + if kind == "file" { + return fmt.Errorf("--kind file requires a file") + } + if kind == "dir" { + return pushDirectory(client, absPath, drivePath) + } + return pushSkillDirectory(client, absPath, drivePath) + } + + // Single file push + if kind == "skill" { + return fmt.Errorf("--kind skill requires a directory containing SKILL.md") + } + if kind == "dir" { + return fmt.Errorf("--kind dir requires a directory") + } + + commitItem, err := uploadAndPrepareCommitItem(client, absPath, drivePath) + if err != nil { + return err + } + + return commitDriveItems(client, []DriveCommitItem{*commitItem}) +} + +func pushDirectory(client StubClient, dirPath string, drivePath string) error { + var items []DriveCommitItem + + err := filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + if info.IsDir() { + if shouldSkipDir(info.Name()) { + return filepath.SkipDir + } + return nil + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("drive push does not support symlinked files: %s", path) + } + + relPath, _ := filepath.Rel(dirPath, path) + driveKey := joinDriveKey(drivePath, filepath.ToSlash(relPath)) + commitItem, err := uploadAndPrepareCommitItem(client, path, driveKey) + if err != nil { + return err + } + items = append(items, *commitItem) + return nil + }) + if err != nil { + return err + } + + if len(items) == 0 { + return fmt.Errorf("directory has no regular files: %s", dirPath) + } + + return commitDriveItems(client, items) +} + +func pushSkillDirectory(client StubClient, dirPath string, drivePath string) error { + skillMDPath := filepath.Join(dirPath, "SKILL.md") + if _, err := os.Stat(skillMDPath); os.IsNotExist(err) { + return fmt.Errorf("--kind skill requires a directory containing SKILL.md") + } + + // Upload SKILL.md + skillMDItem, err := uploadAndPrepareCommitItem(client, skillMDPath, joinDriveKey(drivePath, "SKILL.md")) + if err != nil { + return err + } + + // Build and upload archive + archivePath, err := buildSkillArchive(dirPath) + if err != nil { + return err + } + defer func() { _ = os.Remove(archivePath) }() + + archiveItem, err := uploadAndPrepareCommitItem(client, archivePath, joinDriveKey(drivePath, ".DIFY-SKILL-FULL.zip")) + if err != nil { + return err + } + + return commitDriveItems(client, []DriveCommitItem{*skillMDItem, *archiveItem}) +} + +func uploadAndPrepareCommitItem(client StubClient, filePath string, driveKey string) (*DriveCommitItem, error) { + filename := filepath.Base(filePath) + mimetype := guessMIMEType(filename) + ctx := context.Background() + + // Request upload URL + uploadURL, err := client.CreateFileUploadURL(ctx, filename, mimetype) + if err != nil { + return nil, err + } + + // Upload + uploadBody, err := client.UploadFileToURL(uploadURL, filePath, filename, mimetype) + if err != nil { + return nil, err + } + + var uploadResult map[string]any + if err := json.Unmarshal(uploadBody, &uploadResult); err != nil { + return nil, fmt.Errorf("parse upload result: %w", err) + } + + toolFileID, _ := uploadResult["id"].(string) + if toolFileID == "" { + return nil, fmt.Errorf("upload response is missing id") + } + + return &DriveCommitItem{ + Key: driveKey, + FileRef: DriveFileRef{Kind: "tool_file", ID: toolFileID}, + }, nil +} + +func commitDriveItems(client StubClient, items []DriveCommitItem) error { + body, err := client.CommitDrive(context.Background(), items) + if err != nil { + return err + } + fmt.Println(string(body)) + return nil +} + +func resolveDriveDestination(basePath string, key string) string { + if key == "" { + return basePath + } + return filepath.Join(basePath, filepath.FromSlash(key)) +} + +func joinDriveKey(base string, child string) string { + stripped := strings.TrimRight(base, "/") + child = strings.TrimLeft(child, "/") + if stripped == "" { + return child + } + return stripped + "/" + child +} + +func shouldSkipDir(name string) bool { + skip := map[string]bool{ + ".git": true, "__pycache__": true, ".pytest_cache": true, + ".mypy_cache": true, ".ruff_cache": true, ".venv": true, "node_modules": true, + } + return skip[name] +} + +// buildSkillArchive creates a zip archive of the skill directory. +func buildSkillArchive(dirPath string) (string, error) { + // Create temp file for archive + tmpFile, err := os.CreateTemp("", "skill-archive-*.zip") + if err != nil { + return "", fmt.Errorf("create temp archive: %w", err) + } + archivePath := tmpFile.Name() + _ = tmpFile.Close() + + if err := createZipArchive(archivePath, dirPath); err != nil { + _ = os.Remove(archivePath) + return "", err + } + return archivePath, nil +} diff --git a/dify-agent-runtime/internal/agentcli/env.go b/dify-agent-runtime/internal/agentcli/env.go new file mode 100644 index 00000000000..8e7962f90f8 --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/env.go @@ -0,0 +1,150 @@ +// Package agentcli implements the dify-agent CLI that runs inside the sandbox +// container. It communicates with the Agent Stub server on the host via HTTP +// or gRPC to provide connect, file, drive, and config operations. +package agentcli + +import ( + "errors" + "fmt" + "net/url" + "os" + "strings" +) + +const ( + EnvAPIBaseURL = "DIFY_AGENT_STUB_API_BASE_URL" + EnvAuthJWE = "DIFY_AGENT_STUB_AUTH_JWE" + EnvDriveBase = "DIFY_AGENT_STUB_DRIVE_BASE" + + DefaultDriveBase = "/mnt/drive" +) + +// Environment holds validated Agent Stub connection parameters. +type Environment struct { + URL string + AuthJWE string +} + +// Endpoint represents a parsed Agent Stub endpoint with transport info. +type Endpoint struct { + URL string + Scheme string // "http", "https", or "grpc" + Host string + Port string + IsGRPC bool +} + +var ErrMissingEnvironment = errors.New("missing required Agent Stub environment variables") + +// ReadEnvironment reads and validates the Agent Stub env vars. +func ReadEnvironment() (*Environment, error) { + apiURL := strings.TrimSpace(os.Getenv(EnvAPIBaseURL)) + authJWE := strings.TrimSpace(os.Getenv(EnvAuthJWE)) + + var missing []string + if apiURL == "" { + missing = append(missing, EnvAPIBaseURL) + } + if authJWE == "" { + missing = append(missing, EnvAuthJWE) + } + if len(missing) > 0 { + return nil, fmt.Errorf("%w: %s", ErrMissingEnvironment, strings.Join(missing, ", ")) + } + + endpoint, err := ParseEndpoint(apiURL) + if err != nil { + return nil, fmt.Errorf("invalid %s: %w", EnvAPIBaseURL, err) + } + + return &Environment{ + URL: endpoint.URL, + AuthJWE: authJWE, + }, nil +} + +// HasEnvironment returns whether both required env vars are set. +func HasEnvironment() bool { + return os.Getenv(EnvAPIBaseURL) != "" && os.Getenv(EnvAuthJWE) != "" +} + +// ReadDriveBase returns the configured drive base or the default. +func ReadDriveBase() string { + if v := strings.TrimSpace(os.Getenv(EnvDriveBase)); v != "" { + return v + } + return DefaultDriveBase +} + +// ParseEndpoint parses an Agent Stub URL and normalizes it. +func ParseEndpoint(rawURL string) (*Endpoint, error) { + stripped := strings.TrimSpace(rawURL) + if stripped == "" { + return nil, errors.New("agent stub URL must not be empty") + } + + parsed, err := url.Parse(stripped) + if err != nil { + return nil, fmt.Errorf("invalid URL: %w", err) + } + + switch parsed.Scheme { + case "http", "https": + return parseHTTPEndpoint(parsed) + case "grpc": + return parseGRPCEndpoint(parsed) + default: + return nil, errors.New("agent stub URL must use http, https, or grpc") + } +} + +func parseHTTPEndpoint(parsed *url.URL) (*Endpoint, error) { + if parsed.Host == "" { + return nil, errors.New("agent stub URL must include a host") + } + if parsed.RawQuery != "" || parsed.Fragment != "" { + return nil, errors.New("agent stub URL must not include a query string or fragment") + } + if parsed.User != nil { + return nil, errors.New("agent stub URL must not include user info") + } + + path := strings.TrimRight(parsed.Path, "/") + if path == "" || path == "/" { + path = "/agent-stub" + } else if path != "/agent-stub" { + return nil, errors.New("HTTP agent stub API base URL path must be empty or /agent-stub") + } + + normalizedURL := fmt.Sprintf("%s://%s%s", parsed.Scheme, parsed.Host, path) + return &Endpoint{ + URL: normalizedURL, + Scheme: parsed.Scheme, + Host: parsed.Hostname(), + Port: parsed.Port(), + IsGRPC: false, + }, nil +} + +func parseGRPCEndpoint(parsed *url.URL) (*Endpoint, error) { + if parsed.Host == "" { + return nil, errors.New("gRPC agent stub URL must include a host") + } + path := strings.TrimRight(parsed.Path, "/") + if path != "" && path != "/" { + return nil, errors.New("gRPC agent stub URL must not include a path") + } + port := parsed.Port() + if port == "" { + return nil, errors.New("gRPC agent stub URL must include an explicit port") + } + + normalizedURL := fmt.Sprintf("grpc://%s:%s", parsed.Hostname(), port) + return &Endpoint{ + URL: normalizedURL, + Scheme: "grpc", + Host: parsed.Hostname(), + Port: port, + IsGRPC: true, + }, nil +} diff --git a/dify-agent-runtime/internal/agentcli/env_test.go b/dify-agent-runtime/internal/agentcli/env_test.go new file mode 100644 index 00000000000..6d02f8014a8 --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/env_test.go @@ -0,0 +1,173 @@ +package agentcli + +import ( + "testing" +) + +func TestParseEndpoint_HTTP(t *testing.T) { + tests := []struct { + name string + input string + wantURL string + wantErr bool + }{ + { + name: "bare http host normalizes to /agent-stub", + input: "http://localhost:8080", + wantURL: "http://localhost:8080/agent-stub", + }, + { + name: "bare https host normalizes to /agent-stub", + input: "https://agent.example.com", + wantURL: "https://agent.example.com/agent-stub", + }, + { + name: "explicit /agent-stub path is preserved", + input: "http://localhost:8080/agent-stub", + wantURL: "http://localhost:8080/agent-stub", + }, + { + name: "trailing slash normalized", + input: "http://localhost:8080/", + wantURL: "http://localhost:8080/agent-stub", + }, + { + name: "invalid path rejects", + input: "http://localhost:8080/other-path", + wantErr: true, + }, + { + name: "query string rejects", + input: "http://localhost:8080?foo=bar", + wantErr: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ep, err := ParseEndpoint(tc.input) + if tc.wantErr { + if err == nil { + t.Fatal("expected error") + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ep.URL != tc.wantURL { + t.Errorf("URL = %q, want %q", ep.URL, tc.wantURL) + } + if ep.IsGRPC { + t.Error("expected IsGRPC=false") + } + }) + } +} + +func TestParseEndpoint_GRPC(t *testing.T) { + tests := []struct { + name string + input string + wantURL string + wantErr bool + }{ + { + name: "valid grpc endpoint", + input: "grpc://localhost:50051", + wantURL: "grpc://localhost:50051", + }, + { + name: "grpc without port rejects", + input: "grpc://localhost", + wantErr: true, + }, + { + name: "grpc with path rejects", + input: "grpc://localhost:50051/some-path", + wantErr: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ep, err := ParseEndpoint(tc.input) + if tc.wantErr { + if err == nil { + t.Fatal("expected error") + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ep.URL != tc.wantURL { + t.Errorf("URL = %q, want %q", ep.URL, tc.wantURL) + } + if !ep.IsGRPC { + t.Error("expected IsGRPC=true") + } + }) + } +} + +func TestParseEndpoint_Invalid(t *testing.T) { + tests := []struct { + name string + input string + }{ + {"empty", ""}, + {"spaces only", " "}, + {"unsupported scheme", "ftp://example.com"}, + {"no host", "http://"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := ParseEndpoint(tc.input) + if err == nil { + t.Fatal("expected error") + } + }) + } +} + +func TestReadEnvironment_Missing(t *testing.T) { + t.Setenv(EnvAPIBaseURL, "") + t.Setenv(EnvAuthJWE, "") + + _, err := ReadEnvironment() + if err == nil { + t.Fatal("expected error for missing env vars") + } +} + +func TestReadEnvironment_Valid(t *testing.T) { + t.Setenv(EnvAPIBaseURL, "http://localhost:8080") + t.Setenv(EnvAuthJWE, "test-token") + + env, err := ReadEnvironment() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if env.URL != "http://localhost:8080/agent-stub" { + t.Errorf("URL = %q, want %q", env.URL, "http://localhost:8080/agent-stub") + } + if env.AuthJWE != "test-token" { + t.Errorf("AuthJWE = %q, want %q", env.AuthJWE, "test-token") + } +} + +func TestReadDriveBase_Default(t *testing.T) { + t.Setenv(EnvDriveBase, "") + if got := ReadDriveBase(); got != DefaultDriveBase { + t.Errorf("ReadDriveBase() = %q, want %q", got, DefaultDriveBase) + } +} + +func TestReadDriveBase_Custom(t *testing.T) { + t.Setenv(EnvDriveBase, "/custom/drive") + if got := ReadDriveBase(); got != "/custom/drive" { + t.Errorf("ReadDriveBase() = %q, want %q", got, "/custom/drive") + } +} diff --git a/dify-agent-runtime/internal/agentcli/file.go b/dify-agent-runtime/internal/agentcli/file.go new file mode 100644 index 00000000000..d0f8ab75f1e --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/file.go @@ -0,0 +1,181 @@ +package agentcli + +import ( + "context" + "encoding/json" + "fmt" + "mime" + "os" + "path/filepath" + "strings" +) + +// FileUploadResponse is the JSON output for `dify-agent file upload`. +type FileUploadResponse struct { + TransferMethod string `json:"transfer_method"` + Reference string `json:"reference"` + DownloadURL string `json:"download_url"` +} + +// FileDownloadResponse is the response from a file download request. +type FileDownloadResponse struct { + Filename string `json:"filename"` + MimeType string `json:"mime_type,omitempty"` + Size int64 `json:"size"` + DownloadURL string `json:"download_url"` +} + +// RunFileUpload executes the `file upload` command. +func RunFileUpload(env *Environment, path string) error { + absPath, err := filepath.Abs(path) + if err != nil { + return fmt.Errorf("resolve path: %w", err) + } + info, err := os.Stat(absPath) + if err != nil || info.IsDir() { + return fmt.Errorf("local file not found: %s", absPath) + } + + filename := filepath.Base(absPath) + mimetype := guessMIMEType(filename) + ctx := context.Background() + + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + // Step 1: Request a signed upload URL + uploadURL, err := client.CreateFileUploadURL(ctx, filename, mimetype) + if err != nil { + return err + } + + // Step 2: Upload the file to the signed URL (data-plane) + uploadBody, err := client.UploadFileToURL(uploadURL, absPath, filename, mimetype) + if err != nil { + return err + } + + var uploadResult map[string]any + if err := json.Unmarshal(uploadBody, &uploadResult); err != nil { + return fmt.Errorf("parse upload result: %w", err) + } + + reference, _ := uploadResult["reference"].(string) + if reference == "" { + return fmt.Errorf("signed file upload response is missing reference") + } + + // Step 3: Request download URL for the uploaded file + ref := reference + dlResp, err := client.CreateFileDownloadURL(ctx, "tool_file", &ref, nil, false) + if err != nil { + return err + } + + result := FileUploadResponse{ + TransferMethod: "tool_file", + Reference: reference, + DownloadURL: dlResp.DownloadURL, + } + out, _ := json.Marshal(result) + fmt.Println(string(out)) + return nil +} + +// RunFileDownload executes the `file download` command. +func RunFileDownload(env *Environment, transferMethod string, referenceOrURL string, localDir string) error { + var reference *string + var url *string + if transferMethod == "remote_url" { + url = &referenceOrURL + } else { + reference = &referenceOrURL + } + + ctx := context.Background() + client, err := NewStubClient(env) + if err != nil { + return err + } + defer func() { _ = client.Close() }() + + dlResp, err := client.CreateFileDownloadURL(ctx, transferMethod, reference, url, false) + if err != nil { + return err + } + if dlResp.DownloadURL == "" { + return fmt.Errorf("signed file download response is missing download_url") + } + if dlResp.Filename == "" { + return fmt.Errorf("signed file download response is missing filename") + } + + // Download the file (data-plane) + data, err := client.DownloadFromURL(dlResp.DownloadURL) + if err != nil { + return err + } + + // Determine target directory + targetDir := localDir + if targetDir == "" { + targetDir, _ = os.Getwd() + } + if err := os.MkdirAll(targetDir, 0o755); err != nil { + return fmt.Errorf("create target directory: %w", err) + } + + // Write file + sanitizedName := sanitizeFilename(dlResp.Filename) + destPath := deduplicatePath(filepath.Join(targetDir, sanitizedName)) + if err := os.WriteFile(destPath, data, 0o644); err != nil { + return fmt.Errorf("write file: %w", err) + } + + fmt.Println(destPath) + return nil +} + +func guessMIMEType(filename string) string { + ext := filepath.Ext(filename) + if ext == "" { + return "application/octet-stream" + } + mimeType := mime.TypeByExtension(ext) + if mimeType == "" { + return "application/octet-stream" + } + // Strip parameters (e.g. "; charset=utf-8") so the MIME type matches + // what Flask/Werkzeug returns via file.mimetype during signature + // verification on the Dify API upload endpoint. + if idx := strings.Index(mimeType, ";"); idx >= 0 { + mimeType = strings.TrimSpace(mimeType[:idx]) + } + return mimeType +} + +func sanitizeFilename(filename string) string { + name := filepath.Base(filename) + if name == "" || name == "." || name == ".." { + return "downloaded" + } + return name +} + +func deduplicatePath(path string) string { + if _, err := os.Stat(path); os.IsNotExist(err) { + return path + } + ext := filepath.Ext(path) + stem := strings.TrimSuffix(filepath.Base(path), ext) + dir := filepath.Dir(path) + for counter := 1; ; counter++ { + candidate := filepath.Join(dir, fmt.Sprintf("%s (%d)%s", stem, counter, ext)) + if _, err := os.Stat(candidate); os.IsNotExist(err) { + return candidate + } + } +} diff --git a/dify-agent-runtime/internal/agentcli/httpclient.go b/dify-agent-runtime/internal/agentcli/httpclient.go new file mode 100644 index 00000000000..95c3614bf1b --- /dev/null +++ b/dify-agent-runtime/internal/agentcli/httpclient.go @@ -0,0 +1,230 @@ +package agentcli + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "mime/multipart" + "net/http" + "net/textproto" + "os" + "time" +) + +// HTTPClient wraps HTTP interactions with the Agent Stub server. +type HTTPClient struct { + baseURL string + authJWE string + client *http.Client +} + +// NewHTTPClient creates a new HTTP client for the Agent Stub API. +func NewHTTPClient(env *Environment) *HTTPClient { + return &HTTPClient{ + baseURL: env.URL, + authJWE: env.AuthJWE, + client: &http.Client{Timeout: 30 * time.Second}, + } +} + +// NewHTTPClientWithTimeout creates a client with a custom timeout. +func NewHTTPClientWithTimeout(env *Environment, timeout time.Duration) *HTTPClient { + return &HTTPClient{ + baseURL: env.URL, + authJWE: env.AuthJWE, + client: &http.Client{Timeout: timeout}, + } +} + +// postJSON sends a POST request with JSON body and returns the response body. +func (c *HTTPClient) postJSON(path string, payload any) ([]byte, int, error) { + body, err := json.Marshal(payload) + if err != nil { + return nil, 0, fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequest("POST", c.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, 0, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.authJWE) + + resp, err := c.client.Do(req) + if err != nil { + return nil, 0, fmt.Errorf("request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, resp.StatusCode, fmt.Errorf("read response: %w", err) + } + return respBody, resp.StatusCode, nil +} + +// getJSON sends a GET request and returns the response body. +func (c *HTTPClient) getJSON(path string, params map[string]string) ([]byte, int, error) { + req, err := http.NewRequest("GET", c.baseURL+path, nil) + if err != nil { + return nil, 0, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+c.authJWE) + + if len(params) > 0 { + q := req.URL.Query() + for k, v := range params { + q.Set(k, v) + } + req.URL.RawQuery = q.Encode() + } + + resp, err := c.client.Do(req) + if err != nil { + return nil, 0, fmt.Errorf("request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, resp.StatusCode, fmt.Errorf("read response: %w", err) + } + return respBody, resp.StatusCode, nil +} + +// getRaw sends a GET request and returns raw bytes (for binary downloads). +func (c *HTTPClient) getRaw(path string, params map[string]string) ([]byte, int, error) { + return c.getJSON(path, params) +} + +// patchJSON sends a PATCH request with JSON body. +func (c *HTTPClient) patchJSON(path string, payload any) ([]byte, int, error) { + body, err := json.Marshal(payload) + if err != nil { + return nil, 0, fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequest("PATCH", c.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, 0, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.authJWE) + + resp, err := c.client.Do(req) + if err != nil { + return nil, 0, fmt.Errorf("request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, resp.StatusCode, fmt.Errorf("read response: %w", err) + } + return respBody, resp.StatusCode, nil +} + +// putJSON sends a PUT request with JSON body. +func (c *HTTPClient) putJSON(path string, payload any) ([]byte, int, error) { + body, err := json.Marshal(payload) + if err != nil { + return nil, 0, fmt.Errorf("marshal request: %w", err) + } + + req, err := http.NewRequest("PUT", c.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, 0, fmt.Errorf("create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+c.authJWE) + + resp, err := c.client.Do(req) + if err != nil { + return nil, 0, fmt.Errorf("request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, resp.StatusCode, fmt.Errorf("read response: %w", err) + } + return respBody, resp.StatusCode, nil +} + +// uploadFile uploads a file to a signed URL using multipart form. +func (c *HTTPClient) uploadFile(uploadURL string, filePath string, filename string, mimetype string) ([]byte, error) { + file, err := os.Open(filePath) + if err != nil { + return nil, fmt.Errorf("open file: %w", err) + } + defer func() { _ = file.Close() }() + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + h := make(textproto.MIMEHeader) + h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, filename)) + h.Set("Content-Type", mimetype) + part, err := writer.CreatePart(h) + if err != nil { + return nil, fmt.Errorf("create form file: %w", err) + } + if _, err := io.Copy(part, file); err != nil { + return nil, fmt.Errorf("copy file content: %w", err) + } + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("close multipart writer: %w", err) + } + + uploadClient := &http.Client{Timeout: 120 * time.Second} + req, err := http.NewRequest("POST", uploadURL, &buf) + if err != nil { + return nil, fmt.Errorf("create upload request: %w", err) + } + req.Header.Set("Content-Type", writer.FormDataContentType()) + + resp, err := uploadClient.Do(req) + if err != nil { + return nil, fmt.Errorf("upload request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read upload response: %w", err) + } + if resp.StatusCode >= 400 { + return nil, fmt.Errorf("upload failed with status %d: %s", resp.StatusCode, string(respBody)) + } + return respBody, nil +} + +// downloadFromURL downloads bytes from a signed URL. +func (c *HTTPClient) downloadFromURL(downloadURL string) ([]byte, error) { + dlClient := &http.Client{Timeout: 120 * time.Second} + resp, err := dlClient.Get(downloadURL) + if err != nil { + return nil, fmt.Errorf("download request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("download failed with status %d: %s", resp.StatusCode, string(body)) + } + return io.ReadAll(resp.Body) +} + +// checkHTTPError returns a formatted error if status >= 400. +func checkHTTPError(body []byte, statusCode int, operation string) error { + if statusCode < 400 { + return nil + } + var detail struct { + Detail any `json:"detail"` + } + if json.Unmarshal(body, &detail) == nil && detail.Detail != nil { + return fmt.Errorf("agent stub %s failed (HTTP %d): %v", operation, statusCode, detail.Detail) + } + return fmt.Errorf("agent stub %s failed (HTTP %d): %s", operation, statusCode, string(body)) +} diff --git a/dify-agent-runtime/internal/runner_exit/runner_exit.go b/dify-agent-runtime/internal/runner_exit/runner_exit.go new file mode 100644 index 00000000000..b9057f4ae1b --- /dev/null +++ b/dify-agent-runtime/internal/runner_exit/runner_exit.go @@ -0,0 +1,109 @@ +// Package runner_exit persists a drained shellctl job exit into SQLite. +// +// This runs out-of-process from the main shellctl server, after the tmux +// output pipe reaches EOF. It uses the same CAS semantics as the Python +// shellctl_runtime/runner_exit.py: only non-terminal rows are updated. +package runner_exit + +import ( + "database/sql" + "fmt" + "path/filepath" + + _ "modernc.org/sqlite" +) + +var ( + nonterminalStatuses = []string{"created", "starting", "running"} + terminalStatuses = []string{"exited", "terminated", "failed", "lost"} +) + +// RecordRunnerExit persists exit_code and ended_at into the shellctl SQLite DB. +// The update is idempotent for terminal rows. +func RecordRunnerExit(stateDir, jobID string, exitCode int, endedAt string, busyTimeoutMs int) error { + dbPath := filepath.Join(stateDir, "shellctl.db") + + dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(%d)&_pragma=journal_mode(WAL)", dbPath, busyTimeoutMs) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return fmt.Errorf("open database: %w", err) + } + defer func() { _ = db.Close() }() + db.SetMaxOpenConns(1) + + // Check current status + var status string + err = db.QueryRow("SELECT status FROM jobs WHERE job_id = ?", jobID).Scan(&status) + if err == sql.ErrNoRows { + return fmt.Errorf("unknown job id: %s", jobID) + } + if err != nil { + return fmt.Errorf("query job status: %w", err) + } + + if isTerminal(status) { + return nil + } + + // CAS update: only transition non-terminal rows + result, err := db.Exec(` + UPDATE jobs + SET status = CASE + WHEN status IN (?, ?, ?) THEN ? + ELSE status + END, + exit_code = CASE + WHEN status IN (?, ?, ?) THEN ? + ELSE exit_code + END, + ended_at = CASE + WHEN status IN (?, ?, ?) THEN ? + ELSE ended_at + END, + updated_at = CASE + WHEN status IN (?, ?, ?) THEN ? + ELSE updated_at + END, + reason = CASE + WHEN status IN (?, ?, ?) THEN NULL + ELSE reason + END, + message = CASE + WHEN status IN (?, ?, ?) THEN NULL + ELSE message + END + WHERE job_id = ?`, + // status + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], "exited", + // exit_code + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], exitCode, + // ended_at + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], endedAt, + // updated_at + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], endedAt, + // reason + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], + // message + nonterminalStatuses[0], nonterminalStatuses[1], nonterminalStatuses[2], + // WHERE + jobID, + ) + if err != nil { + return fmt.Errorf("update job: %w", err) + } + + rows, _ := result.RowsAffected() + if rows == 0 { + return fmt.Errorf("unknown job id: %s", jobID) + } + return nil +} + +func isTerminal(status string) bool { + for _, s := range terminalStatuses { + if status == s { + return true + } + } + return false +} diff --git a/dify-agent-runtime/internal/runner_exit/runner_exit_test.go b/dify-agent-runtime/internal/runner_exit/runner_exit_test.go new file mode 100644 index 00000000000..95c8d5557cb --- /dev/null +++ b/dify-agent-runtime/internal/runner_exit/runner_exit_test.go @@ -0,0 +1,182 @@ +package runner_exit + +import ( + "database/sql" + "path/filepath" + "testing" + + _ "modernc.org/sqlite" +) + +func setupTestDB(t *testing.T, dir string) string { + t.Helper() + dbPath := filepath.Join(dir, "shellctl.db") + db, err := sql.Open("sqlite", "file:"+dbPath) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + defer func() { _ = db.Close() }() + + _, err = db.Exec(` + CREATE TABLE 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 + ) + `) + if err != nil { + t.Fatalf("create table: %v", err) + } + + _, err = db.Exec(`INSERT INTO jobs (job_id, script_path, output_path, cwd, status, session_name, pane_target, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`, + "test-job", "s", "o", "/tmp", "running", "sess", "pane", "2025-01-01T00:00:00Z", "2025-01-01T00:00:00Z") + if err != nil { + t.Fatalf("insert job: %v", err) + } + + return dbPath +} + +func TestRecordRunnerExitRunning(t *testing.T) { + dir := t.TempDir() + stateDir := dir + setupTestDB(t, dir) + + err := RecordRunnerExit(stateDir, "test-job", 0, "2025-01-15T12:00:00Z", 5000) + if err != nil { + t.Fatalf("RecordRunnerExit: %v", err) + } + + // Verify the row was updated + db, err := sql.Open("sqlite", "file:"+filepath.Join(dir, "shellctl.db")) + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + + var status string + var exitCode int + if err := db.QueryRow("SELECT status, exit_code FROM jobs WHERE job_id = ?", "test-job").Scan(&status, &exitCode); err != nil { + t.Fatal(err) + } + + if status != "exited" { + t.Errorf("expected status=exited, got %s", status) + } + if exitCode != 0 { + t.Errorf("expected exit_code=0, got %d", exitCode) + } +} + +func TestRecordRunnerExitNonZeroCode(t *testing.T) { + dir := t.TempDir() + setupTestDB(t, dir) + + err := RecordRunnerExit(dir, "test-job", 42, "2025-01-15T12:00:00Z", 5000) + if err != nil { + t.Fatalf("RecordRunnerExit: %v", err) + } + + db, err := sql.Open("sqlite", "file:"+filepath.Join(dir, "shellctl.db")) + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + + var exitCode int + if err := db.QueryRow("SELECT exit_code FROM jobs WHERE job_id = ?", "test-job").Scan(&exitCode); err != nil { + t.Fatal(err) + } + + if exitCode != 42 { + t.Errorf("expected exit_code=42, got %d", exitCode) + } +} + +func TestRecordRunnerExitJobNotFound(t *testing.T) { + dir := t.TempDir() + setupTestDB(t, dir) + + err := RecordRunnerExit(dir, "nonexistent-job", 0, "2025-01-15T12:00:00Z", 5000) + if err == nil { + t.Error("expected error for nonexistent job") + } +} + +func TestRecordRunnerExitDBNotFound(t *testing.T) { + dir := t.TempDir() + + err := RecordRunnerExit(dir, "test-job", 0, "2025-01-15T12:00:00Z", 5000) + if err == nil { + t.Error("expected error when database doesn't exist") + } +} + +func TestRecordRunnerExitTerminalIdempotent(t *testing.T) { + dir := t.TempDir() + dbPath := setupTestDB(t, dir) + + // Manually set job to terminal state + db, err := sql.Open("sqlite", "file:"+dbPath) + if err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`UPDATE jobs SET status='terminated', exit_code=137, ended_at='2025-01-01T00:01:00Z' WHERE job_id='test-job'`); err != nil { + t.Fatal(err) + } + _ = db.Close() + + // Should not overwrite + err = RecordRunnerExit(dir, "test-job", 0, "2025-01-15T12:00:00Z", 5000) + if err != nil { + t.Fatalf("RecordRunnerExit on terminal: %v", err) + } + + db, err = sql.Open("sqlite", "file:"+dbPath) + if err != nil { + t.Fatal(err) + } + defer func() { _ = db.Close() }() + + var status string + var exitCode int + if err := db.QueryRow("SELECT status, exit_code FROM jobs WHERE job_id = ?", "test-job").Scan(&status, &exitCode); err != nil { + t.Fatal(err) + } + + if status != "terminated" { + t.Errorf("expected status=terminated (preserved), got %s", status) + } + if exitCode != 137 { + t.Errorf("expected exit_code=137 (preserved), got %d", exitCode) + } +} + +func TestIsTerminal(t *testing.T) { + terminal := []string{"exited", "terminated", "failed", "lost"} + for _, s := range terminal { + if !isTerminal(s) { + t.Errorf("%s should be terminal", s) + } + } + nonTerminal := []string{"created", "starting", "running"} + for _, s := range nonTerminal { + if isTerminal(s) { + t.Errorf("%s should not be terminal", s) + } + } +} diff --git a/dify-agent-runtime/internal/sanitize/sanitize.go b/dify-agent-runtime/internal/sanitize/sanitize.go new file mode 100644 index 00000000000..f6cebe339e1 --- /dev/null +++ b/dify-agent-runtime/internal/sanitize/sanitize.go @@ -0,0 +1,194 @@ +// Package sanitize implements a streaming PTY output sanitizer. +// +// It strips ANSI escape sequences (CSI, OSC), normalizes carriage-return +// progress updates into the final visible line, and performs incremental +// UTF-8 decoding—mirroring the Python shellctl_runtime/sanitize.py. +package sanitize + +import ( + "bufio" + "io" + "os" + "unicode/utf8" +) + +// escapeState tracks the ANSI escape sequence parser state. +type escapeState int + +const ( + stateNormal escapeState = iota + stateEsc + stateCSI + stateOSC + stateOSCEsc +) + +// PtySanitizer incrementally converts PTY bytes into stable, readable UTF-8. +type PtySanitizer struct { + lineBuffer []byte + pendingCR bool + state escapeState +} + +// New returns a fresh PtySanitizer. +func New() *PtySanitizer { + return &PtySanitizer{} +} + +// Feed consumes one chunk of decoded text and returns newly stable output. +func (s *PtySanitizer) Feed(text []byte) []byte { + var out []byte + for len(text) > 0 { + r, size := utf8.DecodeRune(text) + if r == utf8.RuneError && size <= 1 { + // Replace invalid byte with U+FFFD + text = text[1:] + r = utf8.RuneError + } else { + text = text[size:] + } + out = s.consumeRune(r, out) + } + return out +} + +// Flush returns any remaining buffered content at end-of-stream. +func (s *PtySanitizer) Flush() []byte { + s.state = stateNormal + s.pendingCR = false + result := s.lineBuffer + s.lineBuffer = nil + return result +} + +func (s *PtySanitizer) consumeRune(r rune, out []byte) []byte { + switch s.state { + case stateNormal: + if r == '\x1b' { + s.state = stateEsc + return out + } + return s.consumeVisible(r, out) + + case stateEsc: + switch r { + case '[': + s.state = stateCSI + case ']': + s.state = stateOSC + default: + s.state = stateNormal + if r != '\x1b' && isPrintable(r) { + return s.consumeVisible(r, out) + } + } + return out + + case stateCSI: + // CSI sequence ends at a byte in the range 0x40–0x7E + if r >= '@' && r <= '~' { + s.state = stateNormal + } + return out + + case stateOSC: + switch r { + case '\x07': + s.state = stateNormal + case '\x1b': + s.state = stateOSCEsc + } + return out + + case stateOSCEsc: + if r == '\\' { + s.state = stateNormal + } else { + s.state = stateOSC + } + return out + } + return out +} + +func (s *PtySanitizer) consumeVisible(r rune, out []byte) []byte { + if s.pendingCR { + if r == '\n' { + out = append(out, s.lineBuffer...) + out = append(out, '\n') + s.lineBuffer = nil + s.pendingCR = false + return out + } + // CR without LF: overwrite line buffer (progress update) + s.lineBuffer = nil + s.pendingCR = false + } + + if r == '\r' { + s.pendingCR = true + return out + } + if r == '\n' { + out = append(out, s.lineBuffer...) + out = append(out, '\n') + s.lineBuffer = nil + return out + } + + // Append rune to line buffer + var buf [utf8.UTFMax]byte + n := utf8.EncodeRune(buf[:], r) + s.lineBuffer = append(s.lineBuffer, buf[:n]...) + return out +} + +func isPrintable(r rune) bool { + // Match Python's str.isprintable: exclude C0/C1 control chars + return r >= 0x20 && r != 0x7f +} + +// Run executes the streaming sanitizer as a stdin→stdout filter. +// If readyFile is non-empty, it is touched before reading begins. +func Run(readyFile string, stdin io.Reader, stdout io.Writer) error { + if readyFile != "" { + f, err := os.Create(readyFile) + if err != nil { + return err + } + _ = f.Close() + } + + sanitizer := New() + reader := bufio.NewReaderSize(stdin, 65536) + writer := bufio.NewWriter(stdout) + defer func() { _ = writer.Flush() }() + + buf := make([]byte, 65536) + for { + n, err := reader.Read(buf) + if n > 0 { + out := sanitizer.Feed(buf[:n]) + if len(out) > 0 { + if _, werr := writer.Write(out); werr != nil { + return werr + } + _ = writer.Flush() + } + } + if err != nil { + if err == io.EOF { + break + } + return err + } + } + + tail := sanitizer.Flush() + if len(tail) > 0 { + if _, err := writer.Write(tail); err != nil { + return err + } + } + return writer.Flush() +} diff --git a/dify-agent-runtime/internal/sanitize/sanitize_test.go b/dify-agent-runtime/internal/sanitize/sanitize_test.go new file mode 100644 index 00000000000..4fae6abc3ff --- /dev/null +++ b/dify-agent-runtime/internal/sanitize/sanitize_test.go @@ -0,0 +1,83 @@ +package sanitize + +import ( + "testing" +) + +func TestPlainText(t *testing.T) { + s := New() + out := s.Feed([]byte("hello\nworld\n")) + out = append(out, s.Flush()...) + expected := "hello\nworld\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestStripCSI(t *testing.T) { + // ESC [ 31 m = red color, ESC [ 0 m = reset + input := []byte("\x1b[31mred\x1b[0m\n") + s := New() + out := s.Feed(input) + out = append(out, s.Flush()...) + expected := "red\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestCarriageReturnOverwrite(t *testing.T) { + // Progress: "50%" CR "100%" LF -> only "100%" visible + input := []byte("50%\r100%\n") + s := New() + out := s.Feed(input) + out = append(out, s.Flush()...) + expected := "100%\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestCRLF(t *testing.T) { + input := []byte("line1\r\nline2\r\n") + s := New() + out := s.Feed(input) + out = append(out, s.Flush()...) + expected := "line1\nline2\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestOSCSequence(t *testing.T) { + // OSC: ESC ] ... BEL + input := []byte("\x1b]0;title\x07visible\n") + s := New() + out := s.Feed(input) + out = append(out, s.Flush()...) + expected := "visible\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestFlushUnterminatedLine(t *testing.T) { + s := New() + out := s.Feed([]byte("no newline")) + out = append(out, s.Flush()...) + expected := "no newline" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} + +func TestInvalidUTF8(t *testing.T) { + // 0xFF is not valid UTF-8, should be replaced + s := New() + out := s.Feed([]byte{0xFF, 'a', '\n'}) + out = append(out, s.Flush()...) + expected := "\uFFFDa\n" + if string(out) != expected { + t.Errorf("got %q, want %q", string(out), expected) + } +} diff --git a/dify-agent-runtime/internal/server/api.go b/dify-agent-runtime/internal/server/api.go new file mode 100644 index 00000000000..67c4ee98772 --- /dev/null +++ b/dify-agent-runtime/internal/server/api.go @@ -0,0 +1,285 @@ +package server + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "runtime/debug" + "strconv" + "strings" + "time" +) + +// Handler creates the HTTP handler (mux) for the shellctl API. +func Handler(svc *Service, config *Config) http.Handler { + mux := http.NewServeMux() + + auth := authMiddleware(config.AuthToken) + + mux.HandleFunc("GET /healthz", handleHealthz) + mux.HandleFunc("POST /v1/jobs/run", auth(handleRunJob(svc))) + mux.HandleFunc("POST /v1/jobs/{job_id}/wait", auth(handleWaitJob(svc, config))) + mux.HandleFunc("GET /v1/jobs/{job_id}/log/tail", auth(handleTailJob(svc, config))) + mux.HandleFunc("GET /v1/jobs/{job_id}", auth(handleJobStatus(svc))) + mux.HandleFunc("GET /v1/jobs", auth(handleListJobs(svc, config))) + mux.HandleFunc("POST /v1/jobs/{job_id}/input", auth(handleInputJob(svc, config))) + mux.HandleFunc("POST /v1/jobs/{job_id}/terminate", auth(handleTerminateJob(svc, config))) + mux.HandleFunc("DELETE /v1/jobs/{job_id}", auth(handleDeleteJob(svc, config))) + + return requestLoggingMiddleware(recoveryMiddleware(mux)) +} + +func handleHealthz(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, HealthResponse{Status: HealthStatus}) +} + +func handleRunJob(svc *Service) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + var req RunJobRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeError(w, 400, "invalid_request", "Invalid JSON body") + return + } + if req.Script == "" { + writeError(w, 400, "invalid_request", "script is required") + return + } + // Validate env + if req.Env != nil { + for name, value := range req.Env { + if name == "" { + writeError(w, 422, "validation_error", "env names must be non-empty") + return + } + if strings.Contains(name, "=") { + writeError(w, 422, "validation_error", fmt.Sprintf("env name must not contain '=': %q", name)) + return + } + if strings.Contains(name, "\x00") || strings.Contains(value, "\x00") { + writeError(w, 422, "validation_error", "env entries must not contain NUL") + return + } + } + } + + result, err := svc.RunJob(&req) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +func handleWaitJob(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + var req WaitJobRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeError(w, 400, "invalid_request", "Invalid JSON body") + return + } + if req.IdleFlushSeconds == 0 { + req.IdleFlushSeconds = DefaultIdleFlushSeconds + } + result, err := svc.WaitJob(jobID, &req) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +func handleTailJob(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + outputLimit := config.DefaultOutputLimitBytes + if v := r.URL.Query().Get("output_limit"); v != "" { + if parsed, err := strconv.Atoi(v); err == nil && parsed > 0 { + outputLimit = parsed + } + } + if outputLimit > config.MaxOutputLimitBytes { + outputLimit = config.MaxOutputLimitBytes + } + result, err := svc.TailJob(jobID, outputLimit) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +func handleJobStatus(svc *Service) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + view, err := svc.GetJobStatus(jobID) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, view) + } +} + +func handleListJobs(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + var status *JobStatusName + if v := r.URL.Query().Get("status"); v != "" { + s := JobStatusName(v) + status = &s + } + limit := config.DefaultListLimit + if v := r.URL.Query().Get("limit"); v != "" { + if parsed, err := strconv.Atoi(v); err == nil && parsed > 0 { + limit = parsed + } + } + if limit > config.MaxListLimit { + limit = config.MaxListLimit + } + + result, err := svc.ListJobs(status, limit) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +func handleInputJob(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + var req InputJobRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + writeError(w, 400, "invalid_request", "Invalid JSON body") + return + } + if req.IdleFlushSeconds == 0 { + req.IdleFlushSeconds = DefaultIdleFlushSeconds + } + if req.Timeout == 0 { + req.Timeout = DefaultTimeoutSeconds + } + result, err := svc.SendInput(jobID, &req) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +func handleTerminateJob(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + var req TerminateJobRequest + _ = json.NewDecoder(r.Body).Decode(&req) + graceSeconds := config.DefaultTerminateGraceSeconds + if req.GraceSeconds != nil { + graceSeconds = *req.GraceSeconds + } + view, err := svc.TerminateJob(jobID, graceSeconds) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, view) + } +} + +func handleDeleteJob(svc *Service, config *Config) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + jobID := r.PathValue("job_id") + force := r.URL.Query().Get("force") == "true" + graceSeconds := config.DefaultTerminateGraceSeconds + if v := r.URL.Query().Get("grace_seconds"); v != "" { + if parsed, err := strconv.ParseFloat(v, 64); err == nil { + graceSeconds = parsed + } + } + result, err := svc.DeleteJob(jobID, force, graceSeconds) + if err != nil { + writeServerError(w, err) + return + } + writeJSON(w, http.StatusOK, result) + } +} + +// Middleware + +// statusRecorder wraps ResponseWriter to capture the status code. +type statusRecorder struct { + http.ResponseWriter + statusCode int +} + +func (sr *statusRecorder) WriteHeader(code int) { + sr.statusCode = code + sr.ResponseWriter.WriteHeader(code) +} + +func requestLoggingMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + rec := &statusRecorder{ResponseWriter: w, statusCode: 200} + next.ServeHTTP(rec, r) + log.Printf("%s %s -> %d (%s)", r.Method, r.URL.Path, rec.statusCode, time.Since(start).Round(time.Millisecond)) + }) +} + +func recoveryMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer func() { + if rec := recover(); rec != nil { + stack := debug.Stack() + log.Printf("PANIC %s %s: %v\n%s", r.Method, r.URL.Path, rec, stack) + writeError(w, 500, "internal_panic", fmt.Sprintf("internal server error: %v", rec)) + } + }() + next.ServeHTTP(w, r) + }) +} + +func authMiddleware(token string) func(http.HandlerFunc) http.HandlerFunc { + return func(next http.HandlerFunc) http.HandlerFunc { + if token == "" { + return next // No auth configured + } + expected := "Bearer " + token + return func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != expected { + writeError(w, 401, "unauthorized", "Missing or invalid bearer token") + return + } + next(w, r) + } + } +} + +// Response helpers + +func writeJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(v) +} + +func writeError(w http.ResponseWriter, status int, code, message string) { + writeJSON(w, status, ErrorResponse{Error: ErrorDetail{Code: code, Message: message}}) +} + +func writeServerError(w http.ResponseWriter, err error) { + if se, ok := err.(*ServerError); ok { + log.Printf("ERROR [%d] %s: %s", se.StatusCode, se.Code, se.Message) + writeError(w, se.StatusCode, se.Code, se.Message) + return + } + log.Printf("ERROR [500] internal_error: %v", err) + writeError(w, 500, "internal_error", err.Error()) +} diff --git a/dify-agent-runtime/internal/server/api_test.go b/dify-agent-runtime/internal/server/api_test.go new file mode 100644 index 00000000000..b805143f26a --- /dev/null +++ b/dify-agent-runtime/internal/server/api_test.go @@ -0,0 +1,170 @@ +package server + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// We can't fully test the Service without tmux, but we can test the HTTP +// layer wiring, error handling, and JSON serialization. + +func TestWriteJSON(t *testing.T) { + w := httptest.NewRecorder() + writeJSON(w, 200, HealthResponse{Status: "ok"}) + + if w.Code != 200 { + t.Errorf("expected 200, got %d", w.Code) + } + if w.Header().Get("Content-Type") != "application/json" { + t.Errorf("expected application/json, got %s", w.Header().Get("Content-Type")) + } + + var result HealthResponse + if err := json.NewDecoder(w.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Status != "ok" { + t.Errorf("expected status=ok, got %s", result.Status) + } +} + +func TestWriteError(t *testing.T) { + w := httptest.NewRecorder() + writeError(w, 400, "bad_request", "missing field") + + if w.Code != 400 { + t.Errorf("expected 400, got %d", w.Code) + } + + var result ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Error.Code != "bad_request" { + t.Errorf("expected code=bad_request, got %s", result.Error.Code) + } + if result.Error.Message != "missing field" { + t.Errorf("expected message='missing field', got %s", result.Error.Message) + } +} + +func TestWriteServerError(t *testing.T) { + w := httptest.NewRecorder() + err := NewServerError(404, "job_not_found", "Unknown job id") + writeServerError(w, err) + + if w.Code != 404 { + t.Errorf("expected 404, got %d", w.Code) + } + + var result ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Error.Code != "job_not_found" { + t.Errorf("expected code=job_not_found, got %s", result.Error.Code) + } +} + +func TestWriteServerErrorGeneric(t *testing.T) { + w := httptest.NewRecorder() + writeServerError(w, &json.SyntaxError{Offset: 5}) + + if w.Code != 500 { + t.Errorf("expected 500, got %d", w.Code) + } + + var result ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Error.Code != "internal_error" { + t.Errorf("expected code=internal_error, got %s", result.Error.Code) + } +} + +func TestAuthMiddlewareNoToken(t *testing.T) { + // When no token configured, auth middleware should pass through + handler := authMiddleware("")(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + + req := httptest.NewRequest("GET", "/v1/jobs", nil) + w := httptest.NewRecorder() + handler(w, req) + + if w.Code != 200 { + t.Errorf("expected 200 without auth, got %d", w.Code) + } +} + +func TestAuthMiddlewareWithToken(t *testing.T) { + handler := authMiddleware("secret")(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + + // Without auth header + req := httptest.NewRequest("GET", "/v1/jobs", nil) + w := httptest.NewRecorder() + handler(w, req) + if w.Code != 401 { + t.Errorf("expected 401 without auth, got %d", w.Code) + } + + // With correct auth header + req = httptest.NewRequest("GET", "/v1/jobs", nil) + req.Header.Set("Authorization", "Bearer secret") + w = httptest.NewRecorder() + handler(w, req) + if w.Code != 200 { + t.Errorf("expected 200 with correct auth, got %d", w.Code) + } + + // With wrong auth header + req = httptest.NewRequest("GET", "/v1/jobs", nil) + req.Header.Set("Authorization", "Bearer wrong") + w = httptest.NewRecorder() + handler(w, req) + if w.Code != 401 { + t.Errorf("expected 401 with wrong auth, got %d", w.Code) + } +} + +func TestHealthzHandler(t *testing.T) { + // Create a handler with a nil service (healthz doesn't use it) + mux := http.NewServeMux() + mux.HandleFunc("GET /healthz", handleHealthz) + + req := httptest.NewRequest("GET", "/healthz", nil) + w := httptest.NewRecorder() + mux.ServeHTTP(w, req) + + if w.Code != 200 { + t.Errorf("expected 200, got %d", w.Code) + } + + body := w.Body.String() + if !strings.Contains(body, "ok") { + t.Errorf("expected body to contain 'ok', got %s", body) + } +} + +func TestServerErrorFormat(t *testing.T) { + err := NewServerError(422, "validation_error", "bad input") + expected := "[422] validation_error: bad input" + if err.Error() != expected { + t.Errorf("expected %q, got %q", expected, err.Error()) + } +} + +func TestIsNotFound(t *testing.T) { + if !isNotFound(ErrJobNotFound) { + t.Error("ErrJobNotFound should be detected as not found") + } + if isNotFound(NewServerError(500, "internal_error", "x")) { + t.Error("500 error should not be detected as not found") + } +} diff --git a/dify-agent-runtime/internal/server/config.go b/dify-agent-runtime/internal/server/config.go new file mode 100644 index 00000000000..6b6258a8f54 --- /dev/null +++ b/dify-agent-runtime/internal/server/config.go @@ -0,0 +1,128 @@ +package server + +import ( + "os" + "path/filepath" + "runtime" + "time" +) + +const ( + DefaultListen = "127.0.0.1:8765" + DefaultTimeoutSeconds = 30.0 + DefaultMaxWaitTimeoutSeconds = 600.0 + DefaultIdleFlushSeconds = 0.5 + DefaultTerminalCols = 200 + DefaultTerminalRows = 50 + DefaultOutputLimitBytes = 16 * 1024 + MaxOutputLimitBytes = 512 * 1024 + DefaultListLimit = 50 + MaxListLimit = 200 + DefaultTerminateGraceSeconds = 10.0 + DefaultGCIntervalSeconds = 60.0 + DefaultGCFinishedJobRetentionSeconds = 300.0 + DefaultPollInterval = 50 * time.Millisecond + DefaultPipeMonitorInterval = 1 * time.Second + DefaultPipeReadyTimeout = 10 * time.Second + DefaultSQLiteBusyTimeoutMs = 5000 + DefaultAuthTokenEnv = "SHELLCTL_AUTH_TOKEN" + HealthStatus = "ok" +) + +// Config holds runtime configuration for the shellctl server. +type Config struct { + Listen string + AuthToken string + StateDir string + RuntimeDir string + GCInterval time.Duration + GCFinishedJobRetention time.Duration + DefaultTimeout time.Duration + MaxWaitTimeout time.Duration + IdleFlushDuration time.Duration + DefaultCwd string + DefaultTerminalCols int + DefaultTerminalRows int + DefaultListLimit int + MaxListLimit int + DefaultOutputLimitBytes int + MaxOutputLimitBytes int + DefaultTerminateGraceSeconds float64 + PollInterval time.Duration + PipeMonitorInterval time.Duration + PipeReadyTimeout time.Duration + SQLiteBusyTimeoutMs int + SanitizePtyCommand []string + RunnerExitCommand []string +} + +// DefaultConfig returns a Config with sensible defaults. +func DefaultConfig() *Config { + homeDir, _ := os.UserHomeDir() + stateDir := defaultStateDir() + runtimeDir := filepath.Join(stateDir, "runtime") + + cfg := &Config{ + Listen: DefaultListen, + StateDir: stateDir, + RuntimeDir: runtimeDir, + GCInterval: time.Duration(DefaultGCIntervalSeconds * float64(time.Second)), + GCFinishedJobRetention: time.Duration(DefaultGCFinishedJobRetentionSeconds * float64(time.Second)), + DefaultTimeout: time.Duration(DefaultTimeoutSeconds * float64(time.Second)), + MaxWaitTimeout: time.Duration(DefaultMaxWaitTimeoutSeconds * float64(time.Second)), + IdleFlushDuration: time.Duration(DefaultIdleFlushSeconds * float64(time.Second)), + DefaultCwd: homeDir, + DefaultTerminalCols: DefaultTerminalCols, + DefaultTerminalRows: DefaultTerminalRows, + DefaultListLimit: DefaultListLimit, + MaxListLimit: MaxListLimit, + DefaultOutputLimitBytes: DefaultOutputLimitBytes, + MaxOutputLimitBytes: MaxOutputLimitBytes, + DefaultTerminateGraceSeconds: DefaultTerminateGraceSeconds, + PollInterval: DefaultPollInterval, + PipeMonitorInterval: DefaultPipeMonitorInterval, + PipeReadyTimeout: DefaultPipeReadyTimeout, + SQLiteBusyTimeoutMs: DefaultSQLiteBusyTimeoutMs, + SanitizePtyCommand: []string{"shellctl-sanitize-pty"}, + RunnerExitCommand: []string{"shellctl-runner-exit"}, + } + + // Auth token from environment if not set explicitly + if cfg.AuthToken == "" { + cfg.AuthToken = os.Getenv(DefaultAuthTokenEnv) + } + + return cfg +} + +// JobsDir returns the path to the jobs artifact directory. +func (c *Config) JobsDir() string { + return filepath.Join(c.StateDir, "jobs") +} + +// DBPath returns the path to the SQLite database. +func (c *Config) DBPath() string { + return filepath.Join(c.StateDir, "shellctl.db") +} + +// TmuxSocket returns the path to the dedicated tmux socket. +func (c *Config) TmuxSocket() string { + return filepath.Join(c.RuntimeDir, "tmux.sock") +} + +// RunnerPath returns the path to the installed runner script. +func (c *Config) RunnerPath() string { + return filepath.Join(c.RuntimeDir, "bin", "shellctl-runner") +} + +func defaultStateDir() string { + if runtime.GOOS == "darwin" { + home, _ := os.UserHomeDir() + return filepath.Join(home, ".local", "share", "shellctl") + } + if xdg := os.Getenv("XDG_DATA_HOME"); xdg != "" { + return filepath.Join(xdg, "shellctl") + } + home, _ := os.UserHomeDir() + return filepath.Join(home, ".local", "share", "shellctl") +} diff --git a/dify-agent-runtime/internal/server/config_test.go b/dify-agent-runtime/internal/server/config_test.go new file mode 100644 index 00000000000..0b6aed9449e --- /dev/null +++ b/dify-agent-runtime/internal/server/config_test.go @@ -0,0 +1,56 @@ +package server + +import ( + "testing" +) + +func TestDefaultConfig(t *testing.T) { + cfg := DefaultConfig() + if cfg.Listen != DefaultListen { + t.Errorf("expected Listen=%s, got %s", DefaultListen, cfg.Listen) + } + if cfg.DefaultTerminalCols != DefaultTerminalCols { + t.Errorf("expected cols=%d, got %d", DefaultTerminalCols, cfg.DefaultTerminalCols) + } + if cfg.DefaultTerminalRows != DefaultTerminalRows { + t.Errorf("expected rows=%d, got %d", DefaultTerminalRows, cfg.DefaultTerminalRows) + } + if cfg.SQLiteBusyTimeoutMs != DefaultSQLiteBusyTimeoutMs { + t.Errorf("expected busy_timeout=%d, got %d", DefaultSQLiteBusyTimeoutMs, cfg.SQLiteBusyTimeoutMs) + } +} + +func TestConfigPaths(t *testing.T) { + cfg := DefaultConfig() + cfg.StateDir = "/tmp/shellctl-test" + cfg.RuntimeDir = "/tmp/shellctl-test/runtime" + + if cfg.JobsDir() != "/tmp/shellctl-test/jobs" { + t.Errorf("unexpected JobsDir: %s", cfg.JobsDir()) + } + if cfg.DBPath() != "/tmp/shellctl-test/shellctl.db" { + t.Errorf("unexpected DBPath: %s", cfg.DBPath()) + } + if cfg.TmuxSocket() != "/tmp/shellctl-test/runtime/tmux.sock" { + t.Errorf("unexpected TmuxSocket: %s", cfg.TmuxSocket()) + } + if cfg.RunnerPath() != "/tmp/shellctl-test/runtime/bin/shellctl-runner" { + t.Errorf("unexpected RunnerPath: %s", cfg.RunnerPath()) + } +} + +func TestConfigAuthTokenFromEnv(t *testing.T) { + t.Setenv("SHELLCTL_AUTH_TOKEN", "my-secret-token") + cfg := DefaultConfig() + if cfg.AuthToken != "my-secret-token" { + t.Errorf("expected auth token from env, got %q", cfg.AuthToken) + } +} + +func TestConfigNoAuthToken(t *testing.T) { + t.Setenv("SHELLCTL_AUTH_TOKEN", "") + cfg := DefaultConfig() + if cfg.AuthToken != "" { + t.Errorf("expected empty auth token, got %q", cfg.AuthToken) + } +} diff --git a/dify-agent-runtime/internal/server/db.go b/dify-agent-runtime/internal/server/db.go new file mode 100644 index 00000000000..a7eb6339b52 --- /dev/null +++ b/dify-agent-runtime/internal/server/db.go @@ -0,0 +1,368 @@ +package server + +import ( + "database/sql" + "fmt" + "time" + + _ "modernc.org/sqlite" +) + +// JobStatusName represents the lifecycle states of a shellctl job. +type JobStatusName string + +const ( + StatusCreated JobStatusName = "created" + StatusStarting JobStatusName = "starting" + StatusRunning JobStatusName = "running" + StatusExited JobStatusName = "exited" + StatusTerminated JobStatusName = "terminated" + StatusFailed JobStatusName = "failed" + StatusLost JobStatusName = "lost" +) + +// IsTerminal returns true if the status represents a final job state. +func (s JobStatusName) IsTerminal() bool { + switch s { + case StatusExited, StatusTerminated, StatusFailed, StatusLost: + return true + } + return false +} + +// JobRow represents a row in the jobs SQLite table. +type JobRow struct { + JobID string + ScriptPath string + OutputPath string + Cwd string + TerminalCols int + TerminalRows int + Status JobStatusName + SessionName string + PaneTarget string + ExitCode *int + Reason *string + Message *string + CreatedAt string + StartedAt *string + EndedAt *string + UpdatedAt string +} + +// DB wraps the SQLite database for shellctl job persistence. +type DB struct { + db *sql.DB +} + +// OpenDB opens (or creates) the shellctl SQLite database. +func OpenDB(dbPath string, busyTimeoutMs int) (*DB, error) { + dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(%d)&_pragma=journal_mode(WAL)", dbPath, busyTimeoutMs) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, fmt.Errorf("open sqlite: %w", err) + } + db.SetMaxOpenConns(1) + db.SetConnMaxLifetime(0) + return &DB{db: db}, nil +} + +// Close closes the database connection. +func (d *DB) Close() error { + return d.db.Close() +} + +// InitSchema creates the jobs table if it does not exist. +func (d *DB) InitSchema() error { + _, err := d.db.Exec(` + 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 + ) + `) + return err +} + +// InsertJob inserts a new job row. Returns false if the job_id already exists. +func (d *DB) InsertJob(row *JobRow) (bool, error) { + _, err := d.db.Exec(` + INSERT INTO jobs (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) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + row.JobID, row.ScriptPath, row.OutputPath, row.Cwd, + row.TerminalCols, row.TerminalRows, + string(row.Status), row.SessionName, row.PaneTarget, + row.ExitCode, row.Reason, row.Message, + row.CreatedAt, row.StartedAt, row.EndedAt, row.UpdatedAt, + ) + if err != nil { + // SQLite UNIQUE constraint violation + if isUniqueViolation(err) { + return false, nil + } + return false, err + } + return true, nil +} + +// GetJob retrieves a single job row by ID. +func (d *DB) GetJob(jobID string) (*JobRow, error) { + row := d.db.QueryRow(` + 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 = ?`, jobID) + return scanJobRow(row) +} + +// ListJobs returns all jobs ordered by created_at descending, optionally filtered by status. +func (d *DB) ListJobs(statuses []JobStatusName) ([]*JobRow, error) { + var rows *sql.Rows + var err error + + if len(statuses) == 0 { + rows, err = d.db.Query(`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 ORDER BY created_at DESC`) + } else { + args := make([]any, len(statuses)) + placeholders := "" + for i, s := range statuses { + args[i] = string(s) + if i > 0 { + placeholders += "," + } + placeholders += "?" + } + query := fmt.Sprintf(`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 status IN (%s) ORDER BY created_at DESC`, placeholders) + rows, err = d.db.Query(query, args...) + } + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + + var result []*JobRow + for rows.Next() { + jr, err := scanJobRows(rows) + if err != nil { + return nil, err + } + result = append(result, jr) + } + return result, rows.Err() +} + +// DeleteJob removes a job row by ID. Returns error if not found. +func (d *DB) DeleteJob(jobID string) error { + result, err := d.db.Exec("DELETE FROM jobs WHERE job_id = ?", jobID) + if err != nil { + return err + } + n, _ := result.RowsAffected() + if n == 0 { + return ErrJobNotFound + } + return nil +} + +// TransitionStatus performs a conditional CAS update on a job row. +func (d *DB) TransitionStatus(jobID string, opts TransitionOpts) (*JobRow, error) { + now := FormatTimestamp(time.Now()) + + // Build WHERE clause + allowedList := make([]any, 0, len(opts.AllowedFrom)+len(opts.AllowOverrideFrom)) + placeholders := "" + all := append(opts.AllowedFrom, opts.AllowOverrideFrom...) + for i, s := range all { + allowedList = append(allowedList, string(s)) + if i > 0 { + placeholders += "," + } + placeholders += "?" + } + + setClauses := ` + status = ?, + updated_at = ?, + reason = ?, + message = ?` + + args := []any{string(opts.Target), now, opts.Reason, opts.Message} + + // started_at: set on first transition to starting/running + if opts.Target == StatusStarting || opts.Target == StatusRunning { + setClauses += `, started_at = CASE WHEN started_at IS NULL THEN ? ELSE started_at END` + args = append(args, now) + } + + // ended_at + exit_code: set on terminal transitions + if opts.Target.IsTerminal() { + endedAt := opts.EndedAt + if endedAt == "" { + endedAt = now + } + setClauses += `, ended_at = CASE WHEN ended_at IS NULL THEN ? ELSE ended_at END` + args = append(args, endedAt) + // Set exit_code to 0 if not already set (normal exit without runner-exit callback) + setClauses += `, exit_code = CASE WHEN exit_code IS NULL THEN 0 ELSE exit_code END` + } + + whereClause := fmt.Sprintf("job_id = ? AND status IN (%s)", placeholders) + args = append(args, jobID) + args = append(args, allowedList...) + + if opts.RequireExitCodeNull { + whereClause += " AND exit_code IS NULL" + } + + query := fmt.Sprintf("UPDATE jobs SET %s WHERE %s", setClauses, whereClause) + result, err := d.db.Exec(query, args...) + if err != nil { + return nil, fmt.Errorf("transition status: %w", err) + } + + n, _ := result.RowsAffected() + _ = n // If 0 rows affected, we just return current state + + return d.GetJob(jobID) +} + +// RecordRunnerExit persists exit_code and ended_at for a drained job. +// The update is idempotent for terminal rows: once a job reaches a terminal +// state, this method returns nil without rewriting the row. +func (d *DB) RecordRunnerExit(jobID string, exitCode int, endedAt string) error { + // Check if job exists and is already terminal + var status string + err := d.db.QueryRow("SELECT status FROM jobs WHERE job_id = ?", jobID).Scan(&status) + if err == sql.ErrNoRows { + return ErrJobNotFound + } + if err != nil { + return err + } + if JobStatusName(status).IsTerminal() { + return nil + } + + nonterminal := []any{"created", "starting", "running"} + + result, err := d.db.Exec(` + UPDATE jobs + SET status = CASE WHEN status IN (?, ?, ?) THEN ? ELSE status END, + exit_code = CASE WHEN status IN (?, ?, ?) THEN ? ELSE exit_code END, + ended_at = CASE WHEN status IN (?, ?, ?) AND ended_at IS NULL THEN ? ELSE ended_at END, + updated_at = CASE WHEN status IN (?, ?, ?) THEN ? ELSE updated_at END, + reason = CASE WHEN status IN (?, ?, ?) THEN NULL ELSE reason END, + message = CASE WHEN status IN (?, ?, ?) THEN NULL ELSE message END + WHERE job_id = ? AND status IN (?, ?, ?)`, + nonterminal[0], nonterminal[1], nonterminal[2], string(StatusExited), + // exit_code + nonterminal[0], nonterminal[1], nonterminal[2], exitCode, + // ended_at + nonterminal[0], nonterminal[1], nonterminal[2], endedAt, + // updated_at + nonterminal[0], nonterminal[1], nonterminal[2], endedAt, + // reason + nonterminal[0], nonterminal[1], nonterminal[2], + // message + nonterminal[0], nonterminal[1], nonterminal[2], + // WHERE + jobID, + nonterminal[0], nonterminal[1], nonterminal[2], + ) + if err != nil { + return err + } + n, _ := result.RowsAffected() + if n == 0 { + return ErrJobNotFound + } + return nil +} + +// TransitionOpts holds parameters for a CAS state transition. +type TransitionOpts struct { + AllowedFrom []JobStatusName + AllowOverrideFrom []JobStatusName + Target JobStatusName + RequireExitCodeNull bool + Reason *string + Message *string + EndedAt string +} + +func scanJobRow(row *sql.Row) (*JobRow, error) { + var jr JobRow + var status string + err := row.Scan( + &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &jr.Cwd, + &jr.TerminalCols, &jr.TerminalRows, + &status, &jr.SessionName, &jr.PaneTarget, + &jr.ExitCode, &jr.Reason, &jr.Message, + &jr.CreatedAt, &jr.StartedAt, &jr.EndedAt, &jr.UpdatedAt, + ) + if err == sql.ErrNoRows { + return nil, ErrJobNotFound + } + if err != nil { + return nil, err + } + jr.Status = JobStatusName(status) + return &jr, nil +} + +func scanJobRows(rows *sql.Rows) (*JobRow, error) { + var jr JobRow + var status string + err := rows.Scan( + &jr.JobID, &jr.ScriptPath, &jr.OutputPath, &jr.Cwd, + &jr.TerminalCols, &jr.TerminalRows, + &status, &jr.SessionName, &jr.PaneTarget, + &jr.ExitCode, &jr.Reason, &jr.Message, + &jr.CreatedAt, &jr.StartedAt, &jr.EndedAt, &jr.UpdatedAt, + ) + if err != nil { + return nil, err + } + jr.Status = JobStatusName(status) + return &jr, nil +} + +func isUniqueViolation(err error) bool { + // modernc.org/sqlite returns error messages containing "UNIQUE constraint failed" + return err != nil && contains(err.Error(), "UNIQUE constraint failed") +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(substr) == 0 || + (len(s) > 0 && len(substr) > 0 && searchString(s, substr))) +} + +func searchString(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/dify-agent-runtime/internal/server/db_test.go b/dify-agent-runtime/internal/server/db_test.go new file mode 100644 index 00000000000..c3bbc2c3027 --- /dev/null +++ b/dify-agent-runtime/internal/server/db_test.go @@ -0,0 +1,309 @@ +package server + +import ( + "os" + "path/filepath" + "testing" +) + +func TestJobStatusIsTerminal(t *testing.T) { + terminal := []JobStatusName{StatusExited, StatusTerminated, StatusFailed, StatusLost} + for _, s := range terminal { + if !s.IsTerminal() { + t.Errorf("%s should be terminal", s) + } + } + nonTerminal := []JobStatusName{StatusCreated, StatusStarting, StatusRunning} + for _, s := range nonTerminal { + if s.IsTerminal() { + t.Errorf("%s should not be terminal", s) + } + } +} + +func TestOpenDBAndInitSchema(t *testing.T) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "shellctl.db") + + db, err := OpenDB(dbPath, 5000) + if err != nil { + t.Fatalf("OpenDB: %v", err) + } + defer func() { _ = db.Close() }() + + if err := db.InitSchema(); err != nil { + t.Fatalf("InitSchema: %v", err) + } + + // Verify table exists by inserting and reading back + row := &JobRow{ + JobID: "test-job-1", + ScriptPath: "jobs/test-job-1/script", + OutputPath: "jobs/test-job-1/output.log", + Cwd: "/tmp", + TerminalCols: 80, + TerminalRows: 24, + Status: StatusCreated, + SessionName: "shellctl-test-job-1", + PaneTarget: "shellctl-test-job-1:0.0", + CreatedAt: "2025-01-01T00:00:00Z", + UpdatedAt: "2025-01-01T00:00:00Z", + } + + ok, err := db.InsertJob(row) + if err != nil { + t.Fatalf("InsertJob: %v", err) + } + if !ok { + t.Error("expected insert to succeed (ok=true)") + } + + // Duplicate insert should return ok=false + ok, err = db.InsertJob(row) + if err != nil { + t.Fatalf("InsertJob duplicate: %v", err) + } + if ok { + t.Error("expected duplicate insert to return ok=false") + } +} + +func TestGetJob(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-get-1", StatusCreated) + + row, err := db.GetJob("job-get-1") + if err != nil { + t.Fatalf("GetJob: %v", err) + } + if row.JobID != "job-get-1" { + t.Errorf("expected job_id=job-get-1, got %s", row.JobID) + } + if row.Status != StatusCreated { + t.Errorf("expected status=created, got %s", row.Status) + } + if row.TerminalCols != 80 { + t.Errorf("expected cols=80, got %d", row.TerminalCols) + } +} + +func TestGetJobNotFound(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + _, err := db.GetJob("nonexistent") + if err != ErrJobNotFound { + t.Errorf("expected ErrJobNotFound, got %v", err) + } +} + +func TestListJobs(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-list-1", StatusRunning) + insertTestJob(t, db, "job-list-2", StatusExited) + insertTestJob(t, db, "job-list-3", StatusCreated) + + // List all + rows, err := db.ListJobs(nil) + if err != nil { + t.Fatalf("ListJobs: %v", err) + } + if len(rows) != 3 { + t.Errorf("expected 3 jobs, got %d", len(rows)) + } + + // List by status + rows, err = db.ListJobs([]JobStatusName{StatusRunning}) + if err != nil { + t.Fatalf("ListJobs filtered: %v", err) + } + if len(rows) != 1 { + t.Errorf("expected 1 running job, got %d", len(rows)) + } + if rows[0].JobID != "job-list-1" { + t.Errorf("expected job-list-1, got %s", rows[0].JobID) + } +} + +func TestDeleteJob(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-del-1", StatusExited) + + if err := db.DeleteJob("job-del-1"); err != nil { + t.Fatalf("DeleteJob: %v", err) + } + + _, err := db.GetJob("job-del-1") + if err != ErrJobNotFound { + t.Errorf("expected ErrJobNotFound after delete, got %v", err) + } +} + +func TestDeleteJobNotFound(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + err := db.DeleteJob("nonexistent") + if err != ErrJobNotFound { + t.Errorf("expected ErrJobNotFound, got %v", err) + } +} + +func TestTransitionStatus(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-trans-1", StatusCreated) + + // Transition created → starting + row, err := db.TransitionStatus("job-trans-1", TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated}, + Target: StatusStarting, + }) + if err != nil { + t.Fatalf("TransitionStatus: %v", err) + } + if row.Status != StatusStarting { + t.Errorf("expected starting, got %s", row.Status) + } + + // Transition starting → running + row, err = db.TransitionStatus("job-trans-1", TransitionOpts{ + AllowedFrom: []JobStatusName{StatusStarting}, + Target: StatusRunning, + RequireExitCodeNull: true, + }) + if err != nil { + t.Fatalf("TransitionStatus: %v", err) + } + if row.Status != StatusRunning { + t.Errorf("expected running, got %s", row.Status) + } + + // Transition running → exited + exitCode := 0 + row, err = db.TransitionStatus("job-trans-1", TransitionOpts{ + AllowedFrom: []JobStatusName{StatusRunning}, + Target: StatusExited, + }) + if err != nil { + t.Fatalf("TransitionStatus: %v", err) + } + if row.Status != StatusExited { + t.Errorf("expected exited, got %s", row.Status) + } + if row.ExitCode == nil { + t.Error("expected exit_code to be set after exited transition") + } else if *row.ExitCode != exitCode { + t.Errorf("expected exit_code=0, got %d", *row.ExitCode) + } + if row.EndedAt == nil { + t.Error("expected ended_at to be set after terminal transition") + } +} + +func TestRecordRunnerExit(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-exit-1", StatusRunning) + + if err := db.RecordRunnerExit("job-exit-1", 42, "2025-01-15T12:00:00Z"); err != nil { + t.Fatalf("RecordRunnerExit: %v", err) + } + + row, _ := db.GetJob("job-exit-1") + if row.Status != StatusExited { + t.Errorf("expected exited, got %s", row.Status) + } + if row.ExitCode == nil || *row.ExitCode != 42 { + t.Errorf("expected exit_code=42, got %v", row.ExitCode) + } +} + +func TestRecordRunnerExitIdempotent(t *testing.T) { + dir := t.TempDir() + db := setupTestDB(t, dir) + defer func() { _ = db.Close() }() + + insertTestJob(t, db, "job-exit-2", StatusExited) + exitCode := 10 + row := &JobRow{ + JobID: "job-exit-2", ScriptPath: "x", OutputPath: "y", Cwd: "/tmp", + TerminalCols: 80, TerminalRows: 24, Status: StatusExited, + SessionName: "s", PaneTarget: "p", ExitCode: &exitCode, + CreatedAt: "2025-01-01T00:00:00Z", UpdatedAt: "2025-01-01T00:00:00Z", + EndedAt: strPtr("2025-01-01T00:01:00Z"), + } + _, _ = db.db.Exec(`UPDATE jobs SET exit_code=?, ended_at=? WHERE job_id=?`, + exitCode, "2025-01-01T00:01:00Z", "job-exit-2") + _ = row + + // Should not overwrite existing terminal state + err := db.RecordRunnerExit("job-exit-2", 99, "2025-01-01T00:02:00Z") + if err != nil { + t.Fatalf("RecordRunnerExit on terminal: %v", err) + } + + got, _ := db.GetJob("job-exit-2") + if got.ExitCode != nil && *got.ExitCode != 10 { + t.Errorf("expected exit_code=10 (preserved), got %d", *got.ExitCode) + } +} + +// Helpers + +func setupTestDB(t *testing.T, dir string) *DB { + t.Helper() + dbPath := filepath.Join(dir, "shellctl.db") + db, err := OpenDB(dbPath, 5000) + if err != nil { + t.Fatalf("OpenDB: %v", err) + } + if err := db.InitSchema(); err != nil { + t.Fatalf("InitSchema: %v", err) + } + return db +} + +func insertTestJob(t *testing.T, db *DB, jobID string, status JobStatusName) { + t.Helper() + row := &JobRow{ + JobID: jobID, + ScriptPath: "jobs/" + jobID + "/script", + OutputPath: "jobs/" + jobID + "/output.log", + Cwd: "/tmp", + TerminalCols: 80, + TerminalRows: 24, + Status: status, + SessionName: "shellctl-" + jobID, + PaneTarget: "shellctl-" + jobID + ":0.0", + CreatedAt: "2025-01-01T00:00:00Z", + UpdatedAt: "2025-01-01T00:00:00Z", + } + ok, err := db.InsertJob(row) + if err != nil || !ok { + t.Fatalf("InsertJob(%s): err=%v ok=%v", jobID, err, ok) + } +} + +func strPtr(s string) *string { + return &s +} + +func TestMain(m *testing.M) { + os.Exit(m.Run()) +} diff --git a/dify-agent-runtime/internal/server/errors.go b/dify-agent-runtime/internal/server/errors.go new file mode 100644 index 00000000000..324cc0b375a --- /dev/null +++ b/dify-agent-runtime/internal/server/errors.go @@ -0,0 +1,22 @@ +package server + +import "fmt" + +// ServerError is a structured error with HTTP status code and machine-readable code. +type ServerError struct { + StatusCode int + Code string + Message string +} + +func (e *ServerError) Error() string { + return fmt.Sprintf("[%d] %s: %s", e.StatusCode, e.Code, e.Message) +} + +// Common sentinel errors. +var ErrJobNotFound = &ServerError{StatusCode: 404, Code: "job_not_found", Message: "Unknown job id"} + +// NewServerError creates a new ServerError. +func NewServerError(statusCode int, code, message string) *ServerError { + return &ServerError{StatusCode: statusCode, Code: code, Message: message} +} diff --git a/dify-agent-runtime/internal/server/output.go b/dify-agent-runtime/internal/server/output.go new file mode 100644 index 00000000000..865064bc8df --- /dev/null +++ b/dify-agent-runtime/internal/server/output.go @@ -0,0 +1,167 @@ +package server + +import ( + "fmt" + "os" + "unicode/utf8" +) + +// OutputWindow represents a UTF-8-safe slice of output.log. +type OutputWindow struct { + Output string `json:"output"` + Offset int `json:"offset"` + Truncated bool `json:"truncated"` +} + +// ReadOutputWindow reads a forward UTF-8-safe slice from an output file. +func ReadOutputWindow(path string, offset, limit int) (*OutputWindow, error) { + info, err := os.Stat(path) + if os.IsNotExist(err) { + if offset == 0 { + return &OutputWindow{Output: "", Offset: 0, Truncated: false}, nil + } + return nil, NewServerError(400, "invalid_offset", "offset exceeds current file size 0") + } + if err != nil { + return nil, err + } + + size := int(info.Size()) + if offset > size { + return nil, NewServerError(400, "invalid_offset", + fmt.Sprintf("offset %d exceeds current file size %d", offset, size)) + } + if offset == size { + return &OutputWindow{Output: "", Offset: offset, Truncated: false}, nil + } + + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + + // Read up to limit+4 bytes to handle UTF-8 boundary + readSize := limit + 4 + if offset+readSize > size { + readSize = size - offset + } + + buf := make([]byte, readSize) + _, err = f.ReadAt(buf, int64(offset)) + if err != nil { + return nil, err + } + + // Advance past any UTF-8 continuation bytes at the start + startShift := advanceToUTF8Boundary(buf, 0) + data := buf[startShift:] + budget := limit - startShift + if budget < 0 { + budget = 0 + } + + // Find the longest valid UTF-8 prefix within budget + consumed := validUTF8PrefixLen(data, budget) + if consumed == 0 && len(data) > 0 { + // At least consume one complete rune + consumed = firstCompleteRuneLen(data) + } + + outputBytes := data[:consumed] + newOffset := offset + startShift + consumed + truncated := newOffset < size + + return &OutputWindow{ + Output: string(outputBytes), + Offset: newOffset, + Truncated: truncated, + }, nil +} + +// TailOutputWindow reads a UTF-8-safe tail snapshot from an output file. +func TailOutputWindow(path string, limit int) (*OutputWindow, error) { + info, err := os.Stat(path) + if os.IsNotExist(err) { + return &OutputWindow{Output: "", Offset: 0, Truncated: false}, nil + } + if err != nil { + return nil, err + } + + size := int(info.Size()) + if size == 0 { + return &OutputWindow{Output: "", Offset: 0, Truncated: false}, nil + } + + start := size - limit + if start < 0 { + start = 0 + } + paddedStart := start - 4 + if paddedStart < 0 { + paddedStart = 0 + } + + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer func() { _ = f.Close() }() + + readLen := size - paddedStart + buf := make([]byte, readLen) + _, err = f.ReadAt(buf, int64(paddedStart)) + if err != nil { + return nil, err + } + + relativeStart := advanceToUTF8Boundary(buf, start-paddedStart) + payload := buf[relativeStart:] + consumed := validUTF8PrefixLen(payload, len(payload)) + outputBytes := payload[:consumed] + + return &OutputWindow{ + Output: string(outputBytes), + Offset: paddedStart + relativeStart + consumed, + Truncated: false, + }, nil +} + +// advanceToUTF8Boundary skips continuation bytes (10xxxxxx) at position start. +func advanceToUTF8Boundary(data []byte, start int) int { + for start < len(data) && isUTF8Continuation(data[start]) { + start++ + } + return start +} + +// isUTF8Continuation returns true if byte is a UTF-8 continuation byte. +func isUTF8Continuation(b byte) bool { + return (b & 0xC0) == 0x80 +} + +// validUTF8PrefixLen returns the length of the longest valid UTF-8 prefix up to maxLen. +func validUTF8PrefixLen(data []byte, maxLen int) int { + if maxLen > len(data) { + maxLen = len(data) + } + end := maxLen + for end > 0 { + if utf8.Valid(data[:end]) { + return end + } + end-- + } + return 0 +} + +// firstCompleteRuneLen returns the byte length of the first complete rune in data. +func firstCompleteRuneLen(data []byte) int { + for end := 1; end <= len(data) && end <= 4; end++ { + if utf8.Valid(data[:end]) { + return end + } + } + return 0 +} diff --git a/dify-agent-runtime/internal/server/output_test.go b/dify-agent-runtime/internal/server/output_test.go new file mode 100644 index 00000000000..8b6bbd214ba --- /dev/null +++ b/dify-agent-runtime/internal/server/output_test.go @@ -0,0 +1,275 @@ +package server + +import ( + "os" + "path/filepath" + "testing" +) + +func TestReadOutputWindowEmptyFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + if err := os.WriteFile(path, []byte{}, 0644); err != nil { + t.Fatal(err) + } + + w, err := ReadOutputWindow(path, 0, 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != "" { + t.Errorf("expected empty output, got %q", w.Output) + } + if w.Offset != 0 { + t.Errorf("expected offset=0, got %d", w.Offset) + } + if w.Truncated { + t.Error("expected truncated=false") + } +} + +func TestReadOutputWindowNonexistentFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "nonexistent.log") + + w, err := ReadOutputWindow(path, 0, 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != "" { + t.Errorf("expected empty output, got %q", w.Output) + } +} + +func TestReadOutputWindowFullContent(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "hello\nworld\n" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := ReadOutputWindow(path, 0, 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != content { + t.Errorf("expected %q, got %q", content, w.Output) + } + if w.Truncated { + t.Error("expected truncated=false") + } +} + +func TestReadOutputWindowWithOffset(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "hello\nworld\n" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := ReadOutputWindow(path, 6, 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != "world\n" { + t.Errorf("expected 'world\\n', got %q", w.Output) + } + if w.Offset != len(content) { + t.Errorf("expected offset=%d, got %d", len(content), w.Offset) + } +} + +func TestReadOutputWindowTruncated(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "0123456789" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + // Read only 5 bytes + w, err := ReadOutputWindow(path, 0, 5) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != "01234" { + t.Errorf("expected '01234', got %q", w.Output) + } + if !w.Truncated { + t.Error("expected truncated=true") + } + if w.Offset != 5 { + t.Errorf("expected offset=5, got %d", w.Offset) + } +} + +func TestReadOutputWindowOffsetExceedsSize(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + if err := os.WriteFile(path, []byte("short"), 0644); err != nil { + t.Fatal(err) + } + + _, err := ReadOutputWindow(path, 100, 1024) + if err == nil { + t.Error("expected error for offset > size") + } +} + +func TestReadOutputWindowOffsetAtEnd(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "hello" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := ReadOutputWindow(path, len(content), 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != "" { + t.Errorf("expected empty output at end, got %q", w.Output) + } + if w.Offset != len(content) { + t.Errorf("expected offset=%d, got %d", len(content), w.Offset) + } +} + +func TestReadOutputWindowMultibyteUTF8(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "hello 世界\n" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := ReadOutputWindow(path, 0, 1024) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + if w.Output != content { + t.Errorf("expected %q, got %q", content, w.Output) + } +} + +func TestReadOutputWindowMultibyteTruncated(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + // "世" = 3 bytes (E4 B8 96), "界" = 3 bytes (E7 95 8C) + content := "世界" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + // Read only 4 bytes — should not split a multibyte char + w, err := ReadOutputWindow(path, 0, 4) + if err != nil { + t.Fatalf("ReadOutputWindow: %v", err) + } + // Should get "世" (3 bytes), not partial "界" + if w.Output != "世" { + t.Errorf("expected '世', got %q", w.Output) + } + if !w.Truncated { + t.Error("expected truncated=true") + } +} + +func TestTailOutputWindow(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "line1\nline2\nline3\n" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := TailOutputWindow(path, 100) + if err != nil { + t.Fatalf("TailOutputWindow: %v", err) + } + if w.Output != content { + t.Errorf("expected %q, got %q", content, w.Output) + } +} + +func TestTailOutputWindowLimited(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "output.log") + content := "0123456789" + if err := os.WriteFile(path, []byte(content), 0644); err != nil { + t.Fatal(err) + } + + w, err := TailOutputWindow(path, 5) + if err != nil { + t.Fatalf("TailOutputWindow: %v", err) + } + if w.Output != "56789" { + t.Errorf("expected '56789', got %q", w.Output) + } +} + +func TestTailOutputWindowNonexistent(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "nonexistent.log") + + w, err := TailOutputWindow(path, 100) + if err != nil { + t.Fatalf("TailOutputWindow: %v", err) + } + if w.Output != "" { + t.Errorf("expected empty, got %q", w.Output) + } +} + +func TestTailOutputWindowEmpty(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "empty.log") + if err := os.WriteFile(path, []byte{}, 0644); err != nil { + t.Fatal(err) + } + + w, err := TailOutputWindow(path, 100) + if err != nil { + t.Fatalf("TailOutputWindow: %v", err) + } + if w.Output != "" { + t.Errorf("expected empty, got %q", w.Output) + } +} + +func TestIsUTF8Continuation(t *testing.T) { + // ASCII bytes are not continuation bytes + if isUTF8Continuation('a') { + t.Error("'a' should not be a continuation byte") + } + // 0x80-0xBF are continuation bytes + if !isUTF8Continuation(0x80) { + t.Error("0x80 should be a continuation byte") + } + if !isUTF8Continuation(0xBF) { + t.Error("0xBF should be a continuation byte") + } + // 0xC0 is not a continuation byte (it's a lead byte) + if isUTF8Continuation(0xC0) { + t.Error("0xC0 should not be a continuation byte") + } +} + +func TestAdvanceToUTF8Boundary(t *testing.T) { + // "世" = E4 B8 96, continuation bytes at index 1 and 2 + data := []byte{0xE4, 0xB8, 0x96, 0x41} // 世 + A + // Start at 1, should advance to 3 (the 'A') + result := advanceToUTF8Boundary(data, 1) + if result != 3 { + t.Errorf("expected 3, got %d", result) + } + // Start at 0 (lead byte), should stay at 0 + result = advanceToUTF8Boundary(data, 0) + if result != 0 { + t.Errorf("expected 0, got %d", result) + } +} diff --git a/dify-agent-runtime/internal/server/runtime.go b/dify-agent-runtime/internal/server/runtime.go new file mode 100644 index 00000000000..0cfd5e7cad5 --- /dev/null +++ b/dify-agent-runtime/internal/server/runtime.go @@ -0,0 +1,35 @@ +package server + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "time" +) + +// GenerateJobID produces a short random hex job identifier. +func GenerateJobID() string { + b := make([]byte, 8) + _, _ = rand.Read(b) + return hex.EncodeToString(b) +} + +// FormatTimestamp returns an ISO-8601 UTC timestamp string. +func FormatTimestamp(t time.Time) string { + return t.UTC().Truncate(time.Second).Format("2006-01-02T15:04:05Z") +} + +// ParseTimestamp parses an ISO-8601 UTC timestamp. +func ParseTimestamp(s string) (time.Time, error) { + return time.Parse("2006-01-02T15:04:05Z", s) +} + +// JobSessionName returns the tmux session name for a job. +func JobSessionName(jobID string) string { + return fmt.Sprintf("shellctl-%s", jobID) +} + +// JobPaneTarget returns the tmux pane target for a job. +func JobPaneTarget(jobID string) string { + return fmt.Sprintf("%s:0.0", JobSessionName(jobID)) +} diff --git a/dify-agent-runtime/internal/server/runtime_test.go b/dify-agent-runtime/internal/server/runtime_test.go new file mode 100644 index 00000000000..745a53f715f --- /dev/null +++ b/dify-agent-runtime/internal/server/runtime_test.go @@ -0,0 +1,71 @@ +package server + +import ( + "testing" + "time" +) + +func TestGenerateJobID(t *testing.T) { + id1 := GenerateJobID() + id2 := GenerateJobID() + if id1 == id2 { + t.Error("expected different job IDs") + } + if len(id1) != 16 { + t.Errorf("expected 16-char hex ID, got %d chars: %s", len(id1), id1) + } +} + +func TestGenerateJobIDUniqueness(t *testing.T) { + seen := make(map[string]bool, 1000) + for i := 0; i < 1000; i++ { + id := GenerateJobID() + if seen[id] { + t.Fatalf("collision at iteration %d: %s", i, id) + } + seen[id] = true + } +} + +func TestFormatTimestamp(t *testing.T) { + ts := time.Date(2025, 1, 15, 12, 30, 45, 0, time.UTC) + formatted := FormatTimestamp(ts) + expected := "2025-01-15T12:30:45Z" + if formatted != expected { + t.Errorf("expected %q, got %q", expected, formatted) + } +} + +func TestParseTimestamp(t *testing.T) { + parsed, err := ParseTimestamp("2025-01-15T12:30:45Z") + if err != nil { + t.Fatalf("parse error: %v", err) + } + if parsed.Year() != 2025 || parsed.Month() != 1 || parsed.Day() != 15 { + t.Errorf("unexpected parsed time: %v", parsed) + } + if parsed.Hour() != 12 || parsed.Minute() != 30 || parsed.Second() != 45 { + t.Errorf("unexpected parsed time: %v", parsed) + } +} + +func TestParseTimestampInvalid(t *testing.T) { + _, err := ParseTimestamp("not-a-timestamp") + if err == nil { + t.Error("expected error for invalid timestamp") + } +} + +func TestJobSessionName(t *testing.T) { + name := JobSessionName("abc123") + if name != "shellctl-abc123" { + t.Errorf("expected 'shellctl-abc123', got %q", name) + } +} + +func TestJobPaneTarget(t *testing.T) { + target := JobPaneTarget("abc123") + if target != "shellctl-abc123:0.0" { + t.Errorf("expected 'shellctl-abc123:0.0', got %q", target) + } +} diff --git a/dify-agent-runtime/internal/server/server.go b/dify-agent-runtime/internal/server/server.go new file mode 100644 index 00000000000..d79752f9fa9 --- /dev/null +++ b/dify-agent-runtime/internal/server/server.go @@ -0,0 +1,6 @@ +// Package server implements the shellctl HTTP API and job lifecycle service. +// +// This is the Go equivalent of shellctl/server/ in the Python implementation. +// It manages tmux-backed shell jobs with SQLite persistence and provides +// REST endpoints for run, wait, input, terminate, and delete operations. +package server diff --git a/dify-agent-runtime/internal/server/service.go b/dify-agent-runtime/internal/server/service.go new file mode 100644 index 00000000000..d02e812a3b9 --- /dev/null +++ b/dify-agent-runtime/internal/server/service.go @@ -0,0 +1,943 @@ +package server + +import ( + "context" + "fmt" + "log" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "time" +) + +// Service is the core job lifecycle manager backed by SQLite and tmux. +type Service struct { + config *Config + db *DB + tmux *TmuxController + startingJobs map[string]bool + mu sync.Mutex + cancelGC context.CancelFunc + cancelMon context.CancelFunc +} + +// NewService creates a new shellctl service. +func NewService(config *Config) *Service { + return &Service{ + config: config, + tmux: NewTmuxController(config), + startingJobs: make(map[string]bool), + } +} + +// Initialize performs the full server startup (prepare + reconcile + gc). +func (s *Service) Initialize() error { + if err := s.PrepareRuntime(); err != nil { + return err + } + if err := s.Reconcile(); err != nil { + return err + } + return s.GCOnce() +} + +// PrepareRuntime sets up directories, DB schema, runner script, and tmux server. +func (s *Service) PrepareRuntime() error { + if err := os.MkdirAll(s.config.StateDir, 0700); err != nil { + return err + } + if err := os.MkdirAll(s.config.RuntimeDir, 0700); err != nil { + return err + } + if err := os.MkdirAll(s.config.JobsDir(), 0700); err != nil { + return err + } + runnerDir := filepath.Dir(s.config.RunnerPath()) + if err := os.MkdirAll(runnerDir, 0700); err != nil { + return err + } + + db, err := OpenDB(s.config.DBPath(), s.config.SQLiteBusyTimeoutMs) + if err != nil { + return err + } + s.db = db + + if err := s.db.InitSchema(); err != nil { + return err + } + + s.installRunner() + return s.tmux.StartServer() +} + +// Shutdown stops background goroutines and closes the database. +func (s *Service) Shutdown() { + if s.cancelGC != nil { + s.cancelGC() + } + if s.cancelMon != nil { + s.cancelMon() + } + if s.db != nil { + _ = s.db.Close() + } +} + +// StartBackgroundGC starts the periodic GC goroutine. +func (s *Service) StartBackgroundGC() { + ctx, cancel := context.WithCancel(context.Background()) + s.cancelGC = cancel + go func() { + ticker := time.NewTicker(s.config.GCInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + _ = s.GCOnce() + } + } + }() +} + +// StartBackgroundPipeMonitor starts the periodic pipe health check goroutine. +func (s *Service) StartBackgroundPipeMonitor() { + ctx, cancel := context.WithCancel(context.Background()) + s.cancelMon = cancel + go func() { + ticker := time.NewTicker(s.config.PipeMonitorInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + s.CheckRunningJobsPipeHealth() + } + } + }() +} + +// RunJob creates and starts a new tmux-backed job, then waits for initial output. +func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { + log.Printf("RunJob: script=%d bytes, cwd=%v, env_keys=%d", len(req.Script), req.Cwd, len(req.Env)) + cwd, err := s.resolveCwd(req.Cwd) + if err != nil { + log.Printf("RunJob: resolveCwd failed: %v", err) + return nil, err + } + + cols := s.config.DefaultTerminalCols + rows := s.config.DefaultTerminalRows + if req.Terminal != nil { + cols = req.Terminal.Cols + rows = req.Terminal.Rows + } + + timeout := s.config.DefaultTimeout + if req.Timeout > 0 { + timeout = time.Duration(req.Timeout * float64(time.Second)) + } + outputLimit := s.config.DefaultOutputLimitBytes + if req.OutputLimit > 0 { + outputLimit = req.OutputLimit + } + idleFlush := s.config.IdleFlushDuration + if req.IdleFlushSeconds > 0 { + idleFlush = time.Duration(req.IdleFlushSeconds * float64(time.Second)) + } + + createdAt := FormatTimestamp(time.Now()) + + // Allocate job directory and insert DB row + jobID, jobDir, err := s.allocateJobDir() + if err != nil { + return nil, err + } + + s.mu.Lock() + s.startingJobs[jobID] = true + s.mu.Unlock() + + // Write script and env files + scriptPath := filepath.Join(jobDir, "script") + outputPath := filepath.Join(jobDir, "output.log") + envPath := filepath.Join(jobDir, ".job-env.json") + + if err := os.WriteFile(scriptPath, []byte(req.Script), 0600); err != nil { + s.cleanupStarting(jobID, jobDir) + return nil, err + } + envJSON := "{}" + if req.Env != nil { + pairs := make([]string, 0, len(req.Env)) + for k, v := range req.Env { + pairs = append(pairs, fmt.Sprintf("%q:%q", k, v)) + } + envJSON = "{" + strings.Join(pairs, ",") + "}" + } + if err := os.WriteFile(envPath, []byte(envJSON), 0600); err != nil { + s.cleanupStarting(jobID, jobDir) + return nil, err + } + if err := os.WriteFile(outputPath, []byte{}, 0600); err != nil { + s.cleanupStarting(jobID, jobDir) + return nil, err + } + + row := &JobRow{ + JobID: jobID, + ScriptPath: fmt.Sprintf("jobs/%s/script", jobID), + OutputPath: fmt.Sprintf("jobs/%s/output.log", jobID), + Cwd: cwd, + TerminalCols: cols, + TerminalRows: rows, + Status: StatusCreated, + SessionName: JobSessionName(jobID), + PaneTarget: JobPaneTarget(jobID), + CreatedAt: createdAt, + UpdatedAt: createdAt, + } + + ok, err := s.db.InsertJob(row) + if err != nil { + s.cleanupStarting(jobID, jobDir) + return nil, err + } + if !ok { + s.cleanupStarting(jobID, jobDir) + return nil, NewServerError(500, "job_id_collision", "Failed to allocate a unique job id") + } + + // Transition to starting + _, _ = s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated}, + Target: StatusStarting, + }) + + // Create tmux session and enable output pipe + log.Printf("RunJob [%s]: starting job, cwd=%s", jobID, cwd) + startErr := s.startJob(jobID, jobDir, cwd, cols, rows) + if startErr != nil { + log.Printf("RunJob [%s]: start failed: %v", jobID, startErr) + reason := "start_failed" + msg := startErr.Error() + _, _ = s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting, StatusRunning}, + Target: StatusFailed, + Reason: &reason, + Message: &msg, + }) + s.tmux.CleanupSession(jobID) + } else { + log.Printf("RunJob [%s]: started successfully", jobID) + } + + s.mu.Lock() + delete(s.startingJobs, jobID) + s.mu.Unlock() + + // Wait for initial output + return s.WaitJob(jobID, &WaitJobRequest{ + Offset: 0, + Timeout: timeout.Seconds(), + OutputLimit: outputLimit, + IdleFlushSeconds: idleFlush.Seconds(), + }) +} + +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) + return err + } + + // Wait for pipe ready handshake + if err := s.waitForPipeReady(jobID, pipeReadyPath); err != nil { + log.Printf("startJob [%s]: pipe-ready timeout: %v", jobID, err) + return err + } + + // Open start gate + log.Printf("startJob [%s]: opening start gate", jobID) + gateFile := filepath.Join(jobDir, "start-gate") + if err := os.WriteFile(gateFile, []byte{}, 0600); err != nil { + return err + } + + // Transition to running + _, _ = s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusStarting}, + Target: StatusRunning, + RequireExitCodeNull: true, + }) + + // Clean up ready file + _ = os.Remove(pipeReadyPath) + return nil +} + +func (s *Service) waitForPipeReady(jobID, readyFile string) error { + deadline := time.Now().Add(s.config.PipeReadyTimeout) + for { + if _, err := os.Stat(readyFile); err == nil { + // Ready file exists, verify pipe is active + active, err := s.tmux.IsOutputPipeActive(jobID) + if err != nil { + return err + } + if active == nil { + return NewServerError(500, "pipe_failed", + "tmux pane disappeared after ready-file handshake") + } + if *active { + return nil + } + } + + if time.Now().After(deadline) { + return NewServerError(500, "pipe_failed", + fmt.Sprintf("timed out waiting for pipe ready (%v)", s.config.PipeReadyTimeout)) + } + time.Sleep(s.config.PollInterval) + } +} + +// WaitJob blocks until output, completion, truncation, or timeout. +func (s *Service) WaitJob(jobID string, req *WaitJobRequest) (*JobResult, error) { + row, err := s.db.GetJob(jobID) + if err != nil { + return nil, err + } + + outputPath := s.outputLogPath(row) + timeout := time.Duration(req.Timeout * float64(time.Second)) + outputLimit := req.OutputLimit + if outputLimit <= 0 { + outputLimit = s.config.DefaultOutputLimitBytes + } + idleFlush := time.Duration(req.IdleFlushSeconds * float64(time.Second)) + + deadline := time.Now().Add(timeout) + lastSize := fileSize(outputPath) + sawOutput := lastSize > int64(req.Offset) + var lastGrowthAt *time.Time + if sawOutput { + now := time.Now() + lastGrowthAt = &now + } + + for { + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + + currentSize := fileSize(outputPath) + + // Validate offset against current file size (mirror Python behavior). + if int64(req.Offset) > currentSize { + return nil, NewServerError(400, "invalid_offset", + fmt.Sprintf("offset %d exceeds current file size %d", req.Offset, currentSize)) + } + + // Only update lastGrowthAt when the file actually grows (not just + // when currentSize > offset). Without this, lastGrowthAt resets on + // every poll and idle-flush can never fire. + if currentSize > lastSize { + lastSize = currentSize + if currentSize > int64(req.Offset) { + sawOutput = true + now := time.Now() + lastGrowthAt = &now + } + } + + if view.Done { + window, err := ReadOutputWindow(outputPath, req.Offset, outputLimit) + if err != nil { + return nil, err + } + return s.jobResultFromView(view, row, window), nil + } + + if currentSize > int64(req.Offset) { + window, err := ReadOutputWindow(outputPath, req.Offset, outputLimit) + if err != nil { + return nil, err + } + if window.Truncated { + return s.jobResultFromView(view, row, window), nil + } + if sawOutput && lastGrowthAt != nil { + if time.Since(*lastGrowthAt) >= idleFlush { + return s.jobResultFromView(view, row, window), nil + } + } + } + + if time.Now().After(deadline) { + var window *OutputWindow + if currentSize > int64(req.Offset) { + window, err = ReadOutputWindow(outputPath, req.Offset, outputLimit) + if err != nil { + return nil, err + } + } else { + window = &OutputWindow{Output: "", Offset: req.Offset, Truncated: false} + } + return s.jobResultFromView(view, row, window), nil + } + + time.Sleep(s.config.PollInterval) + } +} + +// TailJob returns the tail of a job's output. +func (s *Service) TailJob(jobID string, outputLimit int) (*JobResult, error) { + row, err := s.db.GetJob(jobID) + if err != nil { + return nil, err + } + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + window, err := TailOutputWindow(s.outputLogPath(row), outputLimit) + if err != nil { + return nil, err + } + return s.jobResultFromView(view, row, window), nil +} + +// GetJobStatus materializes the current status from SQLite + live tmux state. +func (s *Service) GetJobStatus(jobID string) (*JobStatusView, error) { + sessionExists, pipeActive, err := s.liveRuntimeState(jobID) + if err != nil { + return nil, err + } + return s.materializeStatusView(jobID, sessionExists, pipeActive) +} + +// ListJobs returns recent jobs, optionally filtered by status. +func (s *Service) ListJobs(status *JobStatusName, limit int) (*ListJobsResponse, error) { + rows, err := s.db.ListJobs(nil) + if err != nil { + return nil, err + } + + var items []JobInfo + for _, row := range rows { + view, err := s.GetJobStatus(row.JobID) + if err != nil { + if isNotFound(err) { + continue + } + return nil, err + } + if status != nil && view.Status != *status { + continue + } + items = append(items, JobInfo{ + JobID: view.JobID, + Status: view.Status, + CreatedAt: view.CreatedAt, + StartedAt: view.StartedAt, + EndedAt: view.EndedAt, + }) + if len(items) >= limit { + break + } + } + return &ListJobsResponse{Jobs: items}, nil +} + +// SendInput sends input to a running job and waits. +func (s *Service) SendInput(jobID string, req *InputJobRequest) (*JobResult, error) { + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + if view.Done { + return nil, NewServerError(409, "job_not_running", fmt.Sprintf("Job %s is already terminal", jobID)) + } + + if err := s.tmux.SendInput(jobID, req.Text); err != nil { + // Check if job became terminal in the meantime + if se, ok := err.(*ServerError); ok && se.Code == "tmux_target_missing" { + view, _ = s.GetJobStatus(jobID) + if view != nil && view.Done { + return nil, NewServerError(409, "job_not_running", fmt.Sprintf("Job %s is already terminal", jobID)) + } + } + return nil, err + } + + return s.WaitJob(jobID, &WaitJobRequest{ + Offset: req.Offset, + Timeout: req.Timeout, + OutputLimit: req.OutputLimit, + IdleFlushSeconds: req.IdleFlushSeconds, + }) +} + +// TerminateJob terminates a running job. +func (s *Service) TerminateJob(jobID string, graceSeconds float64) (*JobStatusView, error) { + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + if view.Done { + s.tmux.CleanupSession(jobID) + return view, nil + } + + _, _ = s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting, StatusRunning}, + AllowOverrideFrom: []JobStatusName{StatusExited}, + Target: StatusTerminated, + }) + + _ = s.tmux.SendInterrupt(jobID) + if graceSeconds > 0 { + time.Sleep(time.Duration(graceSeconds * float64(time.Second))) + } + s.tmux.CleanupSession(jobID) + + return s.GetJobStatus(jobID) +} + +// DeleteJob deletes a job and its artifacts. +func (s *Service) DeleteJob(jobID string, force bool, graceSeconds float64) (*DeleteJobResponse, error) { + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + if !view.Done { + if !force { + return nil, NewServerError(409, "job_running", fmt.Sprintf("Job %s is still running", jobID)) + } + _, _ = s.TerminateJob(jobID, graceSeconds) + } + + s.tmux.CleanupSession(jobID) + if err := s.db.DeleteJob(jobID); err != nil { + return nil, err + } + _ = os.RemoveAll(filepath.Join(s.config.JobsDir(), jobID)) + return &DeleteJobResponse{JobID: jobID, Deleted: true}, nil +} + +// Reconcile synchronizes SQLite rows with live tmux state. +func (s *Service) Reconcile() error { + rows, err := s.db.ListJobs(nil) + if err != nil { + return err + } + for _, row := range rows { + view, err := s.GetJobStatus(row.JobID) + if err != nil { + if isNotFound(err) { + continue + } + return err + } + if view.Done { + s.tmux.CleanupSession(row.JobID) + } + } + return nil +} + +// GCOnce removes expired terminal jobs. +func (s *Service) GCOnce() error { + cutoff := time.Now().Add(-s.config.GCFinishedJobRetention) + rows, err := s.db.ListJobs(nil) + if err != nil { + return err + } + for _, row := range rows { + view, err := s.GetJobStatus(row.JobID) + if err != nil { + if isNotFound(err) { + continue + } + return err + } + if !view.Done || view.EndedAt == nil { + continue + } + endedAt, err := ParseTimestamp(*view.EndedAt) + if err != nil { + continue + } + if endedAt.After(cutoff) { + continue + } + s.tmux.CleanupSession(row.JobID) + _ = s.db.DeleteJob(row.JobID) + _ = os.RemoveAll(filepath.Join(s.config.JobsDir(), row.JobID)) + } + return nil +} + +// CheckRunningJobsPipeHealth fails running jobs whose pipe died. +func (s *Service) CheckRunningJobsPipeHealth() { + rows, _ := s.db.ListJobs([]JobStatusName{StatusRunning}) + for _, row := range rows { + _, _ = s.GetJobStatus(row.JobID) + } +} + +func (s *Service) materializeStatusView(jobID string, sessionExists bool, pipeActive *bool) (*JobStatusView, error) { + row, err := s.db.GetJob(jobID) + if err != nil { + return nil, err + } + + status := row.Status + + if status.IsTerminal() { + // Ensure ended_at is set + if row.EndedAt == nil { + now := FormatTimestamp(time.Now()) + _, _ = s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{status}, + Target: status, + EndedAt: now, + }) + if r, err := s.db.GetJob(jobID); err == nil { + row = r + } + } + } else if row.ExitCode != nil { + // Already has exit code, transition to exited + if r, err := s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting, StatusRunning}, + Target: StatusExited, + }); err == nil { + row = r + } + } else if exit := s.drainedNormalExitMetadata(jobID); exit != nil { + // Recover from drained exit artifacts + _ = s.db.RecordRunnerExit(jobID, exit.exitCode, exit.endedAt) + if r, err := s.db.GetJob(jobID); err == nil { + row = r + } + } else if sessionExists { + if pipeActive != nil && !*pipeActive { + s.mu.Lock() + isStarting := s.startingJobs[jobID] + s.mu.Unlock() + if !isStarting { + reason := "pipe_failed" + msg := "The tmux output pipe stopped while the job was still running." + if r, err := s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting, StatusRunning}, + Target: StatusFailed, + Reason: &reason, + Message: &msg, + }); err == nil { + row = r + } + } + } else if status == StatusCreated || status == StatusStarting { + s.mu.Lock() + isStarting := s.startingJobs[jobID] + s.mu.Unlock() + if !isStarting { + if r, err := s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting}, + Target: StatusRunning, + RequireExitCodeNull: true, + }); err == nil { + row = r + } + } + } + } else { + // No session + if s.normalExitCommitPending(jobID) { + // Wait for pipe drain finalizer + } else { + s.mu.Lock() + isStarting := s.startingJobs[jobID] + s.mu.Unlock() + if !isStarting || (status != StatusCreated && status != StatusStarting) { + reason := "tmux_session_missing" + msg := "The dedicated tmux session is no longer present." + if r, err := s.db.TransitionStatus(jobID, TransitionOpts{ + AllowedFrom: []JobStatusName{StatusCreated, StatusStarting, StatusRunning}, + Target: StatusLost, + Reason: &reason, + Message: &msg, + }); err == nil { + row = r + } + } + } + } + + return s.statusViewFromRow(row), nil +} + +func (s *Service) liveRuntimeState(jobID string) (bool, *bool, error) { + exists, err := s.tmux.SessionExists(JobSessionName(jobID)) + if err != nil { + return false, nil, err + } + if !exists { + return false, nil, nil + } + active, err := s.tmux.IsOutputPipeActive(jobID) + if err != nil { + return false, nil, err + } + if active == nil { + return false, nil, nil + } + return true, active, nil +} + +type exitMetadata struct { + exitCode int + endedAt string +} + +func (s *Service) drainedNormalExitMetadata(jobID string) *exitMetadata { + jobDir := filepath.Join(s.config.JobsDir(), jobID) + drainedPath := filepath.Join(jobDir, ".pipe-drained") + exitCodePath := filepath.Join(jobDir, "runner-exit-code") + endedAtPath := filepath.Join(jobDir, "runner-ended-at") + + if !fileExists(drainedPath) || !fileExists(exitCodePath) || !fileExists(endedAtPath) { + return nil + } + + exitCodeRaw, err := os.ReadFile(exitCodePath) + if err != nil { + return nil + } + endedAtRaw, err := os.ReadFile(endedAtPath) + if err != nil { + return nil + } + + exitCodeStr := strings.TrimSpace(string(exitCodeRaw)) + endedAtStr := strings.TrimSpace(string(endedAtRaw)) + if exitCodeStr == "" || endedAtStr == "" { + return nil + } + + code, err := strconv.Atoi(exitCodeStr) + if err != nil { + return nil + } + + return &exitMetadata{exitCode: code, endedAt: endedAtStr} +} + +func (s *Service) normalExitCommitPending(jobID string) bool { + jobDir := filepath.Join(s.config.JobsDir(), jobID) + return fileExists(filepath.Join(jobDir, "runner-exit-code")) && + fileExists(filepath.Join(jobDir, "runner-ended-at")) && + !fileExists(filepath.Join(jobDir, ".pipe-drained")) && + !fileExists(filepath.Join(jobDir, ".pipe-failed")) +} + +func (s *Service) outputLogPath(row *JobRow) string { + return filepath.Join(s.config.StateDir, row.OutputPath) +} + +func (s *Service) statusViewFromRow(row *JobRow) *JobStatusView { + outputPath := s.outputLogPath(row) + offset := int(fileSize(outputPath)) + return &JobStatusView{ + JobID: row.JobID, + Status: row.Status, + Done: row.Status.IsTerminal(), + ExitCode: row.ExitCode, + CreatedAt: row.CreatedAt, + StartedAt: row.StartedAt, + EndedAt: row.EndedAt, + Offset: offset, + } +} + +func (s *Service) jobResultFromView(view *JobStatusView, row *JobRow, window *OutputWindow) *JobResult { + outputPath := s.outputLogPath(row) + absPath, _ := filepath.Abs(outputPath) + return &JobResult{ + JobID: view.JobID, + Done: view.Done, + Status: view.Status, + ExitCode: view.ExitCode, + OutputPath: absPath, + Output: window.Output, + Offset: window.Offset, + Truncated: window.Truncated, + } +} + +func (s *Service) resolveCwd(rawCwd *string) (string, error) { + var cwd string + if rawCwd != nil && *rawCwd != "" { + cwd = *rawCwd + } else { + cwd = s.config.DefaultCwd + } + + info, err := os.Stat(cwd) + if err != nil || !info.IsDir() { + return "", NewServerError(400, "invalid_cwd", fmt.Sprintf("cwd is not a directory: %s", cwd)) + } + abs, _ := filepath.Abs(cwd) + return abs, nil +} + +func (s *Service) allocateJobDir() (string, string, error) { + for i := 0; i < 20; i++ { + jobID := GenerateJobID() + jobDir := filepath.Join(s.config.JobsDir(), jobID) + if err := os.Mkdir(jobDir, 0700); err == nil { + return jobID, jobDir, nil + } + } + return "", "", NewServerError(500, "job_id_collision", "Failed to allocate a unique job id") +} + +func (s *Service) cleanupStarting(jobID, jobDir string) { + s.mu.Lock() + delete(s.startingJobs, jobID) + s.mu.Unlock() + _ = os.RemoveAll(jobDir) +} + +func (s *Service) installRunner() { + script := s.runnerScriptSource() + _ = os.WriteFile(s.config.RunnerPath(), []byte(script), 0755) +} + +func (s *Service) runnerScriptSource() string { + // The runner script is a bash wrapper that waits for the start-gate, + // then delegates env loading + exec to a Python bootstrap so that + // JSON env values (including multi-line strings) are applied verbatim. + // The Python helper exec's the target process so SIGINT semantics + // match running the script directly in the tmux pane. + return `#!/usr/bin/env bash +set -uo pipefail + +JOB_DIR="$1" +JOB_ID="$2" +CWD="$3" +SCRIPT_PATH="$JOB_DIR/script" +ENV_PATH="$JOB_DIR/.job-env.json" +START_GATE="$JOB_DIR/start-gate" +RUNNER_EXIT_CODE_PATH="$JOB_DIR/runner-exit-code" +RUNNER_ENDED_AT_PATH="$JOB_DIR/runner-ended-at" + +write_atomic() { + local dest="$1" + local value="$2" + local tmp="${dest}.tmp.$$" + printf '%s\n' "$value" > "$tmp" + mv "$tmp" "$dest" +} + +while [ ! -e "$START_GATE" ]; do + sleep 0.05 +done + +unset TMUX +unset SHELLCTL_STATE_DIR +unset SHELLCTL_RUNTIME_DIR +unset SHELLCTL_TMUX_SOCKET +unset SHELLCTL_RUNNER +unset SHELLCTL_AUTH_TOKEN + +# Use Python bootstrap to load JSON env verbatim and exec the target. +# This avoids depending on jq and correctly handles multi-line values. +python3 -c ' +import json, os, stat, sys +from pathlib import Path + +script_path = Path(sys.argv[1]) +cwd = sys.argv[2] +env_path = Path(sys.argv[3]) + +env = os.environ.copy() +if env_path.exists(): + env.update(json.loads(env_path.read_text(encoding="utf-8"))) + +try: + os.chdir(cwd) +except OSError: + raise SystemExit(111) + +# Ensure HOME directory exists (agent backend may set a per-agent HOME). +home = env.get("HOME", "") +if home and not os.path.isdir(home): + os.makedirs(home, exist_ok=True) + +with script_path.open("r", encoding="utf-8") as handle: + first_line = handle.readline() + +if first_line.startswith("#!"): + script_path.chmod(script_path.stat().st_mode | stat.S_IXUSR) + argv = [str(script_path)] +else: + argv = ["sh", str(script_path)] + +try: + os.execvpe(argv[0], argv, env) +except FileNotFoundError as exc: + print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) + raise SystemExit(127) from exc +except OSError as exc: + print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) + raise SystemExit(126) from exc +' "$SCRIPT_PATH" "$CWD" "$ENV_PATH" +EXIT_CODE=$? + +ENDED_AT="$(date -u +%Y-%m-%dT%H:%M:%SZ)" + +write_atomic "$RUNNER_EXIT_CODE_PATH" "$EXIT_CODE" +write_atomic "$RUNNER_ENDED_AT_PATH" "$ENDED_AT" + +exit "$EXIT_CODE" +` +} + +// Helpers + +func fileExists(path string) bool { + _, err := os.Stat(path) + return err == nil +} + +func fileSize(path string) int64 { + info, err := os.Stat(path) + if err != nil { + return 0 + } + return info.Size() +} + +func isNotFound(err error) bool { + if se, ok := err.(*ServerError); ok { + return se.Code == "job_not_found" + } + return false +} diff --git a/dify-agent-runtime/internal/server/tmux.go b/dify-agent-runtime/internal/server/tmux.go new file mode 100644 index 00000000000..42a05c7de5f --- /dev/null +++ b/dify-agent-runtime/internal/server/tmux.go @@ -0,0 +1,288 @@ +package server + +import ( + "bytes" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" +) + +// TmuxController manages tmux sessions for shellctl jobs via a dedicated socket. +type TmuxController struct { + config *Config +} + +// NewTmuxController creates a new tmux controller. +func NewTmuxController(config *Config) *TmuxController { + return &TmuxController{config: config} +} + +// StartServer ensures the tmux server is running. +func (t *TmuxController) StartServer() error { + _, err := t.runTmux("start-server") + return err +} + +// ListSessions returns the set of active tmux session names. +func (t *TmuxController) ListSessions() (map[string]bool, error) { + result, err := t.runTmuxNoCheck("list-sessions", "-F", "#{session_name}") + if err != nil { + return nil, err + } + if result.exitCode != 0 { + stderr := strings.TrimSpace(result.stderr) + if isTmuxTargetMissing(stderr) { + return make(map[string]bool), nil + } + return nil, NewServerError(500, "tmux_error", stderr) + } + sessions := make(map[string]bool) + for _, line := range strings.Split(strings.TrimSpace(result.stdout), "\n") { + line = strings.TrimSpace(line) + if line != "" { + sessions[line] = true + } + } + return sessions, nil +} + +// SessionExists checks if a tmux session exists by name. +func (t *TmuxController) SessionExists(sessionName string) (bool, error) { + sessions, err := t.ListSessions() + if err != nil { + return false, err + } + return sessions[sessionName], nil +} + +// IsOutputPipeActive checks the #{pane_pipe} format variable for a job pane. +// Returns: true=active, false=inactive, error with nil bool if pane missing. +func (t *TmuxController) IsOutputPipeActive(jobID string) (*bool, error) { + result, err := t.runTmuxNoCheck( + "display-message", "-p", "-t", JobPaneTarget(jobID), "#{pane_pipe}", + ) + if err != nil { + return nil, err + } + if result.exitCode != 0 { + stderr := strings.TrimSpace(result.stderr) + if isTmuxTargetMissing(stderr) { + return nil, nil // pane missing + } + return nil, NewServerError(500, "tmux_error", stderr) + } + active := strings.TrimSpace(result.stdout) == "1" + return &active, nil +} + +// CreateJobSession creates a new tmux session for a job. +func (t *TmuxController) CreateJobSession(jobID, jobDir, cwd string, cols, rows int) error { + runnerCmd := shellJoin([]string{ + t.config.RunnerPath(), jobDir, jobID, cwd, + }) + result, err := t.runTmuxNoCheck( + "-f", "/dev/null", + "new-session", "-d", + "-s", JobSessionName(jobID), + "-x", fmt.Sprintf("%d", cols), + "-y", fmt.Sprintf("%d", rows), + runnerCmd, + ) + if err != nil { + return err + } + if result.exitCode != 0 { + return NewServerError(500, "tmux_new_session_failed", + strings.TrimSpace(result.stderr)) + } + return nil +} + +// 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) + result, err := t.runTmuxNoCheck( + "pipe-pane", "-o", "-t", JobPaneTarget(jobID), pipeCmd, + ) + if err != nil { + return err + } + if result.exitCode != 0 { + return NewServerError(500, "pipe_failed", + strings.TrimSpace(result.stderr)) + } + return nil +} + +// SendInput sends text to a job's tmux pane via load-buffer + paste-buffer. +func (t *TmuxController) SendInput(jobID, text string) error { + bufferName := fmt.Sprintf("shellctl-in-%s", jobID) + + // Write input to a temp file + tmpFile, err := os.CreateTemp(t.config.RuntimeDir, fmt.Sprintf("shellctl-input-%s-", jobID)) + if err != nil { + return err + } + tmpPath := tmpFile.Name() + defer func() { _ = os.Remove(tmpPath) }() + + if _, err := tmpFile.WriteString(text); err != nil { + _ = tmpFile.Close() + return err + } + if err := tmpFile.Close(); err != nil { + return err + } + + // Load buffer + result, err := t.runTmuxNoCheck("load-buffer", "-b", bufferName, tmpPath) + if err != nil { + return err + } + if result.exitCode != 0 { + stderr := strings.TrimSpace(result.stderr) + if isTmuxTargetMissing(stderr) { + return NewServerError(409, "tmux_target_missing", stderr) + } + return NewServerError(500, "tmux_input_failed", stderr) + } + + // Paste buffer + result, err = t.runTmuxNoCheck( + "paste-buffer", "-t", JobPaneTarget(jobID), "-b", bufferName, + ) + // Always clean up buffer + _, _ = t.runTmuxNoCheck("delete-buffer", "-b", bufferName) + + if err != nil { + return err + } + if result.exitCode != 0 { + stderr := strings.TrimSpace(result.stderr) + if isTmuxTargetMissing(stderr) { + return NewServerError(409, "tmux_target_missing", stderr) + } + return NewServerError(500, "tmux_input_failed", stderr) + } + return nil +} + +// SendInterrupt sends C-c to a job's pane. +func (t *TmuxController) SendInterrupt(jobID string) error { + _, _ = t.runTmuxNoCheck("send-keys", "-t", JobPaneTarget(jobID), "C-c") + return nil +} + +// CleanupSession kills the tmux session for a job (best-effort). +func (t *TmuxController) CleanupSession(jobID string) { + _, _ = t.runTmuxNoCheck("kill-session", "-t", JobSessionName(jobID)) +} + +func (t *TmuxController) buildPipeCommand(jobID, jobDir, readyFile string) string { + sanitizeCmd := shellJoin(append(t.config.SanitizePtyCommand, "--ready-file", readyFile)) + runnerExitCmd := shellJoin(append(t.config.RunnerExitCommand, + "--state-dir", t.config.StateDir, + "--job-id", jobID, + "--sqlite-busy-timeout-ms", fmt.Sprintf("%d", t.config.SQLiteBusyTimeoutMs), + )) + + outputPath := shellQuote(filepath.Join(jobDir, "output.log")) + drainedPath := shellQuote(filepath.Join(jobDir, ".pipe-drained")) + errorLogPath := shellQuote(filepath.Join(jobDir, "pipe-error.log")) + failedPath := shellQuote(filepath.Join(jobDir, ".pipe-failed")) + exitCodePath := shellQuote(filepath.Join(jobDir, "runner-exit-code")) + endedAtPath := shellQuote(filepath.Join(jobDir, "runner-ended-at")) + + parts := []string{ + fmt.Sprintf("%s >> %s 2> %s", sanitizeCmd, outputPath, errorLogPath), + "sanitize_status=$?", + "runner_exit_status=0", + fmt.Sprintf( + `if [ "$sanitize_status" -eq 0 ]; then : > %s; if [ -s %s ] && [ -s %s ]; then %s --exit-code "$(cat %s)" --ended-at "$(cat %s)" 2>> %s; runner_exit_status=$?; if [ "$runner_exit_status" -ne 0 ]; then printf 'runner-exit failed with status %%s\n' "$runner_exit_status" >> %s; fi; fi; else : > %s; fi`, + drainedPath, exitCodePath, endedAtPath, runnerExitCmd, + exitCodePath, endedAtPath, errorLogPath, errorLogPath, failedPath, + ), + `if [ "$sanitize_status" -ne 0 ]; then exit "$sanitize_status"; fi`, + `exit "$sanitize_status"`, + } + return strings.Join(parts, " ; ") +} + +type tmuxResult struct { + stdout string + stderr string + exitCode int +} + +func (t *TmuxController) runTmux(args ...string) (string, error) { + result, err := t.runTmuxNoCheck(args...) + if err != nil { + return "", err + } + if result.exitCode != 0 { + stderr := strings.TrimSpace(result.stderr) + if stderr == "" { + stderr = "tmux command failed" + } + return "", NewServerError(500, "tmux_error", stderr) + } + return result.stdout, nil +} + +func (t *TmuxController) runTmuxNoCheck(args ...string) (*tmuxResult, error) { + cmdArgs := append([]string{"-S", t.config.TmuxSocket()}, args...) + cmd := exec.Command("tmux", cmdArgs...) + + // Clear TMUX env to avoid nesting + env := os.Environ() + filtered := env[:0] + for _, e := range env { + if !strings.HasPrefix(e, "TMUX=") { + filtered = append(filtered, e) + } + } + cmd.Env = filtered + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + exitCode := 0 + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + return nil, fmt.Errorf("exec tmux: %w", err) + } + } + + return &tmuxResult{ + stdout: stdout.String(), + stderr: stderr.String(), + exitCode: exitCode, + }, nil +} + +func isTmuxTargetMissing(stderr string) bool { + lower := strings.ToLower(stderr) + return strings.Contains(lower, "can't find pane") || + strings.Contains(lower, "can't find session") || + strings.Contains(lower, "no server running") || + strings.Contains(lower, "failed to connect") || + strings.Contains(lower, "server exited unexpectedly") +} + +func shellJoin(parts []string) string { + quoted := make([]string, len(parts)) + for i, p := range parts { + quoted[i] = shellQuote(p) + } + return strings.Join(quoted, " ") +} + +func shellQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" +} diff --git a/dify-agent-runtime/internal/server/tmux_test.go b/dify-agent-runtime/internal/server/tmux_test.go new file mode 100644 index 00000000000..186b57dbfa7 --- /dev/null +++ b/dify-agent-runtime/internal/server/tmux_test.go @@ -0,0 +1,59 @@ +package server + +import ( + "testing" +) + +func TestShellQuote(t *testing.T) { + tests := []struct { + input string + expected string + }{ + {"hello", "'hello'"}, + {"", "''"}, + {"a b c", "'a b c'"}, + {"it's", "'it'\\''s'"}, + {"/path/to/file", "'/path/to/file'"}, + } + + for _, tc := range tests { + got := shellQuote(tc.input) + if got != tc.expected { + t.Errorf("shellQuote(%q) = %q, want %q", tc.input, got, tc.expected) + } + } +} + +func TestShellJoin(t *testing.T) { + got := shellJoin([]string{"echo", "hello world", "it's"}) + expected := "'echo' 'hello world' 'it'\\''s'" + if got != expected { + t.Errorf("shellJoin = %q, want %q", got, expected) + } +} + +func TestIsTmuxTargetMissing(t *testing.T) { + missingMsgs := []string{ + "can't find pane: shellctl-abc", + "can't find session: shellctl-abc", + "no server running on /tmp/tmux.sock", + "failed to connect to server", + "server exited unexpectedly", + } + for _, msg := range missingMsgs { + if !isTmuxTargetMissing(msg) { + t.Errorf("expected %q to be detected as target missing", msg) + } + } + + validMsgs := []string{ + "some other error", + "permission denied", + "", + } + for _, msg := range validMsgs { + if isTmuxTargetMissing(msg) { + t.Errorf("expected %q to NOT be detected as target missing", msg) + } + } +} diff --git a/dify-agent-runtime/internal/server/types.go b/dify-agent-runtime/internal/server/types.go new file mode 100644 index 00000000000..28eecc37b86 --- /dev/null +++ b/dify-agent-runtime/internal/server/types.go @@ -0,0 +1,100 @@ +package server + +// RunJobRequest is the HTTP request body for POST /v1/jobs/run. +type RunJobRequest struct { + Script string `json:"script"` + Cwd *string `json:"cwd,omitempty"` + Env map[string]string `json:"env,omitempty"` + Terminal *TerminalSize `json:"terminal,omitempty"` + Timeout float64 `json:"timeout,omitempty"` + OutputLimit int `json:"output_limit,omitempty"` + IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` +} + +// TerminalSize specifies the initial PTY geometry. +type TerminalSize struct { + Cols int `json:"cols"` + Rows int `json:"rows"` +} + +// WaitJobRequest is the HTTP request body for POST /v1/jobs/{job_id}/wait. +type WaitJobRequest struct { + Timeout float64 `json:"timeout"` + Offset int `json:"offset"` + OutputLimit int `json:"output_limit,omitempty"` + IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` +} + +// InputJobRequest is the HTTP request body for POST /v1/jobs/{job_id}/input. +type InputJobRequest struct { + Text string `json:"text"` + Timeout float64 `json:"timeout,omitempty"` + Offset int `json:"offset"` + OutputLimit int `json:"output_limit,omitempty"` + IdleFlushSeconds float64 `json:"idle_flush_seconds,omitempty"` +} + +// TerminateJobRequest is the HTTP request body for POST /v1/jobs/{job_id}/terminate. +type TerminateJobRequest struct { + GraceSeconds *float64 `json:"grace_seconds,omitempty"` +} + +// JobResult is the unified response for output-oriented job APIs. +type JobResult struct { + JobID string `json:"job_id"` + Done bool `json:"done"` + Status JobStatusName `json:"status"` + ExitCode *int `json:"exit_code"` + OutputPath string `json:"output_path"` + Output string `json:"output"` + Offset int `json:"offset"` + Truncated bool `json:"truncated"` +} + +// JobStatusView is the materialized lifecycle view. +type JobStatusView struct { + JobID string `json:"job_id"` + Status JobStatusName `json:"status"` + Done bool `json:"done"` + ExitCode *int `json:"exit_code"` + CreatedAt string `json:"created_at"` + StartedAt *string `json:"started_at"` + EndedAt *string `json:"ended_at"` + Offset int `json:"offset"` +} + +// JobInfo is a compact job listing record. +type JobInfo struct { + JobID string `json:"job_id"` + Status JobStatusName `json:"status"` + CreatedAt string `json:"created_at"` + StartedAt *string `json:"started_at,omitempty"` + EndedAt *string `json:"ended_at,omitempty"` +} + +// ListJobsResponse is the response for GET /v1/jobs. +type ListJobsResponse struct { + Jobs []JobInfo `json:"jobs"` +} + +// DeleteJobResponse is the response for DELETE /v1/jobs/{job_id}. +type DeleteJobResponse struct { + JobID string `json:"job_id"` + Deleted bool `json:"deleted"` +} + +// HealthResponse is the health check response. +type HealthResponse struct { + Status string `json:"status"` +} + +// ErrorDetail is the machine-readable API error payload. +type ErrorDetail struct { + Code string `json:"code"` + Message string `json:"message"` +} + +// ErrorResponse is the error envelope. +type ErrorResponse struct { + Error ErrorDetail `json:"error"` +} diff --git a/dify-agent-runtime/internal/stubclient/client.go b/dify-agent-runtime/internal/stubclient/client.go new file mode 100644 index 00000000000..d9b6bdddc8a --- /dev/null +++ b/dify-agent-runtime/internal/stubclient/client.go @@ -0,0 +1,143 @@ +// Package stubclient provides a gRPC client for the Dify Agent Stub service. +// +// The Agent Stub service runs on the host (dify-agent backend) and exposes +// Connect, FileUpload, and FileDownload RPCs to sandbox-resident processes. +// This client is used by the dify-agent CLI binary inside the sandbox container. +package stubclient + +import ( + "context" + "fmt" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + stubv1 "github.com/langgenius/dify/dify-agent-runtime/gen/dify/agent/stub/v1" +) + +// Client wraps the gRPC AgentStubService client with connection lifecycle management. +type Client struct { + conn *grpc.ClientConn + stub stubv1.AgentStubServiceClient + timeout time.Duration +} + +// Option configures a Client. +type Option func(*Client) + +// WithTimeout sets the default per-call timeout. Default is 30s. +func WithTimeout(d time.Duration) Option { + return func(c *Client) { + c.timeout = d + } +} + +// Dial creates a new Client connected to the given target address (host:port). +func Dial(target string, opts ...Option) (*Client, error) { + c := &Client{timeout: 30 * time.Second} + for _, o := range opts { + o(c) + } + + conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + return nil, fmt.Errorf("stubclient: dial %s: %w", target, err) + } + c.conn = conn + c.stub = stubv1.NewAgentStubServiceClient(conn) + return c, nil +} + +// Close closes the underlying gRPC connection. +func (c *Client) Close() error { + if c.conn != nil { + return c.conn.Close() + } + return nil +} + +// ConnectResult holds the response from a Connect RPC. +type ConnectResult struct { + ConnectionID string + Status string +} + +// Connect establishes a logical connection with the Agent Stub server. +func (c *Client) Connect(ctx context.Context, protocolVersion int32, argv []string, metadataJSON string) (*ConnectResult, error) { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + defer cancel() + + resp, err := c.stub.Connect(ctx, &stubv1.ConnectRequest{ + ProtocolVersion: protocolVersion, + Argv: argv, + MetadataJson: metadataJSON, + }) + if err != nil { + return nil, fmt.Errorf("stubclient: connect: %w", err) + } + return &ConnectResult{ + ConnectionID: resp.GetConnectionId(), + Status: resp.GetStatus(), + }, nil +} + +// FileUploadResult holds the response from a CreateFileUploadRequest RPC. +type FileUploadResult struct { + UploadURL string +} + +// CreateFileUpload requests an upload URL from the Agent Stub server. +func (c *Client) CreateFileUpload(ctx context.Context, filename, mimetype string) (*FileUploadResult, error) { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + defer cancel() + + resp, err := c.stub.CreateFileUploadRequest(ctx, &stubv1.FileUploadRequest{ + Filename: filename, + Mimetype: mimetype, + }) + if err != nil { + return nil, fmt.Errorf("stubclient: file upload: %w", err) + } + return &FileUploadResult{ + UploadURL: resp.GetUploadUrl(), + }, nil +} + +// FileDownloadResult holds the response from a CreateFileDownloadRequest RPC. +type FileDownloadResult struct { + Filename string + MimeType string + Size int64 + DownloadURL string +} + +// CreateFileDownload requests a download URL from the Agent Stub server. +func (c *Client) CreateFileDownload(ctx context.Context, transferMethod string, reference, url *string, forExternal bool) (*FileDownloadResult, error) { + ctx, cancel := context.WithTimeout(ctx, c.timeout) + defer cancel() + + fileMapping := &stubv1.FileMapping{ + TransferMethod: transferMethod, + } + if reference != nil { + fileMapping.Reference = reference + } + if url != nil { + fileMapping.Url = url + } + + resp, err := c.stub.CreateFileDownloadRequest(ctx, &stubv1.FileDownloadRequest{ + File: fileMapping, + ForExternal: &forExternal, + }) + if err != nil { + return nil, fmt.Errorf("stubclient: file download: %w", err) + } + return &FileDownloadResult{ + Filename: resp.GetFilename(), + MimeType: resp.GetMimeType(), + Size: resp.GetSize(), + DownloadURL: resp.GetDownloadUrl(), + }, nil +} diff --git a/dify-agent-runtime/internal/stubclient/client_test.go b/dify-agent-runtime/internal/stubclient/client_test.go new file mode 100644 index 00000000000..13499b67499 --- /dev/null +++ b/dify-agent-runtime/internal/stubclient/client_test.go @@ -0,0 +1,114 @@ +package stubclient + +import ( + "context" + "net" + "testing" + + "google.golang.org/grpc" + + stubv1 "github.com/langgenius/dify/dify-agent-runtime/gen/dify/agent/stub/v1" +) + +// fakeServer implements the AgentStubService for testing. +type fakeServer struct { + stubv1.UnimplementedAgentStubServiceServer +} + +func (f *fakeServer) Connect(_ context.Context, req *stubv1.ConnectRequest) (*stubv1.ConnectResponse, error) { + return &stubv1.ConnectResponse{ + ConnectionId: "conn-123", + Status: "connected", + }, nil +} + +func (f *fakeServer) CreateFileUploadRequest(_ context.Context, req *stubv1.FileUploadRequest) (*stubv1.FileUploadResponse, error) { + return &stubv1.FileUploadResponse{ + UploadUrl: "https://upload.example.com/" + req.GetFilename(), + }, nil +} + +func (f *fakeServer) CreateFileDownloadRequest(_ context.Context, req *stubv1.FileDownloadRequest) (*stubv1.FileDownloadResponse, error) { + return &stubv1.FileDownloadResponse{ + Filename: "downloaded.txt", + MimeType: strPtr("text/plain"), + Size: 1024, + DownloadUrl: "https://download.example.com/file", + }, nil +} + +func strPtr(s string) *string { return &s } + +func startFakeServer(t *testing.T) string { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + srv := grpc.NewServer() + stubv1.RegisterAgentStubServiceServer(srv, &fakeServer{}) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + return lis.Addr().String() +} + +func TestConnect(t *testing.T) { + addr := startFakeServer(t) + c, err := Dial(addr) + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer func() { _ = c.Close() }() + + result, err := c.Connect(context.Background(), 1, []string{"hello"}, `{"key":"val"}`) + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + if result.ConnectionID != "conn-123" { + t.Errorf("expected conn-123, got %s", result.ConnectionID) + } + if result.Status != "connected" { + t.Errorf("expected connected, got %s", result.Status) + } +} + +func TestCreateFileUpload(t *testing.T) { + addr := startFakeServer(t) + c, err := Dial(addr) + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer func() { _ = c.Close() }() + + result, err := c.CreateFileUpload(context.Background(), "test.txt", "text/plain") + if err != nil { + t.Fatalf("CreateFileUpload failed: %v", err) + } + if result.UploadURL != "https://upload.example.com/test.txt" { + t.Errorf("unexpected upload URL: %s", result.UploadURL) + } +} + +func TestCreateFileDownload(t *testing.T) { + addr := startFakeServer(t) + c, err := Dial(addr) + if err != nil { + t.Fatalf("Dial failed: %v", err) + } + defer func() { _ = c.Close() }() + + ref := "dify-file-ref:abc123" + result, err := c.CreateFileDownload(context.Background(), "local_file", &ref, nil, false) + if err != nil { + t.Fatalf("CreateFileDownload failed: %v", err) + } + if result.Filename != "downloaded.txt" { + t.Errorf("unexpected filename: %s", result.Filename) + } + if result.Size != 1024 { + t.Errorf("unexpected size: %d", result.Size) + } + if result.DownloadURL != "https://download.example.com/file" { + t.Errorf("unexpected download URL: %s", result.DownloadURL) + } +} diff --git a/dify-agent-runtime/packaging_test.go b/dify-agent-runtime/packaging_test.go new file mode 100644 index 00000000000..acea15db010 --- /dev/null +++ b/dify-agent-runtime/packaging_test.go @@ -0,0 +1,161 @@ +// Package-level packaging tests verify that go.mod, Makefile, and Dockerfile +// stay consistent with the set of binaries the project produces. +// +// These are the Go equivalent of dify-agent/tests/local/test_packaging.py: +// they catch accidental drift in module metadata, build targets, and the +// container image definition. +package packaging_test + +import ( + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +// projectRoot returns the absolute path to the dify-agent-runtime +// module root (the directory containing go.mod). +func projectRoot(t *testing.T) string { + t.Helper() + _, file, _, ok := runtime.Caller(0) + if !ok { + t.Fatal("cannot determine project root via runtime.Caller") + } + return filepath.Dir(file) +} + +func TestGoModDeclares(t *testing.T) { + root := projectRoot(t) + data, err := os.ReadFile(filepath.Join(root, "go.mod")) + if err != nil { + t.Fatalf("read go.mod: %v", err) + } + content := string(data) + + if !strings.Contains(content, "module github.com/langgenius/dify/dify-agent-runtime") { + t.Error("go.mod module path mismatch") + } + if !strings.Contains(content, "modernc.org/sqlite") { + t.Error("go.mod must depend on modernc.org/sqlite") + } +} + +// expectedCmds lists the cmd/ directories that must exist. +// Each one produces a binary via `go build ./cmd/`. +var expectedCmds = []string{ + "shellctl", + "sanitize-pty", + "runner-exit", + "dify-agent-cli", +} + +func TestCmdDirectoriesExist(t *testing.T) { + root := projectRoot(t) + for _, cmd := range expectedCmds { + dir := filepath.Join(root, "cmd", cmd) + info, err := os.Stat(dir) + if err != nil { + t.Errorf("cmd/%s directory missing: %v", cmd, err) + continue + } + if !info.IsDir() { + t.Errorf("cmd/%s is not a directory", cmd) + continue + } + // Each cmd directory must have a main.go + main := filepath.Join(dir, "main.go") + if _, err := os.Stat(main); err != nil { + t.Errorf("cmd/%s/main.go missing: %v", cmd, err) + } + } +} + +func TestMakefileBuildTargets(t *testing.T) { + root := projectRoot(t) + data, err := os.ReadFile(filepath.Join(root, "Makefile")) + if err != nil { + t.Fatalf("read Makefile: %v", err) + } + content := string(data) + + expectedTargets := []string{ + "$(BIN_DIR)/shellctl", + "$(BIN_DIR)/shellctl-sanitize-pty", + "$(BIN_DIR)/shellctl-runner-exit", + "$(BIN_DIR)/dify-agent", + } + for _, target := range expectedTargets { + if !strings.Contains(content, target) { + t.Errorf("Makefile missing build target %s", target) + } + } + + expectedSources := []string{ + "./cmd/shellctl", + "./cmd/sanitize-pty", + "./cmd/runner-exit", + "./cmd/dify-agent-cli", + } + for _, src := range expectedSources { + if !strings.Contains(content, src) { + t.Errorf("Makefile missing build source %s", src) + } + } +} + +func TestDockerfileBuildsAndCopiesAllBinaries(t *testing.T) { + root := projectRoot(t) + data, err := os.ReadFile(filepath.Join(root, "docker", "Dockerfile")) + if err != nil { + t.Fatalf("read Dockerfile: %v", err) + } + content := string(data) + + // Go build stage must produce all binaries + goBuildLines := []string{ + "go build -o /bin/shellctl ./cmd/shellctl", + "go build -o /bin/shellctl-sanitize-pty ./cmd/sanitize-pty", + "go build -o /bin/shellctl-runner-exit ./cmd/runner-exit", + "go build -o /bin/dify-agent ./cmd/dify-agent-cli", + } + for _, line := range goBuildLines { + if !strings.Contains(content, line) { + t.Errorf("Dockerfile missing build line: %s", line) + } + } + + // COPY stage must copy all binaries + copyLines := []string{ + "COPY --from=go-builder /bin/shellctl /usr/local/bin/shellctl", + "COPY --from=go-builder /bin/shellctl-sanitize-pty /usr/local/bin/shellctl-sanitize-pty", + "COPY --from=go-builder /bin/shellctl-runner-exit /usr/local/bin/shellctl-runner-exit", + "COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent", + } + for _, line := range copyLines { + if !strings.Contains(content, line) { + t.Errorf("Dockerfile missing COPY line: %s", line) + } + } + + // Runtime container assertions + assertions := map[string]string{ + "CMD": `CMD ["shellctl", "serve", "--listen", "0.0.0.0:5004"]`, + "EXPOSE": "EXPOSE 5004", + "USER": "USER dify", + "bash": "bash", + "git": "git", + "jq": "jq", + "tmux": "tmux", + } + for name, expected := range assertions { + if !strings.Contains(content, expected) { + t.Errorf("Dockerfile missing %s assertion: %s", name, expected) + } + } + + // Must NOT reference the Python shellctl-server + if strings.Contains(content, "shellctl-server") { + t.Error("Dockerfile must not reference shellctl-server (Python)") + } +} diff --git a/dify-agent-runtime/tests/acceptance_test.go b/dify-agent-runtime/tests/acceptance_test.go new file mode 100644 index 00000000000..cfd5bd1f4be --- /dev/null +++ b/dify-agent-runtime/tests/acceptance_test.go @@ -0,0 +1,587 @@ +//go:build integration + +// Package tests runs the same acceptance test suite against both the Python +// and Go shellctl server implementations to verify API compatibility. +// +// Prerequisites: +// +// docker compose -f tests/docker-compose.yml up --build -d +// go test -tags=integration ./tests/... -v +// docker compose -f tests/docker-compose.yml down +package tests + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strings" + "testing" + "time" +) + +var ( + goURL = envOrDefault("SHELLCTL_GO_URL", "http://localhost:15005") + authToken = envOrDefault("SHELLCTL_TEST_TOKEN", "test-token-123") + + httpClient = &http.Client{Timeout: 120 * time.Second} +) + +// target represents one server under test. +type target struct { + name string + baseURL string +} + +func targets() []target { + return []target{ + {name: "go", baseURL: goURL}, + } +} + +func TestMain(m *testing.M) { + // Warmup: wait for both servers to be ready before running tests + for _, tgt := range targets() { + if !waitForServer(tgt) { + fmt.Fprintf(os.Stderr, "ERROR: %s server not ready at %s\n", tgt.name, tgt.baseURL) + os.Exit(1) + } + } + + for _, tgt := range targets() { + warmupJob(tgt) + } + os.Exit(m.Run()) +} + +func waitForServer(tgt target) bool { + for i := 0; i < 60; i++ { + req, _ := http.NewRequest("GET", tgt.baseURL+"/healthz", nil) + resp, err := httpClient.Do(req) + if err == nil { + resp.Body.Close() + if resp.StatusCode == 200 { + return true + } + } + time.Sleep(time.Second) + } + return false +} + +// warmupJob sends a trivial job to prime the server (tmux bootstrap, lazy init). +// It retries up to 3 times with a 180s timeout per attempt. +func warmupJob(tgt target) { + 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 "+authToken) + 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) + time.Sleep(2 * time.Second) + continue + } + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + if resp.StatusCode == 200 { + return + } + fmt.Fprintf(os.Stderr, "WARN: %s warmup job attempt %d got status %d\n", tgt.name, attempt+1, resp.StatusCode) + time.Sleep(2 * time.Second) + } + fmt.Fprintf(os.Stderr, "WARN: %s warmup job failed after 3 attempts, continuing anyway\n", tgt.name) +} + +// --- Test Cases --- + +func TestHealthz(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + resp := doGet(t, tgt, "/healthz", false) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var result map[string]string + json.Unmarshal(body, &result) + if result["status"] != "ok" { + t.Errorf("expected status=ok, got %q", result["status"]) + } + }) + } +} + +func TestAuthRequired(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Request without auth should fail + req, _ := http.NewRequest("GET", tgt.baseURL+"/v1/jobs", nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != 401 { + t.Errorf("expected 401, got %d", resp.StatusCode) + } + }) + } +} + +func TestRunSimpleScript(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "echo hello-world", + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "hello-world") { + t.Errorf("expected output to contain 'hello-world', got %q", output) + } + }) + } +} + +func TestRunWithEnv(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "echo $MY_VAR", + "env": map[string]string{"MY_VAR": "test-value-42"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "test-value-42") { + t.Errorf("expected output to contain 'test-value-42', got %q", output) + } + }) + } +} + +func TestRunWithCwd(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "pwd", + "cwd": "/tmp", + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "/tmp") { + t.Errorf("expected output to contain '/tmp', got %q", output) + } + }) + } +} + +func TestRunNonZeroExit(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "exit 42", + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 42) + }) + } +} + +func TestWaitForOutput(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Run a script that delays output + result := runJob(t, tgt, map[string]any{ + "script": "sleep 0.5 && echo delayed-output", + "timeout": 5, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "delayed-output") { + t.Errorf("expected 'delayed-output', got %q", output) + } + }) + } +} + +func TestWaitJobWithOffset(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Run a multi-line script + result := runJob(t, tgt, map[string]any{ + "script": "echo line1\necho line2\necho line3", + "timeout": 10, + }) + assertJobDone(t, result) + jobID := result["job_id"].(string) + offset := int(result["offset"].(float64)) + + // Wait with offset should return empty (already at end) + waitResult := waitJob(t, tgt, jobID, map[string]any{ + "offset": offset, + "timeout": 1, + }) + // Should return with empty output since we're already past all data + if waitResult["output"].(string) != "" { + // It might return empty or might not, depending on timing + // Just verify no error occurred + } + }) + } +} + +func TestTailJob(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "echo line1\necho line2\necho final-line", + "timeout": 10, + }) + assertJobDone(t, result) + jobID := result["job_id"].(string) + + // Tail the job + resp := doGet(t, tgt, fmt.Sprintf("/v1/jobs/%s/log/tail", jobID), true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var tailResult map[string]any + json.Unmarshal(body, &tailResult) + output := tailResult["output"].(string) + if !strings.Contains(output, "final-line") { + t.Errorf("tail should contain 'final-line', got %q", output) + } + }) + } +} + +func TestGetJobStatus(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "echo done", + "timeout": 10, + }) + assertJobDone(t, result) + jobID := result["job_id"].(string) + + resp := doGet(t, tgt, fmt.Sprintf("/v1/jobs/%s", jobID), true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var status map[string]any + json.Unmarshal(body, &status) + + if status["job_id"] != jobID { + t.Errorf("expected job_id=%s, got %v", jobID, status["job_id"]) + } + if status["done"] != true { + t.Errorf("expected done=true") + } + if status["status"] != "exited" { + t.Errorf("expected status=exited, got %v", status["status"]) + } + }) + } +} + +func TestListJobs(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Run a job first + runJob(t, tgt, map[string]any{ + "script": "echo for-listing", + "timeout": 10, + }) + + resp := doGet(t, tgt, "/v1/jobs", true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var listResult map[string]any + json.Unmarshal(body, &listResult) + jobs := listResult["jobs"].([]any) + if len(jobs) == 0 { + t.Error("expected at least one job in listing") + } + }) + } +} + +func TestTerminateJob(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Start a long-running job + result := runJob(t, tgt, map[string]any{ + "script": "sleep 60", + "timeout": 1, // short timeout so run returns quickly + }) + jobID := result["job_id"].(string) + + // Terminate it + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/terminate", jobID), + map[string]any{"grace_seconds": 1}, true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var termResult map[string]any + json.Unmarshal(body, &termResult) + + if termResult["done"] != true { + t.Errorf("expected done=true after terminate") + } + status := termResult["status"].(string) + if status != "terminated" && status != "exited" { + t.Errorf("expected terminal status, got %q", status) + } + }) + } +} + +func TestDeleteJob(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "echo to-delete", + "timeout": 10, + }) + assertJobDone(t, result) + jobID := result["job_id"].(string) + + // Delete it + req, _ := http.NewRequest("DELETE", + fmt.Sprintf("%s/v1/jobs/%s", tgt.baseURL, jobID), nil) + req.Header.Set("Authorization", "Bearer "+authToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("delete request failed: %v", err) + } + defer resp.Body.Close() + assertStatus(t, resp, 200) + + // Should be gone now + getResp := doGet(t, tgt, fmt.Sprintf("/v1/jobs/%s", jobID), true) + if getResp.StatusCode != 404 { + t.Errorf("expected 404 after delete, got %d", getResp.StatusCode) + } + getResp.Body.Close() + }) + } +} + +func TestForceDeleteRunningJob(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Start a long-running job + result := runJob(t, tgt, map[string]any{ + "script": "sleep 60", + "timeout": 1, + }) + jobID := result["job_id"].(string) + + // 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) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("delete request failed: %v", err) + } + defer resp.Body.Close() + assertStatus(t, resp, 200) + }) + } +} + +func TestSendInput(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + // Start a script that reads from stdin + result := runJob(t, tgt, map[string]any{ + "script": "read line && echo got:$line", + "timeout": 2, // Will timeout waiting for input + }) + jobID := result["job_id"].(string) + + if result["done"] == true { + // Already finished (possible race), skip + t.Skip("job completed before input could be sent") + } + + // Send input + offset := int(result["offset"].(float64)) + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/input", jobID), + map[string]any{ + "text": "hello-input\n", + "offset": offset, + "timeout": 5, + }, true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var inputResult map[string]any + json.Unmarshal(body, &inputResult) + output := inputResult["output"].(string) + if !strings.Contains(output, "got:hello-input") { + t.Logf("output after input: %q (may need more wait time)", output) + } + }) + } +} + +func TestMultilineOutput(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + script := "for i in $(seq 1 20); do echo \"line $i\"; done" + result := runJob(t, tgt, map[string]any{ + "script": script, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "line 1") { + t.Errorf("missing 'line 1' in output") + } + if !strings.Contains(output, "line 20") { + t.Errorf("missing 'line 20' in output") + } + }) + } +} + +func TestInvalidCwd(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + resp := doPost(t, tgt, "/v1/jobs/run", map[string]any{ + "script": "echo x", + "cwd": "/nonexistent-dir-xyz", + "timeout": 5, + }, true) + if resp.StatusCode != 400 { + body := readBody(t, resp) + t.Errorf("expected 400, got %d: %s", resp.StatusCode, string(body)) + } else { + resp.Body.Close() + } + }) + } +} + +func TestJobNotFound(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + resp := doGet(t, tgt, "/v1/jobs/nonexistent-id-xyz", true) + if resp.StatusCode != 404 { + t.Errorf("expected 404, got %d", resp.StatusCode) + } + resp.Body.Close() + }) + } +} + +// --- Helpers --- + +func runJob(t *testing.T, tgt target, payload map[string]any) map[string]any { + t.Helper() + resp := doPost(t, tgt, "/v1/jobs/run", payload, true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var result map[string]any + if err := json.Unmarshal(body, &result); err != nil { + t.Fatalf("[%s] failed to parse run response: %v\nbody: %s", tgt.name, err, string(body)) + } + return result +} + +func waitJob(t *testing.T, tgt target, jobID string, payload map[string]any) map[string]any { + t.Helper() + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/wait", jobID), payload, true) + assertStatus(t, resp, 200) + body := readBody(t, resp) + var result map[string]any + json.Unmarshal(body, &result) + return result +} + +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) + } + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("[%s] GET %s failed: %v", tgt.name, path, err) + } + return resp +} + +func doPost(t *testing.T, tgt target, path string, payload map[string]any, withAuth bool) *http.Response { + t.Helper() + body, _ := json.Marshal(payload) + 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) + } + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("[%s] POST %s failed: %v", tgt.name, path, err) + } + return resp +} + +func readBody(t *testing.T, resp *http.Response) []byte { + t.Helper() + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("failed to read response body: %v", err) + } + return body +} + +func assertStatus(t *testing.T, resp *http.Response, expected int) { + t.Helper() + if resp.StatusCode != expected { + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + t.Fatalf("expected status %d, got %d: %s", expected, resp.StatusCode, string(body)) + } +} + +func assertJobDone(t *testing.T, result map[string]any) { + t.Helper() + if result["done"] != true { + t.Errorf("expected done=true, got %v (status=%v)", result["done"], result["status"]) + } +} + +func assertExitCode(t *testing.T, result map[string]any, expected int) { + t.Helper() + exitCode, ok := result["exit_code"].(float64) + if !ok { + t.Errorf("exit_code is nil or not a number: %v", result["exit_code"]) + return + } + if int(exitCode) != expected { + t.Errorf("expected exit_code=%d, got %d", expected, int(exitCode)) + } +} + +func envOrDefault(key, defaultVal string) string { + if v := os.Getenv(key); v != "" { + return v + } + return defaultVal +} diff --git a/dify-agent/proto/dify/agent/stub/v1/agent_stub.proto b/dify-agent/proto/dify/agent/stub/v1/agent_stub.proto index 9240c195262..608d6d0006c 100644 --- a/dify-agent/proto/dify/agent/stub/v1/agent_stub.proto +++ b/dify-agent/proto/dify/agent/stub/v1/agent_stub.proto @@ -2,6 +2,8 @@ syntax = "proto3"; package dify.agent.stub.v1; +option go_package = "dify/agent/stub/v1;stubv1"; + service AgentStubService { rpc Connect(ConnectRequest) returns (ConnectResponse); rpc CreateFileUploadRequest(FileUploadRequest) returns (FileUploadResponse); diff --git a/dify-agent/pyproject.toml b/dify-agent/pyproject.toml index fd2fe91a44b..8ce8f2dc05e 100644 --- a/dify-agent/pyproject.toml +++ b/dify-agent/pyproject.toml @@ -5,21 +5,16 @@ description = "Add your description here" readme = "README.md" requires-python = ">=3.12,<4.0" dependencies = [ - "anyio>=4.12.1,<5.0.0", "httpx==0.28.1", "httpx2>=2.5.0,<3.0.0", "pydantic>=2.12.5,<2.13", "pydantic-ai-slim>=1.102.0,<2.0.0", - "typer>=0.16.1,<0.17", "typing-extensions>=4.12.2,<5.0.0", ] [project.scripts] -dify-agent = "dify_agent.agent_stub.cli.main:main" dify-agent-stub-server = "dify_agent.agent_stub.server.cli:main" -shellctl = "shellctl.cli:main" -shellctl-sanitize-pty = "shellctl_runtime.sanitize:main" -shellctl-runner-exit = "shellctl_runtime.runner_exit:main" + [project.optional-dependencies] grpc = ["grpclib[protobuf]>=0.4.9,<0.5.0", "protobuf>=6.33.5,<7.0.0"] @@ -34,16 +29,13 @@ server = [ "redis>=7.4.0,<8.0.0", "uvicorn[standard]==0.46.0", ] -shellctl-server = [ - "aiosqlite>=0.21.0,<1.0.0", - "fastapi==0.136.0", - "sqlmodel>=0.0.24,<0.1.0", - "uvicorn[standard]==0.46.0", -] [tool.setuptools.packages.find] where = ["src"] -include = ["agenton*", "agenton_collections*", "dify_agent*", "shellctl*", "shellctl_runtime*"] +include = ["agenton*", "agenton_collections*", "dify_agent*", "shellctl*"] + +[tool.setuptools.package-data] +"dify_agent.layers" = ["_agent_cli_help.json"] [tool.pyright] include = ["src", "examples", "tests"] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/__init__.py b/dify-agent/src/dify_agent/agent_stub/cli/__init__.py deleted file mode 100644 index 91bd1f4117f..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Client-safe CLI package for the ``dify-agent`` sandbox command.""" - -__all__: list[str] = [] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py b/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py deleted file mode 100644 index 273f17eb917..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py +++ /dev/null @@ -1,20 +0,0 @@ -"""CLI-facing wrapper around the client-safe Agent Stub transport facade.""" - -from __future__ import annotations - -from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.client import connect_agent_stub_sync -from dify_agent.agent_stub.protocol.agent_stub import AgentStubConnectResponse - - -def connect_from_environment(*, argv: list[str]) -> AgentStubConnectResponse: - """Connect to the configured Agent Stub using the current environment.""" - environment = read_agent_stub_environment() - return connect_agent_stub_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - argv=argv, - ) - - -__all__ = ["connect_from_environment"] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_config.py b/dify-agent/src/dify_agent/agent_stub/cli/_config.py deleted file mode 100644 index 5e25de4e6db..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/_config.py +++ /dev/null @@ -1,300 +0,0 @@ -"""CLI helpers for sandbox-visible Agent Stub config commands.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from pathlib import Path -from tempfile import TemporaryDirectory - -from pydantic import BaseModel, ConfigDict - -from dify_agent.agent_stub._drive_materialization import ( - DriveMaterializationTransferError, - DriveMaterializationValidationError, - extract_archive_to_directory, -) -from dify_agent.agent_stub.cli._drive import _build_skill_archive -from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.cli._files import upload_tool_file_resource_from_environment -from dify_agent.agent_stub.client import ( - AgentStubTransferError, - AgentStubValidationError, - request_agent_stub_config_file_pull_sync, - request_agent_stub_config_manifest_sync, - request_agent_stub_config_push_sync, - request_agent_stub_config_skill_pull_sync, -) -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubConfigFileRef, - AgentStubConfigManifestResponse, - AgentStubConfigPushFileItem, - AgentStubConfigPushRequest, - AgentStubConfigPushSkillItem, -) - -_DEFAULT_CONFIG_BASE = Path("./.dify_conf") -_SKILL_MD_FILENAME = "SKILL.md" - - -class ConfigSkillPullResult(BaseModel): - class Item(BaseModel): - name: str - archive_path: str - directory_path: str - skill_md: str - - model_config = ConfigDict(extra="forbid") - - items: list[Item] - - model_config = ConfigDict(extra="forbid") - - -class ConfigFilePullResult(BaseModel): - class Item(BaseModel): - name: str - path: str - - model_config = ConfigDict(extra="forbid") - - items: list[Item] - - model_config = ConfigDict(extra="forbid") - - -@dataclass(frozen=True, slots=True) -class _PreparedPushItem: - name: str - path: Path - - -def manifest_from_environment() -> AgentStubConfigManifestResponse: - environment = read_agent_stub_environment() - return request_agent_stub_config_manifest_sync(url=environment.url, auth_jwe=environment.auth_jwe) - - -def pull_config_skills_from_environment( - names: list[str] | None = None, - local_dir: str | None = None, -) -> ConfigSkillPullResult: - environment = read_agent_stub_environment() - manifest = request_agent_stub_config_manifest_sync(url=environment.url, auth_jwe=environment.auth_jwe) - selected_names = names or [item.name for item in manifest.skills.items] - target_dir = Path(local_dir or (_DEFAULT_CONFIG_BASE / "skills")).expanduser().resolve() - target_dir.mkdir(parents=True, exist_ok=True) - - items: list[ConfigSkillPullResult.Item] = [] - for name in selected_names: - archive_bytes = request_agent_stub_config_skill_pull_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - name=name, - ) - archive_path = target_dir / f"{name}.zip" - skill_dir = target_dir / name - skill_dir.mkdir(parents=True, exist_ok=True) - archive_path.write_bytes(archive_bytes) - try: - extract_archive_to_directory(archive_path, target_dir=skill_dir) - except DriveMaterializationValidationError as exc: - raise AgentStubValidationError(str(exc)) from exc - except DriveMaterializationTransferError as exc: - raise AgentStubTransferError(str(exc)) from exc - skill_md_path = skill_dir / _SKILL_MD_FILENAME - if not skill_md_path.is_file(): - raise AgentStubValidationError(f"pulled config skill is missing {_SKILL_MD_FILENAME}: {name}") - items.append( - ConfigSkillPullResult.Item( - name=name, - archive_path=str(archive_path), - directory_path=str(skill_dir), - skill_md=skill_md_path.read_text(encoding="utf-8"), - ) - ) - return ConfigSkillPullResult(items=items) - - -def pull_config_files_from_environment( - names: list[str] | None = None, - local_dir: str | None = None, -) -> ConfigFilePullResult: - environment = read_agent_stub_environment() - manifest = request_agent_stub_config_manifest_sync(url=environment.url, auth_jwe=environment.auth_jwe) - selected_names = names or [item.name for item in manifest.files.items] - target_dir = Path(local_dir or (_DEFAULT_CONFIG_BASE / "files")).expanduser().resolve() - target_dir.mkdir(parents=True, exist_ok=True) - - items: list[ConfigFilePullResult.Item] = [] - for name in selected_names: - payload = request_agent_stub_config_file_pull_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - name=name, - ) - target_path = target_dir / name - target_path.parent.mkdir(parents=True, exist_ok=True) - target_path.write_bytes(payload) - items.append(ConfigFilePullResult.Item(name=name, path=str(target_path))) - return ConfigFilePullResult(items=items) - - -def pull_config_note_from_environment(local_path: str | None = None) -> Path: - manifest = manifest_from_environment() - target_path = Path(local_path or (_DEFAULT_CONFIG_BASE / "note.md")).expanduser().resolve() - target_path.parent.mkdir(parents=True, exist_ok=True) - target_path.write_text(manifest.note, encoding="utf-8") - return target_path - - -def push_config_note_from_environment(local_path: str | None) -> AgentStubConfigManifestResponse: - note = _read_text_input(local_path, _DEFAULT_CONFIG_BASE / "note.md") - return _push_config_from_environment(note=note) - - -def push_config_env_from_environment(local_path: str) -> AgentStubConfigManifestResponse: - env_text = _read_text_input(local_path, default_path=None) - return _push_config_from_environment(env_text=env_text) - - -def push_config_files_from_environment(paths: list[str]) -> AgentStubConfigManifestResponse: - _require_non_empty_inputs(paths, kind="file path") - items = [ - _PreparedPushItem( - name=_infer_name_from_path(source_path), - path=source_path, - ) - for source_path in _resolve_input_paths(paths) - ] - return _push_config_from_environment(files=[_build_file_push_item(item=item) for item in items]) - - -def delete_config_files_from_environment(names: list[str]) -> AgentStubConfigManifestResponse: - _require_non_empty_inputs(names, kind="file name") - return _push_config_from_environment( - files=[ - AgentStubConfigPushFileItem(name=_require_config_entry_name(name, kind="file"), file_ref=None) - for name in names - ] - ) - - -def push_config_skills_from_environment(paths: list[str]) -> AgentStubConfigManifestResponse: - _require_non_empty_inputs(paths, kind="skill directory") - items = [ - _PreparedPushItem(name=_infer_name_from_path(source_path), path=source_path) - for source_path in _resolve_input_paths(paths) - ] - return _push_config_from_environment(skills=[_build_skill_push_item(item=item) for item in items]) - - -def delete_config_skills_from_environment(names: list[str]) -> AgentStubConfigManifestResponse: - _require_non_empty_inputs(names, kind="skill name") - return _push_config_from_environment( - skills=[ - AgentStubConfigPushSkillItem(name=_require_config_entry_name(name, kind="skill"), file_ref=None) - for name in names - ] - ) - - -def _push_config_from_environment( - *, - files: list[AgentStubConfigPushFileItem] | None = None, - skills: list[AgentStubConfigPushSkillItem] | None = None, - env_text: str | None = None, - note: str | None = None, -) -> AgentStubConfigManifestResponse: - environment = read_agent_stub_environment() - return request_agent_stub_config_push_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - request=AgentStubConfigPushRequest( - files=files or [], - skills=skills or [], - env_text=env_text, - note=note, - ), - ) - - -def _read_text_input(path_value: str | None, default_path: Path | None) -> str: - if path_value == "-": - return os.fdopen(os.dup(0), encoding="utf-8").read() - if path_value is None: - if default_path is None: - raise AgentStubValidationError("local file path or '-' is required") - source_path = default_path.expanduser().resolve() - else: - source_path = Path(path_value).expanduser().resolve() - if not source_path.is_file(): - raise AgentStubValidationError(f"local file not found: {source_path}") - return source_path.read_text(encoding="utf-8") - - -def _resolve_input_paths(paths: list[str]) -> list[Path]: - return [Path(path).expanduser().resolve() for path in paths] - - -def _infer_name_from_path(path: Path) -> str: - return path.name - - -def _build_file_push_item( - *, - item: _PreparedPushItem, -) -> AgentStubConfigPushFileItem: - if not item.path.is_file(): - raise AgentStubValidationError(f"config file path must be a regular file: {item.path}") - uploaded = upload_tool_file_resource_from_environment(path=str(item.path)) - return AgentStubConfigPushFileItem( - name=item.name, - file_ref=AgentStubConfigFileRef(kind="tool_file", id=uploaded.tool_file_id), - ) - - -def _build_skill_push_item( - *, - item: _PreparedPushItem, -) -> AgentStubConfigPushSkillItem: - if not item.path.is_dir(): - raise AgentStubValidationError(f"config skill path must be a directory: {item.path}") - skill_md_path = item.path / _SKILL_MD_FILENAME - if not skill_md_path.is_file(): - raise AgentStubValidationError(f"config skill directory must contain {_SKILL_MD_FILENAME}: {item.path}") - with TemporaryDirectory() as temp_dir: - archive_path = Path(temp_dir) / f"{item.name}.zip" - _build_skill_archive(item.path, archive_path) - uploaded = upload_tool_file_resource_from_environment(path=str(archive_path)) - return AgentStubConfigPushSkillItem( - name=item.name, - file_ref=AgentStubConfigFileRef(kind="tool_file", id=uploaded.tool_file_id), - ) - - -def _require_config_entry_name(name: str, *, kind: str) -> str: - normalized = name.strip() - if not normalized: - raise AgentStubValidationError(f"config {kind} name must not be empty") - return normalized - - -def _require_non_empty_inputs(values: list[str], *, kind: str) -> None: - if not values: - raise AgentStubValidationError(f"at least one {kind} is required") - - -__all__ = [ - "ConfigFilePullResult", - "ConfigSkillPullResult", - "manifest_from_environment", - "pull_config_files_from_environment", - "pull_config_note_from_environment", - "pull_config_skills_from_environment", - "push_config_env_from_environment", - "push_config_files_from_environment", - "push_config_note_from_environment", - "push_config_skills_from_environment", - "delete_config_files_from_environment", - "delete_config_skills_from_environment", -] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py deleted file mode 100644 index 45ae61faabf..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py +++ /dev/null @@ -1,369 +0,0 @@ -"""CLI helpers for sandbox-visible Agent Stub drive commands. - -Drive commands stay in the sandbox-facing CLI because they orchestrate existing -control-plane and signed data-plane helpers. The Agent Stub server authenticates -and injects trusted drive scope; this module only formats manifest output, -downloads signed URLs into a local drive base (including safe auto-extraction of -downloaded skill archives), and uploads local files before committing their -ToolFile ids back into the drive. -""" - -from __future__ import annotations - -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass -from pathlib import Path -from tempfile import TemporaryDirectory -from typing import ClassVar, Literal -from zipfile import ZIP_DEFLATED, ZipFile - -from pydantic import BaseModel, ConfigDict - -from dify_agent.agent_stub._drive_materialization import ( - DriveDownloadPayload, - DriveMaterializationTransferError, - DriveMaterializationValidationError, - SKILL_ARCHIVE_FILENAME, - materialize_drive_downloads, - resolve_drive_destination, -) -from dify_agent.agent_stub._constants import DEFAULT_AGENT_STUB_DRIVE_BASE -from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.cli._files import upload_tool_file_resource_from_environment -from dify_agent.agent_stub.client import ( - AgentStubTransferError, - AgentStubValidationError, - download_file_bytes_from_signed_url_sync, - request_agent_stub_drive_commit_sync, - request_agent_stub_drive_manifest_sync, -) -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubDriveCommitItem, - AgentStubDriveCommitRequest, - AgentStubDriveCommitResponse, - AgentStubDriveFileRef, - AgentStubDriveItem, - AgentStubDriveManifestResponse, -) - -_SKILL_MD_FILENAME = "SKILL.md" -_SKIP_DIR_NAMES = frozenset( - {".git", "__pycache__", ".pytest_cache", ".mypy_cache", ".ruff_cache", ".venv", "node_modules"} -) -_SKIP_FILE_NAMES = frozenset({".DS_Store", SKILL_ARCHIVE_FILENAME}) -DrivePushKind = Literal["file", "skill", "dir"] - - -@dataclass(frozen=True, slots=True) -class _DriveUploadItem: - """Prepared local upload paired with its destination drive key.""" - - local_path: Path - drive_key: str - - -class DrivePullResult(BaseModel): - """Structured JSON result for ``dify-agent drive pull --json``.""" - - class Item(BaseModel): - key: str - local_path: str - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - items: list[Item] - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -def list_drive_manifest_from_environment(prefix: str) -> AgentStubDriveManifestResponse: - """List drive items through the Agent Stub using the current environment. - - Args: - prefix: Optional drive-key prefix forwarded to the manifest request. - - Returns: - The validated manifest response model. - - Side effects: - Calls the Agent Stub drive manifest control-plane endpoint with - ``include_download_url=False`` so list output does not allocate signed - download URLs. - """ - - environment = read_agent_stub_environment() - response = request_agent_stub_drive_manifest_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - prefix=prefix, - include_download_url=False, - ) - return response - - -def pull_drive_from_environment( - targets: list[str] | None = None, - local_base: str | None = None, -) -> DrivePullResult: - """Pull drive files into one local drive base via signed download URLs. - - Args: - targets: Optional drive-key targets or prefixes. An empty list preserves - the historical whole-drive pull by using ``[""]``. - local_base: Local base directory that receives downloaded drive files. - When omitted, the historical Agent Stub drive base is used. - - Returns: - A structured JSON-ready result with requested drive targets/prefixes - that matched at least one manifest item and their local paths. - - Observable behavior: - Requests a manifest with ``include_download_url=True``, requires every - returned item to include ``download_url``, downloads bytes directly from - those signed URLs, blocks path traversal by resolving each destination - under the resolved drive base, writes through a temporary sibling file - before replacing the final path, validates byte length when the manifest - includes ``size``, and automatically extracts - ``.DIFY-SKILL-FULL.zip`` archives into their containing skill - directory with the same path-safety checks. Archive extraction is staged - under a temporary directory and only moved into place after the full - archive validates successfully. Successfully extracted skill archives - are deleted from disk. - - Downloaded files and extracted files are materialized on disk but are - not enumerated in the returned item list; prefix pulls return the local - path corresponding to the requested prefix. - - Raises: - AgentStubValidationError: if a manifest item omits ``download_url``, a - destination would escape the drive base, or a downloaded skill - archive contains unsafe entries such as absolute paths, traversal - entries, or symlink entries. - AgentStubTransferError: if a downloaded payload does not match declared - size metadata or a downloaded skill archive is corrupt / not a valid - zip file. - """ - - environment = read_agent_stub_environment() - manifest_targets = targets or [""] - - def _fetch_manifest(target: str) -> AgentStubDriveManifestResponse: - return request_agent_stub_drive_manifest_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - prefix=target, - include_download_url=True, - ) - - with ThreadPoolExecutor(max_workers=min(len(manifest_targets), 4)) as executor: - responses = list(executor.map(_fetch_manifest, manifest_targets)) - downloads: list[DriveDownloadPayload] = [] - resolved_base_path = Path(local_base or DEFAULT_AGENT_STUB_DRIVE_BASE).expanduser().resolve() - result_items: list[DrivePullResult.Item] = [] - seen_result_targets: set[str] = set() - for target, response in zip(manifest_targets, responses, strict=True): - if not response.items or target in seen_result_targets: - continue - seen_result_targets.add(target) - try: - local_path = resolve_drive_destination(resolved_base_path, target) - except DriveMaterializationValidationError as exc: - raise AgentStubValidationError(str(exc)) from exc - result_items.append(DrivePullResult.Item(key=target, local_path=str(local_path))) - deduplicated_items = {item.key: item for response in responses for item in response.items} - for item in [deduplicated_items[key] for key in sorted(deduplicated_items)]: - download_url = item.download_url - if not isinstance(download_url, str) or not download_url: - raise AgentStubValidationError(f"drive manifest item is missing download_url: {item.key}") - try: - _ = resolve_drive_destination(resolved_base_path, item.key) - except DriveMaterializationValidationError as exc: - raise AgentStubValidationError(str(exc)) from exc - payload = download_file_bytes_from_signed_url_sync(download_url=download_url) - downloads.append(DriveDownloadPayload(key=item.key, payload=payload, size=item.size)) - - try: - _ = materialize_drive_downloads( - base_path=resolved_base_path, - downloads=downloads, - ) - except DriveMaterializationValidationError as exc: - raise AgentStubValidationError(str(exc)) from exc - except DriveMaterializationTransferError as exc: - raise AgentStubTransferError(str(exc)) from exc - - return DrivePullResult(items=result_items) - - -def push_drive_from_environment( - local_path: str, - drive_path: str, - *, - kind: DrivePushKind | None, -) -> AgentStubDriveCommitResponse: - """Upload local files through the Agent Stub and commit them into the drive. - - Args: - local_path: Source file or directory in the sandbox filesystem. - drive_path: Destination drive key or drive-key prefix. - kind: Optional public upload mode. Files infer file mode when omitted, - while directories require explicit ``skill`` or ``dir`` selection. - - Returns: - The validated drive commit response returned by the Agent Stub. - - Mode split: - * If ``local_path`` is a file, upload that file and commit exactly one - ``tool_file`` binding to ``drive_path``. - * If ``local_path`` is a directory and ``kind`` is ``"skill"``, - require ``SKILL.md`` and standardize the upload into - ``/SKILL.md`` plus ``/.DIFY-SKILL-FULL.zip``. - * If ``local_path`` is a directory and ``kind`` is ``"dir"``, upload - each regular file under ``drive_path/`` without skill - standardization. - - Observable safety behavior: - Rejects missing local paths, rejects directory pushes without an - explicit mode, rejects raw directory pushes with no regular files, and - rejects symlinked or escaping paths, including symlinked top-level - ``local_path`` roots, while preparing directory uploads or skill - archives. - """ - - source_path = Path(local_path).expanduser() - if kind not in {None, "file", "skill", "dir"}: - raise AgentStubValidationError(f"invalid drive push kind: {kind}") - if source_path.is_symlink(): - raise AgentStubValidationError(f"drive push does not support symlinked local paths: {source_path}") - source_path = source_path.resolve() - if source_path.is_file(): - if kind == "skill": - raise AgentStubValidationError("--kind skill requires a directory containing SKILL.md") - if kind == "dir": - raise AgentStubValidationError("--kind dir requires a directory") - return _commit_uploaded_items([_prepare_uploaded_file(source_path, drive_path)]) - if not source_path.is_dir(): - raise AgentStubValidationError(f"local path not found: {source_path}") - if kind is None: - raise AgentStubValidationError("directory drive push requires --kind skill or --kind dir") - if kind == "file": - raise AgentStubValidationError("--kind file requires a file") - if kind == "dir": - upload_items = [ - _prepare_uploaded_file(path, _join_drive_key(drive_path, relative_path)) - for path, relative_path in _iter_regular_files(source_path) - ] - if not upload_items: - raise AgentStubValidationError(f"directory has no regular files: {source_path}") - return _commit_uploaded_items(upload_items) - return _push_skill_directory(source_path, drive_path) - - -def _push_skill_directory(source_path: Path, drive_path: str) -> AgentStubDriveCommitResponse: - skill_md_path = source_path / _SKILL_MD_FILENAME - if not skill_md_path.is_file(): - raise AgentStubValidationError("--kind skill requires a directory containing SKILL.md") - with TemporaryDirectory() as temp_dir: - archive_path = Path(temp_dir) / SKILL_ARCHIVE_FILENAME - _build_skill_archive(source_path, archive_path) - return _commit_uploaded_items( - [ - _prepare_uploaded_file(skill_md_path.resolve(), _join_drive_key(drive_path, _SKILL_MD_FILENAME)), - _prepare_uploaded_file(archive_path, _join_drive_key(drive_path, SKILL_ARCHIVE_FILENAME)), - ] - ) - - -def _prepare_uploaded_file(local_path: Path, drive_key: str) -> _DriveUploadItem: - return _DriveUploadItem(local_path=local_path, drive_key=drive_key) - - -def _commit_uploaded_items(items: list[_DriveUploadItem]) -> AgentStubDriveCommitResponse: - environment = read_agent_stub_environment() - commit_items: list[AgentStubDriveCommitItem] = [] - for item in items: - uploaded_file = upload_tool_file_resource_from_environment(path=str(item.local_path)) - commit_items.append( - AgentStubDriveCommitItem( - key=item.drive_key, - file_ref=AgentStubDriveFileRef(kind="tool_file", id=uploaded_file.tool_file_id), - ) - ) - return request_agent_stub_drive_commit_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - request=AgentStubDriveCommitRequest(items=commit_items), - ) - - -def format_drive_manifest(response: AgentStubDriveManifestResponse) -> str: - return "\n".join(_format_manifest_item(item) for item in response.items) - - -def _format_manifest_item(item: AgentStubDriveItem) -> str: - size = str(item.size) if item.size is not None else "-" - mime_type = item.mime_type or "-" - item_hash = item.hash or "-" - return f"{size}\t{mime_type}\t{item_hash}\t{item.key}" - - -def _iter_regular_files(root_path: Path) -> list[tuple[Path, str]]: - """Return all regular files under one directory, rejecting unsafe symlinks.""" - - return _iter_regular_files_with_skip_filter(root_path, skip_filtered=False) - - -def _iter_skill_archive_files(root_path: Path) -> list[tuple[Path, str]]: - """Return regular files for skill packaging, excluding transient content.""" - - return _iter_regular_files_with_skip_filter(root_path, skip_filtered=True) - - -def _iter_regular_files_with_skip_filter(root_path: Path, *, skip_filtered: bool) -> list[tuple[Path, str]]: - root_resolved = root_path.resolve() - collected: list[tuple[Path, str]] = [] - for candidate in sorted(root_path.rglob("*")): - if skip_filtered and _should_skip_path(candidate, root_path): - continue - if candidate.is_symlink(): - raise AgentStubValidationError(f"drive push does not support symlinked files: {candidate}") - if not candidate.is_file(): - continue - resolved_candidate = candidate.resolve() - try: - relative_path = resolved_candidate.relative_to(root_resolved) - except ValueError as exc: - raise AgentStubValidationError( - f"drive push file resolves outside the source directory: {candidate}" - ) from exc - collected.append((resolved_candidate, relative_path.as_posix())) - return collected - - -def _should_skip_path(candidate: Path, root_path: Path) -> bool: - relative_path = candidate.relative_to(root_path) - if any(part in _SKIP_DIR_NAMES for part in relative_path.parts): - return True - return candidate.name in _SKIP_FILE_NAMES - - -def _build_skill_archive(source_path: Path, archive_path: Path) -> None: - with ZipFile(archive_path, mode="w", compression=ZIP_DEFLATED) as archive: - for file_path, relative_path in _iter_skill_archive_files(source_path): - archive.write(file_path, arcname=relative_path) - - -def _join_drive_key(base_key: str, child_key: str) -> str: - stripped_base = base_key.rstrip("/") - stripped_child = child_key.lstrip("/") - return f"{stripped_base}/{stripped_child}" if stripped_base else stripped_child - - -__all__ = [ - "DrivePullResult", - "DrivePushKind", - "format_drive_manifest", - "list_drive_manifest_from_environment", - "pull_drive_from_environment", - "push_drive_from_environment", -] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_env.py b/dify-agent/src/dify_agent/agent_stub/cli/_env.py deleted file mode 100644 index 61f06412262..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/_env.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Environment-variable helpers for the client-safe Agent Stub CLI.""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -import os - -from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE -from dify_agent.agent_stub.protocol.agent_stub import ( - AGENT_STUB_AUTH_JWE_ENV_VAR, - AGENT_STUB_API_BASE_URL_ENV_VAR, - normalize_agent_stub_api_base_url, -) - - -class MissingAgentStubEnvironmentError(RuntimeError): - """Raised when the Agent Stub CLI environment is incomplete.""" - - -@dataclass(slots=True) -class AgentStubEnvironment: - """Validated environment values needed for one CLI forwarding request.""" - - url: str - auth_jwe: str - - -def has_agent_stub_environment(env: Mapping[str, str] | None = None) -> bool: - """Return whether both required Agent Stub environment variables exist.""" - values = env or os.environ - return bool(values.get(AGENT_STUB_API_BASE_URL_ENV_VAR) and values.get(AGENT_STUB_AUTH_JWE_ENV_VAR)) - - -def read_agent_stub_environment(env: Mapping[str, str] | None = None) -> AgentStubEnvironment: - """Read and validate the Agent Stub environment variables.""" - values = env or os.environ - url = (values.get(AGENT_STUB_API_BASE_URL_ENV_VAR) or "").strip() - auth_jwe = (values.get(AGENT_STUB_AUTH_JWE_ENV_VAR) or "").strip() - missing: list[str] = [] - if not url: - missing.append(AGENT_STUB_API_BASE_URL_ENV_VAR) - if not auth_jwe: - missing.append(AGENT_STUB_AUTH_JWE_ENV_VAR) - if missing: - names = ", ".join(missing) - raise MissingAgentStubEnvironmentError(f"missing required Agent Stub environment variables: {names}") - try: - normalized_url = normalize_agent_stub_api_base_url(url) - except ValueError as exc: - raise MissingAgentStubEnvironmentError(f"invalid {AGENT_STUB_API_BASE_URL_ENV_VAR}: {exc}") from exc - return AgentStubEnvironment(url=normalized_url, auth_jwe=auth_jwe) - - -def read_agent_stub_drive_base(env: Mapping[str, str] | None = None) -> str: - """Read the sandbox-local drive base used by ``dify-agent drive pull``. - - The variable is optional because older Agent Stub environments only injected - URL/auth values. Blank values keep the historical ``/mnt/drive`` fallback. - """ - values = env or os.environ - configured_drive_base = (values.get(AGENT_STUB_DRIVE_BASE_ENV_VAR) or "").strip() - return configured_drive_base or DEFAULT_AGENT_STUB_DRIVE_BASE - - -__all__ = [ - "AgentStubEnvironment", - "MissingAgentStubEnvironmentError", - "has_agent_stub_environment", - "read_agent_stub_drive_base", - "read_agent_stub_environment", -] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_files.py b/dify-agent/src/dify_agent/agent_stub/cli/_files.py deleted file mode 100644 index 1ccd57f26da..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/_files.py +++ /dev/null @@ -1,247 +0,0 @@ -"""CLI helpers for sandbox-visible Agent Stub file commands.""" - -from __future__ import annotations - -import mimetypes -from dataclasses import dataclass -from pathlib import Path -from typing import ClassVar, Literal, Protocol, cast - -from pydantic import BaseModel, ConfigDict, ValidationError - -from dify_agent.agent_stub.cli._env import read_agent_stub_environment -from dify_agent.agent_stub.client import ( - AgentStubTransferError, - AgentStubValidationError, - download_file_bytes_from_signed_url_sync, - request_agent_stub_file_download_sync, - request_agent_stub_file_upload_sync, - upload_file_to_signed_url_sync, -) -from dify_agent.agent_stub.protocol.agent_stub import AgentStubFileMapping, is_canonical_dify_file_reference - - -class _AgentStubFileDownloadResponse(Protocol): - download_url: str - - -class UploadedToolFileMapping(BaseModel): - """Canonical Agent output mapping returned by ``dify-agent file upload``.""" - - transfer_method: Literal["tool_file"] = "tool_file" - reference: str - download_url: str - - model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") - - -@dataclass(frozen=True, slots=True) -class DownloadedFileResult: - """Local filesystem result for one CLI download command.""" - - path: Path - - -@dataclass(frozen=True, slots=True) -class UploadedToolFileResource: - """Lower-level upload result carrying the internal mapping and ToolFile id.""" - - mapping: AgentStubFileMapping - tool_file_id: str - - -def upload_file_from_environment(*, path: str) -> UploadedToolFileMapping: - """Upload one sandbox-local file through the Agent Stub control plane. - - The signed upload data-plane response must carry the Dify-generated - ``reference`` for the new ``ToolFile``. The helper then resolves the same - mapping through the signed download-request control plane so the public CLI - output includes a ready-to-share external ``download_url``. - """ - - resource = upload_tool_file_resource_from_environment(path=path) - reference = resource.mapping.reference - if not isinstance(reference, str) or not reference: - raise AgentStubTransferError("signed file upload response is missing reference") - environment = read_agent_stub_environment() - download_url = _request_uploaded_tool_file_download_url( - url=environment.url, - auth_jwe=environment.auth_jwe, - reference=reference, - ) - return UploadedToolFileMapping(reference=reference, download_url=download_url) - - -def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFileResource: - """Upload one sandbox-local file and preserve both reference and ToolFile id. - - This lower-level helper backs ``drive push``. The signed upload data-plane - response must include both the canonical Dify ``reference`` used by public - CLI output and the raw ToolFile ``id`` required by drive commit payloads. - It intentionally stops after the upload allocation so internal flows that - only need the ToolFile identity do not pay for signed download enrichment. - - Raises: - AgentStubValidationError: if ``path`` does not resolve to a local file. - AgentStubTransferError: if the signed upload response omits either the - canonical ``reference`` or the raw ToolFile ``id``, or if the - canonical reference is malformed. - """ - - source_path = Path(path).expanduser().resolve() - if not source_path.is_file(): - raise AgentStubValidationError(f"local file not found: {source_path}") - - environment = read_agent_stub_environment() - filename = source_path.name - mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" - upload_request: object = request_agent_stub_file_upload_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - filename=filename, - mimetype=mime_type, - ) - if not hasattr(upload_request, "upload_url") or not isinstance(upload_request.upload_url, str): - raise AgentStubTransferError("signed file upload response is missing upload_url") - with source_path.open("rb") as file_obj: - payload = upload_file_to_signed_url_sync( - upload_url=upload_request.upload_url, - filename=filename, - file_obj=file_obj, - mimetype=mime_type, - ) - reference, tool_file_id = _normalize_uploaded_tool_file_payload(payload) - return UploadedToolFileResource( - mapping=AgentStubFileMapping(transfer_method="tool_file", reference=reference), - tool_file_id=tool_file_id, - ) - - -def download_file_from_environment( - *, - transfer_method: str | None = None, - reference_or_url: str | None = None, - mapping: str | None = None, - local_dir: str | None = None, -) -> DownloadedFileResult: - """Download one workflow file mapping into the sandbox filesystem. - - Callers may provide either the public positional pair - ``TRANSFER_METHOD REFERENCE_OR_URL`` or one JSON ``--mapping`` payload. - The helper normalizes both forms into ``AgentStubFileMapping`` before - requesting a signed download URL from the Agent Stub. - """ - - file_mapping = _build_download_mapping( - transfer_method=transfer_method, - reference_or_url=reference_or_url, - mapping=mapping, - ) - environment = read_agent_stub_environment() - - download_request: object = request_agent_stub_file_download_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - file=file_mapping, - for_external=False, - ) - if not hasattr(download_request, "filename") or not isinstance(download_request.filename, str): - raise AgentStubTransferError("signed file download response is missing filename") - if not hasattr(download_request, "download_url") or not isinstance(download_request.download_url, str): - raise AgentStubTransferError("signed file download response is missing download_url") - target_dir = Path(local_dir).expanduser().resolve() if local_dir else Path.cwd() - target_dir.mkdir(parents=True, exist_ok=True) - destination = _deduplicate_destination_path(target_dir / _sanitize_download_filename(download_request.filename)) - _ = destination.write_bytes(download_file_bytes_from_signed_url_sync(download_url=download_request.download_url)) - return DownloadedFileResult(path=destination) - - -def _build_download_mapping( - *, - transfer_method: str | None, - reference_or_url: str | None, - mapping: str | None, -) -> AgentStubFileMapping: - if mapping is not None: - if transfer_method is not None or reference_or_url is not None: - raise AgentStubValidationError("--mapping cannot be combined with TRANSFER_METHOD or REFERENCE_OR_URL") - try: - return AgentStubFileMapping.model_validate_json(mapping) - except ValidationError as exc: - raise AgentStubValidationError("invalid file download mapping") from exc - - if transfer_method is None or reference_or_url is None: - raise AgentStubValidationError("file download requires either --mapping or TRANSFER_METHOD REFERENCE_OR_URL") - - normalized_transfer_method = cast( - Literal["local_file", "tool_file", "datasource_file", "remote_url"], - transfer_method, - ) - try: - return AgentStubFileMapping( - transfer_method=normalized_transfer_method, - url=reference_or_url if normalized_transfer_method == "remote_url" else None, - reference=reference_or_url if normalized_transfer_method != "remote_url" else None, - ) - except ValidationError as exc: - raise AgentStubValidationError("invalid file download arguments") from exc - - -def _normalize_uploaded_tool_file_payload(payload: dict[str, object]) -> tuple[str, str]: - reference = payload.get("reference") - if not isinstance(reference, str) or not reference: - raise AgentStubTransferError("signed file upload response is missing reference") - if not is_canonical_dify_file_reference(reference): - raise AgentStubTransferError("signed file upload response has invalid canonical reference") - tool_file_id = payload.get("id") - if not isinstance(tool_file_id, str) or not tool_file_id: - raise AgentStubTransferError("signed file upload response is missing id") - return reference, tool_file_id - - -def _request_uploaded_tool_file_download_url(*, url: str, auth_jwe: str, reference: str) -> str: - download_request = cast( - _AgentStubFileDownloadResponse, - request_agent_stub_file_download_sync( - url=url, - auth_jwe=auth_jwe, - file=AgentStubFileMapping(transfer_method="tool_file", reference=reference), - ), - ) - if not hasattr(download_request, "download_url") or not isinstance(download_request.download_url, str): - raise AgentStubTransferError("signed file download response is missing download_url") - download_url = download_request.download_url - if not download_url: - raise AgentStubTransferError("signed file download response is missing download_url") - return download_url - - -def _deduplicate_destination_path(path: Path) -> Path: - if not path.exists(): - return path - - suffix = "".join(path.suffixes) - stem = path.name[: -len(suffix)] if suffix else path.name - counter = 1 - while True: - candidate = path.with_name(f"{stem} ({counter}){suffix}") - if not candidate.exists(): - return candidate - counter += 1 - - -def _sanitize_download_filename(filename: str) -> str: - sanitized = Path(filename).name - if sanitized in {"", ".", ".."}: - raise AgentStubTransferError("signed file download response has invalid filename") - return sanitized - - -__all__ = [ - "DownloadedFileResult", - "UploadedToolFileMapping", - "UploadedToolFileResource", - "download_file_from_environment", - "upload_file_from_environment", - "upload_tool_file_resource_from_environment", -] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/main.py b/dify-agent/src/dify_agent/agent_stub/cli/main.py deleted file mode 100644 index 592ff35253e..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/cli/main.py +++ /dev/null @@ -1,667 +0,0 @@ -"""Typer entry point for the client-safe ``dify-agent`` console script. - -The CLI supports explicit ``connect``, ``file``, and ``drive`` commands and -treats unknown bare commands as Agent Stub forwards. When the injected Agent -Stub environment variables are missing, that path intentionally surfaces a -clear missing-env error instead of Typer's generic unknown-command message. The -module depends only on client-safe code so importing the console entry point -does not pull in FastAPI, Redis, shellctl, or JWE runtime dependencies. -""" - -from __future__ import annotations - -from functools import cache -from importlib import import_module -import sys - -import click -import typer -from typer.main import get_command - -from dify_agent.agent_stub._constants import AGENT_STUB_DRIVE_BASE_ENV_VAR, DEFAULT_AGENT_STUB_DRIVE_BASE - -_CONFIG_MANIFEST_STDOUT_EXCLUDE = { - "skills": {"items": {"__all__": {"hash"}}}, - "files": {"items": {"__all__": {"hash"}}}, -} -_FILE_DOWNLOAD_TRANSFER_METHOD_HELP = ( - "File mapping transfer_method: local_file, tool_file, datasource_file, or remote_url." -) -_FILE_DOWNLOAD_REFERENCE_OR_URL_HELP = ( - "File mapping reference or URL. Use dify-file-ref:... with local_file, tool_file, or datasource_file; " - "use https://... with remote_url." -) - - -app = typer.Typer( - add_completion=False, - help="Forward shell-visible dify-agent commands to the Dify Agent Stub server.", - no_args_is_help=True, - rich_markup_mode=None, -) -file_app = typer.Typer(help="Upload or download workflow files through the Agent Stub.", rich_markup_mode=None) -config_app = typer.Typer( - help="Inspect or update Agent Soul-backed config assets through the Agent Stub.", - rich_markup_mode=None, -) -config_skills_app = typer.Typer(help="Pull or update config skills through the Agent Stub.", rich_markup_mode=None) -config_files_app = typer.Typer(help="Pull or update config files through the Agent Stub.", rich_markup_mode=None) -config_skill_pull_alias_app = typer.Typer(help="Pull config skills through the Agent Stub.", rich_markup_mode=None) -config_file_pull_alias_app = typer.Typer(help="Pull config files through the Agent Stub.", rich_markup_mode=None) -config_env_app = typer.Typer(help="Update config env variables visible to the current run.", rich_markup_mode=None) -config_note_app = typer.Typer(help="Pull or update the current config note.", rich_markup_mode=None) -drive_app = typer.Typer(help="List, pull, or push agent drive files through the Agent Stub.", rich_markup_mode=None) -app.add_typer(file_app, name="file") -app.add_typer(config_app, name="config") -config_app.add_typer(config_skills_app, name="skills") -# Keep hidden singular aliases on separate pull-only apps. Reusing the plural -# apps here would also expose push/delete through `config skill|file ...`, -# which is intentionally not part of the compatibility surface. -config_app.add_typer(config_skill_pull_alias_app, name="skill", hidden=True) -config_app.add_typer(config_files_app, name="files") -config_app.add_typer(config_file_pull_alias_app, name="file", hidden=True) -config_app.add_typer(config_env_app, name="env") -config_app.add_typer(config_note_app, name="note") -app.add_typer(drive_app, name="drive") -_KNOWN_ROOT_COMMANDS = frozenset({"config", "connect", "drive", "file"}) - - -@app.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def connect( - json_output: bool = typer.Option(False, "--json", help="Emit the connection response as JSON."), - argv: list[str] = typer.Argument(default_factory=list, metavar="ARGV"), -) -> None: - """Establish one Agent Stub connection using the current environment.""" - _run_connect(argv=list(argv), json_output=json_output) - - -@file_app.command("upload") -def upload(path: str = typer.Argument(..., metavar="PATH")) -> None: - """Upload one sandbox-local file as a ToolFile output reference.""" - _run_file_upload(path=path) - - -@file_app.command("download") -def download( - transfer_method: str = typer.Argument(..., metavar="TRANSFER_METHOD", help=_FILE_DOWNLOAD_TRANSFER_METHOD_HELP), - reference_or_url: str = typer.Argument( - ..., - metavar="REFERENCE_OR_URL", - help=_FILE_DOWNLOAD_REFERENCE_OR_URL_HELP, - ), - local_dir: str | None = typer.Option(None, "--to", help="Local directory for the downloaded file."), -) -> None: - """Download one workflow file mapping into the local sandbox directory.""" - _run_file_download( - transfer_method=transfer_method, - reference_or_url=reference_or_url, - local_dir=local_dir, - ) - - -@config_app.command("manifest") -def config_manifest() -> None: - """Show the current visible Agent config manifest as JSON.""" - _run_config_manifest() - - -@config_skills_app.command("pull") -def config_skills_pull( - names: list[str] = typer.Argument(None, metavar="NAME"), - local_dir: str | None = typer.Option(None, "--to", help="Local directory for pulled config skills."), - json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), -) -> None: - """Pull one or all visible config skills into ./.dify_conf/skills by default.""" - _run_config_skill_pull(names=names or None, local_dir=local_dir, json_output=json_output) - - -@config_skill_pull_alias_app.command("pull") -def config_skill_pull_alias( - names: list[str] = typer.Argument(None, metavar="NAME"), - local_dir: str | None = typer.Option(None, "--to", help="Local directory for pulled config skills."), - json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), -) -> None: - """Pull one or all visible config skills into ./.dify_conf/skills by default.""" - _run_config_skill_pull(names=names or None, local_dir=local_dir, json_output=json_output) - - -@config_skills_app.command("push") -def config_skills_push( - paths: list[str] = typer.Argument(..., metavar="PATH..."), -) -> None: - """Upload one or more local skill directories into the current config manifest. - - Pass a directory such as ./skills/researcher that contains SKILL.md. Other files in that directory are - archived with the skill. Pushing a skill with an existing name replaces that config skill. - - Skill directory requirements: - - - Each PATH must be one skill directory; the directory basename is the config skill name. - - - The directory must contain a top-level SKILL.md. - - - SKILL.md must be non-empty UTF-8 Markdown. - - - SKILL.md must start with YAML frontmatter matching this schema: - - \b - --- - name: - description: - --- - - - Symlinked files are rejected. - - - Dependency/cache folders such as .git, __pycache__, .venv and node_modules should be manually cleared before push. - """ - _run_config_skills_push(paths=paths) - - -@config_skills_app.command("delete") -def config_skills_delete( - names: list[str] = typer.Argument(..., metavar="NAME"), -) -> None: - """Delete one or more config skills by name without touching local directories.""" - _run_config_skills_delete(names=names) - - -@config_files_app.command("pull") -def config_files_pull( - names: list[str] = typer.Argument(None, metavar="NAME"), - local_dir: str | None = typer.Option(None, "--to", help="Local directory for pulled config files."), - json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), -) -> None: - """Pull one or all visible config files into ./.dify_conf/files by default.""" - _run_config_file_pull(names=names or None, local_dir=local_dir, json_output=json_output) - - -@config_file_pull_alias_app.command("pull") -def config_file_pull_alias( - names: list[str] = typer.Argument(None, metavar="NAME"), - local_dir: str | None = typer.Option(None, "--to", help="Local directory for pulled config files."), - json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), -) -> None: - """Pull one or all visible config files into ./.dify_conf/files by default.""" - _run_config_file_pull(names=names or None, local_dir=local_dir, json_output=json_output) - - -@config_files_app.command("push") -def config_files_push( - paths: list[str] = typer.Argument(..., metavar="PATH..."), -) -> None: - """Upload one or more local files into the current config manifest. - - Pass one or more regular local files such as ./guide.md and ./policy.txt. To upload a folder, compress it - first and pass the archive file. The config file name is the local basename. Pushing a file with an existing config - name replaces that config file. - """ - _run_config_files_push(paths=paths) - - -@config_files_app.command("delete") -def config_files_delete( - names: list[str] = typer.Argument(..., metavar="NAME"), -) -> None: - """Delete one or more config files by name without touching local files.""" - _run_config_files_delete(names=names) - - -@config_env_app.command("push") -def config_env_push( - local_path: str = typer.Argument(..., metavar="PATH|-"), -) -> None: - """Set or delete config env entries from one local dotenv file or stdin. - - Pass dotenv text with lines like API_KEY=value and OLD_KEY=. KEY=value sets or updates an entry, KEY= deletes - it, and omitted keys are unchanged. - """ - _run_config_env_push(local_path=local_path) - - -@config_note_app.command("pull") -def config_note_pull( - local_path: str | None = typer.Option(None, "--to", help="Local markdown file path."), -) -> None: - """Export the current config note into ./.dify_conf/note.md by default.""" - _run_config_note_pull(local_path=local_path) - - -@config_note_app.command("push") -def config_note_push( - local_path: str | None = typer.Argument(None, metavar="[PATH|-]"), -) -> None: - """Replace the current config note from one local text file or stdin. - - Pass a plain text or Markdown note file, or use - to read the note from stdin. When PATH is omitted, the - command reads ./.dify_conf/note.md. - """ - _run_config_note_push(local_path=local_path) - - -@drive_app.command("list") -def drive_list( - path_prefix: str = typer.Argument("", metavar="REMOTE_PREFIX"), - json_output: bool = typer.Option(False, "--json", help="Emit the drive manifest as JSON."), -) -> None: - """List drive files visible to the current sandbox execution.""" - _run_drive_list(path_prefix=path_prefix, json_output=json_output) - - -@drive_app.command("pull") -def drive_pull( - targets: list[str] = typer.Argument(None, metavar="REMOTE"), - local_base: str | None = typer.Option( - None, - "--to", - help=( - f"Local base directory for pulled drive files. Defaults to ${AGENT_STUB_DRIVE_BASE_ENV_VAR} " - f"or {DEFAULT_AGENT_STUB_DRIVE_BASE}." - ), - ), - json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), -) -> None: - """Pull one or more drive keys/prefixes into one local directory tree. - - Passing no ``TARGET`` preserves the historical whole-drive behavior by - pulling from the empty prefix. - """ - _run_drive_pull(targets=targets or None, local_base=local_base, json_output=json_output) - - -@drive_app.command("push") -def drive_push( - local_path: str = typer.Argument(..., metavar="LOCAL_PATH"), - drive_path: str = typer.Argument(..., metavar="REMOTE_PATH"), - kind: str | None = typer.Option(None, "--kind", help="Directory upload kind: skill or dir."), - json_output: bool = typer.Option( - False, - "--json", - help="Accepted for consistency; drive push output is already emitted as JSON.", - ), -) -> None: - """Upload one local file or directory into the agent drive.""" - del json_output - _run_drive_push(local_path=local_path, drive_path=drive_path, kind=kind) - - -def main(argv: list[str] | None = None) -> None: - """Run the ``dify-agent`` CLI with optional argv injection for tests.""" - args = list(sys.argv[1:] if argv is None else argv) - if args[:1] == ["connect"] and not _is_help_request(args[1:]): - json_output, forwarded_args = _parse_connect_args(args[1:]) - _run_connect(argv=forwarded_args, json_output=json_output) - return - json_output, forwarded_args = _extract_root_json_flag(args) - if _is_unknown_bare_command(forwarded_args): - if not _env_module().has_agent_stub_environment(): - _show_root_help() - _run_connect(argv=forwarded_args, json_output=json_output) - return - app(prog_name="dify-agent", args=args) - - -def _extract_root_json_flag(argv: list[str]) -> tuple[bool, list[str]]: - if len(argv) >= 2 and argv[0] == "--json" and argv[1] not in _KNOWN_ROOT_COMMANDS: - return True, argv[1:] - return False, argv - - -def _is_unknown_bare_command(argv: list[str]) -> bool: - if not argv: - return False - first = argv[0] - return first not in _KNOWN_ROOT_COMMANDS and not first.startswith("-") - - -def _parse_connect_args(argv: list[str]) -> tuple[bool, list[str]]: - json_output = False - remaining = list(argv) - if remaining[:1] == ["--json"]: - json_output = True - remaining = remaining[1:] - if remaining[:1] == ["--"]: - remaining = remaining[1:] - return json_output, remaining - - -def _is_help_request(argv: list[str]) -> bool: - return any(value in {"--help", "-h"} for value in argv) - - -def _show_root_help() -> None: - """Render root CLI guidance before reporting unknown-command env errors.""" - command = get_command(app) - context = command.make_context("dify-agent", [], resilient_parsing=True) - typer.echo(command.get_help(context)) - - -def render_agent_stub_cli_help(args: tuple[str, ...]) -> str: - """Render Click help for one known ``dify-agent`` subcommand without executing a shell.""" - command: click.Command = get_command(app) - parent_context: click.Context | None = None - command_path = ["dify-agent"] - for name in args: - if not isinstance(command, click.Group): - raise ValueError(f"dify-agent {' '.join(args)} is not a command group") - next_command = command.commands.get(name) - if next_command is None: - raise ValueError(f"unknown dify-agent command path: {' '.join(args)}") - current_context = click.Context(command, info_name=command_path[-1], parent=parent_context) - parent_context = current_context - command = next_command - command_path.append(name) - context = click.Context(command, info_name=command_path[-1], parent=parent_context) - return command.get_help(context).strip() - - -def _run_connect(*, argv: list[str], json_output: bool) -> None: - env_module = _env_module() - client_module = _client_module() - agent_stub_module = _agent_stub_module() - try: - response = agent_stub_module.connect_from_environment(argv=argv) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - - if json_output: - typer.echo(response.model_dump_json()) - return - typer.echo(f"connected {response.connection_id}") - - -def _run_file_upload(*, path: str) -> None: - env_module = _env_module() - client_module = _client_module() - files_module = _files_module() - try: - response = files_module.upload_file_from_environment(path=path) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_file_download( - *, - transfer_method: str, - reference_or_url: str, - local_dir: str | None, -) -> None: - env_module = _env_module() - client_module = _client_module() - files_module = _files_module() - try: - response = files_module.download_file_from_environment( - transfer_method=transfer_method, - reference_or_url=reference_or_url, - local_dir=local_dir, - ) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(str(response.path)) - - -def _run_config_manifest() -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.manifest_from_environment() - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json(exclude=_CONFIG_MANIFEST_STDOUT_EXCLUDE)) - - -def _run_config_skill_pull(*, names: list[str] | None, local_dir: str | None, json_output: bool) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.pull_config_skills_from_environment(names=names, local_dir=local_dir) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - if json_output: - typer.echo(response.model_dump_json()) - return - for index, item in enumerate(response.items): - if index: - typer.echo("") - typer.echo(item.directory_path) - typer.echo(item.skill_md, nl=False) - - -def _run_config_file_pull(*, names: list[str] | None, local_dir: str | None, json_output: bool) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.pull_config_files_from_environment(names=names, local_dir=local_dir) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - if json_output: - typer.echo(response.model_dump_json()) - return - for item in response.items: - typer.echo(item.path) - - -def _run_config_note_pull(*, local_path: str | None) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - path = config_module.pull_config_note_from_environment(local_path=local_path) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(str(path)) - - -def _run_config_note_push(*, local_path: str | None) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.push_config_note_from_environment(local_path=local_path) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_config_env_push(*, local_path: str) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.push_config_env_from_environment(local_path=local_path) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_config_files_push(*, paths: list[str]) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.push_config_files_from_environment(paths=paths) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_config_files_delete(*, names: list[str]) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.delete_config_files_from_environment(names=names) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_config_skills_push(*, paths: list[str]) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.push_config_skills_from_environment(paths=paths) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_config_skills_delete(*, names: list[str]) -> None: - env_module = _env_module() - client_module = _client_module() - config_module = _config_module() - try: - response = config_module.delete_config_skills_from_environment(names=names) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -def _run_drive_list(*, path_prefix: str, json_output: bool) -> None: - env_module = _env_module() - client_module = _client_module() - drive_module = _drive_module() - try: - response = drive_module.list_drive_manifest_from_environment(prefix=path_prefix) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - if json_output: - typer.echo(response.model_dump_json()) - return - typer.echo(drive_module.format_drive_manifest(response)) - - -def _run_drive_pull(*, targets: list[str] | None, local_base: str | None, json_output: bool) -> None: - env_module = _env_module() - client_module = _client_module() - drive_module = _drive_module() - try: - response = drive_module.pull_drive_from_environment( - targets=targets, - local_base=local_base or env_module.read_agent_stub_drive_base(), - ) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - if json_output: - typer.echo(response.model_dump_json()) - return - for item in response.items: - typer.echo(item.local_path) - - -def _run_drive_push(*, local_path: str, drive_path: str, kind: str | None) -> None: - env_module = _env_module() - client_module = _client_module() - drive_module = _drive_module() - try: - response = drive_module.push_drive_from_environment( - local_path=local_path, - drive_path=drive_path, - kind=kind, - ) - except env_module.MissingAgentStubEnvironmentError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(2) from exc - except client_module.AgentStubClientError as exc: - typer.echo(str(exc), err=True) - raise SystemExit(1) from exc - typer.echo(response.model_dump_json()) - - -# Keep helper imports on demand so importing CLI/help stays free of server-side -# and unrelated heavy runtime dependencies. -@cache -def _agent_stub_module(): - return import_module("dify_agent.agent_stub.cli._agent_stub") - - -@cache -def _config_module(): - return import_module("dify_agent.agent_stub.cli._config") - - -@cache -def _drive_module(): - return import_module("dify_agent.agent_stub.cli._drive") - - -@cache -def _files_module(): - return import_module("dify_agent.agent_stub.cli._files") - - -@cache -def _env_module(): - return import_module("dify_agent.agent_stub.cli._env") - - -@cache -def _client_module(): - return import_module("dify_agent.agent_stub.client") - - -__all__ = ["app", "main"] diff --git a/dify-agent/src/dify_agent/agent_stub/client/__init__.py b/dify-agent/src/dify_agent/agent_stub/client/__init__.py deleted file mode 100644 index 422f4a27301..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/client/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Client-safe helpers for the Dify Agent Stub control plane.""" - -from ._agent_stub import ( - connect_agent_stub_sync, - download_file_bytes_from_signed_url_sync, - request_agent_stub_config_env_update_sync, - request_agent_stub_config_file_pull_sync, - request_agent_stub_config_manifest_sync, - request_agent_stub_config_note_update_sync, - request_agent_stub_config_push_sync, - request_agent_stub_config_skill_inspect_sync, - request_agent_stub_config_skill_pull_sync, - request_agent_stub_drive_commit_sync, - request_agent_stub_drive_manifest_sync, - request_agent_stub_file_download_sync, - request_agent_stub_file_upload_sync, - upload_file_to_signed_url_sync, -) -from ._errors import ( - AgentStubClientError, - AgentStubGRPCError, - AgentStubHTTPError, - AgentStubMissingGRPCDependencyError, - AgentStubTransferError, - AgentStubValidationError, -) - -__all__ = [ - "AgentStubClientError", - "AgentStubGRPCError", - "AgentStubHTTPError", - "AgentStubMissingGRPCDependencyError", - "AgentStubTransferError", - "AgentStubValidationError", - "connect_agent_stub_sync", - "download_file_bytes_from_signed_url_sync", - "request_agent_stub_config_env_update_sync", - "request_agent_stub_config_file_pull_sync", - "request_agent_stub_config_manifest_sync", - "request_agent_stub_config_note_update_sync", - "request_agent_stub_config_push_sync", - "request_agent_stub_config_skill_inspect_sync", - "request_agent_stub_config_skill_pull_sync", - "request_agent_stub_drive_commit_sync", - "request_agent_stub_drive_manifest_sync", - "request_agent_stub_file_download_sync", - "request_agent_stub_file_upload_sync", - "upload_file_to_signed_url_sync", -] diff --git a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub.py b/dify-agent/src/dify_agent/agent_stub/client/_agent_stub.py deleted file mode 100644 index 1f5150b429e..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub.py +++ /dev/null @@ -1,352 +0,0 @@ -"""Transport-dispatch facade for Agent Stub control-plane calls.""" - -from __future__ import annotations - -import httpx -from pydantic import JsonValue - -from dify_agent.agent_stub.client._agent_stub_http import ( - connect_agent_stub_http_sync, - download_file_bytes_from_signed_url_sync, - request_agent_stub_config_env_update_http_sync, - request_agent_stub_config_file_pull_http_sync, - request_agent_stub_config_manifest_http_sync, - request_agent_stub_config_note_update_http_sync, - request_agent_stub_config_push_http_sync, - request_agent_stub_config_skill_inspect_http_sync, - request_agent_stub_config_skill_pull_http_sync, - request_agent_stub_drive_commit_http_sync, - request_agent_stub_drive_manifest_http_sync, - request_agent_stub_file_download_http_sync, - request_agent_stub_file_upload_http_sync, - upload_file_to_signed_url_sync, -) -from dify_agent.agent_stub.client._errors import AgentStubValidationError -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubConfigPushRequest, - AgentStubDriveCommitRequest, - parse_agent_stub_endpoint, - AgentStubFileMapping, -) - - -def connect_agent_stub_sync( - *, - url: str, - auth_jwe: str, - argv: list[str], - metadata: dict[str, JsonValue] | None = None, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Connect through HTTP or gRPC based on the configured Agent Stub URL.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - from dify_agent.agent_stub.client._agent_stub_grpc import connect_agent_stub_grpc_sync - - return connect_agent_stub_grpc_sync( - url=endpoint.url, - auth_jwe=auth_jwe, - argv=argv, - metadata=metadata, - timeout=timeout, - ) - return connect_agent_stub_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - argv=argv, - metadata=metadata, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_file_upload_sync( - *, - url: str, - auth_jwe: str, - filename: str, - mimetype: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Request a signed upload URL through the selected Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - from dify_agent.agent_stub.client._agent_stub_grpc import request_agent_stub_file_upload_grpc_sync - - return request_agent_stub_file_upload_grpc_sync( - url=endpoint.url, - auth_jwe=auth_jwe, - filename=filename, - mimetype=mimetype, - timeout=timeout, - ) - return request_agent_stub_file_upload_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - filename=filename, - mimetype=mimetype, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_file_download_sync( - *, - url: str, - auth_jwe: str, - file: AgentStubFileMapping, - for_external: bool = True, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Request a signed download URL through the selected Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - from dify_agent.agent_stub.client._agent_stub_grpc import request_agent_stub_file_download_grpc_sync - - return request_agent_stub_file_download_grpc_sync( - url=endpoint.url, - auth_jwe=auth_jwe, - file=file, - for_external=for_external, - timeout=timeout, - ) - return request_agent_stub_file_download_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - file=file, - for_external=for_external, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_drive_manifest_sync( - *, - url: str, - auth_jwe: str, - prefix: str, - include_download_url: bool, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Request one drive manifest through the HTTP Agent Stub transport. - - Drive operations are intentionally HTTP-only in this stage. Callers must - provide an ``http://`` or ``https://`` Agent Stub URL; ``grpc://`` endpoints - raise ``AgentStubValidationError`` instead of attempting transport fallback. - """ - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub drive operations require an HTTP Agent Stub URL") - return request_agent_stub_drive_manifest_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - prefix=prefix, - include_download_url=include_download_url, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_drive_commit_sync( - *, - url: str, - auth_jwe: str, - request: AgentStubDriveCommitRequest, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Commit one drive batch through the HTTP Agent Stub transport. - - Drive operations are intentionally HTTP-only in this stage. Callers must - provide an ``http://`` or ``https://`` Agent Stub URL; ``grpc://`` endpoints - raise ``AgentStubValidationError`` instead of attempting transport fallback. - """ - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub drive operations require an HTTP Agent Stub URL") - return request_agent_stub_drive_commit_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - request=request, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_manifest_sync( - *, - url: str, - auth_jwe: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Fetch the current config manifest through the HTTP Agent Stub transport. - - Config operations are HTTP-only in this stage. ``grpc://`` endpoints raise - ``AgentStubValidationError`` instead of attempting transport fallback. - """ - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_manifest_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_skill_pull_sync( - *, - url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Download one config skill archive through the HTTP Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_skill_pull_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - name=name, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_skill_inspect_sync( - *, - url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Fetch one config skill inspect view through the HTTP Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_skill_inspect_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - name=name, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_file_pull_sync( - *, - url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Download one config file payload through the HTTP Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_file_pull_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - name=name, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_push_sync( - *, - url: str, - auth_jwe: str, - request: AgentStubConfigPushRequest, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Push config file/skill/env/note mutations through the HTTP transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_push_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - request=request, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_env_update_sync( - *, - url: str, - auth_jwe: str, - env_text: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Set or delete config env entries through the HTTP transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_env_update_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - env_text=env_text, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def request_agent_stub_config_note_update_sync( - *, - url: str, - auth_jwe: str, - note: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -): - """Update the config note text through the HTTP Agent Stub transport.""" - endpoint = _parse_endpoint(url) - if endpoint.is_grpc: - raise AgentStubValidationError("Agent Stub config operations require an HTTP Agent Stub URL") - return request_agent_stub_config_note_update_http_sync( - base_url=endpoint.url, - auth_jwe=auth_jwe, - note=note, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def _parse_endpoint(url: str): - """Normalize one Agent Stub base URL and map parse failures to client errors.""" - try: - return parse_agent_stub_endpoint(url) - except ValueError as exc: - raise AgentStubValidationError("invalid Agent Stub base URL") from exc - - -__all__ = [ - "connect_agent_stub_sync", - "download_file_bytes_from_signed_url_sync", - "request_agent_stub_config_env_update_sync", - "request_agent_stub_config_file_pull_sync", - "request_agent_stub_config_manifest_sync", - "request_agent_stub_config_note_update_sync", - "request_agent_stub_config_push_sync", - "request_agent_stub_config_skill_inspect_sync", - "request_agent_stub_config_skill_pull_sync", - "request_agent_stub_drive_commit_sync", - "request_agent_stub_drive_manifest_sync", - "request_agent_stub_file_download_sync", - "request_agent_stub_file_upload_sync", - "upload_file_to_signed_url_sync", -] diff --git a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_grpc.py b/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_grpc.py deleted file mode 100644 index 7c5f5c4fbbb..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_grpc.py +++ /dev/null @@ -1,243 +0,0 @@ -"""Client-safe gRPC helpers for Agent Stub control-plane endpoints. - -These entrypoints mirror the HTTP helpers in ``_agent_stub_http.py`` but keep -the gRPC dependency chain lazy so default installs remain import-safe. Callers -only reach this module after choosing a ``grpc://`` Agent Stub URL. At that -point the public failure contract is: - -- missing optional gRPC/protobuf dependencies raise - ``AgentStubMissingGRPCDependencyError``; -- non-OK server statuses raise ``AgentStubGRPCError`` with the surfaced gRPC - status name and detail text; -- transport/runtime failures such as bad targets, terminated streams, or socket - errors raise ``AgentStubClientError``. - -Authentication follows gRPC conventions rather than HTTP headers: the compact -JWE bearer token is sent as ``authorization: Bearer `` metadata on each -unary RPC. -""" - -from __future__ import annotations - -import asyncio -from dataclasses import dataclass -from types import ModuleType -from typing import TYPE_CHECKING - -import httpx -from pydantic import JsonValue - -from dify_agent.agent_stub.client._errors import ( - AgentStubClientError, - AgentStubGRPCError, - AgentStubMissingGRPCDependencyError, -) -from dify_agent.agent_stub.protocol.agent_stub import AgentStubFileMapping, parse_agent_stub_endpoint - -if TYPE_CHECKING: - from grpclib.client import Channel - from grpclib.exceptions import GRPCError, StreamTerminatedError - - from dify_agent.agent_stub.grpc._generated.agent_stub_grpc import AgentStubServiceStub - - -@dataclass(frozen=True, slots=True) -class _GRPCRuntime: - Channel: type[Channel] - AgentStubServiceStub: type[AgentStubServiceStub] - agent_stub_pb2: ModuleType - GRPCError: type[GRPCError] - StreamTerminatedError: type[StreamTerminatedError] - - -def connect_agent_stub_grpc_sync( - *, - url: str, - auth_jwe: str, - argv: list[str], - metadata: dict[str, JsonValue] | None = None, - timeout: float | httpx.Timeout = 30.0, -): - """Create one gRPC Agent Stub connection using the provided bearer JWE. - - Raises: - AgentStubMissingGRPCDependencyError: if the optional gRPC runtime or - generated protobuf support is not installed. - AgentStubGRPCError: if the server returns a non-OK gRPC status. - AgentStubClientError: if the target URL is invalid for gRPC or the - runtime/transport fails before a valid gRPC response is received. - """ - return asyncio.run( - _call_grpc( - url=url, - auth_jwe=auth_jwe, - method_name="Connect", - request_factory=lambda runtime: _require_conversions().proto_connect_request( - runtime.agent_stub_pb2, - argv=argv, - metadata=metadata, - ), - response_parser=lambda response: _require_conversions().connect_response_from_proto(response), - timeout=timeout, - ) - ) - - -def request_agent_stub_file_upload_grpc_sync( - *, - url: str, - auth_jwe: str, - filename: str, - mimetype: str, - timeout: float | httpx.Timeout = 30.0, -): - """Request one signed upload URL through the gRPC Agent Stub endpoint. - - The compact-JWE bearer token is sent via ``authorization`` metadata on the - unary RPC rather than an HTTP header. - - Raises: - AgentStubMissingGRPCDependencyError: if the optional gRPC runtime or - generated protobuf support is not installed. - AgentStubGRPCError: if the server returns a non-OK gRPC status. - AgentStubClientError: if the target URL is invalid for gRPC or the - runtime/transport fails before a valid gRPC response is received. - """ - return asyncio.run( - _call_grpc( - url=url, - auth_jwe=auth_jwe, - method_name="CreateFileUploadRequest", - request_factory=lambda runtime: _require_conversions().proto_file_upload_request( - runtime.agent_stub_pb2, - filename=filename, - mimetype=mimetype, - ), - response_parser=lambda response: _require_conversions().file_upload_response_from_proto(response), - timeout=timeout, - ) - ) - - -def request_agent_stub_file_download_grpc_sync( - *, - url: str, - auth_jwe: str, - file: AgentStubFileMapping, - for_external: bool = True, - timeout: float | httpx.Timeout = 30.0, -): - """Request one signed download URL through the gRPC Agent Stub endpoint. - - The compact-JWE bearer token is sent via ``authorization`` metadata on the - unary RPC rather than an HTTP header. - - Raises: - AgentStubMissingGRPCDependencyError: if the optional gRPC runtime or - generated protobuf support is not installed. - AgentStubGRPCError: if the server returns a non-OK gRPC status. - AgentStubClientError: if the target URL is invalid for gRPC or the - runtime/transport fails before a valid gRPC response is received. - """ - return asyncio.run( - _call_grpc( - url=url, - auth_jwe=auth_jwe, - method_name="CreateFileDownloadRequest", - request_factory=lambda runtime: _require_conversions().proto_file_download_request( - runtime.agent_stub_pb2, file=file, for_external=for_external - ), - response_parser=lambda response: _require_conversions().file_download_response_from_proto(response), - timeout=timeout, - ) - ) - - -async def _call_grpc[TProto, TResult]( - *, - url: str, - auth_jwe: str, - method_name: str, - request_factory, - response_parser, - timeout: float | httpx.Timeout, -) -> TResult: - """Execute one unary Agent Stub gRPC call with shared error mapping. - - This helper attaches ``authorization`` metadata for every RPC, normalizes - ``httpx.Timeout`` into one gRPC timeout value, and translates grpclib status - failures into ``AgentStubGRPCError`` while keeping lower-level transport - failures in the broader ``AgentStubClientError`` family. - """ - runtime = _require_runtime() - endpoint = parse_agent_stub_endpoint(url) - if not endpoint.is_grpc or endpoint.port is None: - raise AgentStubClientError("gRPC Agent Stub requests require a grpc://host:port URL") - - channel = runtime.Channel(host=endpoint.host, port=endpoint.port, ssl=False) - try: - stub = runtime.AgentStubServiceStub(channel) - method = getattr(stub, method_name) - response = await method( - request_factory(runtime), - metadata=(("authorization", f"Bearer {auth_jwe}"),), - timeout=_grpc_timeout_seconds(timeout), - ) - return response_parser(response) - except runtime.GRPCError as exc: - detail = getattr(exc, "message", "") or getattr(exc, "details", "") or "request failed" - status = getattr(getattr(exc, "status", None), "name", str(getattr(exc, "status", "UNKNOWN"))) - raise AgentStubGRPCError(status, detail) from exc - except runtime.StreamTerminatedError as exc: - raise AgentStubClientError(f"Agent Stub gRPC {method_name} request terminated unexpectedly") from exc - except OSError as exc: - raise AgentStubClientError(f"Agent Stub gRPC {method_name} request failed: {exc}") from exc - finally: - channel.close() - - -def _grpc_timeout_seconds(timeout: float | httpx.Timeout) -> float | None: - if isinstance(timeout, httpx.Timeout): - values = [timeout.read, timeout.connect, timeout.write, timeout.pool] - resolved = next((value for value in values if value is not None), None) - return float(resolved) if resolved is not None else None - return float(timeout) - - -def _require_runtime() -> _GRPCRuntime: - """Import grpclib and generated protobuf runtime only when gRPC is selected.""" - try: - from grpclib.client import Channel - from grpclib.exceptions import GRPCError, StreamTerminatedError - - from dify_agent.agent_stub.grpc._generated import agent_stub_pb2 - from dify_agent.agent_stub.grpc._generated.agent_stub_grpc import AgentStubServiceStub - except ImportError as exc: - raise AgentStubMissingGRPCDependencyError( - "Agent Stub gRPC transport requires the optional dify-agent[grpc] dependencies" - ) from exc - return _GRPCRuntime( - Channel=Channel, - AgentStubServiceStub=AgentStubServiceStub, - agent_stub_pb2=agent_stub_pb2, - GRPCError=GRPCError, - StreamTerminatedError=StreamTerminatedError, - ) - - -def _require_conversions(): - """Import protobuf conversion helpers lazily for HTTP-only installations.""" - try: - from dify_agent.agent_stub.grpc import conversions - except ImportError as exc: - raise AgentStubMissingGRPCDependencyError( - "Agent Stub gRPC transport requires the optional dify-agent[grpc] dependencies" - ) from exc - return conversions - - -__all__ = [ - "connect_agent_stub_grpc_sync", - "request_agent_stub_file_download_grpc_sync", - "request_agent_stub_file_upload_grpc_sync", -] diff --git a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_http.py b/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_http.py deleted file mode 100644 index 543446ad894..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/client/_agent_stub_http.py +++ /dev/null @@ -1,597 +0,0 @@ -"""Client-safe HTTP helpers for Agent Stub control-plane endpoints. - -The main ``Client`` class stays focused on run APIs. Sandbox-visible CLI -commands only need a narrow synchronous subset of the stub server contract and -must stay safe to import in default installations, so these helpers live under -``dify_agent.agent_stub.client`` rather than the standard run client package. -""" - -from __future__ import annotations - -import json -from collections.abc import Callable -from typing import BinaryIO -from typing import cast - -import httpx -from pydantic import BaseModel, JsonValue, ValidationError - -from dify_agent.agent_stub.client._errors import ( - AgentStubClientError, - AgentStubHTTPError, - AgentStubTransferError, - AgentStubValidationError, -) -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubConnectRequest, - AgentStubConnectResponse, - AgentStubConfigEnvUpdateRequest, - AgentStubConfigManifestResponse, - AgentStubConfigNoteUpdateRequest, - AgentStubConfigPushRequest, - AgentStubConfigPushResponse, - AgentStubDriveCommitRequest, - AgentStubDriveCommitResponse, - AgentStubDriveManifestResponse, - AgentStubFileDownloadRequest, - AgentStubFileDownloadResponse, - AgentStubFileMapping, - AgentStubFileUploadRequest, - AgentStubFileUploadResponse, - agent_stub_config_env_url, - agent_stub_config_file_pull_url, - agent_stub_config_manifest_url, - agent_stub_config_note_url, - agent_stub_config_push_url, - agent_stub_config_skill_inspect_url, - agent_stub_config_skill_pull_url, - agent_stub_connections_url, - agent_stub_drive_commit_url, - agent_stub_drive_manifest_url, - agent_stub_file_download_request_url, - agent_stub_file_upload_request_url, -) - - -def connect_agent_stub_http_sync( - *, - base_url: str, - auth_jwe: str, - argv: list[str], - metadata: dict[str, JsonValue] | None = None, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubConnectResponse: - """Create one HTTP Agent Stub connection using the provided bearer JWE. - - Raises: - AgentStubValidationError: if the base URL is invalid, the request DTO is - invalid, or the success response body does not match the public - Agent Stub response schema. - AgentStubHTTPError: if the server returns a non-2xx HTTP response. - AgentStubClientError: if the request times out, the transport fails, or - the response body cannot be parsed as JSON. - """ - request_model = _validate_request(argv=argv, metadata=metadata) - response = _post_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="connect", - endpoint_url_factory=agent_stub_connections_url, - request_body=request_model.model_dump_json(), - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response(response=response, response_model=AgentStubConnectResponse, label="connection") - - -def request_agent_stub_file_upload_http_sync( - *, - base_url: str, - auth_jwe: str, - filename: str, - mimetype: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubFileUploadResponse: - """Request one signed upload URL from the HTTP Agent Stub endpoint.""" - - try: - request_model = AgentStubFileUploadRequest(filename=filename, mimetype=mimetype) - except ValidationError as exc: - raise AgentStubValidationError("invalid Agent Stub file upload request") from exc - response = _post_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="file upload request", - endpoint_url_factory=agent_stub_file_upload_request_url, - request_body=request_model.model_dump_json(), - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response(response=response, response_model=AgentStubFileUploadResponse, label="file upload") - - -def request_agent_stub_file_download_http_sync( - *, - base_url: str, - auth_jwe: str, - file: AgentStubFileMapping, - for_external: bool = True, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubFileDownloadResponse: - """Request one signed download URL from the HTTP Agent Stub endpoint.""" - - try: - request_model = AgentStubFileDownloadRequest(file=file, for_external=for_external) - except ValidationError as exc: - raise AgentStubValidationError("invalid Agent Stub file download request") from exc - response = _post_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="file download request", - endpoint_url_factory=agent_stub_file_download_request_url, - request_body=request_model.model_dump_json(exclude_none=True, exclude_defaults=True), - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response( - response=response, - response_model=AgentStubFileDownloadResponse, - label="file download", - ) - - -def request_agent_stub_drive_manifest_http_sync( - *, - base_url: str, - auth_jwe: str, - prefix: str, - include_download_url: bool, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubDriveManifestResponse: - """Request one drive manifest from the HTTP Agent Stub endpoint.""" - - response = _get_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="drive manifest request", - endpoint_url_factory=agent_stub_drive_manifest_url, - params={ - "prefix": prefix, - "include_download_url": str(include_download_url).lower(), - }, - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response( - response=response, response_model=AgentStubDriveManifestResponse, label="drive manifest" - ) - - -def request_agent_stub_drive_commit_http_sync( - *, - base_url: str, - auth_jwe: str, - request: AgentStubDriveCommitRequest, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubDriveCommitResponse: - """Commit one drive batch through the HTTP Agent Stub endpoint.""" - - response = _post_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="drive commit request", - endpoint_url_factory=agent_stub_drive_commit_url, - request_body=_dump_drive_commit_request(request), - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response(response=response, response_model=AgentStubDriveCommitResponse, label="drive commit") - - -def request_agent_stub_config_manifest_http_sync( - *, - base_url: str, - auth_jwe: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubConfigManifestResponse: - """Fetch the current config manifest from the HTTP Agent Stub endpoint.""" - response = _get_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config manifest request", - endpoint_url_factory=agent_stub_config_manifest_url, - params={}, - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response( - response=response, - response_model=AgentStubConfigManifestResponse, - label="config manifest", - ) - - -def request_agent_stub_config_skill_pull_http_sync( - *, - base_url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> bytes: - """Download one config skill archive from the HTTP Agent Stub endpoint.""" - response = _get_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config skill pull request", - endpoint_url_factory=lambda resolved_base_url: agent_stub_config_skill_pull_url(resolved_base_url, name), - params={}, - timeout=timeout, - sync_http_client=sync_http_client, - ) - if response.is_error: - payload = _parse_json_payload( - response, invalid_json_message="Agent Stub returned invalid JSON for config skill pull" - ) - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - return response.content - - -def request_agent_stub_config_skill_inspect_http_sync( - *, - base_url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> dict[str, object]: - """Fetch the JSON inspect view for one config skill.""" - response = _get_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config skill inspect request", - endpoint_url_factory=lambda resolved_base_url: agent_stub_config_skill_inspect_url(resolved_base_url, name), - params={}, - timeout=timeout, - sync_http_client=sync_http_client, - ) - payload = _parse_json_payload( - response, invalid_json_message="Agent Stub returned invalid JSON for config skill inspect" - ) - if response.is_error: - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - if not isinstance(payload, dict): - raise AgentStubValidationError("invalid Agent Stub config skill inspect response") - return cast(dict[str, object], payload) - - -def request_agent_stub_config_file_pull_http_sync( - *, - base_url: str, - auth_jwe: str, - name: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> bytes: - """Download one config file payload from the HTTP Agent Stub endpoint.""" - response = _get_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config file pull request", - endpoint_url_factory=lambda resolved_base_url: agent_stub_config_file_pull_url(resolved_base_url, name), - params={}, - timeout=timeout, - sync_http_client=sync_http_client, - ) - if response.is_error: - payload = _parse_json_payload( - response, invalid_json_message="Agent Stub returned invalid JSON for config file pull" - ) - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - return response.content - - -def request_agent_stub_config_push_http_sync( - *, - base_url: str, - auth_jwe: str, - request: AgentStubConfigPushRequest, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> AgentStubConfigPushResponse: - """Push config file/skill/env/note mutations to the HTTP Agent Stub endpoint.""" - response = _post_agent_stub_json( - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config push request", - endpoint_url_factory=agent_stub_config_push_url, - request_body=request.model_dump_json(exclude_none=True, exclude_defaults=True), - timeout=timeout, - sync_http_client=sync_http_client, - ) - return _parse_success_response(response=response, response_model=AgentStubConfigPushResponse, label="config push") - - -def _dump_drive_commit_request(request: AgentStubDriveCommitRequest) -> str: - payload = cast(dict[str, object], request.model_dump(mode="json", exclude_none=True)) - items = payload.get("items") - if isinstance(items, list): - for item in items: - if isinstance(item, dict) and item.get("is_skill") is False: - item.pop("is_skill") - return json.dumps(payload, separators=(",", ":")) - - -def request_agent_stub_config_env_update_http_sync( - *, - base_url: str, - auth_jwe: str, - env_text: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> dict[str, object]: - """Set or delete config env entries through the HTTP Agent Stub endpoint.""" - request_model = AgentStubConfigEnvUpdateRequest(env_text=env_text) - response = _send_agent_stub_json( - method="PATCH", - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config env update request", - endpoint_url_factory=agent_stub_config_env_url, - request_body=request_model.model_dump_json(), - timeout=timeout, - sync_http_client=sync_http_client, - ) - payload = _parse_json_payload( - response, invalid_json_message="Agent Stub returned invalid JSON for config env update" - ) - if response.is_error: - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - if not isinstance(payload, dict): - raise AgentStubValidationError("invalid Agent Stub config env update response") - return cast(dict[str, object], payload) - - -def request_agent_stub_config_note_update_http_sync( - *, - base_url: str, - auth_jwe: str, - note: str, - timeout: float | httpx.Timeout = 30.0, - sync_http_client: httpx.Client | None = None, -) -> dict[str, object]: - """Update the config note text through the HTTP Agent Stub endpoint.""" - request_model = AgentStubConfigNoteUpdateRequest(note=note) - response = _send_agent_stub_json( - method="PUT", - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name="config note update request", - endpoint_url_factory=agent_stub_config_note_url, - request_body=request_model.model_dump_json(), - timeout=timeout, - sync_http_client=sync_http_client, - ) - payload = _parse_json_payload( - response, invalid_json_message="Agent Stub returned invalid JSON for config note update" - ) - if response.is_error: - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - if not isinstance(payload, dict): - raise AgentStubValidationError("invalid Agent Stub config note update response") - return cast(dict[str, object], payload) - - -def upload_file_to_signed_url_sync( - *, - upload_url: str, - filename: str, - file_obj: BinaryIO, - mimetype: str, - timeout: float | httpx.Timeout = 120.0, - sync_http_client: httpx.Client | None = None, -) -> dict[str, object]: - """Upload one local file directly to a signed Dify API data-plane URL.""" - - owns_client = sync_http_client is None - client = sync_http_client or httpx.Client(timeout=timeout, follow_redirects=True) - try: - response = client.post( - upload_url, - files={"file": (filename, file_obj, mimetype)}, - timeout=timeout, - ) - except httpx.TimeoutException as exc: - raise AgentStubTransferError("signed file upload timed out") from exc - except httpx.RequestError as exc: - raise AgentStubTransferError(f"signed file upload failed: {exc}") from exc - finally: - if owns_client: - client.close() - - payload = _parse_json_payload(response, invalid_json_message="signed file upload returned invalid JSON") - if response.is_error: - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - if not isinstance(payload, dict): - raise AgentStubValidationError("invalid signed file upload response") - return cast(dict[str, object], payload) - - -def download_file_bytes_from_signed_url_sync( - *, - download_url: str, - timeout: float | httpx.Timeout = 120.0, - sync_http_client: httpx.Client | None = None, -) -> bytes: - """Download one file directly from a signed Dify API data-plane URL.""" - - owns_client = sync_http_client is None - client = sync_http_client or httpx.Client(timeout=timeout, follow_redirects=True) - try: - response = client.get(download_url, timeout=timeout) - except httpx.TimeoutException as exc: - raise AgentStubTransferError("signed file download timed out") from exc - except httpx.RequestError as exc: - raise AgentStubTransferError(f"signed file download failed: {exc}") from exc - finally: - if owns_client: - client.close() - - if response.is_error: - payload = _parse_json_payload(response, invalid_json_message="signed file download returned invalid JSON") - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - return response.content - - -def _validate_request(*, argv: list[str], metadata: dict[str, JsonValue] | None) -> AgentStubConnectRequest: - try: - return AgentStubConnectRequest(argv=argv, metadata=cast(dict[str, JsonValue], metadata or {})) - except ValidationError as exc: - raise AgentStubValidationError("invalid Agent Stub connection request") from exc - - -def _post_agent_stub_json( - *, - base_url: str, - auth_jwe: str, - endpoint_name: str, - endpoint_url_factory: Callable[[str], str], - request_body: str, - timeout: float | httpx.Timeout, - sync_http_client: httpx.Client | None, -) -> httpx.Response: - return _send_agent_stub_json( - method="POST", - base_url=base_url, - auth_jwe=auth_jwe, - endpoint_name=endpoint_name, - endpoint_url_factory=endpoint_url_factory, - request_body=request_body, - timeout=timeout, - sync_http_client=sync_http_client, - ) - - -def _send_agent_stub_json( - *, - method: str, - base_url: str, - auth_jwe: str, - endpoint_name: str, - endpoint_url_factory: Callable[[str], str], - request_body: str, - timeout: float | httpx.Timeout, - sync_http_client: httpx.Client | None, -) -> httpx.Response: - try: - endpoint_url = endpoint_url_factory(base_url) - except ValueError as exc: - raise AgentStubValidationError("invalid Agent Stub base URL") from exc - owns_client = sync_http_client is None - client = sync_http_client or httpx.Client(timeout=timeout, follow_redirects=True) - try: - return client.request( - method, - endpoint_url, - content=request_body, - headers={ - "Authorization": f"Bearer {auth_jwe}", - "Content-Type": "application/json", - }, - timeout=timeout, - ) - except httpx.TimeoutException as exc: - raise AgentStubClientError(f"Agent Stub {endpoint_name} timed out") from exc - except httpx.RequestError as exc: - raise AgentStubClientError(f"Agent Stub {endpoint_name} request failed: {exc}") from exc - finally: - if owns_client: - client.close() - - -def _get_agent_stub_json( - *, - base_url: str, - auth_jwe: str, - endpoint_name: str, - endpoint_url_factory: Callable[[str], str], - params: dict[str, str], - timeout: float | httpx.Timeout, - sync_http_client: httpx.Client | None, -) -> httpx.Response: - try: - endpoint_url = endpoint_url_factory(base_url) - except ValueError as exc: - raise AgentStubValidationError("invalid Agent Stub base URL") from exc - owns_client = sync_http_client is None - client = sync_http_client or httpx.Client(timeout=timeout, follow_redirects=True) - try: - return client.get( - endpoint_url, - params=params, - headers={"Authorization": f"Bearer {auth_jwe}"}, - timeout=timeout, - ) - except httpx.TimeoutException as exc: - raise AgentStubClientError(f"Agent Stub {endpoint_name} timed out") from exc - except httpx.RequestError as exc: - raise AgentStubClientError(f"Agent Stub {endpoint_name} request failed: {exc}") from exc - finally: - if owns_client: - client.close() - - -def _parse_success_response[T: BaseModel]( - *, - response: httpx.Response, - response_model: type[T], - label: str, -) -> T: - payload = _parse_json_payload(response, invalid_json_message=f"Agent Stub returned invalid JSON for {label}") - - if response.is_error: - detail = payload.get("detail", payload) if isinstance(payload, dict) else payload - raise AgentStubHTTPError(response.status_code, detail) - - try: - return response_model.model_validate(payload) - except ValidationError as exc: - raise AgentStubValidationError(f"invalid Agent Stub {label} response") from exc - - -def _parse_json_payload(response: httpx.Response, *, invalid_json_message: str) -> object: - try: - return response.json() - except ValueError as exc: - raise AgentStubClientError(invalid_json_message) from exc - - -__all__ = [ - "connect_agent_stub_http_sync", - "download_file_bytes_from_signed_url_sync", - "request_agent_stub_config_env_update_http_sync", - "request_agent_stub_config_file_pull_http_sync", - "request_agent_stub_config_manifest_http_sync", - "request_agent_stub_config_note_update_http_sync", - "request_agent_stub_config_push_http_sync", - "request_agent_stub_config_skill_inspect_http_sync", - "request_agent_stub_config_skill_pull_http_sync", - "request_agent_stub_drive_commit_http_sync", - "request_agent_stub_drive_manifest_http_sync", - "request_agent_stub_file_download_http_sync", - "request_agent_stub_file_upload_http_sync", - "upload_file_to_signed_url_sync", -] diff --git a/dify-agent/src/dify_agent/agent_stub/client/_errors.py b/dify-agent/src/dify_agent/agent_stub/client/_errors.py deleted file mode 100644 index b495728ee9c..00000000000 --- a/dify-agent/src/dify_agent/agent_stub/client/_errors.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Shared client-safe exception types for Agent Stub transports.""" - -from __future__ import annotations - - -class AgentStubClientError(RuntimeError): - """Base class for client-safe Agent Stub control-plane failures.""" - - -class AgentStubHTTPError(AgentStubClientError): - """Raised when the HTTP transport returns a non-success response.""" - - status_code: int - detail: object - - def __init__(self, status_code: int, detail: object) -> None: - self.status_code = status_code - self.detail = detail - super().__init__(f"Agent Stub HTTP {status_code}: {detail}") - - -class AgentStubGRPCError(AgentStubClientError): - """Raised when the gRPC transport returns a non-OK status.""" - - status: str - detail: str - - def __init__(self, status: str, detail: str) -> None: - self.status = status - self.detail = detail - super().__init__(f"Agent Stub gRPC {status}: {detail}") - - -class AgentStubValidationError(AgentStubClientError): - """Raised when request or response DTO validation fails.""" - - -class AgentStubTransferError(AgentStubClientError): - """Raised when a signed upload/download data-plane request fails.""" - - -class AgentStubMissingGRPCDependencyError(AgentStubClientError): - """Raised when the optional gRPC extra is required but unavailable.""" - - -__all__ = [ - "AgentStubClientError", - "AgentStubGRPCError", - "AgentStubHTTPError", - "AgentStubMissingGRPCDependencyError", - "AgentStubTransferError", - "AgentStubValidationError", -] diff --git a/dify-agent/src/dify_agent/layers/_agent_cli_help.json b/dify-agent/src/dify_agent/layers/_agent_cli_help.json new file mode 100644 index 00000000000..1361d1a4170 --- /dev/null +++ b/dify-agent/src/dify_agent/layers/_agent_cli_help.json @@ -0,0 +1,25 @@ +{ + "config": "Inspect or update Agent Soul-backed config assets through the Agent Stub.\n\nUsage:\n dify-agent config [command]\n\nAvailable Commands:\n env Update config env variables visible to the current run.\n files Pull or update config files through the Agent Stub.\n manifest Show the current visible Agent config manifest as JSON.\n note Pull or update the current config note.\n skills Pull or update config skills through the Agent Stub.\n\nFlags:\n -h, --help help for config\n\nUse \"dify-agent config [command] --help\" for more information about a command.", + "config env": "Update config env variables visible to the current run.\n\nUsage:\n dify-agent config env [command]\n\nAvailable Commands:\n push Set or delete config env entries from one local dotenv file or stdin.\n\nFlags:\n -h, --help help for env\n\nUse \"dify-agent config env [command] --help\" for more information about a command.", + "config env push": "Set or delete config env entries from one local dotenv file or stdin.\n\nUsage:\n dify-agent config env push PATH|- [flags]\n\nFlags:\n -h, --help help for push", + "config files": "Pull or update config files through the Agent Stub.\n\nUsage:\n dify-agent config files [command]\n\nAvailable Commands:\n delete Delete one or more config files by name without touching local files.\n pull Pull one or all visible config files into ./.dify_conf/files by default.\n push Upload one or more local files into the current config manifest.\n\nFlags:\n -h, --help help for files\n\nUse \"dify-agent config files [command] --help\" for more information about a command.", + "config files delete": "Delete one or more config files by name without touching local files.\n\nUsage:\n dify-agent config files delete NAME... [flags]\n\nFlags:\n -h, --help help for delete", + "config files pull": "Pull one or all visible config files into ./.dify_conf/files by default.\n\nUsage:\n dify-agent config files pull [NAME]... [flags]\n\nFlags:\n -h, --help help for pull\n --json Emit the pull result as JSON.\n --to string Local directory for pulled config files.", + "config files push": "Upload one or more local files into the current config manifest.\n\nUsage:\n dify-agent config files push PATH... [flags]\n\nFlags:\n -h, --help help for push", + "config manifest": "Show the current visible Agent config manifest as JSON.\n\nUsage:\n dify-agent config manifest [flags]\n\nFlags:\n -h, --help help for manifest", + "config note": "Pull or update the current config note.\n\nUsage:\n dify-agent config note [command]\n\nAvailable Commands:\n pull Export the current config note into ./.dify_conf/note.md by default.\n push Replace the current config note from one local text file or stdin.\n\nFlags:\n -h, --help help for note\n\nUse \"dify-agent config note [command] --help\" for more information about a command.", + "config note pull": "Export the current config note into ./.dify_conf/note.md by default.\n\nUsage:\n dify-agent config note pull [flags]\n\nFlags:\n -h, --help help for pull\n --to string Local markdown file path.", + "config note push": "Replace the current config note from one local text file or stdin.\n\nUsage:\n dify-agent config note push [PATH|-] [flags]\n\nFlags:\n -h, --help help for push", + "config skills": "Pull or update config skills through the Agent Stub.\n\nUsage:\n dify-agent config skills [command]\n\nAvailable Commands:\n delete Delete one or more config skills by name without touching local directories.\n pull Pull one or all visible config skills into ./.dify_conf/skills by default.\n push Upload one or more local skill directories into the current config manifest.\n\nFlags:\n -h, --help help for skills\n\nUse \"dify-agent config skills [command] --help\" for more information about a command.", + "config skills delete": "Delete one or more config skills by name without touching local directories.\n\nUsage:\n dify-agent config skills delete NAME... [flags]\n\nFlags:\n -h, --help help for delete", + "config skills pull": "Pull one or all visible config skills into ./.dify_conf/skills by default.\n\nUsage:\n dify-agent config skills pull [NAME]... [flags]\n\nFlags:\n -h, --help help for pull\n --json Emit the pull result as JSON.\n --to string Local directory for pulled config skills.", + "config skills push": "Upload one or more local skill directories into the current config manifest.\n\nUsage:\n dify-agent config skills push PATH... [flags]\n\nFlags:\n -h, --help help for push", + "connect": "Establish one Agent Stub connection using the current environment.\n\nUsage:\n dify-agent connect [ARGV]... [flags]\n\nFlags:\n -h, --help help for connect\n --json Emit the connection response as JSON.", + "drive": "List, pull, or push agent drive files through the Agent Stub.\n\nUsage:\n dify-agent drive [command]\n\nAvailable Commands:\n list List drive files visible to the current sandbox execution.\n pull Pull one or more drive keys/prefixes into one local directory tree.\n push Upload one local file or directory into the agent drive.\n\nFlags:\n -h, --help help for drive\n\nUse \"dify-agent drive [command] --help\" for more information about a command.", + "drive list": "List drive files visible to the current sandbox execution.\n\nUsage:\n dify-agent drive list [REMOTE_PREFIX] [flags]\n\nFlags:\n -h, --help help for list\n --json Emit the drive manifest as JSON.", + "drive pull": "Pull one or more drive keys/prefixes into one local directory tree.\n\nUsage:\n dify-agent drive pull [REMOTE]... [flags]\n\nFlags:\n -h, --help help for pull\n --json Emit the pull result as JSON.\n --to string Local base directory for pulled drive files.", + "drive push": "Upload one local file or directory into the agent drive.\n\nUsage:\n dify-agent drive push LOCAL_PATH REMOTE_PATH [flags]\n\nFlags:\n -h, --help help for push\n --json Accepted for consistency; drive push output is already emitted as JSON.\n --kind string Directory upload kind: skill or dir.", + "file": "Upload or download workflow files through the Agent Stub.\n\nUsage:\n dify-agent file [command]\n\nAvailable Commands:\n download Download one workflow file mapping into the local sandbox directory.\n upload Upload one sandbox-local file as a ToolFile output reference.\n\nFlags:\n -h, --help help for file\n\nUse \"dify-agent file [command] --help\" for more information about a command.", + "file download": "Download one workflow file mapping into the local sandbox directory.\n\nUsage:\n dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [flags]\n\nFlags:\n -h, --help help for download\n --to string Local directory for the downloaded file.", + "file upload": "Upload one sandbox-local file as a ToolFile output reference.\n\nUsage:\n dify-agent file upload PATH [flags]\n\nFlags:\n -h, --help help for upload" +} diff --git a/dify-agent/src/dify_agent/layers/_agent_cli_help.py b/dify-agent/src/dify_agent/layers/_agent_cli_help.py new file mode 100644 index 00000000000..db7deb6884e --- /dev/null +++ b/dify-agent/src/dify_agent/layers/_agent_cli_help.py @@ -0,0 +1,28 @@ +"""Loader for the generated ``dify-agent`` CLI help snapshot.""" + +from __future__ import annotations + +import json +from functools import cache +from importlib import resources + +_HELP_RESOURCE = "_agent_cli_help.json" + + +@cache +def _help_table() -> dict[str, str]: + """Load and cache the command-path -> help-text map from the JSON snapshot.""" + raw = resources.files(__package__).joinpath(_HELP_RESOURCE).read_text(encoding="utf-8") + return json.loads(raw) + + +def render_agent_stub_cli_help(args: tuple[str, ...]) -> str: + """Return the stored ``dify-agent --help`` output for one command path.""" + key = " ".join(args) + try: + return _help_table()[key] + except KeyError as exc: + raise ValueError(f"unknown dify-agent command path: {key}") from exc + + +__all__ = ["render_agent_stub_cli_help"] diff --git a/dify-agent/src/dify_agent/layers/config/layer.py b/dify-agent/src/dify_agent/layers/config/layer.py index 6a4a5dea4db..7088179996f 100644 --- a/dify-agent/src/dify_agent/layers/config/layer.py +++ b/dify-agent/src/dify_agent/layers/config/layer.py @@ -10,7 +10,7 @@ from typing import ClassVar from typing_extensions import Self, override from agenton.layers import LayerDeps, PlainLayer -from dify_agent.agent_stub.cli.main import render_agent_stub_cli_help +from dify_agent.layers._agent_cli_help import render_agent_stub_cli_help from dify_agent.layers._agent_file_cli_help import AGENT_FILE_UPLOAD_REPLY_HINT as _AGENT_FILE_UPLOAD_REPLY_HINT from dify_agent.layers.config.configs import ( DIFY_CONFIG_LAYER_TYPE_ID, diff --git a/dify-agent/src/shellctl/cli.py b/dify-agent/src/shellctl/cli.py deleted file mode 100644 index 4dac2e8e9c2..00000000000 --- a/dify-agent/src/shellctl/cli.py +++ /dev/null @@ -1,587 +0,0 @@ -"""Typer CLI for network-backed shellctl commands. - -Job-management commands in this module intentionally stay on the SDK side of -the boundary: they parse CLI options, call `ShellctlClient`, and render compact -JSON. That keeps `shellctl --help` and `shellctl run --help` free of FastAPI, -SQLAlchemy, tmux, and local runtime bootstrap imports. - -Only `serve` lazily imports server-side modules when that subcommand is -actually invoked. -""" - -from __future__ import annotations - -import json -from collections.abc import Awaitable, Callable -from pathlib import Path -from typing import NoReturn - -import anyio -import httpx2 as httpx -import typer -from pydantic import BaseModel, ValidationError - -from shellctl.client import ShellctlClient, ShellctlClientError -from shellctl.shared.constants import ( - DEFAULT_AUTH_TOKEN_ENV, - DEFAULT_BASE_URL, - DEFAULT_BASE_URL_ENV, - DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS, - DEFAULT_GC_INTERVAL_SECONDS, - DEFAULT_IDLE_FLUSH_SECONDS, - DEFAULT_LIST_LIMIT, - DEFAULT_OUTPUT_LIMIT_BYTES, - DEFAULT_TERMINAL_COLS, - DEFAULT_TERMINAL_ROWS, - DEFAULT_TERMINATE_GRACE_SECONDS, - DEFAULT_TIMEOUT_SECONDS, - MAX_LIST_LIMIT, - MAX_OUTPUT_LIMIT_BYTES, -) -from shellctl.shared.schemas import ( - DeleteJobResponse, - HealthResponse, - JobInfo, - JobResult, - JobStatusName, - JobStatusView, - RunJobRequest, - TerminalSize, -) - -cli = typer.Typer( - no_args_is_help=True, - pretty_exceptions_enable=False, - rich_markup_mode=None, -) - - -@cli.command("health") -def health_command( - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help="Accepted for CLI consistency but ignored because /healthz is public.", - ), -) -> None: - """Call the public health endpoint and report JSON.""" - - del auth_token - - async def action(client: ShellctlClient) -> HealthResponse: - return await client.health() - - _run_client_action( - base_url=base_url, - auth_token=None, - action=action, - emit=_emit_model, - ) - - -@cli.command("run") -def run_command( - script: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - cwd: Path | None = typer.Option(None, "--cwd"), - env: list[str] | None = typer.Option(None, "--env"), - timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"), - output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"), - idle_flush_seconds: float = typer.Option( - DEFAULT_IDLE_FLUSH_SECONDS, - "--idle-flush-seconds", - ), - cols: int | None = typer.Option(None, "--cols"), - rows: int | None = typer.Option(None, "--rows"), -) -> None: - """Create a job through the running shellctl server.""" - - request = _build_model( - RunJobRequest, - script=script, - cwd=str(cwd) if cwd is not None else None, - env=_parse_env(env), - terminal=_terminal_size(cols=cols, rows=rows), - timeout=timeout, - output_limit=output_limit, - idle_flush_seconds=idle_flush_seconds, - ) - - async def action(client: ShellctlClient) -> JobResult: - return await client.run( - request.script, - cwd=request.cwd, - env=request.env, - timeout=request.timeout, - terminal=request.terminal, - ) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - output_limit=output_limit, - idle_flush_seconds=idle_flush_seconds, - action=action, - emit=_emit_model, - ) - - -@cli.command("wait") -def wait_command( - job_id: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - offset: int = typer.Option(..., "--offset"), - timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"), - output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"), - idle_flush_seconds: float = typer.Option( - DEFAULT_IDLE_FLUSH_SECONDS, - "--idle-flush-seconds", - ), -) -> None: - """Wait for incremental output, completion, truncation, or timeout.""" - - async def action(client: ShellctlClient) -> JobResult: - return await client.wait(job_id, offset=offset, timeout=timeout) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - output_limit=output_limit, - idle_flush_seconds=idle_flush_seconds, - action=action, - emit=_emit_model, - ) - - -@cli.command("status") -def status_command( - job_id: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), -) -> None: - """Materialize the current status view for one job.""" - - async def action(client: ShellctlClient) -> JobStatusView: - return await client.status(job_id) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - action=action, - emit=_emit_model, - ) - - -@cli.command("list") -def list_command( - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - status: JobStatusName | None = typer.Option(None, "--status"), - limit: int = typer.Option(DEFAULT_LIST_LIMIT, "--limit", min=1, max=MAX_LIST_LIMIT), -) -> None: - """List recent jobs, optionally filtered by lifecycle status.""" - - async def action(client: ShellctlClient) -> list[JobInfo]: - return await client.list_jobs(status=status, limit=limit) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - action=action, - emit=_emit_job_list, - ) - - -@cli.command("input") -def input_command( - job_id: str = typer.Argument(...), - text: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - offset: int = typer.Option(..., "--offset"), - timeout: float = typer.Option(DEFAULT_TIMEOUT_SECONDS, "--timeout"), - output_limit: int = typer.Option(DEFAULT_OUTPUT_LIMIT_BYTES, "--output-limit"), - idle_flush_seconds: float = typer.Option( - DEFAULT_IDLE_FLUSH_SECONDS, - "--idle-flush-seconds", - ), -) -> None: - """Send text input to a running job and wait for the next result window.""" - - async def action(client: ShellctlClient) -> JobResult: - return await client.input(job_id, text, offset=offset, timeout=timeout) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - output_limit=output_limit, - idle_flush_seconds=idle_flush_seconds, - action=action, - emit=_emit_model, - ) - - -@cli.command("tail") -def tail_command( - job_id: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - output_limit: int = typer.Option( - DEFAULT_OUTPUT_LIMIT_BYTES, - "--output-limit", - min=1, - max=MAX_OUTPUT_LIMIT_BYTES, - ), -) -> None: - """Read a UTF-8-safe output tail for one job.""" - - async def action(client: ShellctlClient) -> JobResult: - return await client.tail(job_id) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - output_limit=output_limit, - action=action, - emit=_emit_model, - ) - - -@cli.command("terminate") -def terminate_command( - job_id: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - grace_seconds: float = typer.Option( - DEFAULT_TERMINATE_GRACE_SECONDS, - "--grace-seconds", - ), -) -> None: - """Terminate a job and return its materialized status.""" - - async def action(client: ShellctlClient) -> JobStatusView: - return await client.terminate(job_id, grace_seconds=grace_seconds) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - action=action, - emit=_emit_model, - ) - - -@cli.command("delete") -def delete_command( - job_id: str = typer.Argument(...), - base_url: str = typer.Option( - DEFAULT_BASE_URL, - "--base-url", - envvar=DEFAULT_BASE_URL_ENV, - help="shellctl server base URL. You can also set SHELLCTL_BASE_URL.", - ), - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty when the server does not require auth." - ), - ), - force: bool = typer.Option(False, "--force"), - grace_seconds: float = typer.Option( - DEFAULT_TERMINATE_GRACE_SECONDS, - "--grace-seconds", - ), -) -> None: - """Delete a job row and artifacts, optionally terminating first.""" - - async def action(client: ShellctlClient) -> DeleteJobResponse: - return await client.delete( - job_id, - force=force, - grace_seconds=grace_seconds, - ) - - _run_client_action( - base_url=base_url, - auth_token=auth_token, - action=action, - emit=_emit_model, - ) - - -@cli.command("serve") -def serve_command( - listen: str = "127.0.0.1:8765", - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty to disable HTTP bearer auth." - ), - ), - state_dir: Path | None = None, - runtime_dir: Path | None = None, - gc_interval_seconds: float = typer.Option( - DEFAULT_GC_INTERVAL_SECONDS, - "--gc-interval-seconds", - ), - gc_finished_job_retention_seconds: float = typer.Option( - DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS, - "--gc-finished-job-retention-seconds", - ), -) -> None: - """Run the local shellctl FastAPI server via uvicorn.""" - - from shellctl.server.serve import ( - serve_command as server_serve_command, - ) - - server_serve_command( - listen=listen, - auth_token=auth_token, - state_dir=state_dir, - runtime_dir=runtime_dir, - gc_interval_seconds=gc_interval_seconds, - gc_finished_job_retention_seconds=gc_finished_job_retention_seconds, - ) - - -def main() -> None: - """CLI entrypoint used by the console script and `python -m` invocations.""" - - cli() - - -def _parse_env(values: list[str] | None) -> dict[str, str] | None: - if not values: - return None - - parsed: dict[str, str] = {} - for value in values: - if "=" not in value: - raise typer.BadParameter( - "env entries must use NAME=VALUE format", - param_hint="--env", - ) - name, env_value = value.split("=", 1) - if not name: - raise typer.BadParameter( - "env names must be non-empty", - param_hint="--env", - ) - parsed[name] = env_value - return parsed - - -def _terminal_size(*, cols: int | None, rows: int | None) -> TerminalSize | None: - if cols is None and rows is None: - return None - return _build_model( - TerminalSize, - cols=cols if cols is not None else DEFAULT_TERMINAL_COLS, - rows=rows if rows is not None else DEFAULT_TERMINAL_ROWS, - ) - - -def _build_model[ModelT: BaseModel](model_type: type[ModelT], /, **data: object) -> ModelT: - try: - return model_type(**data) - except ValidationError as exc: - raise typer.BadParameter(_validation_error_message(exc)) from exc - - -async def _with_client[ResponseT]( - base_url: str, - auth_token: str | None, - output_limit: int, - idle_flush_seconds: float, - action: Callable[[ShellctlClient], Awaitable[ResponseT]], -) -> ResponseT: - async with ShellctlClient( - base_url, - output_limit=output_limit, - idle_flush_seconds=idle_flush_seconds, - token=auth_token, - ) as client: - return await action(client) - - -def _run_client_action[ResponseT]( - *, - base_url: str, - auth_token: str | None, - action: Callable[[ShellctlClient], Awaitable[ResponseT]], - emit: Callable[[ResponseT], None], - output_limit: int = DEFAULT_OUTPUT_LIMIT_BYTES, - idle_flush_seconds: float = DEFAULT_IDLE_FLUSH_SECONDS, -) -> None: - try: - payload = anyio.run( - _with_client, - base_url, - auth_token, - output_limit, - idle_flush_seconds, - action, - ) - except ShellctlClientError as exc: - _emit_error_and_exit(exc.code, exc.message) - except httpx.TimeoutException: - _emit_error_and_exit("request_timeout", "request timed out") - except httpx.TransportError as exc: - _emit_error_and_exit("connection_error", str(exc)) - - emit(payload) - - -def _emit_model(model: BaseModel) -> None: - typer.echo(model.model_dump_json(exclude_none=True), color=False) - - -def _emit_job_list(jobs: list[JobInfo]) -> None: - typer.echo( - json.dumps( - [item.model_dump(mode="json", exclude_none=True) for item in jobs], - separators=(",", ":"), - ), - color=False, - ) - - -def _emit_error_and_exit(code: str, message: str) -> NoReturn: - typer.echo( - json.dumps( - {"error": {"code": code, "message": message}}, - separators=(",", ":"), - ), - err=True, - color=False, - ) - raise typer.Exit(code=1) - - -def _validation_error_message(exc: ValidationError) -> str: - detail = exc.errors(include_url=False)[0] - location = ".".join(str(part) for part in detail.get("loc", ())) - message = detail["msg"] - return f"{location}: {message}" if location else str(message) - - -__all__ = [ - "cli", - "delete_command", - "health_command", - "input_command", - "list_command", - "main", - "run_command", - "serve_command", - "status_command", - "tail_command", - "terminate_command", - "wait_command", -] diff --git a/dify-agent/src/shellctl/server/__init__.py b/dify-agent/src/shellctl/server/__init__.py deleted file mode 100644 index fcb034f3e30..00000000000 --- a/dify-agent/src/shellctl/server/__init__.py +++ /dev/null @@ -1,58 +0,0 @@ -"""shellctl server package. - -The server stack sits behind lazy exports so importing the network CLI does not -pull in FastAPI, SQLAlchemy, tmux, or the local service runtime unless a -server-side symbol is actually used. -""" - -from __future__ import annotations - -from importlib import import_module -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from shellctl.server.api import create_app - from shellctl.server.cli import ( - cli, - main, - ) - from shellctl.server.config import ShellctlConfig - from shellctl.server.db import JobRow - from shellctl.server.errors import ShellctlServerError - from shellctl.server.serve import serve_command - from shellctl.server.service import ShellctlService - -__all__ = [ - "JobRow", - "ShellctlConfig", - "ShellctlServerError", - "ShellctlService", - "cli", - "create_app", - "main", - "serve_command", -] - -_EXPORTS = { - "JobRow": "shellctl.server.db", - "ShellctlConfig": "shellctl.server.config", - "ShellctlServerError": "shellctl.server.errors", - "ShellctlService": "shellctl.server.service", - "cli": "shellctl.server.cli", - "create_app": "shellctl.server.api", - "main": "shellctl.server.cli", - "serve_command": "shellctl.server.serve", -} - - -def __getattr__(name: str) -> Any: - if name not in _EXPORTS: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - module = import_module(_EXPORTS[name]) - value = getattr(module, name) # noqa: no-new-getattr lazy export proxy - globals()[name] = value - return value - - -def __dir__() -> list[str]: - return sorted(set(globals()) | set(__all__)) diff --git a/dify-agent/src/shellctl/server/__main__.py b/dify-agent/src/shellctl/server/__main__.py deleted file mode 100644 index b5925643528..00000000000 --- a/dify-agent/src/shellctl/server/__main__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Support `python -m shellctl.server`.""" - -from shellctl.cli import main - -if __name__ == "__main__": - main() diff --git a/dify-agent/src/shellctl/server/api.py b/dify-agent/src/shellctl/server/api.py deleted file mode 100644 index 07299e2db63..00000000000 --- a/dify-agent/src/shellctl/server/api.py +++ /dev/null @@ -1,198 +0,0 @@ -"""FastAPI wiring for shellctl server endpoints.""" - -from __future__ import annotations - -from contextlib import asynccontextmanager -from typing import Annotated, cast - -from fastapi import Depends, FastAPI, Header, Query, Request -from fastapi.responses import JSONResponse - -from shellctl.server.config import ShellctlConfig -from shellctl.server.errors import ShellctlServerError -from shellctl.server.service import ShellctlService -from shellctl.shared.constants import ( - DEFAULT_HEALTH_STATUS, - DEFAULT_LIST_LIMIT, - DEFAULT_OUTPUT_LIMIT_BYTES, - DEFAULT_TERMINATE_GRACE_SECONDS, - MAX_LIST_LIMIT, - MAX_OUTPUT_LIMIT_BYTES, -) -from shellctl.shared.schemas import ( - DeleteJobResponse, - ErrorDetail, - ErrorResponse, - HealthResponse, - InputJobRequest, - JobResult, - JobStatusName, - JobStatusView, - ListJobsResponse, - RunJobRequest, - TerminateJobRequest, - WaitJobRequest, -) - - -def create_app( - config: ShellctlConfig | None = None, - *, - service: ShellctlService | None = None, -) -> FastAPI: - """Create the FastAPI application used by `shellctl serve`.""" - - resolved_config = config or ShellctlConfig() - resolved_service = service or ShellctlService(resolved_config) - - @asynccontextmanager - async def lifespan(_app: FastAPI): - await resolved_service.initialize() - resolved_service.start_background_gc() - resolved_service.start_background_pipe_monitor() - try: - yield - finally: - await resolved_service.shutdown() - - app = FastAPI(title="shellctl", version="0.1.0", lifespan=lifespan) - app.state.shellctl_service = resolved_service - - @app.exception_handler(ShellctlServerError) - async def handle_shellctl_error( - _request: Request, - exc: ShellctlServerError, - ) -> JSONResponse: - return JSONResponse( - status_code=exc.status_code, - content=ErrorResponse(error=ErrorDetail(code=exc.code, message=exc.message)).model_dump(mode="json"), - ) - - @app.exception_handler(RuntimeError) - async def handle_runtime_error(_request: Request, exc: RuntimeError) -> JSONResponse: - return JSONResponse( - status_code=500, - content=ErrorResponse( - error=ErrorDetail( - code="internal_error", - message=str(exc) or "internal server error", - ) - ).model_dump(mode="json"), - ) - - def get_service() -> ShellctlService: - return cast(ShellctlService, app.state.shellctl_service) - - def verify_auth( - authorization: Annotated[str | None, Header()] = None, - ) -> None: - token = resolved_config.auth_token - if token is None: - return - expected = f"Bearer {token}" - if authorization != expected: - raise ShellctlServerError(401, "unauthorized", "Missing or invalid bearer token") - - @app.get("/healthz", response_model=HealthResponse) - async def healthz() -> HealthResponse: - return HealthResponse(status=DEFAULT_HEALTH_STATUS) - - @app.post( - "/v1/jobs/run", - response_model=JobResult, - dependencies=[Depends(verify_auth)], - ) - async def run_job( - payload: RunJobRequest, - svc: ShellctlService = Depends(get_service), - ) -> JobResult: - return await svc.run_job(payload) - - @app.post( - "/v1/jobs/{job_id}/wait", - response_model=JobResult, - dependencies=[Depends(verify_auth)], - ) - async def wait_job( - job_id: str, - payload: WaitJobRequest, - svc: ShellctlService = Depends(get_service), - ) -> JobResult: - return await svc.wait_job(job_id, payload) - - @app.get( - "/v1/jobs/{job_id}/log/tail", - response_model=JobResult, - dependencies=[Depends(verify_auth)], - ) - async def tail_job( - job_id: str, - output_limit: Annotated[int, Query(ge=1, le=MAX_OUTPUT_LIMIT_BYTES)] = DEFAULT_OUTPUT_LIMIT_BYTES, - svc: ShellctlService = Depends(get_service), - ) -> JobResult: - return await svc.tail_job(job_id, output_limit=output_limit) - - @app.get( - "/v1/jobs/{job_id}", - response_model=JobStatusView, - dependencies=[Depends(verify_auth)], - ) - async def job_status( - job_id: str, - svc: ShellctlService = Depends(get_service), - ) -> JobStatusView: - return await svc.get_job_status(job_id) - - @app.get( - "/v1/jobs", - response_model=ListJobsResponse, - dependencies=[Depends(verify_auth)], - ) - async def list_jobs( - status: Annotated[JobStatusName | None, Query()] = None, - limit: Annotated[int, Query(ge=1, le=MAX_LIST_LIMIT)] = DEFAULT_LIST_LIMIT, - svc: ShellctlService = Depends(get_service), - ) -> ListJobsResponse: - return await svc.list_jobs(status=status, limit=limit) - - @app.post( - "/v1/jobs/{job_id}/input", - response_model=JobResult, - dependencies=[Depends(verify_auth)], - ) - async def input_job( - job_id: str, - payload: InputJobRequest, - svc: ShellctlService = Depends(get_service), - ) -> JobResult: - return await svc.send_input(job_id, payload) - - @app.post( - "/v1/jobs/{job_id}/terminate", - response_model=JobStatusView, - dependencies=[Depends(verify_auth)], - ) - async def terminate_job( - job_id: str, - payload: TerminateJobRequest, - svc: ShellctlService = Depends(get_service), - ) -> JobStatusView: - return await svc.terminate_job(job_id, payload) - - @app.delete( - "/v1/jobs/{job_id}", - response_model=DeleteJobResponse, - dependencies=[Depends(verify_auth)], - ) - async def delete_job( - job_id: str, - force: bool = False, - grace_seconds: float = DEFAULT_TERMINATE_GRACE_SECONDS, - svc: ShellctlService = Depends(get_service), - ) -> DeleteJobResponse: - return await svc.delete_job(job_id, force=force, grace_seconds=grace_seconds) - - return app - - -__all__ = ["create_app"] diff --git a/dify-agent/src/shellctl/server/artifacts.py b/dify-agent/src/shellctl/server/artifacts.py deleted file mode 100644 index ca013f4a36a..00000000000 --- a/dify-agent/src/shellctl/server/artifacts.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Per-job artifact names used by shellctl server/runtime code. - -Normal job completion is coordinated through small marker files inside each -`jobs//` directory so the tmux output-pipe finalizer can publish the -SQLite `exited(exit_code, ended_at)` state only after PTY output is fully -drained into `output.log`. The same artifact directory also stores the request's -environment overlay so the runner can merge arbitrary key/value pairs without -shell-escaping them into the generated script. Separate failure markers and a -dedicated `pipe-error.log` stderr capture keep startup diagnostics available -when the sanitizer never reaches its ready-file handshake. -""" - -from __future__ import annotations - -from pathlib import Path - -RUNNER_EXIT_CODE_FILENAME = ".runner-exit-code" -RUNNER_ENDED_AT_FILENAME = ".runner-ended-at" -JOB_ENV_FILENAME = ".job-env.json" -PIPE_DRAINED_FILENAME = ".pipe-drained" -PIPE_FAILED_FILENAME = ".pipe-failed" -PIPE_ERROR_LOG_FILENAME = "pipe-error.log" - - -def runner_exit_code_path(job_dir: Path) -> Path: - return job_dir / RUNNER_EXIT_CODE_FILENAME - - -def runner_ended_at_path(job_dir: Path) -> Path: - return job_dir / RUNNER_ENDED_AT_FILENAME - - -def job_env_path(job_dir: Path) -> Path: - return job_dir / JOB_ENV_FILENAME - - -def pipe_drained_path(job_dir: Path) -> Path: - return job_dir / PIPE_DRAINED_FILENAME - - -def pipe_failed_path(job_dir: Path) -> Path: - return job_dir / PIPE_FAILED_FILENAME - - -def pipe_error_log_path(job_dir: Path) -> Path: - return job_dir / PIPE_ERROR_LOG_FILENAME - - -__all__ = [ - "JOB_ENV_FILENAME", - "PIPE_DRAINED_FILENAME", - "PIPE_ERROR_LOG_FILENAME", - "PIPE_FAILED_FILENAME", - "RUNNER_ENDED_AT_FILENAME", - "RUNNER_EXIT_CODE_FILENAME", - "job_env_path", - "pipe_drained_path", - "pipe_error_log_path", - "pipe_failed_path", - "runner_ended_at_path", - "runner_exit_code_path", -] diff --git a/dify-agent/src/shellctl/server/cli.py b/dify-agent/src/shellctl/server/cli.py deleted file mode 100644 index 306981fb696..00000000000 --- a/dify-agent/src/shellctl/server/cli.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Compatibility shim for historical `shellctl.server.cli` imports. - -Job-management commands now live in `shellctl.cli`. This -module is intentionally a thin legacy re-export for callers that still import -CLI symbols from `shellctl.server.cli`. -""" - -from __future__ import annotations - -from shellctl.cli import cli, main -from shellctl.server.serve import serve_command - -__all__ = [ - "cli", - "main", - "serve_command", -] diff --git a/dify-agent/src/shellctl/server/config.py b/dify-agent/src/shellctl/server/config.py deleted file mode 100644 index 979d61c59b8..00000000000 --- a/dify-agent/src/shellctl/server/config.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Configuration objects for shellctl server/runtime modules.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from pathlib import Path -from typing import cast - -from shellctl.shared.constants import ( - DEFAULT_AUTH_TOKEN_ENV, - DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS, - DEFAULT_GC_INTERVAL_SECONDS, - DEFAULT_IDLE_FLUSH_SECONDS, - DEFAULT_LIST_LIMIT, - DEFAULT_OUTPUT_LIMIT_BYTES, - DEFAULT_TERMINAL_COLS, - DEFAULT_TERMINAL_ROWS, - DEFAULT_TERMINATE_GRACE_SECONDS, - DEFAULT_TIMEOUT_SECONDS, - MAX_LIST_LIMIT, - MAX_OUTPUT_LIMIT_BYTES, - MAX_WAIT_TIMEOUT_SECONDS, -) -from shellctl.shared.runtime import ( - default_runtime_dir, - default_state_dir, -) -from shellctl_runtime.paths import DEFAULT_SQLITE_BUSY_TIMEOUT_MS - - -@dataclass(slots=True, frozen=True) -class ShellctlConfig: - """Runtime configuration for the shellctl service and CLI. - - The tmux subprocess hooks use dedicated console scripts instead of - `python -m shellctl...` entrypoints so every job does not pay - the shellctl client/server import cost just to sanitize PTY bytes or record - an exit row. `sqlite_busy_timeout_ms` still applies to the out-of-process - `runner-exit` callback, so non-default deployments keep one SQLite timeout - policy across the service and tmux finalizer. - - Bearer auth is opt-in: if the explicit `auth_token` and the fallback - `SHELLCTL_AUTH_TOKEN` environment variable are both missing or empty, - `shellctl serve` accepts requests without checking an Authorization header. - """ - - listen: str = "127.0.0.1:8765" - auth_token: str | None = None - state_dir: Path = field(default_factory=default_state_dir) - runtime_dir: Path | None = None - gc_interval_seconds: float = DEFAULT_GC_INTERVAL_SECONDS - gc_finished_job_retention_seconds: float = DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS - default_timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS - max_wait_timeout_seconds: float = MAX_WAIT_TIMEOUT_SECONDS - idle_flush_seconds: float = DEFAULT_IDLE_FLUSH_SECONDS - default_cwd: Path = field(default_factory=Path.home) - default_terminal_cols: int = DEFAULT_TERMINAL_COLS - default_terminal_rows: int = DEFAULT_TERMINAL_ROWS - default_list_limit: int = DEFAULT_LIST_LIMIT - max_list_limit: int = MAX_LIST_LIMIT - default_output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES - max_output_limit_bytes: int = MAX_OUTPUT_LIMIT_BYTES - default_terminate_grace_seconds: float = DEFAULT_TERMINATE_GRACE_SECONDS - poll_interval_seconds: float = 0.05 - pipe_monitor_interval_seconds: float = 1.0 - pipe_ready_timeout_seconds: float = 10.0 - sqlite_busy_timeout_ms: int = DEFAULT_SQLITE_BUSY_TIMEOUT_MS - sanitize_pty_command: tuple[str, ...] = ("shellctl-sanitize-pty",) - runner_exit_command: tuple[str, ...] = ("shellctl-runner-exit",) - - def __post_init__(self) -> None: - if self.runtime_dir is None: - object.__setattr__(self, "runtime_dir", default_runtime_dir(self.state_dir)) - token = self.auth_token - if token is None: - token = os.environ.get(DEFAULT_AUTH_TOKEN_ENV) - if not token: - token = None - object.__setattr__(self, "auth_token", token) - - @property - def jobs_dir(self) -> Path: - return self.state_dir / "jobs" - - @property - def db_path(self) -> Path: - return self.state_dir / "shellctl.db" - - @property - def database_url(self) -> str: - return f"sqlite+aiosqlite:///{self.db_path}" - - @property - def tmux_socket(self) -> Path: - runtime_dir = cast(Path, self.runtime_dir) - return runtime_dir / "tmux.sock" - - @property - def runner_path(self) -> Path: - runtime_dir = cast(Path, self.runtime_dir) - return runtime_dir / "bin" / "shellctl-runner" - - -__all__ = ["ShellctlConfig"] diff --git a/dify-agent/src/shellctl/server/db.py b/dify-agent/src/shellctl/server/db.py deleted file mode 100644 index 88ac516b9fc..00000000000 --- a/dify-agent/src/shellctl/server/db.py +++ /dev/null @@ -1,52 +0,0 @@ -"""SQLite models and engine helpers for shellctl.""" - -from __future__ import annotations - -from typing import Any, cast - -from sqlalchemy import event -from sqlalchemy.ext.asyncio import AsyncEngine -from sqlmodel import Field, SQLModel - - -class JobRow(SQLModel, table=True): - """SQLite source-of-truth row for shellctl jobs. - - The table intentionally combines immutable metadata, mutable lifecycle state, - and exit facts so state transitions can be expressed as single conditional - `UPDATE` statements without synchronizing separate records. - """ - - __tablename__ = cast(Any, "jobs") - - job_id: str = Field(primary_key=True) - script_path: str - output_path: str - cwd: str - terminal_cols: int - terminal_rows: int - status: str = Field(index=True) - session_name: str - pane_target: str - exit_code: int | None = Field(default=None, nullable=True) - reason: str | None = Field(default=None, nullable=True) - message: str | None = Field(default=None, nullable=True) - created_at: str = Field(index=True) - started_at: str | None = Field(default=None, nullable=True) - ended_at: str | None = Field(default=None, nullable=True, index=True) - updated_at: str - - -def configure_sqlite_engine(engine: AsyncEngine, *, busy_timeout_ms: int) -> None: - """Install SQLite pragmas required by the proposal's concurrency model.""" - - @event.listens_for(engine.sync_engine, "connect") - def _set_pragmas(dbapi_connection: Any, _connection_record: Any) -> None: - cursor = dbapi_connection.cursor() - cursor.execute("PRAGMA journal_mode=WAL") - cursor.execute("PRAGMA foreign_keys=ON") - cursor.execute(f"PRAGMA busy_timeout={busy_timeout_ms}") - cursor.close() - - -__all__ = ["JobRow", "configure_sqlite_engine"] diff --git a/dify-agent/src/shellctl/server/errors.py b/dify-agent/src/shellctl/server/errors.py deleted file mode 100644 index 3e673860bc1..00000000000 --- a/dify-agent/src/shellctl/server/errors.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Shared exception types for shellctl server-side modules.""" - - -class ShellctlServerError(RuntimeError): - """Structured server-side error that maps directly to API responses.""" - - def __init__(self, status_code: int, code: str, message: str) -> None: - super().__init__(message) - self.status_code = status_code - self.code = code - self.message = message - - -__all__ = ["ShellctlServerError"] diff --git a/dify-agent/src/shellctl/server/serve.py b/dify-agent/src/shellctl/server/serve.py deleted file mode 100644 index de5957254d0..00000000000 --- a/dify-agent/src/shellctl/server/serve.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Server-side `shellctl serve` implementation. - -This module stays separate from the top-level CLI so ordinary client commands -do not import the FastAPI or uvicorn stack. `serve_command()` performs those -imports lazily because only the long-running server path needs them. -""" - -from __future__ import annotations - -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import typer - -from shellctl.shared.constants import ( - DEFAULT_AUTH_TOKEN_ENV, - DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS, - DEFAULT_GC_INTERVAL_SECONDS, -) -from shellctl.shared.runtime import ( - default_state_dir, -) - -if TYPE_CHECKING: - from shellctl.server.config import ShellctlConfig - - -def serve_command( - listen: str = "127.0.0.1:8765", - auth_token: str | None = typer.Option( - None, - "--auth-token", - envvar=DEFAULT_AUTH_TOKEN_ENV, - help=( - "Bearer token value. You can also set SHELLCTL_AUTH_TOKEN. " - "Leave it unset or empty to disable HTTP bearer auth." - ), - ), - state_dir: Path | None = None, - runtime_dir: Path | None = None, - gc_interval_seconds: float = DEFAULT_GC_INTERVAL_SECONDS, - gc_finished_job_retention_seconds: float = DEFAULT_GC_FINISHED_JOB_RETENTION_SECONDS, -) -> None: - """Build `ShellctlConfig` from CLI inputs and run the local HTTP server. - - Args: - listen: Host/port pair for the uvicorn listener in `host:port` form. - auth_token: Optional bearer token value. An explicit empty string - disables HTTP auth, an explicit non-empty token enables it, and an - omitted/`None` value may still resolve from `SHELLCTL_AUTH_TOKEN` - through the Typer env var binding or `ShellctlConfig` fallback. - state_dir: Persistent shellctl state directory; defaults to the shared - XDG-style state path when omitted. - runtime_dir: Optional runtime directory override for tmux/runtime - artifacts. - gc_interval_seconds: Background GC wake-up cadence for finished jobs. - gc_finished_job_retention_seconds: Retention window before finished jobs - are eligible for GC. - - This entrypoint is the only CLI path that should pull in the FastAPI and - uvicorn stack. It parses the listener, constructs `ShellctlConfig`, and - then hands the configured app to uvicorn. - """ - - from shellctl.server.config import ShellctlConfig - - host, port = _parse_listen(listen) - config = ShellctlConfig( - listen=listen, - auth_token=auth_token, - state_dir=state_dir or default_state_dir(), - runtime_dir=runtime_dir, - gc_interval_seconds=gc_interval_seconds, - gc_finished_job_retention_seconds=gc_finished_job_retention_seconds, - ) - _uvicorn_run(_create_app(config), host=host, port=port, log_level="info") - - -def _parse_listen(value: str) -> tuple[str, int]: - if ":" not in value: - raise typer.BadParameter("listen must use host:port format") - host, raw_port = value.rsplit(":", 1) - host = host.strip("[]") - try: - port = int(raw_port) - except ValueError as exc: - raise typer.BadParameter(f"invalid port: {raw_port}") from exc - return host, port - - -def _create_app(config: ShellctlConfig) -> Any: - from shellctl.server.api import create_app - - return create_app(config) - - -def _uvicorn_run(app: Any, *, host: str, port: int, log_level: str) -> None: - import uvicorn - - uvicorn.run(app, host=host, port=port, log_level=log_level) - - -__all__ = ["serve_command"] diff --git a/dify-agent/src/shellctl/server/service.py b/dify-agent/src/shellctl/server/service.py deleted file mode 100644 index 741cbb6406d..00000000000 --- a/dify-agent/src/shellctl/server/service.py +++ /dev/null @@ -1,1184 +0,0 @@ -"""SQLite-backed shellctl job service. - -This module owns the long-lived job lifecycle rules: artifact creation, -SQLite compare-and-swap transitions, tmux reconciliation, runner exit -materialization, and GC. API and local server entrypoints should call into -`ShellctlService` rather than duplicating lifecycle logic, choosing either the -lightweight `prepare_runtime()` bootstrap or the full `initialize()` startup -depending on whether they need bounded setup work or long-running maintenance. -""" - -from __future__ import annotations - -import asyncio -import json -import shlex -import shutil -import stat -import sys -from contextlib import suppress -from pathlib import Path -from typing import Any, cast - -import anyio -from sqlalchemy import case, delete, select, update -from sqlalchemy.exc import IntegrityError -from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine -from sqlmodel import SQLModel - -from shellctl.server.artifacts import ( - JOB_ENV_FILENAME, - RUNNER_ENDED_AT_FILENAME, - RUNNER_EXIT_CODE_FILENAME, - job_env_path, - pipe_drained_path, - pipe_error_log_path, - pipe_failed_path, - runner_ended_at_path, - runner_exit_code_path, -) -from shellctl.server.config import ShellctlConfig -from shellctl.server.db import JobRow, configure_sqlite_engine -from shellctl.server.errors import ShellctlServerError -from shellctl.server.tmux import ( - TmuxController, - TmuxControllerProtocol, -) -from shellctl.shared.constants import DEFAULT_AUTH_TOKEN_ENV -from shellctl.shared.output import ( - OutputWindow, - read_output_window, - tail_output_window, -) -from shellctl.shared.runtime import ( - format_timestamp, - generate_job_id, - is_terminal_status, - job_pane_target, - job_session_name, - parse_timestamp, - utc_now, -) -from shellctl.shared.schemas import ( - DeleteJobResponse, - InputJobRequest, - JobInfo, - JobResult, - JobStatusName, - JobStatusView, - ListJobsResponse, - RunJobRequest, - TerminalSize, - TerminateJobRequest, - WaitJobRequest, -) - - -class ShellctlService: - """SQLite-backed shellctl job service used by HTTP and local server entrypoints. - - The service keeps only minimal in-memory coordination state: - - - `_starting_jobs` prevents `created`/`starting` rows from being marked - `lost` while the current startup path is still installing tmux state. - - Everything else uses SQLite conditional updates so concurrent readers and - lifecycle paths converge through the DB rather than through mutable JSON - artifacts. - """ - - def __init__( - self, - config: ShellctlConfig, - *, - tmux: TmuxControllerProtocol | None = None, - ) -> None: - self.config = config - self._tmux = tmux or TmuxController(config) - self._gc_task: asyncio.Task[None] | None = None - self._monitor_task: asyncio.Task[None] | None = None - self._starting_jobs: set[str] = set() - self._engine = create_async_engine(self.config.database_url, future=True) - configure_sqlite_engine( - self._engine, - busy_timeout_ms=self.config.sqlite_busy_timeout_ms, - ) - self._session_factory = async_sessionmaker( - self._engine, - expire_on_commit=False, - class_=AsyncSession, - ) - - async def initialize_database(self) -> None: - """Ensure the SQLite database exists and its schema is initialized.""" - - self._ensure_dir(self.config.state_dir) - async with self._engine.begin() as connection: - await connection.run_sync(SQLModel.metadata.create_all) - - async def prepare_runtime(self) -> None: - """Prepare the minimal local runtime needed for bounded server-side work. - - This startup path intentionally skips reconciliation, GC, and background - tasks. Entry points that only need local runtime artifacts should stay - bounded by the requested operation instead of scanning or maintaining - the full historical job set on every invocation. - """ - - await self.initialize_database() - self._ensure_dir(cast(Path, self.config.runtime_dir)) - self._ensure_dir(self.config.jobs_dir) - self._ensure_dir(self.config.runner_path.parent) - self._install_runner() - await self._tmux.start_server() - - async def initialize(self) -> None: - """Run the full server startup path used by `shellctl serve`. - - Auth is configured at the API layer, so service startup must also work - when `shellctl serve` is intentionally running without a bearer token. - Unlike `prepare_runtime()`, this path also reconciles persisted jobs and - performs one foreground GC pass before the long-running server starts. - """ - - await self.prepare_runtime() - await self.reconcile() - await self.gc_once() - - async def shutdown(self) -> None: - """Stop background work and close the SQLite engine.""" - - for task_name in ("_gc_task", "_monitor_task"): - task = getattr(self, task_name) # noqa: no-new-getattr dynamic task cleanup - if task is not None: - task.cancel() - with anyio.CancelScope(shield=True): - with suppress(asyncio.CancelledError): - await task - setattr(self, task_name, None) - await self._engine.dispose() - - def start_background_gc(self) -> None: - if self._gc_task is None: - self._gc_task = asyncio.create_task(self._gc_loop(), name="shellctl-gc") - - def start_background_pipe_monitor(self) -> None: - if self._monitor_task is None: - self._monitor_task = asyncio.create_task(self._pipe_monitor_loop(), name="shellctl-pipe-monitor") - - async def run_job(self, request: RunJobRequest) -> JobResult: - """Create a tmux-backed job and wait for its initial result window. - - Side effects: - - allocates `jobs//` and writes `script`, `.job-env.json`, and - `output.log` - - inserts a `jobs` row into SQLite, then conditionally transitions it - through `created -> starting -> running` - - creates a dedicated tmux session, installs `pipe-pane`, waits for the - sanitize/output pipeline ready-file handshake, and only then opens the - `start-gate` - - on startup failure, marks the job `failed` in SQLite and cleans up the - tmux session without deleting the job artifacts - - Returns: - The same `JobResult` shape as `wait()`, starting from byte offset 0. - - Notes: - The startup reservation in `_starting_jobs` exists so concurrent - status/list/reconcile calls can observe a fresh row as a legitimate - `created`/`starting` job instead of materializing it to `lost` before - tmux becomes visible. - """ - - cwd = self._resolve_cwd(request.cwd) - env = self._resolve_env(request.env) - terminal = request.terminal or TerminalSize( - cols=self.config.default_terminal_cols, - rows=self.config.default_terminal_rows, - ) - created_at = format_timestamp() - - job_id: str | None = None - job_dir: Path | None = None - while True: - candidate_job_id, candidate_job_dir = self._allocate_job_dir() - try: - self._starting_jobs.add(candidate_job_id) - script_path = candidate_job_dir / "script" - output_path = candidate_job_dir / "output.log" - script_path.write_text(request.script, encoding="utf-8") - job_env_path(candidate_job_dir).write_text( - json.dumps(env, ensure_ascii=False), - encoding="utf-8", - ) - output_path.touch() - row = JobRow( - job_id=candidate_job_id, - script_path=f"jobs/{candidate_job_id}/script", - output_path=f"jobs/{candidate_job_id}/output.log", - cwd=str(cwd), - terminal_cols=terminal.cols, - terminal_rows=terminal.rows, - status=JobStatusName.CREATED.value, - session_name=job_session_name(candidate_job_id), - pane_target=job_pane_target(candidate_job_id), - exit_code=None, - reason=None, - message=None, - created_at=created_at, - started_at=None, - ended_at=None, - updated_at=created_at, - ) - if await self._insert_job_row(row): - job_id = candidate_job_id - job_dir = candidate_job_dir - break - self._starting_jobs.discard(candidate_job_id) - except BaseException: - self._starting_jobs.discard(candidate_job_id) - shutil.rmtree(candidate_job_dir, ignore_errors=True) - raise - shutil.rmtree(candidate_job_dir, ignore_errors=True) - - assert job_id is not None and job_dir is not None - await self._transition_status( - job_id, - allowed_from={JobStatusName.CREATED}, - target=JobStatusName.STARTING, - ) - - try: - pipe_ready_path = job_dir / ".pipe-ready" - await self._tmux.create_job_session( - job_id=job_id, - job_dir=job_dir, - cwd=cwd, - terminal=terminal, - ) - await self._tmux.enable_output_pipe( - job_id=job_id, - job_dir=job_dir, - ready_file=pipe_ready_path, - ) - await self._wait_for_output_pipe_ready( - job_id=job_id, - ready_file=pipe_ready_path, - ) - (job_dir / "start-gate").touch() - await self._transition_status( - job_id, - allowed_from={JobStatusName.STARTING}, - target=JobStatusName.RUNNING, - require_exit_code_null=True, - ) - except Exception as exc: - await self._transition_status( - job_id, - allowed_from={ - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - }, - target=JobStatusName.FAILED, - reason=getattr(exc, "code", "start_failed"), # noqa: no-new-getattr exception code fallback - message=str(exc), - ) - await self._tmux.cleanup_session(job_id=job_id) - finally: - (job_dir / ".pipe-ready").unlink(missing_ok=True) - self._starting_jobs.discard(job_id) - - return await self.wait_job( - job_id, - WaitJobRequest( - offset=0, - timeout=request.timeout, - output_limit=request.output_limit, - idle_flush_seconds=request.idle_flush_seconds, - ), - ) - - async def wait_job(self, job_id: str, request: WaitJobRequest) -> JobResult: - """Block until output, completion, truncation, or timeout.""" - - row = await self._get_job_row(job_id) - output_path = self._output_log_path(row) - self._validate_offset(output_path, request.offset) - deadline = anyio.current_time() + request.timeout - last_size = output_path.stat().st_size if output_path.exists() else 0 - saw_output = last_size > request.offset - last_growth_at: float | None = anyio.current_time() if saw_output else None - - while True: - view = await self.get_job_status(job_id) - current_size = output_path.stat().st_size if output_path.exists() else 0 - if request.offset > current_size: - raise ShellctlServerError( - 400, - "invalid_offset", - f"offset {request.offset} exceeds current file size {current_size}", - ) - - if current_size > last_size: - last_size = current_size - if current_size > request.offset: - saw_output = True - last_growth_at = anyio.current_time() - - if view.done: - window = read_output_window( - output_path, - offset=request.offset, - limit=request.output_limit, - ) - return self._job_result_from_view(view, row, window) - - if current_size > request.offset: - window = read_output_window( - output_path, - offset=request.offset, - limit=request.output_limit, - ) - if window.truncated: - return self._job_result_from_view(view, row, window) - if saw_output and last_growth_at is not None: - idle_elapsed = anyio.current_time() - last_growth_at - if idle_elapsed >= request.idle_flush_seconds: - return self._job_result_from_view(view, row, window) - - if anyio.current_time() >= deadline: - window = ( - read_output_window( - output_path, - offset=request.offset, - limit=request.output_limit, - ) - if current_size > request.offset - else OutputWindow(output="", offset=request.offset, truncated=False) - ) - return self._job_result_from_view(view, row, window) - - await anyio.sleep(self.config.poll_interval_seconds) - - async def tail_job(self, job_id: str, *, output_limit: int) -> JobResult: - row = await self._get_job_row(job_id) - view = await self.get_job_status(job_id) - tail = tail_output_window(self._output_log_path(row), limit=output_limit) - return self._job_result_from_view(view, row, tail, truncated=False) - - async def get_job_status(self, job_id: str) -> JobStatusView: - session_exists, pipe_active = await self._live_runtime_state(job_id) - view = await self._materialize_status_view( - job_id, - session_exists=session_exists, - pipe_active=pipe_active, - ) - if view.done: - await self._tmux.cleanup_session(job_id=job_id) - return view - - async def list_jobs( - self, - *, - status: JobStatusName | None = None, - limit: int, - ) -> ListJobsResponse: - """List recent jobs ordered by `created_at` descending. - - This intentionally uses live per-job runtime queries instead of a single - tmux snapshot so concurrent startup cannot turn a healthy job into `lost`. - Rows may also disappear concurrently because delete removes the DB record - before artifact cleanup; such races are treated as benign skips rather - than failing the whole list request. - """ - - rows = await self._list_job_rows() - items: list[JobInfo] = [] - for row in rows: - try: - view = await self.get_job_status(row.job_id) - except ShellctlServerError as exc: - if exc.code == "job_not_found": - continue - raise - if status is not None and view.status != status: - continue - items.append( - JobInfo( - job_id=view.job_id, - status=view.status, - created_at=view.created_at, - started_at=view.started_at, - ended_at=view.ended_at, - ) - ) - if len(items) >= limit: - break - return ListJobsResponse(jobs=items) - - async def send_input(self, job_id: str, request: InputJobRequest) -> JobResult: - """Send input to a running job and then wait using the same semantics as `wait()`. - - If the pane disappears after the initial running-state check, this method - re-materializes the job state before deciding whether to return terminal - `409` semantics or a true internal tmux failure. - """ - - view = await self.get_job_status(job_id) - if view.done: - raise ShellctlServerError(409, "job_not_running", f"Job {job_id} is already terminal") - try: - await self._tmux.send_input(job_id=job_id, text=request.text) - except ShellctlServerError as exc: - if exc.code != "tmux_target_missing": - raise - view = await self.get_job_status(job_id) - if view.done: - raise ShellctlServerError(409, "job_not_running", f"Job {job_id} is already terminal") from exc - raise ShellctlServerError(500, "tmux_input_failed", exc.message) from exc - return await self.wait_job( - job_id, - WaitJobRequest( - offset=request.offset, - timeout=request.timeout, - output_limit=request.output_limit, - idle_flush_seconds=request.idle_flush_seconds, - ), - ) - - async def terminate_job(self, job_id: str, request: TerminateJobRequest) -> JobStatusView: - """Terminate a job, reserving `terminated` before tmux cleanup. - - The initial conditional DB update allows an `exited` row to be overridden - only when it raced after an earlier non-terminal read in this method. That - preserves proposal semantics: once a user terminates a non-terminal job, - the final API status stays `terminated` even if the runner exits during the - same window. - """ - - view = await self.get_job_status(job_id) - if view.done: - await self._tmux.cleanup_session(job_id=job_id) - return view - - await self._transition_status( - job_id, - allowed_from={ - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - }, - allow_override_from={JobStatusName.EXITED}, - target=JobStatusName.TERMINATED, - ) - await self._tmux.send_interrupt(job_id=job_id) - if request.grace_seconds > 0: - await anyio.sleep(request.grace_seconds) - await self._tmux.cleanup_session(job_id=job_id) - return await self.get_job_status(job_id) - - async def delete_job( - self, - job_id: str, - *, - force: bool, - grace_seconds: float, - ) -> DeleteJobResponse: - """Delete a job row and its artifacts. - - Behavior: - - if the job is still non-terminal and `force` is false, raises `409` - - if `force` is true, first applies normal terminate semantics - - always best-effort cleans the tmux session before removing persisted - state - - Deletion ordering matters: the SQLite row is removed before best-effort - artifact-directory cleanup. That ordering intentionally avoids exposing a - visible half-deleted state where API listing/status can still see a row - whose `script` / `output.log` artifacts are already gone. - """ - - view = await self.get_job_status(job_id) - if not view.done: - if not force: - raise ShellctlServerError(409, "job_running", f"Job {job_id} is still running") - await self.terminate_job(job_id, TerminateJobRequest(grace_seconds=grace_seconds)) - await self._tmux.cleanup_session(job_id=job_id) - await self._delete_job_row(job_id) - shutil.rmtree(self._artifact_dir(job_id), ignore_errors=True) - return DeleteJobResponse(job_id=job_id) - - async def reconcile(self) -> None: - """Reconcile SQLite rows against live tmux state using per-job probes. - - Reconciliation also uses a live per-job runtime query to avoid turning a - job that is still completing startup into `lost` based on an earlier tmux - snapshot. - Jobs may disappear between the initial row query and the later status - materialization because delete removes the SQLite row first. Reconciliation - treats that race as benign and continues with the remaining jobs. - """ - - for row in await self._list_job_rows(): - try: - view = await self.get_job_status(row.job_id) - except ShellctlServerError as exc: - if exc.code == "job_not_found": - continue - raise - if view.done: - await self._tmux.cleanup_session(job_id=row.job_id) - - async def gc_once(self) -> None: - """Delete expired terminal rows and their artifact directories. - - Like reconciliation, GC tolerates rows disappearing mid-pass due to a - concurrent explicit delete. - - Deletion ordering matches explicit delete semantics: remove the SQLite - row first, then best-effort delete the artifact directory. That keeps the - DB/API source of truth from exposing a row whose filesystem artifacts have - already been removed. - """ - - cutoff = utc_now().timestamp() - self.config.gc_finished_job_retention_seconds - for row in await self._list_job_rows(): - try: - view = await self.get_job_status(row.job_id) - except ShellctlServerError as exc: - if exc.code == "job_not_found": - continue - raise - if not view.done or view.ended_at is None: - continue - if parse_timestamp(view.ended_at).timestamp() >= cutoff: - continue - await self._tmux.cleanup_session(job_id=row.job_id) - try: - await self._delete_job_row(row.job_id) - except ShellctlServerError as exc: - if exc.code == "job_not_found": - continue - raise - shutil.rmtree(self._artifact_dir(row.job_id), ignore_errors=True) - - async def record_runner_exit(self, job_id: str, exit_code: int, ended_at: str) -> None: - """Persist a drained normal-exit fact into SQLite. - - The usual caller is the tmux pipe finalizer after the lightweight PTY - sanitizer reaches EOF and flushes `output.log`. `_materialize_status_view()` - may also call this as a recovery path when `.pipe-drained` plus the - normal exit metadata files exist but SQLite is still non-terminal. The - statement always records `exit_code` and `ended_at`, but it only changes - `status` to `exited` when the current row is still non-terminal. - """ - - nonterminal = [ - JobStatusName.CREATED.value, - JobStatusName.STARTING.value, - JobStatusName.RUNNING.value, - ] - job_id_col = cast(Any, JobRow.job_id) - status_col = cast(Any, JobRow.status) - ended_at_col = cast(Any, JobRow.ended_at) - reason_col = cast(Any, JobRow.reason) - message_col = cast(Any, JobRow.message) - async with self._session_factory() as session: - stmt = ( - update(JobRow) - .where(job_id_col == job_id) - .values( - exit_code=exit_code, - ended_at=case( - (ended_at_col.is_(None), ended_at), - else_=ended_at_col, - ), - updated_at=ended_at, - status=case( - (status_col.in_(nonterminal), JobStatusName.EXITED.value), - else_=status_col, - ), - reason=case( - (status_col.in_(nonterminal), None), - else_=reason_col, - ), - message=case( - (status_col.in_(nonterminal), None), - else_=message_col, - ), - ) - ) - result = await session.execute(stmt) - if cast(Any, result).rowcount == 0: - raise ShellctlServerError(404, "job_not_found", f"Unknown job id: {job_id}") - await session.commit() - - async def check_running_jobs_pipe_health(self) -> None: - """Fail running jobs whose tmux output pipe disappeared after startup.""" - - for row in await self._list_job_rows(statuses={JobStatusName.RUNNING}): - session_exists, pipe_active = await self._live_runtime_state(row.job_id) - view = await self._materialize_status_view( - row.job_id, - session_exists=session_exists, - pipe_active=pipe_active, - ) - if view.done: - await self._tmux.cleanup_session(job_id=row.job_id) - - async def _gc_loop(self) -> None: - while True: - await anyio.sleep(self.config.gc_interval_seconds) - await self.gc_once() - - async def _pipe_monitor_loop(self) -> None: - while True: - await anyio.sleep(self.config.pipe_monitor_interval_seconds) - await self.check_running_jobs_pipe_health() - - async def _wait_for_output_pipe_ready(self, *, job_id: str, ready_file: Path) -> None: - """Confirm the sanitize/output pipeline is live before opening start-gate. - - Timeout failures include the ready-file path, current `#{pane_pipe}` - state, and a summary of `pipe-error.log` so operators can tell the - difference between slow startup, tmux pipe loss, and sanitizer crashes. - """ - - started_at = anyio.current_time() - deadline = started_at + self.config.pipe_ready_timeout_seconds - while True: - waited_seconds = anyio.current_time() - started_at - if ready_file.exists(): - pipe_state = await self._tmux.is_output_pipe_active(job_id=job_id) - if pipe_state is True: - return - if pipe_state is None: - raise ShellctlServerError( - 500, - "pipe_failed", - self._pipe_ready_failure_message( - job_id=job_id, - ready_file=ready_file, - waited_seconds=waited_seconds, - pipe_state=pipe_state, - cause="tmux pane disappeared after the ready-file handshake", - ), - ) - else: - if not await self._tmux.session_exists(job_session_name(job_id)): - pipe_state = await self._tmux.is_output_pipe_active(job_id=job_id) - raise ShellctlServerError( - 500, - "pipe_failed", - self._pipe_ready_failure_message( - job_id=job_id, - ready_file=ready_file, - waited_seconds=waited_seconds, - pipe_state=pipe_state, - cause="tmux session exited before the ready-file handshake", - ), - ) - if anyio.current_time() >= deadline: - pipe_state = await self._tmux.is_output_pipe_active(job_id=job_id) - raise ShellctlServerError( - 500, - "pipe_failed", - self._pipe_ready_failure_message( - job_id=job_id, - ready_file=ready_file, - waited_seconds=waited_seconds, - pipe_state=pipe_state, - cause="timed out waiting for the sanitize/output pipeline handshake", - ), - ) - await anyio.sleep(self.config.poll_interval_seconds) - - def _pipe_ready_failure_message( - self, - *, - job_id: str, - ready_file: Path, - waited_seconds: float, - pipe_state: bool | None, - cause: str, - ) -> str: - """Build a startup error message with enough pipe diagnostics to debug.""" - - return ( - f"Output pipe never became ready for {job_id}: {cause}; " - f"waited {waited_seconds:.3f}s; " - f"ready-file={ready_file}; " - f"tmux #{{pane_pipe}}={self._format_pane_pipe_state(pipe_state)}; " - f"pipe-error.log={self._pipe_error_log_summary(ready_file.parent)}" - ) - - def _format_pane_pipe_state(self, pipe_state: bool | None) -> str: - if pipe_state is True: - return "1" - if pipe_state is False: - return "0" - return "pane-missing" - - def _pipe_error_log_summary(self, job_dir: Path) -> str: - """Summarize sanitizer stderr without flooding API error responses.""" - - error_log = pipe_error_log_path(job_dir) - if not error_log.exists(): - return f"{error_log} missing" - stderr_text = error_log.read_text(encoding="utf-8", errors="replace").strip() - if not stderr_text: - return f"{error_log} empty" - summary = " | ".join(line.strip() for line in stderr_text.splitlines() if line.strip()) - if not summary: - return f"{error_log} contains only whitespace" - if len(summary) > 240: - summary = f"{summary[:237]}..." - return f"{error_log}: {summary}" - - async def _materialize_status_view( - self, - job_id: str, - *, - session_exists: bool, - pipe_active: bool | None, - ) -> JobStatusView: - """Materialize the current API status view from SQLite + live tmux state. - - Priority/invariants: - - terminal SQLite status wins and is preserved - - if SQLite already has an `exit_code` on a non-terminal row, materialize - `exited` - - if normal-exit metadata and `.pipe-drained` already exist, recover the - drained normal exit immediately even when the finalizer has not yet - committed SQLite state - - if normal-exit metadata exists and neither `.pipe-drained` nor - `.pipe-failed` exists, keep the non-terminal row instead of guessing - `lost` while the pipe finalizer is still responsible for the eventual - `runner-exit` commit - - if a live tmux session exists but the output pipe is known-dead, - conditionally materialize `failed(reason=pipe_failed)` - - if no live session exists and the row is not protected by the local - startup reservation, conditionally materialize `lost` - - This helper is the main reconciliation point between live tmux probes and - persisted SQLite state. It must never revive terminal rows and must never - convert an ambiguous startup race into `lost` while the current `run_job` - path still owns that startup window. - """ - - row = await self._get_job_row(job_id) - status = JobStatusName(row.status) - - if is_terminal_status(status): - if row.ended_at is None: - row = await self._transition_status( - job_id, - allowed_from={status}, - target=status, - ended_at=format_timestamp(), - ) - elif row.exit_code is not None: - row = await self._transition_status( - job_id, - allowed_from={ - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - }, - target=JobStatusName.EXITED, - ended_at=row.ended_at or format_timestamp(), - ) - elif (drained_exit := self._drained_normal_exit_metadata(job_id)) is not None: - exit_code, ended_at = drained_exit - await self.record_runner_exit(job_id, exit_code, ended_at) - row = await self._get_job_row(job_id) - elif session_exists: - if pipe_active is False: - if ( - status - in { - JobStatusName.CREATED, - JobStatusName.STARTING, - } - and job_id in self._starting_jobs - ): - pass - else: - row = await self._transition_status( - job_id, - allowed_from={ - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - }, - target=JobStatusName.FAILED, - reason="pipe_failed", - message="The tmux output pipe stopped while the job was still running.", - ended_at=format_timestamp(), - ) - elif status in {JobStatusName.CREATED, JobStatusName.STARTING} and job_id not in self._starting_jobs: - row = await self._transition_status( - job_id, - allowed_from={JobStatusName.CREATED, JobStatusName.STARTING}, - target=JobStatusName.RUNNING, - require_exit_code_null=True, - ) - else: - if self._normal_exit_commit_pending(job_id): - pass - elif not (status in {JobStatusName.CREATED, JobStatusName.STARTING} and job_id in self._starting_jobs): - row = await self._transition_status( - job_id, - allowed_from={ - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - }, - target=JobStatusName.LOST, - reason="tmux_session_missing", - message="The dedicated tmux session is no longer present.", - ended_at=format_timestamp(), - ) - - return self._status_view_from_row(row) - - async def _transition_status( - self, - job_id: str, - *, - allowed_from: set[JobStatusName], - target: JobStatusName, - allow_override_from: set[JobStatusName] | None = None, - require_exit_code_null: bool = False, - reason: str | None = None, - message: str | None = None, - ended_at: str | None = None, - ) -> JobRow: - """Conditionally transition a job row through a SQLite CAS update. - - This is the SQLite replacement for the old JSON+lock compare-and-swap - flow. The method performs a single conditional `UPDATE ... WHERE ...` - against the current row and returns the post-update materialized row. - - Invariants: - - non-terminal states may only advance from explicitly allowed source - states - - callers can require `exit_code IS NULL` so late runner exits do not let - startup paths overwrite a concurrently materialized terminal outcome - - terminal transitions preserve first-writer semantics via the WHERE - clause and fill `ended_at` only when needed - - when the conditional update affects zero rows, the method returns the - current persisted row instead of forcing a stale overwrite - """ - - now = format_timestamp() - allowed_statuses = {status.value for status in allowed_from} - if allow_override_from is not None: - allowed_statuses.update(status.value for status in allow_override_from) - - values: dict[str, Any] = { - "status": target.value, - "updated_at": now, - "reason": reason, - "message": message, - } - status_col = cast(Any, JobRow.status) - job_id_col = cast(Any, JobRow.job_id) - started_at_col = cast(Any, JobRow.started_at) - ended_at_col = cast(Any, JobRow.ended_at) - exit_code_col = cast(Any, JobRow.exit_code) - if target in {JobStatusName.STARTING, JobStatusName.RUNNING}: - values["started_at"] = case( - (started_at_col.is_(None), now), - else_=started_at_col, - ) - if is_terminal_status(target): - values["ended_at"] = case( - (ended_at_col.is_(None), ended_at or now), - else_=ended_at_col, - ) - - async with self._session_factory() as session: - stmt = update(JobRow).where(job_id_col == job_id, status_col.in_(allowed_statuses)) - if require_exit_code_null: - stmt = stmt.where(exit_code_col.is_(None)) - result = await session.execute(stmt.values(**values)) - await session.commit() - if cast(Any, result).rowcount == 0: - return await self._get_job_row(job_id) - return await self._get_job_row(job_id) - - async def _live_runtime_state(self, job_id: str) -> tuple[bool, bool | None]: - """Read tmux session/pipe liveness for one job. - - If the pane disappears between the session probe and the pipe probe, this - helper degrades to `(False, None)` so reconciliation treats the job as a - missing-session case rather than a `pipe_failed` case. - """ - - session_exists = await self._tmux.session_exists(job_session_name(job_id)) - if not session_exists: - return False, None - pipe_active = await self._tmux.is_output_pipe_active(job_id=job_id) - if pipe_active is None: - return False, None - return True, pipe_active - - async def _insert_job_row(self, row: JobRow) -> bool: - async with self._session_factory() as session: - session.add(row) - try: - await session.commit() - except IntegrityError: - await session.rollback() - return False - return True - - async def _delete_job_row(self, job_id: str) -> None: - async with self._session_factory() as session: - result = await session.execute(delete(JobRow).where(cast(Any, JobRow.job_id) == job_id)) - if cast(Any, result).rowcount == 0: - raise ShellctlServerError(404, "job_not_found", f"Unknown job id: {job_id}") - await session.commit() - - async def _get_job_row(self, job_id: str) -> JobRow: - async with self._session_factory() as session: - row = await session.get(JobRow, job_id) - if row is None: - raise ShellctlServerError(404, "job_not_found", f"Unknown job id: {job_id}") - return row - - async def _list_job_rows(self, *, statuses: set[JobStatusName] | None = None) -> list[JobRow]: - created_at_col = cast(Any, JobRow.created_at) - status_col = cast(Any, JobRow.status) - async with self._session_factory() as session: - stmt = select(JobRow).order_by(created_at_col.desc()) - if statuses is not None: - stmt = stmt.where(status_col.in_([status.value for status in statuses])) - rows = (await session.execute(stmt)).scalars().all() - return list(rows) - - def _status_view_from_row(self, row: JobRow) -> JobStatusView: - status = JobStatusName(row.status) - output_path = self._output_log_path(row) - offset = output_path.stat().st_size if output_path.exists() else 0 - return JobStatusView( - job_id=row.job_id, - status=status, - done=is_terminal_status(status), - exit_code=row.exit_code, - created_at=row.created_at, - started_at=row.started_at, - ended_at=row.ended_at, - offset=offset, - ) - - def _job_result_from_view( - self, - view: JobStatusView, - row: JobRow, - window: OutputWindow, - *, - truncated: bool | None = None, - ) -> JobResult: - return JobResult( - job_id=view.job_id, - done=view.done, - status=view.status, - exit_code=view.exit_code, - output_path=str(self._output_log_path(row).resolve()), - output=window.output, - offset=window.offset, - truncated=window.truncated if truncated is None else truncated, - ) - - def _output_log_path(self, row: JobRow) -> Path: - return self.config.state_dir / row.output_path - - def _artifact_dir(self, job_id: str) -> Path: - return self.config.jobs_dir / job_id - - def _normal_exit_commit_pending(self, job_id: str) -> bool: - """Check whether a normal exit is still waiting for pipe drain + commit.""" - - job_dir = self._artifact_dir(job_id) - return ( - runner_exit_code_path(job_dir).exists() - and runner_ended_at_path(job_dir).exists() - and not pipe_drained_path(job_dir).exists() - and not pipe_failed_path(job_dir).exists() - ) - - def _drained_normal_exit_metadata(self, job_id: str) -> tuple[int, str] | None: - """Read drained normal-exit metadata after successful sanitize drain.""" - - job_dir = self._artifact_dir(job_id) - drained_path = pipe_drained_path(job_dir) - exit_code_path = runner_exit_code_path(job_dir) - ended_at_path = runner_ended_at_path(job_dir) - if not drained_path.exists() or not exit_code_path.exists() or not ended_at_path.exists(): - return None - raw_exit_code = exit_code_path.read_text(encoding="utf-8").strip() - raw_ended_at = ended_at_path.read_text(encoding="utf-8").strip() - if not raw_exit_code or not raw_ended_at: - return None - return int(raw_exit_code), raw_ended_at - - def _allocate_job_dir(self) -> tuple[str, Path]: - """Atomically allocate a unique artifact directory for a new job. - - The proposal requires `mkdir jobs/` itself to be the collision - detector. We therefore create the directory inside this helper and retry - on `FileExistsError` instead of using a separate `exists()` pre-check. - """ - - for _ in range(20): - job_id = generate_job_id() - job_dir = self.config.jobs_dir / job_id - try: - job_dir.mkdir(mode=0o700) - return job_id, job_dir - except FileExistsError: - continue - raise ShellctlServerError(500, "job_id_collision", "Failed to allocate a unique job id") - - def _resolve_cwd(self, raw_cwd: str | None) -> Path: - cwd = Path(raw_cwd).expanduser() if raw_cwd is not None else self.config.default_cwd.expanduser() - cwd = cwd.resolve() - if not cwd.exists() or not cwd.is_dir(): - raise ShellctlServerError(400, "invalid_cwd", f"cwd is not a directory: {cwd}") - return cwd - - def _resolve_env(self, raw_env: dict[str, str] | None) -> dict[str, str]: - """Return a detached environment overlay for the generated runner. - - `RunJobRequest` validates names/values before the service sees them. - The overlay augments the runner's inherited environment after shellctl - scrubs its own control variables, so explicit request values can be used - without replacing ambient entries like `PATH`. - """ - - return dict(raw_env or {}) - - def _validate_offset(self, output_path: Path, offset: int) -> None: - size = output_path.stat().st_size if output_path.exists() else 0 - if offset > size: - raise ShellctlServerError( - 400, - "invalid_offset", - f"offset {offset} exceeds current file size {size}", - ) - - @staticmethod - def _ensure_dir(path: Path) -> None: - path.mkdir(mode=0o700, parents=True, exist_ok=True) - path.chmod(0o700) - - def _install_runner(self) -> None: - self.config.runner_path.write_text(self._runner_script_source(), encoding="utf-8") - mode = self.config.runner_path.stat().st_mode - self.config.runner_path.chmod(mode | stat.S_IXUSR) - - def _runner_script_source(self) -> str: - """Build the bash runner installed into the shellctl runtime directory. - - The wrapper still handles gate waiting and exit metadata in bash, but it - delegates launch setup to Python so per-job env overlays can be loaded - from JSON without shell-escaping arbitrary values into `export` - statements. The Python helper then `exec`s the target process instead of - waiting on it, which keeps `SIGINT` behavior identical to running the - script directly in the tmux pane. The bootstrap uses `python -c ...` - instead of a stdin-fed heredoc so the exec'd job still inherits the tmux - pane's PTY on fd 0 and can receive later `send_input()` data. - """ - - auth_env = shlex.quote(DEFAULT_AUTH_TOKEN_ENV) - bootstrap_source = shlex.quote( - """ -import json -import os -import stat -import sys -from pathlib import Path - -script_path = Path(sys.argv[1]) -cwd = sys.argv[2] -env_path = Path(sys.argv[3]) - -env = os.environ.copy() -if env_path.exists(): - env.update(json.loads(env_path.read_text(encoding="utf-8"))) - -try: - os.chdir(cwd) -except OSError: - raise SystemExit(111) - -with script_path.open("r", encoding="utf-8") as handle: - first_line = handle.readline() - -if first_line.startswith("#!"): - script_path.chmod(script_path.stat().st_mode | stat.S_IXUSR) - argv = [str(script_path)] -else: - argv = ["sh", str(script_path)] - -try: - os.execvpe(argv[0], argv, env) -except FileNotFoundError as exc: - print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) - raise SystemExit(127) from exc -except OSError as exc: - print(f"{argv[0]}: {exc.strerror}", file=sys.stderr) - raise SystemExit(126) from exc - """.strip() - ) - return f"""#!/usr/bin/env bash -set -uo pipefail - -JOB_DIR="$1" -JOB_ID="$2" -CWD="$3" -SCRIPT_PATH="$JOB_DIR/script" -ENV_PATH="$JOB_DIR/{JOB_ENV_FILENAME}" -START_GATE="$JOB_DIR/start-gate" -RUNNER_EXIT_CODE_PATH="$JOB_DIR/{RUNNER_EXIT_CODE_FILENAME}" -RUNNER_ENDED_AT_PATH="$JOB_DIR/{RUNNER_ENDED_AT_FILENAME}" - -write_atomic() {{ - local dest="$1" - local value="$2" - local tmp="${{dest}}.tmp.$$" - printf '%s\n' "$value" > "$tmp" - mv "$tmp" "$dest" -}} - -while [ ! -e "$START_GATE" ]; do - sleep 0.05 -done - -unset TMUX -unset SHELLCTL_STATE_DIR -unset SHELLCTL_RUNTIME_DIR -unset SHELLCTL_TMUX_SOCKET -unset SHELLCTL_RUNNER -unset {auth_env} - -{shlex.quote(sys.executable)} -c {bootstrap_source} "$SCRIPT_PATH" "$CWD" "$ENV_PATH" -EXIT_CODE=$? - -ENDED_AT="$({shlex.quote(sys.executable)} - <<'PY' -from datetime import UTC, datetime -print(datetime.now(UTC).replace(microsecond=0).isoformat().replace('+00:00', 'Z')) -PY -)" - -write_atomic "$RUNNER_EXIT_CODE_PATH" "$EXIT_CODE" -write_atomic "$RUNNER_ENDED_AT_PATH" "$ENDED_AT" - -exit "$EXIT_CODE" -""" - - -__all__ = ["ShellctlService", "anyio", "shutil"] diff --git a/dify-agent/src/shellctl/server/tmux.py b/dify-agent/src/shellctl/server/tmux.py deleted file mode 100644 index e412f812207..00000000000 --- a/dify-agent/src/shellctl/server/tmux.py +++ /dev/null @@ -1,343 +0,0 @@ -"""tmux control layer for shellctl jobs. - -The rest of shellctl treats this module as the only place that knows tmux CLI - command shapes. Tests can replace `TmuxControllerProtocol` with a fake to - exercise SQLite/output semantics without depending on a local tmux daemon. -""" - -from __future__ import annotations - -import os -import shlex -import subprocess -import tempfile -from pathlib import Path -from typing import Protocol, cast - -import anyio - -from shellctl.server.artifacts import ( - pipe_drained_path, - pipe_error_log_path, - pipe_failed_path, - runner_ended_at_path, - runner_exit_code_path, -) -from shellctl.server.config import ShellctlConfig -from shellctl.server.errors import ShellctlServerError -from shellctl.shared.runtime import ( - job_pane_target, - job_session_name, -) -from shellctl.shared.schemas import TerminalSize - - -class TmuxControllerProtocol(Protocol): - """Protocol used by `ShellctlService` for tmux interactions.""" - - async def start_server(self) -> None: ... - - async def list_sessions(self) -> set[str]: ... - - async def session_exists(self, session_name: str) -> bool: ... - - async def is_output_pipe_active(self, *, job_id: str) -> bool | None: ... - - async def create_job_session( - self, - *, - job_id: str, - job_dir: Path, - cwd: Path, - terminal: TerminalSize, - ) -> None: ... - - async def enable_output_pipe(self, *, job_id: str, job_dir: Path, ready_file: Path) -> None: ... - - async def send_input(self, *, job_id: str, text: str) -> None: ... - - async def send_interrupt(self, *, job_id: str) -> None: ... - - async def cleanup_session(self, *, job_id: str) -> None: ... - - -class TmuxController: - """Best-effort wrapper around a dedicated tmux socket. - - The controller always clears `TMUX` from the child environment and always - passes `-S ` so shellctl sessions stay isolated from the user's - default tmux server. - """ - - def __init__(self, config: ShellctlConfig) -> None: - self._config = config - - async def start_server(self) -> None: - await self._run_tmux("start-server") - - async def list_sessions(self) -> set[str]: - result = await self._run_tmux("list-sessions", "-F", "#{session_name}", check=False) - if result.returncode != 0: - stderr = result.stderr.decode("utf-8", errors="replace") - if _tmux_target_missing(stderr): - return set() - raise ShellctlServerError(500, "tmux_error", stderr.strip() or "tmux list-sessions failed") - output = result.stdout.decode("utf-8", errors="replace") - return {line.strip() for line in output.splitlines() if line.strip()} - - async def session_exists(self, session_name: str) -> bool: - return session_name in await self.list_sessions() - - async def is_output_pipe_active(self, *, job_id: str) -> bool | None: - result = await self._run_tmux( - "display-message", - "-p", - "-t", - job_pane_target(job_id), - "#{pane_pipe}", - check=False, - ) - if result.returncode != 0: - stderr = result.stderr.decode("utf-8", errors="replace") - if _tmux_target_missing(stderr): - return None - raise ShellctlServerError( - 500, - "tmux_error", - stderr.strip() or f"Failed to inspect output pipe for {job_id}", - ) - return result.stdout.decode("utf-8", errors="replace").strip() == "1" - - async def create_job_session( - self, - *, - job_id: str, - job_dir: Path, - cwd: Path, - terminal: TerminalSize, - ) -> None: - runner_command = self._shell_join( - [ - str(self._config.runner_path), - str(job_dir), - job_id, - str(cwd), - ] - ) - result = await self._run_tmux( - "-f", - "/dev/null", - "new-session", - "-d", - "-s", - job_session_name(job_id), - "-x", - str(terminal.cols), - "-y", - str(terminal.rows), - runner_command, - check=False, - ) - if result.returncode != 0: - raise ShellctlServerError( - 500, - "tmux_new_session_failed", - result.stderr.decode("utf-8", errors="replace").strip() - or f"Failed to create tmux session for {job_id}", - ) - - async def enable_output_pipe(self, *, job_id: str, job_dir: Path, ready_file: Path) -> None: - output_command = self._pipe_command_source( - job_id=job_id, - job_dir=job_dir, - ready_file=ready_file, - ) - result = await self._run_tmux( - "pipe-pane", - "-o", - "-t", - job_pane_target(job_id), - output_command, - check=False, - ) - if result.returncode != 0: - raise ShellctlServerError( - 500, - "pipe_failed", - result.stderr.decode("utf-8", errors="replace").strip() or f"Failed to attach output pipe for {job_id}", - ) - - def _pipe_command_source(self, *, job_id: str, job_dir: Path, ready_file: Path) -> str: - """Build the tmux `pipe-pane` command that drains and finalizes output. - - For normal exits, the runner now records completion metadata into job - artifacts and the pipe finalizer commits `runner-exit` only after - the lightweight sanitizer reaches EOF and flushes `output.log` - successfully. Sanitizer stderr is captured into `pipe-error.log` so - startup timeouts can distinguish slow imports from subprocess crashes. - If the follow-up `runner-exit` write fails, the drain marker remains in - place and stderr is appended to the same log with an explicit status - line. The pipe still exits with the sanitizer status so a drained job is - not misclassified as `pipe_failed` before reconciliation can recover the - SQLite write from the drained artifacts. - """ - - sanitize_command = self._shell_join( - ( - *self._config.sanitize_pty_command, - "--ready-file", - str(ready_file), - ) - ) - runner_exit_command = self._shell_join( - ( - *self._config.runner_exit_command, - "--state-dir", - str(self._config.state_dir), - "--job-id", - job_id, - "--sqlite-busy-timeout-ms", - str(self._config.sqlite_busy_timeout_ms), - ) - ) - output_path = shlex.quote(str(job_dir / "output.log")) - drained_path = shlex.quote(str(pipe_drained_path(job_dir))) - error_log_path = shlex.quote(str(pipe_error_log_path(job_dir))) - failed_path = shlex.quote(str(pipe_failed_path(job_dir))) - exit_code_path = shlex.quote(str(runner_exit_code_path(job_dir))) - ended_at_path = shlex.quote(str(runner_ended_at_path(job_dir))) - return " ; ".join( - [ - f"{sanitize_command} >> {output_path} 2> {error_log_path}", - "sanitize_status=$?", - "runner_exit_status=0", - ( - 'if [ "$sanitize_status" -eq 0 ]; then ' - f": > {drained_path}; " - f"if [ -s {exit_code_path} ] && [ -s {ended_at_path} ]; then " - f'{runner_exit_command} --exit-code "$(cat {exit_code_path})" ' - f'--ended-at "$(cat {ended_at_path})" 2>> {error_log_path}; ' - "runner_exit_status=$?; " - 'if [ "$runner_exit_status" -ne 0 ]; then ' - f"printf 'runner-exit failed with status %s\\n' \"$runner_exit_status\" >> {error_log_path}; " - "fi; fi; " - f"else : > {failed_path}; fi" - ), - 'if [ "$sanitize_status" -ne 0 ]; then exit "$sanitize_status"; fi', - 'exit "$sanitize_status"', - ] - ) - - async def send_input(self, *, job_id: str, text: str) -> None: - buffer_name = f"shellctl-in-{job_id}" - runtime_dir = cast(Path, self._config.runtime_dir) - runtime_dir.mkdir(parents=True, exist_ok=True) - fd, tmp_name = tempfile.mkstemp(prefix=f"shellctl-input-{job_id}-", dir=runtime_dir) - tmp_path = Path(tmp_name) - try: - with os.fdopen(fd, "w", encoding="utf-8") as handle: - handle.write(text) - load_result = await self._run_tmux( - "load-buffer", - "-b", - buffer_name, - str(tmp_path), - check=False, - ) - if load_result.returncode != 0: - stderr = load_result.stderr.decode("utf-8", errors="replace").strip() - if _tmux_target_missing(stderr): - raise ShellctlServerError( - 409, - "tmux_target_missing", - stderr or f"The tmux pane for {job_id} is no longer available", - ) - raise ShellctlServerError( - 500, - "tmux_input_failed", - stderr or f"Failed to load input buffer for {job_id}", - ) - paste_result = await self._run_tmux( - "paste-buffer", - "-t", - job_pane_target(job_id), - "-b", - buffer_name, - check=False, - ) - if paste_result.returncode != 0: - stderr = paste_result.stderr.decode("utf-8", errors="replace").strip() - if _tmux_target_missing(stderr): - raise ShellctlServerError( - 409, - "tmux_target_missing", - stderr or f"The tmux pane for {job_id} is no longer available", - ) - raise ShellctlServerError( - 500, - "tmux_input_failed", - stderr or f"Failed to paste input buffer for {job_id}", - ) - finally: - await self._run_tmux("delete-buffer", "-b", buffer_name, check=False) - tmp_path.unlink(missing_ok=True) - - async def send_interrupt(self, *, job_id: str) -> None: - await self._run_tmux( - "send-keys", - "-t", - job_pane_target(job_id), - "C-c", - check=False, - ) - - async def cleanup_session(self, *, job_id: str) -> None: - await self._run_tmux( - "kill-session", - "-t", - job_session_name(job_id), - check=False, - ) - - async def _run_tmux(self, *args: str, check: bool = True) -> subprocess.CompletedProcess[bytes]: - env = dict(os.environ) - env.pop("TMUX", None) - try: - result = await anyio.run_process( - ["tmux", "-S", str(self._config.tmux_socket), *args], - env=env, - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - except FileNotFoundError as exc: - raise ShellctlServerError(500, "tmux_not_installed", "tmux executable was not found") from exc - if check and result.returncode != 0: - raise ShellctlServerError( - 500, - "tmux_error", - result.stderr.decode("utf-8", errors="replace").strip() or "tmux command failed", - ) - return result - - @staticmethod - def _shell_join(parts: tuple[str, ...] | list[str]) -> str: - return " ".join(shlex.quote(part) for part in parts) - - -def _tmux_target_missing(stderr: str) -> bool: - normalized = stderr.lower() - return ( - "can't find pane" in normalized - or "can't find session" in normalized - or "no server running" in normalized - or "failed to connect" in normalized - or "server exited unexpectedly" in normalized - ) - - -__all__ = [ - "TmuxController", - "TmuxControllerProtocol", - "_tmux_target_missing", -] diff --git a/dify-agent/src/shellctl_runtime/__init__.py b/dify-agent/src/shellctl_runtime/__init__.py deleted file mode 100644 index daa28c47ea5..00000000000 --- a/dify-agent/src/shellctl_runtime/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Stdlib-only shellctl subprocess entrypoints. - -This package is reserved for tmux hot-path helpers that must start quickly for -every job. Keep `__init__` import-free so `shellctl-sanitize-pty` and -`shellctl-runner-exit` do not accidentally pull in the main shellctl client or -server stacks. -""" diff --git a/dify-agent/src/shellctl_runtime/paths.py b/dify-agent/src/shellctl_runtime/paths.py deleted file mode 100644 index 4ba2295a658..00000000000 --- a/dify-agent/src/shellctl_runtime/paths.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Stdlib-only filesystem helpers shared by shellctl runtime commands.""" - -from __future__ import annotations - -from pathlib import Path - -DEFAULT_SQLITE_BUSY_TIMEOUT_MS = 5000 - - -def db_path_from_state_dir(state_dir: Path) -> Path: - """Resolve the SQLite database path for a shellctl state directory.""" - - return state_dir / "shellctl.db" - - -__all__ = ["DEFAULT_SQLITE_BUSY_TIMEOUT_MS", "db_path_from_state_dir"] diff --git a/dify-agent/src/shellctl_runtime/py.typed b/dify-agent/src/shellctl_runtime/py.typed deleted file mode 100644 index 8b137891791..00000000000 --- a/dify-agent/src/shellctl_runtime/py.typed +++ /dev/null @@ -1 +0,0 @@ - diff --git a/dify-agent/src/shellctl_runtime/runner_exit.py b/dify-agent/src/shellctl_runtime/runner_exit.py deleted file mode 100644 index b78f1758a73..00000000000 --- a/dify-agent/src/shellctl_runtime/runner_exit.py +++ /dev/null @@ -1,158 +0,0 @@ -"""Lightweight SQLite updater for drained shellctl job exits. - -This command runs after the tmux output pipe reaches EOF and flushes -`output.log`. It must stay stdlib-only so exit-state materialization does not -pay the FastAPI/SQLAlchemy import cost on every job completion. The CLI accepts -the busy-timeout override from `ShellctlConfig` so non-default deployments keep -the same SQLite locking behavior in and out of process. -""" - -from __future__ import annotations - -import argparse -import sqlite3 -from pathlib import Path - -from shellctl_runtime.paths import ( - DEFAULT_SQLITE_BUSY_TIMEOUT_MS, - db_path_from_state_dir, -) - -TERMINAL_STATUSES = frozenset({"exited", "terminated", "failed", "lost"}) -NONTERMINAL_STATUSES = frozenset({"created", "starting", "running"}) - - -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - """Parse the tmux finalizer CLI contract.""" - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--state-dir", required=True, type=Path) - parser.add_argument("--job-id", required=True) - parser.add_argument("--exit-code", required=True, type=int) - parser.add_argument("--ended-at", required=True) - parser.add_argument( - "--sqlite-busy-timeout-ms", - type=int, - default=DEFAULT_SQLITE_BUSY_TIMEOUT_MS, - ) - return parser.parse_args(argv) - - -def record_runner_exit( - *, - state_dir: Path, - job_id: str, - exit_code: int, - ended_at: str, - busy_timeout_ms: int = DEFAULT_SQLITE_BUSY_TIMEOUT_MS, -) -> None: - """Persist a drained runner exit directly into the shellctl SQLite DB. - - The update is idempotent for terminal rows: once a job reaches a terminal - state, this helper returns successfully without rewriting the row. The - terminal/non-terminal decision is repeated inside the `UPDATE` itself so a - concurrent writer cannot be clobbered by a stale pre-read. - """ - - db_path = db_path_from_state_dir(state_dir) - if not db_path.exists(): - raise FileNotFoundError(f"shellctl database not found: {db_path}") - - connection = sqlite3.connect(db_path, timeout=busy_timeout_ms / 1000) - try: - connection.execute(f"PRAGMA busy_timeout={busy_timeout_ms}") - row = connection.execute( - "SELECT status FROM jobs WHERE job_id = ?", - (job_id,), - ).fetchone() - if row is None: - raise LookupError(f"Unknown job id: {job_id}") - status = str(row[0]) - if status in TERMINAL_STATUSES: - return - if status not in NONTERMINAL_STATUSES: - raise ValueError(f"Unsupported shellctl job status: {status}") - nonterminal = tuple(NONTERMINAL_STATUSES) - cursor = connection.execute( - """ - UPDATE jobs - SET status = CASE - WHEN status IN (?, ?, ?) THEN ? - ELSE status - END, - exit_code = CASE - WHEN status IN (?, ?, ?) THEN ? - ELSE exit_code - END, - ended_at = CASE - WHEN status IN (?, ?, ?) THEN ? - ELSE ended_at - END, - updated_at = CASE - WHEN status IN (?, ?, ?) THEN ? - ELSE updated_at - END, - reason = CASE - WHEN status IN (?, ?, ?) THEN NULL - ELSE reason - END, - message = CASE - WHEN status IN (?, ?, ?) THEN NULL - ELSE message - END - WHERE job_id = ? - """, - ( - *nonterminal, - "exited", - *nonterminal, - exit_code, - *nonterminal, - ended_at, - *nonterminal, - ended_at, - *nonterminal, - *nonterminal, - job_id, - ), - ) - if cursor.rowcount == 0: - raise LookupError(f"Unknown job id: {job_id}") - connection.commit() - finally: - connection.close() - - -def main(argv: list[str] | None = None) -> None: - """Run the standalone runner-exit command.""" - - args = parse_args(argv) - try: - record_runner_exit( - state_dir=args.state_dir, - job_id=args.job_id, - exit_code=args.exit_code, - ended_at=args.ended_at, - busy_timeout_ms=args.sqlite_busy_timeout_ms, - ) - except ( - FileNotFoundError, - LookupError, - OSError, - sqlite3.DatabaseError, - ValueError, - ) as exc: - raise SystemExit(str(exc)) from exc - - -if __name__ == "__main__": - main() - - -__all__ = [ - "NONTERMINAL_STATUSES", - "TERMINAL_STATUSES", - "main", - "parse_args", - "record_runner_exit", -] diff --git a/dify-agent/src/shellctl_runtime/sanitize.py b/dify-agent/src/shellctl_runtime/sanitize.py deleted file mode 100644 index f536e9629b1..00000000000 --- a/dify-agent/src/shellctl_runtime/sanitize.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Lightweight PTY sanitizer used by tmux `pipe-pane`. - -This module stays stdlib-only because shellctl starts it once per job to drain -PTY output into `output.log`. The ready-file handshake must happen before any -stdin reads so the server can distinguish slow startup from a stuck tmux pipe. -""" - -from __future__ import annotations - -import argparse -import codecs -import sys -from pathlib import Path -from typing import BinaryIO - - -class PtySanitizer: - """Incrementally convert PTY bytes into stable, readable UTF-8 text. - - Responsibilities: - - preserve UTF-8 decoder state across chunk boundaries - - strip common ANSI control sequences without leaking partial escape state - - normalize carriage-return progress updates into the final visible line - - This adapter intentionally keeps only minimal terminal state. It aims for a - practical log representation rather than a full terminal emulator. - """ - - def __init__(self) -> None: - self._decoder = codecs.getincrementaldecoder("utf-8")("replace") - self._line_buffer = "" - self._pending_cr = False - self._escape_state = "normal" - - def feed(self, raw: bytes) -> str: - """Consume one PTY byte chunk and return newly stable text.""" - - return self._consume_text(self._decoder.decode(raw, final=False), final=False) - - def flush(self) -> str: - """Flush buffered decoder/line state at end-of-stream.""" - - tail = self._consume_text(self._decoder.decode(b"", final=True), final=True) - if self._line_buffer: - tail += self._line_buffer - self._line_buffer = "" - self._pending_cr = False - self._escape_state = "normal" - return tail - - def _consume_text(self, text: str, *, final: bool) -> str: - parts: list[str] = [] - for char in text: - self._consume_char(char, parts) - if final and self._escape_state != "normal": - self._escape_state = "normal" - return "".join(parts) - - def _consume_char(self, char: str, parts: list[str]) -> None: - state = self._escape_state - if state == "normal": - if char == "\x1b": - self._escape_state = "esc" - return - self._consume_visible_char(char, parts) - return - if state == "esc": - if char == "[": - self._escape_state = "csi" - return - if char == "]": - self._escape_state = "osc" - return - self._escape_state = "normal" - if char.isprintable() and char != "\x1b": - self._consume_visible_char(char, parts) - return - if state == "csi": - if "@" <= char <= "~": - self._escape_state = "normal" - return - if state == "osc": - if char == "\x07": - self._escape_state = "normal" - return - if char == "\x1b": - self._escape_state = "osc_esc" - return - if state == "osc_esc": - self._escape_state = "normal" if char == "\\" else "osc" - - def _consume_visible_char(self, char: str, parts: list[str]) -> None: - if self._pending_cr: - if char == "\n": - parts.append(self._line_buffer) - parts.append("\n") - self._line_buffer = "" - self._pending_cr = False - return - self._line_buffer = "" - self._pending_cr = False - if char == "\r": - self._pending_cr = True - return - if char == "\n": - parts.append(self._line_buffer) - parts.append("\n") - self._line_buffer = "" - return - self._line_buffer += char - - -def sanitize_pty_output(raw: bytes) -> str: - """Sanitize a complete PTY byte string into readable UTF-8 text.""" - - sanitizer = PtySanitizer() - return sanitizer.feed(raw) + sanitizer.flush() - - -def sanitize_pty_stream( - stdin: BinaryIO, - stdout: BinaryIO, - *, - chunk_size: int = 65536, -) -> None: - """Run the streaming PTY sanitizer as a Unix-style filter.""" - - sanitizer = PtySanitizer() - while True: - chunk = stdin.read(chunk_size) - if not chunk: - break - output = sanitizer.feed(chunk) - if output: - stdout.write(output.encode("utf-8")) - if hasattr(stdout, "flush"): - stdout.flush() - tail = sanitizer.flush() - if tail: - stdout.write(tail.encode("utf-8")) - if hasattr(stdout, "flush"): - stdout.flush() - if hasattr(stdout, "flush"): - stdout.flush() - - -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - """Parse the tiny CLI contract used by tmux `pipe-pane`.""" - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--ready-file", type=Path) - return parser.parse_args(argv) - - -def run_sanitize_pty( - ready_file: Path | None, - *, - stdin: BinaryIO, - stdout: BinaryIO, -) -> None: - """Touch the ready file, then sanitize stdin into stdout.""" - - if ready_file is not None: - ready_file.touch() - sanitize_pty_stream(stdin, stdout) - - -def main(argv: list[str] | None = None) -> None: - """Run the standalone PTY sanitizer module.""" - - args = parse_args(argv) - run_sanitize_pty(args.ready_file, stdin=sys.stdin.buffer, stdout=sys.stdout.buffer) - - -if __name__ == "__main__": - main() - - -__all__ = [ - "PtySanitizer", - "main", - "parse_args", - "run_sanitize_pty", - "sanitize_pty_output", - "sanitize_pty_stream", -] diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_config.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_config.py deleted file mode 100644 index 0032d9c6142..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_config.py +++ /dev/null @@ -1,337 +0,0 @@ -from __future__ import annotations - -import io -import zipfile -from pathlib import Path -from types import SimpleNamespace -from typing import cast - -import pytest - -from dify_agent.agent_stub.cli import _config as config_cli -from dify_agent.agent_stub.cli._env import AgentStubEnvironment -from dify_agent.agent_stub.client._errors import AgentStubValidationError -from dify_agent.agent_stub.protocol.agent_stub import AgentStubConfigManifestResponse, AgentStubConfigPushRequest - - -def _manifest_payload() -> AgentStubConfigManifestResponse: - return AgentStubConfigManifestResponse.model_validate( - { - "agent_id": "agent-1", - "config_version": {"id": "cfg-1", "kind": "build_draft", "writable": True}, - "skills": {"items": [{"name": "alpha", "description": "Alpha skill"}]}, - "files": {"items": [{"name": "guide.txt"}]}, - "env_keys": ["API_KEY", "JSON_VALUE"], - "note": "Use carefully.", - } - ) - - -def _zip_bytes(members: dict[str, bytes]) -> bytes: - buffer = io.BytesIO() - with zipfile.ZipFile(buffer, "w") as archive: - for name, payload in members.items(): - archive.writestr(name, payload) - return buffer.getvalue() - - -def _environment() -> AgentStubEnvironment: - return AgentStubEnvironment(url="https://agent.example.com/agent-stub", auth_jwe="test-jwe") - - -def test_pull_config_skills_from_environment_downloads_and_extracts_archives( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - monkeypatch.setattr(config_cli, "request_agent_stub_config_manifest_sync", lambda **_kwargs: _manifest_payload()) - monkeypatch.setattr( - config_cli, - "request_agent_stub_config_skill_pull_sync", - lambda **_kwargs: _zip_bytes({"SKILL.md": b"# Alpha\n", "refs/spec.md": b"spec"}), - ) - - result = config_cli.pull_config_skills_from_environment(local_dir=str(tmp_path)) - - assert [item.name for item in result.items] == ["alpha"] - assert Path(result.items[0].archive_path).read_bytes() - assert Path(result.items[0].directory_path, "SKILL.md").read_text(encoding="utf-8") == "# Alpha\n" - assert result.items[0].skill_md == "# Alpha\n" - - -def test_pull_config_files_from_environment_downloads_visible_files( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - monkeypatch.setattr(config_cli, "request_agent_stub_config_manifest_sync", lambda **_kwargs: _manifest_payload()) - monkeypatch.setattr( - config_cli, - "request_agent_stub_config_file_pull_sync", - lambda **_kwargs: b"guide-bytes", - ) - - result = config_cli.pull_config_files_from_environment(local_dir=str(tmp_path)) - - assert result.items == [config_cli.ConfigFilePullResult.Item(name="guide.txt", path=str(tmp_path / "guide.txt"))] - assert (tmp_path / "guide.txt").read_bytes() == b"guide-bytes" - - -def test_pull_config_note_uses_hidden_default_dir( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.chdir(tmp_path) - monkeypatch.setattr(config_cli, "manifest_from_environment", _manifest_payload) - - note_path = config_cli.pull_config_note_from_environment() - - assert note_path == (tmp_path / ".dify_conf" / "note.md").resolve() - assert note_path.read_text(encoding="utf-8") == "Use carefully." - - -def test_push_config_note_from_environment_reads_default_file( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.chdir(tmp_path) - note_path = tmp_path / ".dify_conf" / "note.md" - note_path.parent.mkdir() - note_path.write_text("hello", encoding="utf-8") - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - response = config_cli.push_config_note_from_environment(None) - - assert response.agent_id == "agent-1" - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.note == "hello" - assert request.env_text is None - assert request.files == [] - assert request.skills == [] - - -def test_push_config_note_from_environment_reads_explicit_file( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - note_path = tmp_path / "note.md" - note_path.write_text("explicit", encoding="utf-8") - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_note_from_environment(str(note_path)) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.note == "explicit" - - -def test_push_config_note_from_environment_reads_stdin(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - monkeypatch.setattr(config_cli.os, "dup", lambda _fd: 99) - monkeypatch.setattr(config_cli.os, "fdopen", lambda _fd, encoding: io.StringIO("from-stdin")) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_note_from_environment("-") - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.note == "from-stdin" - - -def test_push_config_env_from_environment_reads_explicit_file( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - env_path = tmp_path / "custom.env" - env_path.write_text("API_KEY=custom\n", encoding="utf-8") - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_env_from_environment(str(env_path)) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.env_text == "API_KEY=custom\n" - - -def test_push_config_env_from_environment_reads_stdin(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - monkeypatch.setattr(config_cli.os, "dup", lambda _fd: 99) - monkeypatch.setattr(config_cli.os, "fdopen", lambda _fd, encoding: io.StringIO("API_KEY=stdin\n")) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_env_from_environment("-") - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.env_text == "API_KEY=stdin\n" - - -def test_push_config_files_from_environment_builds_upload_items( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - file_path = tmp_path / "guide.txt" - file_path.write_text("guide", encoding="utf-8") - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - monkeypatch.setattr( - config_cli, - "upload_tool_file_resource_from_environment", - lambda *, path: SimpleNamespace(tool_file_id=f"tool-file:{Path(path).name}"), - ) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_files_from_environment([str(file_path)]) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.files[0].name == "guide.txt" - assert request.files[0].file_ref is not None - assert request.files[0].file_ref.id == "tool-file:guide.txt" - assert request.skills == [] - - -def test_push_config_files_from_environment_rejects_empty_paths() -> None: - with pytest.raises(AgentStubValidationError, match="at least one file path is required"): - config_cli.push_config_files_from_environment([]) - - -def test_delete_config_files_from_environment_builds_delete_items(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.delete_config_files_from_environment(["old.txt", "legacy.txt"]) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert [item.name for item in request.files] == ["old.txt", "legacy.txt"] - assert all(item.file_ref is None for item in request.files) - - -def test_delete_config_files_from_environment_rejects_empty_names() -> None: - with pytest.raises(AgentStubValidationError, match="at least one file name is required"): - config_cli.delete_config_files_from_environment([]) - - -def test_push_config_skills_from_environment_builds_archive_upload_items( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "alpha" - skill_dir.mkdir() - (skill_dir / "SKILL.md").write_text("# Alpha\n", encoding="utf-8") - uploaded_paths: list[str] = [] - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - def fake_upload_tool_file_resource_from_environment(*, path: str): - uploaded_paths.append(path) - return SimpleNamespace(tool_file_id=f"tool-file-{len(uploaded_paths)}") - - monkeypatch.setattr( - config_cli, - "upload_tool_file_resource_from_environment", - fake_upload_tool_file_resource_from_environment, - ) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.push_config_skills_from_environment([str(skill_dir)]) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert request.skills[0].name == "alpha" - assert request.skills[0].file_ref is not None - assert request.skills[0].file_ref.id == "tool-file-1" - assert len(uploaded_paths) == 1 - assert uploaded_paths[0].endswith("/alpha.zip") - - -def test_push_config_skills_from_environment_rejects_empty_paths() -> None: - with pytest.raises(AgentStubValidationError, match="at least one skill directory is required"): - config_cli.push_config_skills_from_environment([]) - - -def test_delete_config_skills_from_environment_builds_delete_items(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(config_cli, "read_agent_stub_environment", lambda: _environment()) - - captured: dict[str, object] = {} - - def fake_push_sync(**kwargs): - captured.update(kwargs) - return _manifest_payload() - - monkeypatch.setattr(config_cli, "request_agent_stub_config_push_sync", fake_push_sync) - - config_cli.delete_config_skills_from_environment(["alpha", "beta"]) - - request = cast(AgentStubConfigPushRequest, captured["request"]) - assert [item.name for item in request.skills] == ["alpha", "beta"] - assert all(item.file_ref is None for item in request.skills) - - -def test_delete_config_skills_from_environment_rejects_empty_names() -> None: - with pytest.raises(AgentStubValidationError, match="at least one skill name is required"): - config_cli.delete_config_skills_from_environment([]) - - -def test_build_file_push_item_rejects_non_regular_files(tmp_path: Path) -> None: - with pytest.raises(AgentStubValidationError, match="regular file"): - config_cli._build_file_push_item(item=config_cli._PreparedPushItem(name="bad", path=tmp_path)) - - -def test_build_skill_push_item_rejects_missing_skill_md(tmp_path: Path) -> None: - skill_dir = tmp_path / "alpha" - skill_dir.mkdir() - - with pytest.raises(AgentStubValidationError, match="must contain SKILL.md"): - config_cli._build_skill_push_item(item=config_cli._PreparedPushItem(name="alpha", path=skill_dir)) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py deleted file mode 100644 index 5dbf0bafe5a..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py +++ /dev/null @@ -1,870 +0,0 @@ -from __future__ import annotations - -import base64 -import json -from io import BytesIO -from pathlib import Path -import stat -from zipfile import ZipFile, ZipInfo - -import pytest - -from dify_agent.agent_stub.cli._drive import ( - DrivePullResult, - format_drive_manifest, - list_drive_manifest_from_environment, - pull_drive_from_environment, - push_drive_from_environment, -) -from dify_agent.agent_stub.cli._files import UploadedToolFileResource -from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubDriveCommitRequest, - AgentStubDriveCommitResponse, - AgentStubDriveItem, - AgentStubFileMapping, - AgentStubDriveManifestResponse, -) - - -def _reference(record_id: str) -> str: - payload = base64.urlsafe_b64encode(json.dumps({"record_id": record_id}, separators=(",", ":")).encode()).decode() - return f"dify-file-ref:{payload}" - - -def test_list_drive_manifest_from_environment_returns_manifest_model(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured: dict[str, object] = {} - - def fake_manifest(**kwargs): - captured.update(kwargs) - return AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/example/SKILL.md", - size=12, - hash="sha256:abc", - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - ) - ] - ) - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - fake_manifest, - ) - - result = list_drive_manifest_from_environment(prefix="skills/") - - assert isinstance(result, AgentStubDriveManifestResponse) - assert result.items[0].key == "skills/example/SKILL.md" - assert captured["prefix"] == "skills/" - assert captured["include_download_url"] is False - - -def test_format_drive_manifest_returns_human_readable_listing(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured: dict[str, object] = {} - - def fake_manifest(**kwargs): - captured.update(kwargs) - return AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/example/SKILL.md", - size=12, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - ), - AgentStubDriveItem( - key="skills/example/helper.py", - size=None, - hash="sha256:abc", - mime_type=None, - file_kind="tool_file", - file_id="tool-file-2", - ), - ] - ) - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - fake_manifest, - ) - - result = format_drive_manifest(list_drive_manifest_from_environment(prefix="skills/")) - - assert result == ("12\ttext/markdown\t-\tskills/example/SKILL.md\n-\t-\tsha256:abc\tskills/example/helper.py") - assert captured["prefix"] == "skills/" - assert captured["include_download_url"] is False - - -def test_pull_drive_from_environment_writes_files_under_drive_base( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured: dict[str, object] = {} - - def fake_manifest(**kwargs): - captured.update(kwargs) - return AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/example/SKILL.md", - size=11, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ) - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - fake_manifest, - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"hello world", - ) - - result = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) - - assert result.model_dump() == {"items": [{"key": "skills/", "local_path": str(tmp_path / "skills")}]} - assert (tmp_path / "skills" / "example" / "SKILL.md").read_bytes() == b"hello world" - assert captured["prefix"] == "skills/" - assert captured["include_download_url"] is True - - -def test_pull_drive_from_environment_auto_extracts_skill_archive( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - archive_buffer = BytesIO() - with ZipFile(archive_buffer, mode="w") as archive: - archive.writestr("SKILL.md", "# Example\n") - archive.writestr("nested/helper.py", "print('x')\n") - archive_bytes = archive_buffer.getvalue() - - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/.DIFY-SKILL-FULL.zip", - size=len(archive_bytes), - hash=None, - mime_type="application/zip", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: archive_bytes, - ) - - result = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) - - archive_path = tmp_path / "skills" / "foo" / ".DIFY-SKILL-FULL.zip" - assert result.model_dump() == {"items": [{"key": "skills/foo", "local_path": str(tmp_path / "skills" / "foo")}]} - assert not archive_path.exists() - assert (tmp_path / "skills" / "foo" / "SKILL.md").read_text(encoding="utf-8") == "# Example\n" - assert (tmp_path / "skills" / "foo" / "nested" / "helper.py").read_text(encoding="utf-8") == "print('x')\n" - - -def test_pull_drive_from_environment_rejects_traversal_keys( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="../escape.txt", - size=4, - hash=None, - mime_type="text/plain", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - - with pytest.raises(AgentStubValidationError, match="outside the drive base"): - _ = pull_drive_from_environment(targets=[""], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_rejects_skill_archive_path_traversal( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - archive_buffer = BytesIO() - with ZipFile(archive_buffer, mode="w") as archive: - archive.writestr("SKILL.md", "# Example\n") - archive.writestr("../escape.txt", "escape") - archive_bytes = archive_buffer.getvalue() - - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/.DIFY-SKILL-FULL.zip", - size=len(archive_bytes), - hash=None, - mime_type="application/zip", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: archive_bytes, - ) - - with pytest.raises(AgentStubValidationError, match="path traversal"): - _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) - assert not (tmp_path / "skills" / "foo" / "SKILL.md").exists() - - -def test_pull_drive_from_environment_rejects_skill_archive_absolute_entry( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - archive_buffer = BytesIO() - with ZipFile(archive_buffer, mode="w") as archive: - archive.writestr("/escape.txt", "escape") - archive_bytes = archive_buffer.getvalue() - - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/.DIFY-SKILL-FULL.zip", - size=len(archive_bytes), - hash=None, - mime_type="application/zip", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: archive_bytes, - ) - - with pytest.raises(AgentStubValidationError, match="absolute path"): - _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_rejects_skill_archive_symlink_entry( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - archive_buffer = BytesIO() - with ZipFile(archive_buffer, mode="w") as archive: - symlink_info = ZipInfo("linked.txt") - symlink_info.external_attr = (stat.S_IFLNK | 0o777) << 16 - archive.writestr(symlink_info, "outside.txt") - archive_bytes = archive_buffer.getvalue() - - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/.DIFY-SKILL-FULL.zip", - size=len(archive_bytes), - hash=None, - mime_type="application/zip", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: archive_bytes, - ) - - with pytest.raises(AgentStubValidationError, match="symlink entry"): - _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_rejects_invalid_skill_archive( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - archive_bytes = b"not-a-zip" - - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/.DIFY-SKILL-FULL.zip", - size=len(archive_bytes), - hash=None, - mime_type="application/zip", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: archive_bytes, - ) - - with pytest.raises(AgentStubTransferError, match="downloaded skill archive is invalid"): - _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_rejects_missing_download_url( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/example/SKILL.md", - size=11, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - ) - ] - ), - ) - - with pytest.raises(AgentStubValidationError, match="missing download_url"): - _ = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_rejects_size_mismatch( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/example/SKILL.md", - size=99, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/download", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"hello world", - ) - - with pytest.raises(AgentStubTransferError, match="size mismatch"): - _ = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) - - -def test_pull_drive_from_environment_requests_multiple_targets_and_deduplicates_overlaps( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured_prefixes: list[str] = [] - - def fake_manifest(**kwargs): - captured_prefixes.append(kwargs["prefix"]) - if kwargs["prefix"] == "skills/foo": - return AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/SKILL.md", - size=5, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/skill-md", - ) - ] - ) - return AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="skills/foo/SKILL.md", - size=5, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/skill-md", - ), - AgentStubDriveItem( - key="files/a.txt", - size=1, - hash=None, - mime_type="text/plain", - file_kind="tool_file", - file_id="tool-file-2", - download_url="https://files.example.com/a-txt", - ), - ] - ) - - downloaded_urls: list[str] = [] - monkeypatch.setattr("dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", fake_manifest) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda *, download_url: ( - downloaded_urls.append(download_url) or (b"hello" if download_url.endswith("skill-md") else b"a") - ), - ) - - result = pull_drive_from_environment(targets=["skills/foo", "files/a.txt"], local_base=str(tmp_path)) - - assert set(captured_prefixes) == {"skills/foo", "files/a.txt"} - assert len(captured_prefixes) == 2 - assert {(item.key, item.local_path) for item in result.items} == { - ("files/a.txt", str(tmp_path / "files" / "a.txt")), - ("skills/foo", str(tmp_path / "skills" / "foo")), - } - assert set(downloaded_urls) == {"https://files.example.com/a-txt", "https://files.example.com/skill-md"} - assert len(downloaded_urls) == 2 - - -def test_pull_drive_from_environment_without_targets_preserves_whole_drive_pull( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured_prefixes: list[str] = [] - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **kwargs: captured_prefixes.append(kwargs["prefix"]) or AgentStubDriveManifestResponse(items=[]), - ) - - assert pull_drive_from_environment(local_base=str(tmp_path)).model_dump() == {"items": []} - assert captured_prefixes == [""] - - -def test_pull_drive_from_environment_returns_json_result( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", - lambda **_kwargs: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key="files/a.txt", - size=1, - hash=None, - mime_type="text/plain", - file_kind="tool_file", - file_id="tool-file-1", - download_url="https://files.example.com/a-txt", - ) - ] - ), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"a", - ) - - result = pull_drive_from_environment(targets=["files/a.txt"], local_base=str(tmp_path)) - - assert isinstance(result, DrivePullResult) - assert result.model_dump() == {"items": [{"key": "files/a.txt", "local_path": str(tmp_path / "files" / "a.txt")}]} - assert (tmp_path / "files" / "a.txt").read_bytes() == b"a" - - -def test_push_drive_from_environment_commits_single_file(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.upload_tool_file_resource_from_environment", - lambda *, path: UploadedToolFileResource( - mapping=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")), - tool_file_id="tool-file-1", - ), - ) - captured: dict[str, object] = {} - - def fake_commit(**kwargs): - captured.update(kwargs) - return AgentStubDriveCommitResponse( - items=[ - AgentStubDriveItem( - key="files/report.pdf", - size=6, - hash=None, - mime_type="application/pdf", - file_kind="tool_file", - file_id="tool-file-1", - value_owned_by_drive=True, - ) - ] - ) - - monkeypatch.setattr("dify_agent.agent_stub.cli._drive.request_agent_stub_drive_commit_sync", fake_commit) - - response = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind=None) - - assert response.items[0].key == "files/report.pdf" - request = captured["request"] - assert isinstance(request, AgentStubDriveCommitRequest) - assert request.items[0].key == "files/report.pdf" - assert request.items[0].file_ref is not None - assert request.items[0].file_ref.kind == "tool_file" - assert request.items[0].file_ref.id == "tool-file-1" - - -def test_push_drive_from_environment_rejects_file_with_kind_skill( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="--kind skill requires a directory containing SKILL.md"): - _ = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind="skill") - - -def test_push_drive_from_environment_rejects_symlinked_file_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report") - symlink_path = tmp_path / "report-link.pdf" - symlink_path.symlink_to(source) - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(symlink_path), drive_path="files/report.pdf", kind=None) - - -def test_push_drive_from_environment_requires_kind_for_directory( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="requires --kind skill or --kind dir"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind=None) - - -def test_push_drive_from_environment_kind_skill_requires_skill_md( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="requires a directory containing SKILL.md"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") - - -def test_push_drive_from_environment_kind_skill_standardizes_skill_directory( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - (skill_dir / "SKILL.md").write_text("# Example\n", encoding="utf-8") - (skill_dir / "helper.py").write_text("print('x')\n", encoding="utf-8") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - uploaded_paths: list[str] = [] - - def fake_upload(*, path: str) -> UploadedToolFileResource: - uploaded_paths.append(Path(path).name) - return UploadedToolFileResource( - mapping=AgentStubFileMapping( - transfer_method="tool_file", - reference=_reference(Path(path).name), - ), - tool_file_id=Path(path).name, - ) - - monkeypatch.setattr("dify_agent.agent_stub.cli._drive.upload_tool_file_resource_from_environment", fake_upload) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_commit_sync", - lambda **kwargs: AgentStubDriveCommitResponse( - items=[ - AgentStubDriveItem( - key=item.key, - size=None, - hash=None, - mime_type=None, - file_kind=item.file_ref.kind, - file_id=item.file_ref.id, - value_owned_by_drive=item.value_owned_by_drive, - ) - for item in kwargs["request"].items - ] - ), - ) - - response = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") - - assert set(uploaded_paths) == {"SKILL.md", ".DIFY-SKILL-FULL.zip"} - assert {item.key for item in response.items} == { - "skills/example/SKILL.md", - "skills/example/.DIFY-SKILL-FULL.zip", - } - - -def test_push_drive_from_environment_kind_skill_archive_excludes_transient_entries( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - (skill_dir / "SKILL.md").write_text("# Example\n", encoding="utf-8") - (skill_dir / "helper.py").write_text("print('x')\n", encoding="utf-8") - (skill_dir / ".DIFY-SKILL-FULL.zip").write_bytes(b"old-archive") - git_dir = skill_dir / ".git" - git_dir.mkdir() - (git_dir / "config").write_text("[core]\n", encoding="utf-8") - pycache_dir = skill_dir / "__pycache__" - pycache_dir.mkdir() - (pycache_dir / "helper.pyc").write_bytes(b"compiled") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - archive_entries: list[str] = [] - - def fake_upload(*, path: str) -> UploadedToolFileResource: - if Path(path).name == ".DIFY-SKILL-FULL.zip": - with ZipFile(path) as archive: - archive_entries.extend(sorted(archive.namelist())) - return UploadedToolFileResource( - mapping=AgentStubFileMapping( - transfer_method="tool_file", - reference=_reference(Path(path).name), - ), - tool_file_id=Path(path).name, - ) - - monkeypatch.setattr("dify_agent.agent_stub.cli._drive.upload_tool_file_resource_from_environment", fake_upload) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_commit_sync", - lambda **kwargs: AgentStubDriveCommitResponse( - items=[ - AgentStubDriveItem( - key=item.key, - size=None, - hash=None, - mime_type=None, - file_kind=item.file_ref.kind, - file_id=item.file_ref.id, - value_owned_by_drive=item.value_owned_by_drive, - ) - for item in kwargs["request"].items - ] - ), - ) - - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") - - assert {"SKILL.md", "helper.py"}.issubset(archive_entries) - assert ".git/config" not in archive_entries - assert "__pycache__/helper.pyc" not in archive_entries - assert ".DIFY-SKILL-FULL.zip" not in archive_entries - - -def test_push_drive_from_environment_kind_skill_rejects_symlinked_archive_entries( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - (skill_dir / "SKILL.md").write_text("# Example\n", encoding="utf-8") - outside = tmp_path / "outside.txt" - outside.write_text("outside", encoding="utf-8") - (skill_dir / "linked.txt").symlink_to(outside) - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") - - -def test_push_drive_from_environment_kind_dir_requires_directory( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="--kind dir requires a directory"): - _ = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind="dir") - - -def test_push_drive_from_environment_kind_file_requires_file( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - skill_dir = tmp_path / "skill" - skill_dir.mkdir() - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="--kind file requires a file"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="file") - - -def test_push_drive_from_environment_rejects_symlinked_directory_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source_dir = tmp_path / "skill" - source_dir.mkdir() - (source_dir / "SKILL.md").write_text("# Example\n", encoding="utf-8") - symlink_path = tmp_path / "skill-link" - symlink_path.symlink_to(source_dir, target_is_directory=True) - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(symlink_path), drive_path="skills/example", kind="skill") - - -def test_push_drive_from_environment_kind_dir_rejects_symlinked_files( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - root = tmp_path / "skill" - root.mkdir() - outside = tmp_path / "outside.txt" - outside.write_text("outside", encoding="utf-8") - (root / "linked.txt").symlink_to(outside) - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(root), drive_path="skills/example", kind="dir") - - -def test_push_drive_from_environment_kind_dir_keeps_user_files_that_skill_packaging_skips( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - root = tmp_path / "skill" - root.mkdir() - (root / ".DIFY-SKILL-FULL.zip").write_bytes(b"archive") - node_modules_dir = root / "node_modules" - node_modules_dir.mkdir() - (node_modules_dir / "module.js").write_text("export default 1\n", encoding="utf-8") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - uploaded_paths: list[str] = [] - - def fake_upload(*, path: str) -> UploadedToolFileResource: - uploaded_paths.append(Path(path).relative_to(root).as_posix()) - return UploadedToolFileResource( - mapping=AgentStubFileMapping( - transfer_method="tool_file", - reference=_reference(Path(path).name), - ), - tool_file_id=Path(path).name, - ) - - monkeypatch.setattr("dify_agent.agent_stub.cli._drive.upload_tool_file_resource_from_environment", fake_upload) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_commit_sync", - lambda **kwargs: AgentStubDriveCommitResponse( - items=[ - AgentStubDriveItem( - key=item.key, - size=None, - hash=None, - mime_type=None, - file_kind=item.file_ref.kind, - file_id=item.file_ref.id, - value_owned_by_drive=item.value_owned_by_drive, - ) - for item in kwargs["request"].items - ] - ), - ) - - response = push_drive_from_environment(local_path=str(root), drive_path="skills/example", kind="dir") - - assert set(uploaded_paths) == {".DIFY-SKILL-FULL.zip", "node_modules/module.js"} - assert {item.key for item in response.items} == { - "skills/example/.DIFY-SKILL-FULL.zip", - "skills/example/node_modules/module.js", - } diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py deleted file mode 100644 index 3cf22e04d8e..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py +++ /dev/null @@ -1,365 +0,0 @@ -from __future__ import annotations - -import base64 -import json -from pathlib import Path - -import pytest - -from dify_agent.agent_stub.cli._files import ( - download_file_from_environment, - upload_file_from_environment, - upload_tool_file_resource_from_environment, -) -from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError -from dify_agent.agent_stub.protocol.agent_stub import AgentStubFileMapping - - -def _reference(record_id: str) -> str: - payload = base64.urlsafe_b64encode(json.dumps({"record_id": record_id}, separators=(",", ":")).encode()).decode() - return f"dify-file-ref:{payload}" - - -def test_upload_file_from_environment_requests_signed_url_and_normalizes_output( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report-bytes") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", - lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), - ) - captured = {} - - def fake_upload_file_to_signed_url_sync(**kwargs): - captured["filename"] = kwargs["filename"] - captured["mimetype"] = kwargs["mimetype"] - captured["file_bytes"] = kwargs["file_obj"].read() - kwargs["file_obj"].seek(0) - return { - "id": "tool-file-1", - "reference": _reference("tool-file-1"), - } - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", - fake_upload_file_to_signed_url_sync, - ) - captured_download_request: dict[str, object] = {} - - def fake_request_agent_stub_file_download_sync(**kwargs): - captured_download_request["file"] = kwargs["file"] - captured_download_request["for_external"] = kwargs.get("for_external", True) - return type( - "Response", - (), - { - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 12, - "download_url": "https://files.example.com/download", - }, - )() - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", - fake_request_agent_stub_file_download_sync, - ) - - result = upload_file_from_environment(path=str(source)) - - assert result.model_dump() == { - "transfer_method": "tool_file", - "reference": _reference("tool-file-1"), - "download_url": "https://files.example.com/download", - } - assert captured == { - "filename": "report.pdf", - "mimetype": "application/pdf", - "file_bytes": b"report-bytes", - } - assert captured_download_request["file"] == AgentStubFileMapping( - transfer_method="tool_file", - reference=_reference("tool-file-1"), - ) - assert captured_download_request["for_external"] is True - - -def test_upload_tool_file_resource_from_environment_preserves_tool_file_id( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report-bytes") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", - lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", - lambda **_kwargs: {"id": "tool-file-1", "reference": _reference("tool-file-1")}, - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", - lambda **_kwargs: pytest.fail("resource helper must not request download_url"), - ) - - result = upload_tool_file_resource_from_environment(path=str(source)) - - assert result.mapping.model_dump() == { - "transfer_method": "tool_file", - "reference": _reference("tool-file-1"), - "url": None, - } - assert result.tool_file_id == "tool-file-1" - - -def test_download_file_from_environment_saves_bytes_and_renames_on_collision( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - target_dir = tmp_path / "downloads" - target_dir.mkdir() - (target_dir / "report.pdf").write_bytes(b"existing") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", - lambda **_kwargs: type( - "Response", - (), - { - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 12, - "download_url": "https://files.example.com/download", - }, - )(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"downloaded", - ) - - result = download_file_from_environment( - transfer_method="tool_file", - reference_or_url=_reference("tool-file-1"), - local_dir=str(target_dir), - ) - - assert result.path.name == "report (1).pdf" - assert result.path.read_bytes() == b"downloaded" - - -def test_download_file_from_environment_sanitizes_server_filename( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - target_dir = tmp_path / "downloads" - target_dir.mkdir() - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", - lambda **_kwargs: type( - "Response", - (), - { - "filename": "../nested/evil.txt", - "mime_type": "text/plain", - "size": 12, - "download_url": "https://files.example.com/download", - }, - )(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"downloaded", - ) - - result = download_file_from_environment( - transfer_method="tool_file", - reference_or_url=_reference("tool-file-1"), - local_dir=str(target_dir), - ) - - assert result.path.parent == target_dir - assert result.path.name == "evil.txt" - assert result.path.read_bytes() == b"downloaded" - - -def test_upload_file_from_environment_rejects_non_canonical_reference( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report-bytes") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", - lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", - lambda **_kwargs: {"id": "tool-file-1", "reference": "raw-tool-file-uuid"}, - ) - - with pytest.raises(AgentStubTransferError, match="invalid canonical reference"): - _ = upload_file_from_environment(path=str(source)) - - -def test_upload_file_from_environment_rejects_missing_download_url( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report-bytes") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", - lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", - lambda **_kwargs: {"id": "tool-file-1", "reference": _reference("tool-file-1")}, - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", - lambda **_kwargs: type( - "Response", (), {"filename": "report.pdf", "mime_type": "application/pdf", "size": 12} - )(), - ) - - with pytest.raises(AgentStubTransferError, match="missing download_url"): - _ = upload_file_from_environment(path=str(source)) - - -def test_download_file_from_environment_supports_mapping_json( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - target_dir = tmp_path / "inputs" - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured: dict[str, object] = {} - - def fake_request_download(**kwargs): - captured["file"] = kwargs["file"] - captured["for_external"] = kwargs["for_external"] - return type( - "Response", - (), - { - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 12, - "download_url": "https://files.example.com/download", - }, - )() - - monkeypatch.setattr("dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", fake_request_download) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"downloaded", - ) - - result = download_file_from_environment( - mapping=json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), - local_dir=str(target_dir), - ) - - assert captured["file"].model_dump() == { - "transfer_method": "tool_file", - "reference": _reference("tool-file-1"), - "url": None, - } - assert captured["for_external"] is False - assert result.path == target_dir / "report.pdf" - assert result.path.read_bytes() == b"downloaded" - - -def test_download_file_from_environment_requests_internal_download_url( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - target_dir = tmp_path / "downloads" - target_dir.mkdir() - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - captured: dict[str, object] = {} - - def fake_request_download(**kwargs): - captured.update(kwargs) - return type( - "Response", - (), - { - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 12, - "download_url": "http://internal-files/report.pdf", - }, - )() - - monkeypatch.setattr("dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", fake_request_download) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync", - lambda **_kwargs: b"downloaded", - ) - - _ = download_file_from_environment( - transfer_method="tool_file", - reference_or_url=_reference("tool-file-1"), - local_dir=str(target_dir), - ) - - assert captured["for_external"] is False - - -def test_download_file_from_environment_requires_mapping_or_positional_pair() -> None: - with pytest.raises(AgentStubValidationError, match="requires either --mapping or TRANSFER_METHOD REFERENCE_OR_URL"): - _ = download_file_from_environment() - - -def test_download_file_from_environment_rejects_mapping_mixed_with_positionals() -> None: - with pytest.raises(AgentStubValidationError, match="cannot be combined"): - _ = download_file_from_environment( - transfer_method="tool_file", - reference_or_url=_reference("tool-file-1"), - mapping=json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), - ) - - -def test_upload_tool_file_resource_from_environment_rejects_missing_id( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - source = tmp_path / "report.pdf" - source.write_bytes(b"report-bytes") - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.request_agent_stub_file_upload_sync", - lambda **_kwargs: type("Response", (), {"upload_url": "https://files.example.com/upload"})(), - ) - monkeypatch.setattr( - "dify_agent.agent_stub.cli._files.upload_file_to_signed_url_sync", - lambda **_kwargs: {"reference": _reference("tool-file-1")}, - ) - - with pytest.raises(AgentStubTransferError, match="missing id"): - _ = upload_tool_file_resource_from_environment(path=str(source)) diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py deleted file mode 100644 index 02d505b5552..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py +++ /dev/null @@ -1,924 +0,0 @@ -from __future__ import annotations - -import base64 -import json -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from dify_agent.agent_stub.cli._drive import DrivePullResult -from dify_agent.agent_stub.cli.main import main -from dify_agent.agent_stub.client._errors import AgentStubTransferError -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubConfigFileItemsResponse, - AgentStubConfigFileItem, - AgentStubConfigManifestResponse, - AgentStubConfigSkillItemsResponse, - AgentStubConfigSkillItem, - AgentStubConfigVersionInfo, - AgentStubDriveCommitResponse, - AgentStubDriveItem, - AgentStubDriveManifestResponse, -) -from dify_agent.agent_stub.protocol.agent_stub import AgentStubConnectResponse - - -def _reference(record_id: str) -> str: - payload = base64.urlsafe_b64encode(json.dumps({"record_id": record_id}, separators=(",", ":")).encode()).decode() - return f"dify-file-ref:{payload}" - - -def _config_manifest_response() -> AgentStubConfigManifestResponse: - return AgentStubConfigManifestResponse( - agent_id="agent-1", - config_version=AgentStubConfigVersionInfo(id="cfg-1", kind="build_draft", writable=True), - skills=AgentStubConfigSkillItemsResponse(items=[]), - files=AgentStubConfigFileItemsResponse(items=[]), - env_keys=[], - note="Runtime note.", - ) - - -def _patch_cli_module(monkeypatch: pytest.MonkeyPatch, accessor_name: str, **attrs: object) -> None: - monkeypatch.setattr( - f"dify_agent.agent_stub.cli.main.{accessor_name}", - lambda: SimpleNamespace(**attrs), - ) - - -def test_cli_connect_reports_missing_environment_variables(capsys: pytest.CaptureFixture[str]) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["connect"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "DIFY_AGENT_STUB_API_BASE_URL" in captured.err - assert "DIFY_AGENT_STUB_AUTH_JWE" in captured.err - - -def test_cli_connect_supports_json_output( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - def fake_connect_from_environment(*, argv: list[str]) -> AgentStubConnectResponse: - assert argv == ["echo", "hello"] - return AgentStubConnectResponse(connection_id="conn-1", status="connected") - - _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) - - main(["connect", "--json", "--", "echo", "hello"]) - - captured = capsys.readouterr() - assert json.loads(captured.out) == {"connection_id": "conn-1", "status": "connected"} - - -def test_cli_unknown_command_auto_forwards_when_agent_stub_env_is_present( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - def fake_connect_from_environment(*, argv: list[str]) -> AgentStubConnectResponse: - assert argv == ["run", "--target", "prod"] - return AgentStubConnectResponse(connection_id="conn-1", status="connected") - - _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) - - main(["run", "--target", "prod"]) - - captured = capsys.readouterr() - assert captured.out.strip() == "connected conn-1" - - -def test_cli_unknown_command_reports_missing_environment_variables( - capsys: pytest.CaptureFixture[str], -) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["run", "--target", "prod"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "Usage: dify-agent" in captured.out - assert "connect" in captured.out - assert "DIFY_AGENT_STUB_API_BASE_URL" in captured.err - assert "DIFY_AGENT_STUB_AUTH_JWE" in captured.err - - -def test_cli_connect_help_routes_to_typer_help(capsys: pytest.CaptureFixture[str]) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["connect", "--help"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert "Establish one Agent Stub connection" in captured.out - assert "--json" in captured.out - assert "╭" not in captured.out - assert "─" not in captured.out - - -def test_cli_config_push_is_not_a_valid_command(capsys: pytest.CaptureFixture[str]) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["config", "push"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "No such command 'push'" in captured.err - assert "╭" not in captured.err - assert "─" not in captured.err - - -def test_cli_config_help_lists_plural_groups_and_no_root_push(capsys: pytest.CaptureFixture[str]) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["config", "--help"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert "Usage: dify-agent config" in captured.out - assert "manifest" in captured.out - assert "files" in captured.out - assert "skills" in captured.out - assert "env" in captured.out - assert "note" in captured.out - - -@pytest.mark.parametrize("argv", [["config", "files", "--help"], ["config", "skills", "--help"]]) -def test_cli_plural_config_groups_expose_pull_push_and_delete( - capsys: pytest.CaptureFixture[str], - argv: list[str], -) -> None: - with pytest.raises(SystemExit) as exc_info: - main(argv) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert "pull" in captured.out - assert "push" in captured.out - assert "delete" in captured.out - - -def test_cli_config_skills_push_help_describes_skill_directory_requirements( - capsys: pytest.CaptureFixture[str], -) -> None: - with pytest.raises(SystemExit) as exc_info: - main(["config", "skills", "push", "--help"]) - - captured = capsys.readouterr() - normalized_output = " ".join(captured.out.split()) - assert exc_info.value.code == 0 - assert "Skill directory requirements:" in captured.out - assert "top-level SKILL.md" in captured.out - assert "UTF-8 Markdown" in captured.out - assert "YAML frontmatter matching this schema" in captured.out - assert "--- name: description: ---" in normalized_output - assert "Symlinked files are rejected" in captured.out - assert ".venv and node_modules should be manually cleared before push" in normalized_output - - -@pytest.mark.parametrize("argv", [["config", "file", "--help"], ["config", "skill", "--help"]]) -def test_cli_hidden_singular_alias_help_exposes_pull_only( - capsys: pytest.CaptureFixture[str], - argv: list[str], -) -> None: - with pytest.raises(SystemExit) as exc_info: - main(argv) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert "pull" in captured.out - assert "push" not in captured.out - assert "delete" not in captured.out - - -def test_cli_reports_invalid_agent_stub_api_base_url_environment_value( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub?x=1") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(SystemExit) as exc_info: - main(["connect"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "invalid DIFY_AGENT_STUB_API_BASE_URL" in captured.err - assert "query string or fragment" in captured.err - - -@pytest.mark.parametrize( - ("invalid_url", "expected_message"), - [ - ("not-a-url", "http, https, or grpc"), - ("ftp://agent.example.com/agent-stub", "http, https, or grpc"), - ("https:///agent-stub", "include a host"), - ("grpc://agent.example.com", "explicit port"), - ], -) -def test_cli_reports_structurally_invalid_agent_stub_api_base_url_environment_value( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], - invalid_url: str, - expected_message: str, -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", invalid_url) - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - with pytest.raises(SystemExit) as exc_info: - main(["connect"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "invalid DIFY_AGENT_STUB_API_BASE_URL" in captured.err - assert expected_message in captured.err - - -def test_cli_connect_accepts_grpc_agent_stub_api_base_url( - monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "grpc://agent.example.com:9091") - monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - - def fake_connect_from_environment(*, argv: list[str]) -> AgentStubConnectResponse: - assert argv == ["echo", "hello"] - return AgentStubConnectResponse(connection_id="conn-1", status="connected") - - _patch_cli_module(monkeypatch, "_agent_stub_module", connect_from_environment=fake_connect_from_environment) - - main(["connect", "echo", "hello"]) - - captured = capsys.readouterr() - assert captured.out.strip() == "connected conn-1" - - -def test_cli_config_manifest_omits_hash_fields( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_config_module", - manifest_from_environment=lambda: AgentStubConfigManifestResponse( - agent_id="agent-1", - config_version=AgentStubConfigVersionInfo(id="cfg-1", kind="build_draft", writable=True), - skills=AgentStubConfigSkillItemsResponse( - items=[ - AgentStubConfigSkillItem( - name="alpha", - description="Alpha skill.", - size=12, - hash="sha256:skill", - mime_type="application/zip", - ) - ] - ), - files=AgentStubConfigFileItemsResponse( - items=[ - AgentStubConfigFileItem( - name="guide.txt", - size=34, - hash="sha256:file", - mime_type="text/plain", - ) - ] - ), - env_keys=["RUNTIME_KEY"], - note="Runtime note.", - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["config", "manifest"]) - - captured = capsys.readouterr() - payload = json.loads(captured.out) - assert exc_info.value.code == 0 - assert payload["skills"] == { - "items": [ - { - "name": "alpha", - "description": "Alpha skill.", - "size": 12, - "mime_type": "application/zip", - } - ] - } - assert payload["files"] == { - "items": [ - { - "name": "guide.txt", - "size": 34, - "mime_type": "text/plain", - } - ] - } - - -@pytest.mark.parametrize( - ("argv", "helper_name", "expected_kwargs"), - [ - (["config", "note", "push"], "push_config_note_from_environment", {"local_path": None}), - ( - ["config", "note", "push", "/tmp/note.md"], - "push_config_note_from_environment", - {"local_path": "/tmp/note.md"}, - ), - (["config", "env", "push", "/tmp/.env"], "push_config_env_from_environment", {"local_path": "/tmp/.env"}), - ( - ["config", "files", "push", "/tmp/guide.txt"], - "push_config_files_from_environment", - {"paths": ["/tmp/guide.txt"]}, - ), - ( - ["config", "files", "delete", "old.txt", "legacy.txt"], - "delete_config_files_from_environment", - {"names": ["old.txt", "legacy.txt"]}, - ), - ( - ["config", "skills", "push", "/tmp/alpha", "/tmp/beta"], - "push_config_skills_from_environment", - {"paths": ["/tmp/alpha", "/tmp/beta"]}, - ), - ( - ["config", "skills", "delete", "alpha", "beta"], - "delete_config_skills_from_environment", - {"names": ["alpha", "beta"]}, - ), - ], -) -def test_cli_config_mutation_commands_forward_and_print_manifest_json( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], - argv: list[str], - helper_name: str, - expected_kwargs: dict[str, object], -) -> None: - captured_kwargs: dict[str, object] = {} - - def fake_helper(**kwargs): - captured_kwargs.update(kwargs) - return _config_manifest_response() - - _patch_cli_module(monkeypatch, "_config_module", **{helper_name: fake_helper}) - - with pytest.raises(SystemExit) as exc_info: - main(argv) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out) == json.loads(_config_manifest_response().model_dump_json()) - assert captured_kwargs == expected_kwargs - - -@pytest.mark.parametrize( - ("argv", "helper_name"), - [ - (["config", "skills", "pull", "alpha", "--json"], "pull_config_skills_from_environment"), - (["config", "skill", "pull", "alpha", "--json"], "pull_config_skills_from_environment"), - (["config", "files", "pull", "guide.txt", "--json"], "pull_config_files_from_environment"), - (["config", "file", "pull", "guide.txt", "--json"], "pull_config_files_from_environment"), - ], -) -def test_cli_config_pull_commands_support_plural_and_hidden_singular_aliases( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], - argv: list[str], - helper_name: str, -) -> None: - if "skills" in helper_name: - expected_json = { - "items": [ - { - "name": "alpha", - "archive_path": "/tmp/alpha.zip", - "directory_path": "/tmp/alpha", - "skill_md": "# Alpha\n", - } - ] - } - response = type( - "Response", - (), - {"model_dump_json": lambda self: json.dumps(expected_json)}, - )() - expected_kwargs = {"names": ["alpha"], "local_dir": None} - else: - expected_json = {"items": [{"name": "guide.txt", "path": "/tmp/guide.txt"}]} - response = type( - "Response", - (), - {"model_dump_json": lambda self: json.dumps(expected_json)}, - )() - expected_kwargs = {"names": ["guide.txt"], "local_dir": None} - - captured_kwargs: dict[str, object] = {} - - def fake_helper(*, names, local_dir): - captured_kwargs["names"] = names - captured_kwargs["local_dir"] = local_dir - return response - - _patch_cli_module(monkeypatch, "_config_module", **{helper_name: fake_helper}) - - with pytest.raises(SystemExit) as exc_info: - main(argv) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out) == expected_json - assert captured_kwargs == expected_kwargs - - -@pytest.mark.parametrize( - "argv", - [ - ["config", "skill", "push", "/tmp/alpha"], - ["config", "skill", "delete", "alpha"], - ["config", "file", "push", "/tmp/guide.txt"], - ["config", "file", "delete", "guide.txt"], - ], -) -def test_cli_hidden_singular_aliases_do_not_expose_mutation_commands( - capsys: pytest.CaptureFixture[str], - argv: list[str], -) -> None: - with pytest.raises(SystemExit) as exc_info: - main(argv) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert "No such command" in captured.err - - -def test_cli_file_upload_prints_uploaded_tool_file_json( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_files_module", - upload_file_from_environment=lambda *, path: type( - "Response", - (), - { - "model_dump_json": lambda self: json.dumps( - { - "transfer_method": "tool_file", - "reference": _reference(Path(path).name), - "download_url": f"https://files.example.com/{Path(path).name}", - } - ) - }, - )(), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["file", "upload", "/tmp/report.pdf"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out) == { - "transfer_method": "tool_file", - "reference": _reference("report.pdf"), - "download_url": "https://files.example.com/report.pdf", - } - - -def test_cli_file_upload_exits_non_zero_without_partial_json_when_download_lookup_fails( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_files_module", - upload_file_from_environment=lambda *, path: (_ for _ in ()).throw( - AgentStubTransferError("signed file download request failed") - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["file", "upload", "/tmp/report.pdf"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 1 - assert captured.out == "" - assert "signed file download request failed" in captured.err - - -def test_cli_file_download_prints_saved_path( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_files_module", - download_file_from_environment=lambda **_kwargs: type("Response", (), {"path": Path("/tmp/report.pdf")})(), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["file", "download", "tool_file", _reference("tool-file-1"), "--to", "/tmp"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured.out.strip() == "/tmp/report.pdf" - - -def test_cli_file_download_rejects_mapping_option( - monkeypatch: pytest.MonkeyPatch, -) -> None: - called = False - - def fake_download_file_from_environment(**_kwargs): - nonlocal called - called = True - return type("Response", (), {"path": Path("/tmp/inputs/report.pdf")})() - - _patch_cli_module( - monkeypatch, - "_files_module", - download_file_from_environment=fake_download_file_from_environment, - ) - - with pytest.raises(SystemExit) as exc_info: - main( - [ - "file", - "download", - "--mapping", - json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), - "--to", - "/tmp/inputs", - ] - ) - - assert exc_info.value.code == 2 - assert called is False - - -def test_cli_file_download_rejects_legacy_positional_directory( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - called = False - - def fake_download_file_from_environment(**_kwargs): - nonlocal called - called = True - return type("Response", (), {"path": Path("/tmp/report.pdf")})() - - _patch_cli_module( - monkeypatch, - "_files_module", - download_file_from_environment=fake_download_file_from_environment, - ) - - with pytest.raises(SystemExit) as exc_info: - main(["file", "download", "tool_file", _reference("tool-file-1"), "/tmp"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert called is False - assert "/tmp" in captured.err - - -def test_cli_drive_list_prints_manifest_json( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_drive_module", - list_drive_manifest_from_environment=lambda *, prefix: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key=prefix + "example/SKILL.md", - size=12, - hash="sha256:abc", - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - ) - ] - ), - format_drive_manifest=lambda response: ( - f"{response.items[0].size}\t{response.items[0].mime_type}\t{response.items[0].hash or '-'}\t" - f"{response.items[0].key}" - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "list", "skills/", "--json"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out)["items"][0]["key"] == "skills/example/SKILL.md" - - -def test_cli_drive_list_prints_human_readable_listing( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_drive_module", - list_drive_manifest_from_environment=lambda *, prefix: AgentStubDriveManifestResponse( - items=[ - AgentStubDriveItem( - key=f"{prefix}example/SKILL.md", - size=12, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id="tool-file-1", - ) - ] - ), - format_drive_manifest=lambda response: ( - f"{response.items[0].size}\t{response.items[0].mime_type}\t{response.items[0].hash or '-'}\t" - f"{response.items[0].key}" - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "list", "skills/"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured.out.strip() == "12\ttext/markdown\t-\tskills/example/SKILL.md" - - -def test_cli_drive_pull_prints_downloaded_paths( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_drive_module", - pull_drive_from_environment=lambda *, targets, local_base: DrivePullResult( - items=[ - DrivePullResult.Item( - key=f"{targets[0]}/SKILL.md", local_path=str(Path(local_base) / targets[0] / "SKILL.md") - ), - DrivePullResult.Item( - key=f"{targets[0]}/helper.py", local_path=str(Path(local_base) / targets[0] / "helper.py") - ), - ] - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/example", "--to", "/tmp/drive"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured.out.strip().splitlines() == [ - "/tmp/drive/skills/example/SKILL.md", - "/tmp/drive/skills/example/helper.py", - ] - - -def test_cli_drive_pull_prints_json_result( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_drive_module", - pull_drive_from_environment=lambda *, targets, local_base: DrivePullResult( - items=[ - DrivePullResult.Item(key="files/a.txt", local_path=f"{local_base}/files/a.txt"), - DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=f"{local_base}/skills/foo/SKILL.md"), - ] - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "files/a.txt", "--to", "/tmp/drive", "--json"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out) == { - "items": [ - {"key": "files/a.txt", "local_path": "/tmp/drive/files/a.txt"}, - {"key": "skills/foo/SKILL.md", "local_path": "/tmp/drive/skills/foo/SKILL.md"}, - ] - } - - -def test_cli_drive_pull_forwards_multiple_targets( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - captured_kwargs: dict[str, object] = {} - - def fake_pull_drive_from_environment(*, targets, local_base): - captured_kwargs["targets"] = targets - captured_kwargs["local_base"] = local_base - return DrivePullResult( - items=[ - DrivePullResult.Item( - key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md") - ) - ] - ) - - _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/foo", "files/a.txt", "--to", "/tmp/drive"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo", "files/a.txt"], "local_base": "/tmp/drive"} - assert captured.out.strip() == "/tmp/drive/skills/foo/SKILL.md" - - -def test_cli_drive_pull_uses_environment_drive_base_default( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.setenv("DIFY_AGENT_STUB_DRIVE_BASE", "/env/drive") - captured_kwargs: dict[str, object] = {} - - def fake_pull_drive_from_environment(*, targets, local_base): - captured_kwargs["targets"] = targets - captured_kwargs["local_base"] = local_base - return DrivePullResult( - items=[ - DrivePullResult.Item( - key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md") - ) - ] - ) - - _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/foo"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo"], "local_base": "/env/drive"} - assert captured.out.strip() == "/env/drive/skills/foo/SKILL.md" - - -def test_cli_drive_pull_keeps_historical_drive_base_when_env_is_missing( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - monkeypatch.delenv("DIFY_AGENT_STUB_DRIVE_BASE", raising=False) - captured_kwargs: dict[str, object] = {} - - def fake_pull_drive_from_environment(*, targets, local_base): - captured_kwargs["targets"] = targets - captured_kwargs["local_base"] = local_base - return DrivePullResult( - items=[ - DrivePullResult.Item( - key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md") - ) - ] - ) - - _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/foo"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo"], "local_base": "/mnt/drive"} - assert captured.out.strip() == "/mnt/drive/skills/foo/SKILL.md" - - -def test_cli_drive_pull_without_targets_pulls_whole_visible_drive( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - captured_kwargs: dict[str, object] = {} - - def fake_pull_drive_from_environment(*, targets, local_base): - captured_kwargs["targets"] = targets - captured_kwargs["local_base"] = local_base - return DrivePullResult( - items=[DrivePullResult.Item(key="files/a.txt", local_path=str(Path(local_base) / "files" / "a.txt"))] - ) - - _patch_cli_module(monkeypatch, "_drive_module", pull_drive_from_environment=fake_pull_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "--to", "/tmp/drive"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": None, "local_base": "/tmp/drive"} - assert captured.out.strip() == "/tmp/drive/files/a.txt" - - -def test_cli_drive_push_prints_commit_json( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - _patch_cli_module( - monkeypatch, - "_drive_module", - push_drive_from_environment=lambda *, local_path, drive_path, kind: AgentStubDriveCommitResponse( - items=[ - AgentStubDriveItem( - key=drive_path, - size=12, - hash=None, - mime_type="text/markdown", - file_kind="tool_file", - file_id=Path(local_path).name, - value_owned_by_drive=kind != "dir", - ) - ] - ), - ) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "push", "/tmp/report.md", "skills/example/SKILL.md"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out)["items"][0]["key"] == "skills/example/SKILL.md" - - -def test_cli_drive_push_forwards_kind( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - captured_kwargs: dict[str, object] = {} - - def fake_push_drive_from_environment(*, local_path, drive_path, kind): - captured_kwargs["local_path"] = local_path - captured_kwargs["drive_path"] = drive_path - captured_kwargs["kind"] = kind - return AgentStubDriveCommitResponse(items=[]) - - _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "push", "/tmp/skill", "skills/example", "--kind", "skill"]) - - capsys.readouterr() - assert exc_info.value.code == 0 - assert captured_kwargs == { - "local_path": "/tmp/skill", - "drive_path": "skills/example", - "kind": "skill", - } - - -def test_cli_drive_push_accepts_json_flag( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - captured_kwargs: dict[str, object] = {} - - def fake_push_drive_from_environment(*, local_path, drive_path, kind): - captured_kwargs["local_path"] = local_path - captured_kwargs["drive_path"] = drive_path - captured_kwargs["kind"] = kind - return AgentStubDriveCommitResponse(items=[]) - - _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "push", "/tmp/report.md", "files/report.md", "--json"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 0 - assert json.loads(captured.out) == {"items": []} - assert captured_kwargs == { - "local_path": "/tmp/report.md", - "drive_path": "files/report.md", - "kind": None, - } - - -def test_cli_drive_push_rejects_recursive_option( - monkeypatch: pytest.MonkeyPatch, - capsys: pytest.CaptureFixture[str], -) -> None: - called = False - - def fake_push_drive_from_environment(*, local_path, drive_path, kind): - nonlocal called - called = True - return AgentStubDriveCommitResponse(items=[]) - - _patch_cli_module(monkeypatch, "_drive_module", push_drive_from_environment=fake_push_drive_from_environment) - - with pytest.raises(SystemExit) as exc_info: - main(["drive", "push", "/tmp/dir", "files/dir", "--recursive"]) - - captured = capsys.readouterr() - assert exc_info.value.code == 2 - assert called is False - assert "--recursive" in captured.err diff --git a/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_client.py b/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_client.py deleted file mode 100644 index be4fb1def39..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_client.py +++ /dev/null @@ -1,667 +0,0 @@ -from __future__ import annotations - -import base64 -import builtins -import json -from types import SimpleNamespace - -import httpx -import pytest - -from dify_agent.agent_stub.client._agent_stub import ( - connect_agent_stub_sync, - download_file_bytes_from_signed_url_sync, - request_agent_stub_drive_commit_sync, - request_agent_stub_drive_manifest_sync, - request_agent_stub_file_download_sync, - request_agent_stub_file_upload_sync, - upload_file_to_signed_url_sync, -) -from dify_agent.agent_stub.client._errors import ( - AgentStubClientError, - AgentStubGRPCError, - AgentStubHTTPError, - AgentStubMissingGRPCDependencyError, - AgentStubTransferError, - AgentStubValidationError, -) -from dify_agent.agent_stub.protocol.agent_stub import ( - AgentStubDriveCommitItem, - AgentStubDriveCommitRequest, - AgentStubDriveFileRef, - AgentStubFileMapping, -) - - -def _reference(record_id: str) -> str: - payload = base64.urlsafe_b64encode(json.dumps({"record_id": record_id}, separators=(",", ":")).encode()).decode() - return f"dify-file-ref:{payload}" - - -def test_connect_agent_stub_sync_posts_connections_request_with_authorization() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/connections" - assert request.headers["Authorization"] == "Bearer test-jwe" - assert json.loads(request.content) == { - "protocol_version": 1, - "argv": ["connect", "--", "echo", "hello"], - "metadata": {}, - } - return httpx.Response(200, json={"connection_id": "conn-1", "status": "connected"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = connect_agent_stub_sync( - url="https://agent.example.com/agent-stub/", - auth_jwe="test-jwe", - argv=["connect", "--", "echo", "hello"], - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.connection_id == "conn-1" - assert response.status == "connected" - - -def test_connect_agent_stub_sync_rejects_invalid_base_url() -> None: - with pytest.raises( - AgentStubValidationError, match="invalid DIFY_AGENT_STUB_API_BASE_URL|invalid Agent Stub base URL" - ): - _ = connect_agent_stub_sync( - url="https://agent.example.com/agent-stub?x=1", - auth_jwe="test-jwe", - argv=["connect"], - ) - - -def test_connect_agent_stub_sync_maps_non_2xx_response_to_http_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(401, json={"detail": "invalid token"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubHTTPError, match="401") as exc_info: - _ = connect_agent_stub_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - argv=["connect"], - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert exc_info.value.status_code == 401 - assert exc_info.value.detail == "invalid token" - - -def test_connect_agent_stub_sync_maps_malformed_json_response_to_client_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, text="not-json", headers={"Content-Type": "application/json"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubClientError, match="invalid JSON"): - _ = connect_agent_stub_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - argv=["connect"], - sync_http_client=http_client, - ) - finally: - http_client.close() - - -def test_connect_agent_stub_sync_maps_schema_invalid_success_payload_to_validation_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json={"connection_id": 123, "status": "unexpected"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubValidationError, match="invalid Agent Stub connection response"): - _ = connect_agent_stub_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - argv=["connect"], - sync_http_client=http_client, - ) - finally: - http_client.close() - - -def test_request_agent_stub_file_upload_sync_posts_upload_request() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/files/upload-request" - assert json.loads(request.content) == {"filename": "report.pdf", "mimetype": "application/pdf"} - return httpx.Response(200, json={"upload_url": "https://files.example.com/upload"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_file_upload_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - filename="report.pdf", - mimetype="application/pdf", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.upload_url == "https://files.example.com/upload" - - -def test_request_agent_stub_file_download_sync_posts_download_request() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/files/download-request" - assert json.loads(request.content) == { - "file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")}, - } - return httpx.Response( - 200, - json={ - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 123, - "download_url": "https://files.example.com/download", - }, - ) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_file_download_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")), - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.download_url == "https://files.example.com/download" - - -def test_request_agent_stub_file_download_sync_posts_internal_download_request() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/files/download-request" - assert json.loads(request.content) == { - "file": {"transfer_method": "tool_file", "reference": _reference("tool-file-1")}, - "for_external": False, - } - return httpx.Response( - 200, - json={ - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 123, - "download_url": "http://internal-files/report.pdf", - }, - ) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_file_download_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")), - for_external=False, - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.download_url == "http://internal-files/report.pdf" - - -def test_request_agent_stub_drive_manifest_sync_gets_manifest_request() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == ( - "https://agent.example.com/agent-stub/drive/manifest?prefix=skills%2F&include_download_url=true" - ) - assert request.headers["Authorization"] == "Bearer test-jwe" - return httpx.Response( - 200, - json={ - "items": [ - { - "key": "skills/example/SKILL.md", - "name": "SKILL.md", - "size": 12, - "hash": "sha256:abc", - "mime_type": "text/markdown", - "file_kind": "tool_file", - "file_id": "tool-file-1", - "created_at": 123, - } - ] - }, - ) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_drive_manifest_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - prefix="skills/", - include_download_url=True, - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.items[0].key == "skills/example/SKILL.md" - assert response.items[0].model_extra == {"name": "SKILL.md"} - - -def test_request_agent_stub_drive_commit_sync_posts_commit_request() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/drive/commit" - assert request.headers["Authorization"] == "Bearer test-jwe" - assert json.loads(request.content) == { - "items": [ - { - "key": "skills/example/SKILL.md", - "file_ref": {"kind": "tool_file", "id": "tool-file-1"}, - "value_owned_by_drive": True, - } - ] - } - return httpx.Response( - 200, - json={ - "items": [ - { - "key": "skills/example/SKILL.md", - "size": 12, - "mime_type": "text/markdown", - "file_kind": "tool_file", - "file_id": "tool-file-1", - "value_owned_by_drive": True, - } - ] - }, - ) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_drive_commit_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - request=AgentStubDriveCommitRequest( - items=[ - AgentStubDriveCommitItem( - key="skills/example/SKILL.md", - file_ref=AgentStubDriveFileRef(kind="tool_file", id="tool-file-1"), - ) - ] - ), - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.items[0].file_id == "tool-file-1" - - -def test_request_agent_stub_drive_manifest_sync_maps_invalid_json_to_client_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, text="not-json", headers={"Content-Type": "application/json"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubClientError, match="invalid JSON"): - _ = request_agent_stub_drive_manifest_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - prefix="", - include_download_url=False, - sync_http_client=http_client, - ) - finally: - http_client.close() - - -def test_request_agent_stub_drive_commit_sync_maps_non_2xx_to_http_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(403, json={"detail": "forbidden"}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubHTTPError, match="403") as exc_info: - _ = request_agent_stub_drive_commit_sync( - url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - request=AgentStubDriveCommitRequest( - items=[ - AgentStubDriveCommitItem( - key="skills/example/SKILL.md", - file_ref=AgentStubDriveFileRef(kind="tool_file", id="tool-file-1"), - ) - ] - ), - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert exc_info.value.detail == "forbidden" - - -def test_upload_file_to_signed_url_sync_posts_multipart_file() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://files.example.com/upload" - assert b"report.pdf" in request.content - return httpx.Response(201, json={"id": "tool-file-1", "name": "report.pdf", "size": 9}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - from io import BytesIO - - payload = upload_file_to_signed_url_sync( - upload_url="https://files.example.com/upload", - filename="report.pdf", - file_obj=BytesIO(b"contents!"), - mimetype="application/pdf", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert payload["id"] == "tool-file-1" - - -def test_download_file_bytes_from_signed_url_sync_returns_bytes() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, content=b"payload") - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - payload = download_file_bytes_from_signed_url_sync( - download_url="https://files.example.com/download", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert payload == b"payload" - - -def test_download_file_bytes_from_signed_url_sync_maps_timeout_to_transfer_error() -> None: - def handler(_request: httpx.Request) -> httpx.Response: - raise httpx.ReadTimeout("boom") - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(AgentStubTransferError, match="download timed out"): - _ = download_file_bytes_from_signed_url_sync( - download_url="https://files.example.com/download", - sync_http_client=http_client, - ) - finally: - http_client.close() - - -def test_connect_agent_stub_sync_dispatches_grpc_urls(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - - monkeypatch.setattr( - "dify_agent.agent_stub.client._agent_stub_grpc.connect_agent_stub_grpc_sync", - lambda **kwargs: ( - captured.update(kwargs) or type("Response", (), {"connection_id": "conn-1", "status": "connected"})() - ), - ) - - response = connect_agent_stub_sync(url="grpc://agent.example.com:9091", auth_jwe="token", argv=["connect"]) - - assert captured["url"] == "grpc://agent.example.com:9091" - assert response.connection_id == "conn-1" - - -def test_connect_agent_stub_sync_reports_missing_grpc_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: - original_import = builtins.__import__ - - def fake_import(name: str, globals=None, locals=None, fromlist=(), level: int = 0): - if name in {"grpclib.client", "grpclib.exceptions"}: - raise ImportError("grpclib is not installed") - return original_import(name, globals, locals, fromlist, level) - - monkeypatch.setattr(builtins, "__import__", fake_import) - - with pytest.raises(AgentStubMissingGRPCDependencyError, match=r"optional dify-agent\[grpc\] dependencies"): - _ = connect_agent_stub_sync(url="grpc://agent.example.com:9091", auth_jwe="token", argv=["connect"]) - - -def test_request_agent_stub_file_upload_sync_dispatches_grpc_urls(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - - monkeypatch.setattr( - "dify_agent.agent_stub.client._agent_stub_grpc.request_agent_stub_file_upload_grpc_sync", - lambda **kwargs: ( - captured.update(kwargs) or type("Response", (), {"upload_url": "https://files.example.com/upload"})() - ), - ) - - response = request_agent_stub_file_upload_sync( - url="grpc://agent.example.com:9091", - auth_jwe="token", - filename="report.pdf", - mimetype="application/pdf", - ) - - assert captured == { - "url": "grpc://agent.example.com:9091", - "auth_jwe": "token", - "filename": "report.pdf", - "mimetype": "application/pdf", - "timeout": 30.0, - } - assert response.upload_url == "https://files.example.com/upload" - - -def test_request_agent_stub_file_download_sync_dispatches_grpc_urls(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, object] = {} - file_mapping = AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")) - - monkeypatch.setattr( - "dify_agent.agent_stub.client._agent_stub_grpc.request_agent_stub_file_download_grpc_sync", - lambda **kwargs: ( - captured.update(kwargs) - or type( - "Response", - (), - { - "filename": "report.pdf", - "mime_type": "application/pdf", - "size": 123, - "download_url": "https://files.example.com/download", - }, - )() - ), - ) - - response = request_agent_stub_file_download_sync( - url="grpc://agent.example.com:9091", - auth_jwe="token", - file=file_mapping, - for_external=False, - ) - - assert captured == { - "url": "grpc://agent.example.com:9091", - "auth_jwe": "token", - "file": file_mapping, - "for_external": False, - "timeout": 30.0, - } - assert response.download_url == "https://files.example.com/download" - - -def test_request_agent_stub_drive_manifest_sync_rejects_grpc_urls() -> None: - with pytest.raises(AgentStubValidationError, match="require an HTTP Agent Stub URL"): - _ = request_agent_stub_drive_manifest_sync( - url="grpc://agent.example.com:9091", - auth_jwe="token", - prefix="skills/", - include_download_url=False, - ) - - -def test_request_agent_stub_drive_commit_sync_rejects_grpc_urls() -> None: - with pytest.raises(AgentStubValidationError, match="require an HTTP Agent Stub URL"): - _ = request_agent_stub_drive_commit_sync( - url="grpc://agent.example.com:9091", - auth_jwe="token", - request=AgentStubDriveCommitRequest( - items=[ - AgentStubDriveCommitItem( - key="skills/example/SKILL.md", - file_ref=AgentStubDriveFileRef(kind="tool_file", id="tool-file-1"), - ) - ] - ), - ) - - -def test_request_agent_stub_file_upload_grpc_sync_attaches_bearer_metadata(monkeypatch: pytest.MonkeyPatch) -> None: - import dify_agent.agent_stub.client._agent_stub_grpc as grpc_module - - captured: dict[str, object] = {} - - class FakeChannel: - def __init__(self, *, host: str, port: int, ssl: bool) -> None: - captured.update(host=host, port=port, ssl=ssl) - - def close(self) -> None: - captured["closed"] = True - - class FakeMethod: - async def __call__(self, request, *, metadata, timeout): - captured.update(request=request, metadata=metadata, timeout=timeout) - return {"upload_url": "https://files.example.com/upload"} - - class FakeStub: - def __init__(self, channel) -> None: - captured["channel"] = channel - self.CreateFileUploadRequest = FakeMethod() - - class FakeGRPCError(Exception): - def __init__(self, status, message: str) -> None: - self.status = status - self.message = message - super().__init__(message) - - class FakeStreamTerminatedError(Exception): - pass - - monkeypatch.setattr( - grpc_module, - "_require_runtime", - lambda: SimpleNamespace( - Channel=FakeChannel, - AgentStubServiceStub=FakeStub, - agent_stub_pb2=object(), - GRPCError=FakeGRPCError, - StreamTerminatedError=FakeStreamTerminatedError, - ), - ) - monkeypatch.setattr( - grpc_module, - "_require_conversions", - lambda: SimpleNamespace( - proto_file_upload_request=lambda _pb2, *, filename, mimetype: { - "filename": filename, - "mimetype": mimetype, - }, - file_upload_response_from_proto=lambda response: type( - "Response", - (), - {"upload_url": response["upload_url"]}, - )(), - ), - ) - - response = grpc_module.request_agent_stub_file_upload_grpc_sync( - url="grpc://agent.example.com:9091", - auth_jwe="test-jwe", - filename="report.pdf", - mimetype="application/pdf", - timeout=12.0, - ) - - assert captured["host"] == "agent.example.com" - assert captured["port"] == 9091 - assert captured["ssl"] is False - assert captured["request"] == {"filename": "report.pdf", "mimetype": "application/pdf"} - assert captured["metadata"] == (("authorization", "Bearer test-jwe"),) - assert captured["timeout"] == 12.0 - assert captured["closed"] is True - assert response.upload_url == "https://files.example.com/upload" - - -def test_request_agent_stub_file_download_grpc_sync_maps_grpc_errors(monkeypatch: pytest.MonkeyPatch) -> None: - import dify_agent.agent_stub.client._agent_stub_grpc as grpc_module - - class FakeChannel: - def __init__(self, *, host: str, port: int, ssl: bool) -> None: - del host, port, ssl - - def close(self) -> None: - return None - - class FakeMethod: - async def __call__(self, request, *, metadata, timeout): - del request, metadata, timeout - raise FakeGRPCError(SimpleNamespace(name="RESOURCE_EXHAUSTED"), "rate limited") - - class FakeStub: - def __init__(self, channel) -> None: - del channel - self.CreateFileDownloadRequest = FakeMethod() - - class FakeGRPCError(Exception): - def __init__(self, status, message: str) -> None: - self.status = status - self.message = message - super().__init__(message) - - class FakeStreamTerminatedError(Exception): - pass - - monkeypatch.setattr( - grpc_module, - "_require_runtime", - lambda: SimpleNamespace( - Channel=FakeChannel, - AgentStubServiceStub=FakeStub, - agent_stub_pb2=object(), - GRPCError=FakeGRPCError, - StreamTerminatedError=FakeStreamTerminatedError, - ), - ) - monkeypatch.setattr( - grpc_module, - "_require_conversions", - lambda: SimpleNamespace( - proto_file_download_request=lambda _pb2, *, file, for_external: { - "file": file, - "for_external": for_external, - }, - file_download_response_from_proto=lambda response: response, - ), - ) - - with pytest.raises(AgentStubGRPCError, match="RESOURCE_EXHAUSTED") as exc_info: - _ = grpc_module.request_agent_stub_file_download_grpc_sync( - url="grpc://agent.example.com:9091", - auth_jwe="test-jwe", - file=AgentStubFileMapping(transfer_method="tool_file", reference=_reference("tool-file-1")), - ) - - assert exc_info.value.status == "RESOURCE_EXHAUSTED" - assert exc_info.value.detail == "rate limited" diff --git a/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_config_client.py b/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_config_client.py deleted file mode 100644 index ed2fbf5a678..00000000000 --- a/dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_config_client.py +++ /dev/null @@ -1,310 +0,0 @@ -from __future__ import annotations - -import json - -import httpx -import pytest - -from dify_agent.agent_stub.client._agent_stub import ( - request_agent_stub_config_env_update_sync, - request_agent_stub_config_file_pull_sync, - request_agent_stub_config_manifest_sync, - request_agent_stub_config_note_update_sync, - request_agent_stub_config_push_sync, - request_agent_stub_config_skill_inspect_sync, - request_agent_stub_config_skill_pull_sync, -) -from dify_agent.agent_stub.client._agent_stub_http import ( - request_agent_stub_config_env_update_http_sync, - request_agent_stub_config_file_pull_http_sync, - request_agent_stub_config_manifest_http_sync, - request_agent_stub_config_note_update_http_sync, - request_agent_stub_config_push_http_sync, - request_agent_stub_config_skill_inspect_http_sync, - request_agent_stub_config_skill_pull_http_sync, -) -from dify_agent.agent_stub.client._errors import AgentStubClientError, AgentStubHTTPError, AgentStubValidationError -from dify_agent.agent_stub.protocol.agent_stub import AgentStubConfigPushRequest - - -def _manifest_payload() -> dict[str, object]: - return { - "agent_id": "agent-1", - "config_version": {"id": "cfg-1", "kind": "build_draft", "writable": True}, - "skills": {"items": [{"name": "alpha", "description": "Alpha skill"}]}, - "files": {"items": [{"name": "guide.txt"}]}, - "env_keys": ["API_KEY"], - "note": "Use carefully.", - } - - -def test_config_manifest_sync_normalizes_url_and_uses_http_transport() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == "https://agent.example.com/agent-stub/config/manifest" - assert request.headers["Authorization"] == "Bearer test-jwe" - return httpx.Response(200, json=_manifest_payload()) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_config_manifest_sync( - url="https://agent.example.com/agent-stub/", - auth_jwe="test-jwe", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.agent_id == "agent-1" - - -@pytest.mark.parametrize( - ("func", "kwargs"), - [ - (request_agent_stub_config_manifest_sync, {}), - (request_agent_stub_config_skill_pull_sync, {"name": "alpha"}), - (request_agent_stub_config_skill_inspect_sync, {"name": "alpha"}), - (request_agent_stub_config_file_pull_sync, {"name": "guide.txt"}), - (request_agent_stub_config_push_sync, {"request": AgentStubConfigPushRequest(note="hello")}), - (request_agent_stub_config_env_update_sync, {"env_text": "API_KEY=value\n"}), - (request_agent_stub_config_note_update_sync, {"note": "hello"}), - ], -) -def test_config_sync_entrypoints_reject_grpc_urls(func, kwargs: dict[str, object]) -> None: - with pytest.raises(AgentStubValidationError, match="require an HTTP Agent Stub URL"): - func( - url="grpc://agent.example.com:9091", - auth_jwe="test-jwe", - **kwargs, - ) - - -def test_request_agent_stub_config_manifest_http_sync_gets_manifest() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == "https://agent.example.com/agent-stub/config/manifest" - assert request.headers["Authorization"] == "Bearer test-jwe" - return httpx.Response(200, json=_manifest_payload()) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_config_manifest_http_sync( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.agent_id == "agent-1" - assert response.config_version.kind == "build_draft" - - -def test_request_agent_stub_config_skill_pull_http_sync_returns_binary_payload() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == "https://agent.example.com/agent-stub/config/skills/alpha/pull" - return httpx.Response(200, content=b"zip-bytes") - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - payload = request_agent_stub_config_skill_pull_http_sync( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - name="alpha", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert payload == b"zip-bytes" - - -def test_request_agent_stub_config_skill_inspect_http_sync_returns_json_dict() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == "https://agent.example.com/agent-stub/config/skills/alpha/inspect" - return httpx.Response(200, json={"name": "alpha", "files": ["SKILL.md"]}) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - payload = request_agent_stub_config_skill_inspect_http_sync( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - name="alpha", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert payload == {"name": "alpha", "files": ["SKILL.md"]} - - -def test_request_agent_stub_config_file_pull_http_sync_returns_binary_payload() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "GET" - assert str(request.url) == "https://agent.example.com/agent-stub/config/files/guide.txt/pull" - return httpx.Response(200, content=b"file-bytes") - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - payload = request_agent_stub_config_file_pull_http_sync( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - name="guide.txt", - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert payload == b"file-bytes" - - -def test_request_agent_stub_config_push_http_sync_posts_json_and_validates_response() -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert request.method == "POST" - assert str(request.url) == "https://agent.example.com/agent-stub/config/push" - assert json.loads(request.content) == {"note": "hello"} - return httpx.Response(200, json=_manifest_payload()) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - response = request_agent_stub_config_push_http_sync( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - request=AgentStubConfigPushRequest(note="hello"), - sync_http_client=http_client, - ) - finally: - http_client.close() - - assert response.note == "Use carefully." - - -@pytest.mark.parametrize( - ("func", "kwargs", "expected_path", "expected_body", "expected_response"), - [ - ( - request_agent_stub_config_env_update_http_sync, - {"env_text": "API_KEY=value\n"}, - "https://agent.example.com/agent-stub/config/env", - {"env_text": "API_KEY=value\n"}, - {"env_keys": ["API_KEY"]}, - ), - ( - request_agent_stub_config_note_update_http_sync, - {"note": "hello"}, - "https://agent.example.com/agent-stub/config/note", - {"note": "hello"}, - {"note": "hello"}, - ), - ], -) -def test_config_update_http_entrypoints_round_trip_json( - func, - kwargs: dict[str, object], - expected_path: str, - expected_body: dict[str, object], - expected_response: dict[str, object], -) -> None: - def handler(request: httpx.Request) -> httpx.Response: - assert str(request.url) == expected_path - assert json.loads(request.content) == expected_body - return httpx.Response(200, json=expected_response) - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - payload = func( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - sync_http_client=http_client, - **kwargs, - ) - finally: - http_client.close() - - assert payload == expected_response - - -@pytest.mark.parametrize( - ("func", "kwargs", "response", "expected_error", "expected_match"), - [ - ( - request_agent_stub_config_manifest_http_sync, - {}, - httpx.Response(401, json={"detail": "denied"}), - AgentStubHTTPError, - "401", - ), - ( - request_agent_stub_config_manifest_http_sync, - {}, - httpx.Response(200, text="not-json", headers={"Content-Type": "application/json"}), - AgentStubClientError, - "invalid JSON", - ), - ( - request_agent_stub_config_push_http_sync, - {"request": AgentStubConfigPushRequest(note="hello")}, - httpx.Response(200, json={"bad": "shape"}), - AgentStubValidationError, - "config push response", - ), - ( - request_agent_stub_config_skill_pull_http_sync, - {"name": "alpha"}, - httpx.Response(404, json={"detail": "missing"}), - AgentStubHTTPError, - "404", - ), - ( - request_agent_stub_config_skill_inspect_http_sync, - {"name": "alpha"}, - httpx.Response(200, json=["bad"]), - AgentStubValidationError, - "config skill inspect response", - ), - ( - request_agent_stub_config_file_pull_http_sync, - {"name": "guide.txt"}, - httpx.Response(404, json={"detail": "missing"}), - AgentStubHTTPError, - "404", - ), - ( - request_agent_stub_config_env_update_http_sync, - {"env_text": "API_KEY=value\n"}, - httpx.Response(200, json=["bad"]), - AgentStubValidationError, - "config env update response", - ), - ( - request_agent_stub_config_note_update_http_sync, - {"note": "hello"}, - httpx.Response(200, text="not-json", headers={"Content-Type": "application/json"}), - AgentStubClientError, - "invalid JSON", - ), - ], -) -def test_config_http_entrypoints_map_http_json_and_validation_errors( - func, - kwargs: dict[str, object], - response: httpx.Response, - expected_error: type[Exception], - expected_match: str, -) -> None: - def handler(_request: httpx.Request) -> httpx.Response: - return response - - http_client = httpx.Client(transport=httpx.MockTransport(handler)) - try: - with pytest.raises(expected_error, match=expected_match): - func( - base_url="https://agent.example.com/agent-stub", - auth_jwe="test-jwe", - sync_http_client=http_client, - **kwargs, - ) - finally: - http_client.close() diff --git a/dify-agent/tests/local/dify_agent/test_client_safe_exports.py b/dify-agent/tests/local/dify_agent/test_client_safe_exports.py index 28b21ff52c7..175eb61aa5e 100644 --- a/dify-agent/tests/local/dify_agent/test_client_safe_exports.py +++ b/dify-agent/tests/local/dify_agent/test_client_safe_exports.py @@ -68,9 +68,8 @@ def test_client_public_exports_work_with_default_dependencies_only(tmp_path: Pat dify_agent = importlib.import_module("dify_agent") client_module = importlib.import_module("dify_agent.client") protocol_module = importlib.import_module("dify_agent.protocol") - agent_stub_client_module = importlib.import_module("dify_agent.agent_stub.client") agent_stub_protocol_module = importlib.import_module("dify_agent.agent_stub.protocol") - agent_stub_cli_main_module = importlib.import_module("dify_agent.agent_stub.cli.main") + agent_cli_help_module = importlib.import_module("dify_agent.layers._agent_cli_help") agent_stub_shell_env_module = importlib.import_module("dify_agent.agent_stub.shell_env") shell_module = importlib.import_module("dify_agent.layers.shell") drive_module = importlib.import_module("dify_agent.layers.drive") @@ -89,11 +88,10 @@ def test_client_public_exports_work_with_default_dependencies_only(tmp_path: Pat assert protocol_module.CreateRunRequest is not None assert protocol_module.RunComposition is not None assert protocol_module.RunLayerSpec is not None - assert agent_stub_client_module.connect_agent_stub_sync is not None - assert agent_stub_client_module.request_agent_stub_config_manifest_sync is not None - assert agent_stub_client_module.request_agent_stub_drive_manifest_sync is not None assert agent_stub_protocol_module.AgentStubConnectRequest is not None - assert agent_stub_cli_main_module.main is not None + assert agent_cli_help_module.render_agent_stub_cli_help is not None + # Exercises the generated JSON snapshot to confirm it ships in the wheel. + assert "Usage:" in agent_cli_help_module.render_agent_stub_cli_help(("config",)) assert agent_stub_shell_env_module.build_shell_agent_stub_env is not None assert shell_module.DifyShellLayerConfig is not None assert drive_module.DifyDriveLayerConfig is not None @@ -102,18 +100,6 @@ def test_client_public_exports_work_with_default_dependencies_only(tmp_path: Pat assert ask_human_module.DifyAskHumanLayerConfig is not None assert output_module.DifyOutputLayerConfig is not None - grpc_error = importlib.import_module("dify_agent.agent_stub.client._errors").AgentStubMissingGRPCDependencyError - try: - agent_stub_client_module.connect_agent_stub_sync( - url="grpc://agent.example.com:9091", - auth_jwe="test-jwe", - argv=["connect"], - ) - except grpc_error: - pass - else: - raise AssertionError("grpc:// dispatch should fail with AgentStubMissingGRPCDependencyError without grpc extras") - unexpectedly_installed = [] for dependency_name in sorted(server_only_dependency_names): try: diff --git a/dify-agent/tests/local/dify_agent/test_import_boundaries.py b/dify-agent/tests/local/dify_agent/test_import_boundaries.py index 82fbda61667..961dab1b8d2 100644 --- a/dify-agent/tests/local/dify_agent/test_import_boundaries.py +++ b/dify-agent/tests/local/dify_agent/test_import_boundaries.py @@ -122,7 +122,6 @@ def test_protocol_and_dify_plugin_exports_do_not_import_server_only_modules() -> "pydantic_settings", "redis", "shellctl.client", - "shellctl.server", ], imports=[ "dify_agent.protocol", @@ -147,7 +146,7 @@ def test_protocol_and_dify_plugin_exports_do_not_import_server_only_modules() -> ) -def test_agent_stub_cli_main_import_is_client_safe() -> None: +def test_agent_cli_help_import_is_client_safe() -> None: _run_import_check( blocked_imports=[ "dify_agent.server", @@ -158,20 +157,17 @@ def test_agent_stub_cli_main_import_is_client_safe() -> None: "jwcrypto", "pydantic_settings", "redis", - "shellctl.server", ], imports=[ - "dify_agent.agent_stub.client", "dify_agent.agent_stub.protocol", - "dify_agent.agent_stub.cli.main", + "dify_agent.layers._agent_cli_help", "dify_agent.agent_stub.shell_env", "dify_agent.layers.shell.layer", "dify_agent.runtime.compositor_factory", ], assertions=[ - "assert hasattr(dify_agent_agent_stub_client, 'request_agent_stub_drive_manifest_sync')", "assert hasattr(dify_agent_agent_stub_protocol, 'AgentStubConnectRequest')", - "assert hasattr(dify_agent_agent_stub_cli_main, 'main')", + "assert hasattr(dify_agent_layers__agent_cli_help, 'render_agent_stub_cli_help')", "assert hasattr(dify_agent_agent_stub_shell_env, 'build_shell_agent_stub_env')", "assert hasattr(dify_agent_layers_shell_layer, 'DifyShellLayer')", "assert hasattr(dify_agent_runtime_compositor_factory, 'create_default_layer_providers')", @@ -228,7 +224,7 @@ def test_agent_stub_cli_main_import_is_client_safe() -> None: ) -def test_agent_stub_cli_help_render_does_not_load_server_modules() -> None: +def test_agent_cli_help_render_does_not_load_server_or_cli_modules() -> None: blocked_modules = [ "dify_agent.server", "dify_agent.agent_stub.server", @@ -238,26 +234,17 @@ def test_agent_stub_cli_help_render_does_not_load_server_modules() -> None: "jwcrypto", "pydantic_settings", "redis", - "shellctl.server", ] script = "\n".join( [ - "import click", "import importlib", - "import os", "import sys", - "from typer.main import get_command", f"blocked_modules = {blocked_modules!r}", - 'original_disable_plugins = os.environ.get("PYDANTIC_DISABLE_PLUGINS")', - 'original_disable_plugins_present = "PYDANTIC_DISABLE_PLUGINS" in os.environ', - 'module = importlib.import_module("dify_agent.agent_stub.cli.main")', - "command = get_command(module.app)", - "help_text = command.get_help(click.Context(command))", - 'assert "Forward shell-visible dify-agent commands" in help_text', - "if original_disable_plugins_present:", - ' assert os.environ.get("PYDANTIC_DISABLE_PLUGINS") == original_disable_plugins', - "else:", - ' assert "PYDANTIC_DISABLE_PLUGINS" not in os.environ', + 'module = importlib.import_module("dify_agent.layers._agent_cli_help")', + 'help_text = module.render_agent_stub_cli_help(("config",))', + 'assert "Inspect or update Agent Soul-backed config assets" in help_text', + 'assert "typer" not in sys.modules', + 'assert "click" not in sys.modules', "loaded_blocked = sorted(", " name", " for name in sys.modules", @@ -276,17 +263,13 @@ def test_shellctl_client_imports_do_not_import_server_modules() -> None: "fastapi", "sqlalchemy", "sqlmodel", - "shellctl.server.api", - "shellctl.server.service", - "shellctl.server.tmux", "uvicorn", ], - imports=["shellctl", "shellctl.client", "shellctl.shared", "shellctl.cli"], + imports=["shellctl", "shellctl.client", "shellctl.shared"], assertions=[ "assert hasattr(shellctl, 'ShellctlClient')", "assert hasattr(shellctl_client, 'ShellctlClient')", "assert hasattr(shellctl_shared, 'JobResult')", - "assert hasattr(shellctl_cli, 'cli')", ], ) diff --git a/dify-agent/tests/local/shellctl/golden_shellctl_sanitize.json b/dify-agent/tests/local/shellctl/golden_shellctl_sanitize.json deleted file mode 100644 index 10eba3678ef..00000000000 --- a/dify-agent/tests/local/shellctl/golden_shellctl_sanitize.json +++ /dev/null @@ -1,32 +0,0 @@ -[ - { - "name": "utf8_passthrough", - "chunks_hex": ["68656c6c6f20e4b896e7958cf09f99820a"], - "expected": "hello 世界🙂\n" - }, - { - "name": "split_utf8_chunks", - "chunks_hex": ["e4bda0e5a5bd20f09f", "99820a"], - "expected": "你好 🙂\n" - }, - { - "name": "invalid_utf8_replacement", - "chunks_hex": ["666f80ff6f0a"], - "expected": "fo��o\n" - }, - { - "name": "ansi_sequences_removed", - "chunks_hex": ["1b5b33316d7265641b5b306d1b5b4b0a"], - "expected": "red\n" - }, - { - "name": "carriage_return_progress", - "chunks_hex": ["30250d3530250d313030250a"], - "expected": "100%\n" - }, - { - "name": "crlf_normalized", - "chunks_hex": ["6f6b0d0a6e6578740d0a"], - "expected": "ok\nnext\n" - } -] diff --git a/dify-agent/tests/local/shellctl/test_shellctl_cli.py b/dify-agent/tests/local/shellctl/test_shellctl_cli.py deleted file mode 100644 index ac8eed51843..00000000000 --- a/dify-agent/tests/local/shellctl/test_shellctl_cli.py +++ /dev/null @@ -1,688 +0,0 @@ -from __future__ import annotations - -import importlib -import importlib.util -import json -import os -import subprocess -import sys -from pathlib import Path -from typing import ClassVar, TypeVar, cast - -import click -import httpx2 as httpx -import pytest -import typer.main -from typer.testing import CliRunner - -from shellctl.client import ShellctlClientError -from shellctl.shared.constants import DEFAULT_BASE_URL -from shellctl.shared.schemas import ( - DeleteJobResponse, - HealthResponse, - JobInfo, - JobResult, - JobStatusName, - JobStatusView, - TerminalSize, -) - -cli_module = importlib.import_module("shellctl.cli") -cli = cli_module.cli -runner = CliRunner() -ResultT = TypeVar("ResultT") - - -def _click_command(command_name: str) -> click.Command: - return cast(click.Group, typer.main.get_command(cli)).commands[command_name] - - -def _command_option_names(command_name: str) -> set[str]: - command = _click_command(command_name) - return { - option - for parameter in command.params - for option in getattr(parameter, "opts", []) # noqa: no-new-getattr click option introspection - } - - -def _command_option(command_name: str, option_name: str) -> click.Option: - command = _click_command(command_name) - for parameter in command.params: - if isinstance(parameter, click.Option) and option_name in parameter.opts: - return parameter - raise AssertionError(f"{option_name} not found on {command_name}") - - -class RecordingShellctlClient: - init_calls: ClassVar[list[dict[str, object]]] = [] - calls: ClassVar[list[tuple[str, tuple[object, ...], dict[str, object]]]] = [] - results: ClassVar[dict[str, object]] = {} - error: ClassVar[BaseException | None] = None - - def __init__( - self, - base_url: str, - *, - output_limit: int, - idle_flush_seconds: float, - token: str | None = None, - client: object | None = None, - transport: object | None = None, - ) -> None: - del client, transport - type(self).init_calls.append( - { - "base_url": base_url, - "output_limit": output_limit, - "idle_flush_seconds": idle_flush_seconds, - "token": token, - } - ) - - @classmethod - def reset(cls) -> None: - cls.init_calls = [] - cls.calls = [] - cls.results = {} - cls.error = None - - async def __aenter__(self) -> RecordingShellctlClient: - return self - - async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: - return None - - def _result(self, method: str, result_type: type[ResultT]) -> ResultT: - del result_type - error = type(self).error - if error is not None: - raise error - return cast(ResultT, type(self).results[method]) - - async def health(self) -> HealthResponse: - type(self).calls.append(("health", (), {})) - return self._result("health", HealthResponse) - - async def run( - self, - script: str, - *, - cwd: str | None = None, - env: dict[str, str] | None = None, - timeout: float, - terminal: TerminalSize | None = None, - ) -> JobResult: - type(self).calls.append( - ( - "run", - (script,), - { - "cwd": cwd, - "env": env, - "timeout": timeout, - "terminal": terminal, - }, - ) - ) - return self._result("run", JobResult) - - async def wait(self, job_id: str, *, offset: int, timeout: float) -> JobResult: - type(self).calls.append(("wait", (job_id,), {"offset": offset, "timeout": timeout})) - return self._result("wait", JobResult) - - async def status(self, job_id: str) -> JobStatusView: - type(self).calls.append(("status", (job_id,), {})) - return self._result("status", JobStatusView) - - async def list_jobs( - self, - *, - status: JobStatusName | None = None, - limit: int, - ) -> list[JobInfo]: - type(self).calls.append(("list_jobs", (), {"status": status, "limit": limit})) - return cast(list[JobInfo], self._result("list_jobs", list)) - - async def input( - self, - job_id: str, - text: str, - *, - offset: int, - timeout: float, - ) -> JobResult: - type(self).calls.append( - ( - "input", - (job_id, text), - {"offset": offset, "timeout": timeout}, - ) - ) - return self._result("input", JobResult) - - async def tail(self, job_id: str) -> JobResult: - type(self).calls.append(("tail", (job_id,), {})) - return self._result("tail", JobResult) - - async def terminate(self, job_id: str, grace_seconds: float) -> JobStatusView: - type(self).calls.append(("terminate", (job_id,), {"grace_seconds": grace_seconds})) - return self._result("terminate", JobStatusView) - - async def delete( - self, - job_id: str, - *, - force: bool = False, - grace_seconds: float | None = None, - ) -> DeleteJobResponse: - type(self).calls.append( - ( - "delete", - (job_id,), - {"force": force, "grace_seconds": grace_seconds}, - ) - ) - return self._result("delete", DeleteJobResponse) - - -@pytest.fixture -def patched_client( - monkeypatch: pytest.MonkeyPatch, -) -> type[RecordingShellctlClient]: - RecordingShellctlClient.reset() - monkeypatch.setattr(cli_module, "ShellctlClient", RecordingShellctlClient) - return RecordingShellctlClient - - -def _package_env() -> dict[str, str]: - package_root = Path(__file__).resolve().parents[3] - src_path = package_root / "src" - env = dict(os.environ) - current_pythonpath = env.get("PYTHONPATH") - env["PYTHONPATH"] = f"{src_path}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(src_path) - return env - - -def test_shellctl_help_lists_network_commands() -> None: - result = runner.invoke(cli, ["--help"]) - - assert result.exit_code == 0, result.stderr - assert "health" in result.stdout - assert "run" in result.stdout - assert "wait" in result.stdout - assert "status" in result.stdout - assert "list" in result.stdout - assert "input" in result.stdout - assert "tail" in result.stdout - assert "terminate" in result.stdout - assert "delete" in result.stdout - assert "serve" in result.stdout - - -def test_shellctl_run_and_serve_help_show_the_new_option_boundaries() -> None: - run_result = runner.invoke(cli, ["run", "--help"]) - serve_result = runner.invoke(cli, ["serve", "--help"]) - - assert run_result.exit_code == 0, run_result.stderr - assert "--base-url" in run_result.stdout - assert "--auth-token" in run_result.stdout - assert "--state-dir" not in run_result.stdout - assert "--runtime-dir" not in run_result.stdout - - assert serve_result.exit_code == 0, serve_result.stderr - assert "--state-dir" in serve_result.stdout - assert "--runtime-dir" in serve_result.stdout - assert "--auth-token" in serve_result.stdout - - -def test_shellctl_health_uses_sdk_health_and_ignores_auth_token_input( - monkeypatch: pytest.MonkeyPatch, - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.results["health"] = HealthResponse(status="ok") - monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "from-env") - - result = runner.invoke(cli, ["health", "--auth-token", "flag-token"]) - - assert result.exit_code == 0, result.stderr - assert json.loads(result.stdout) == {"status": "ok"} - assert patched_client.calls == [("health", (), {})] - assert patched_client.init_calls == [ - { - "base_url": DEFAULT_BASE_URL, - "output_limit": 8192, - "idle_flush_seconds": 0.5, - "token": None, - } - ] - - -def test_shellctl_run_builds_sdk_request_and_emits_json( - patched_client: type[RecordingShellctlClient], - tmp_path: Path, -) -> None: - patched_client.results["run"] = JobResult( - job_id="job-run", - done=False, - status=JobStatusName.RUNNING, - output_path="/tmp/output.log", - output="hello\n", - offset=6, - truncated=False, - ) - - result = runner.invoke( - cli, - [ - "run", - "printf hello\\n", - "--cwd", - str(tmp_path / "workspace"), - "--env", - "A=1", - "--env", - "EMPTY=", - "--timeout", - "12", - "--output-limit", - "4096", - "--idle-flush-seconds", - "0.25", - "--cols", - "90", - ], - ) - - assert result.exit_code == 0, result.stderr - assert json.loads(result.stdout) == { - "job_id": "job-run", - "done": False, - "status": "running", - "output_path": "/tmp/output.log", - "output": "hello\n", - "offset": 6, - "truncated": False, - } - assert patched_client.init_calls == [ - { - "base_url": DEFAULT_BASE_URL, - "output_limit": 4096, - "idle_flush_seconds": 0.25, - "token": None, - } - ] - assert patched_client.calls[0][0] == "run" - assert patched_client.calls[0][1] == ("printf hello\\n",) - assert patched_client.calls[0][2]["cwd"] == str(tmp_path / "workspace") - assert patched_client.calls[0][2]["env"] == {"A": "1", "EMPTY": ""} - assert patched_client.calls[0][2]["timeout"] == 12.0 - terminal = patched_client.calls[0][2]["terminal"] - assert isinstance(terminal, TerminalSize) - assert terminal.cols == 90 - assert terminal.rows == 80 - - -def test_shellctl_wait_and_input_require_offset() -> None: - assert _command_option("wait", "--offset").required is True - assert _command_option("input", "--offset").required is True - - -def test_shellctl_wait_and_input_map_requests( - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.results["wait"] = JobResult( - job_id="job-1", - done=False, - status=JobStatusName.RUNNING, - output_path="/tmp/wait.log", - output="chunk", - offset=5, - truncated=False, - ) - - wait_result = runner.invoke( - cli, - [ - "wait", - "job-1", - "--offset", - "3", - "--timeout", - "9", - "--output-limit", - "2048", - "--idle-flush-seconds", - "0.1", - ], - ) - - assert wait_result.exit_code == 0, wait_result.stderr - assert patched_client.calls[0] == ( - "wait", - ("job-1",), - {"offset": 3, "timeout": 9.0}, - ) - assert patched_client.init_calls[0]["output_limit"] == 2048 - assert patched_client.init_calls[0]["idle_flush_seconds"] == 0.1 - - RecordingShellctlClient.reset() - patched_client.results["input"] = JobResult( - job_id="job-1", - done=False, - status=JobStatusName.RUNNING, - output_path="/tmp/input.log", - output="reply", - offset=8, - truncated=False, - ) - - input_result = runner.invoke( - cli, - [ - "input", - "job-1", - "hello\n", - "--offset", - "5", - "--timeout", - "4", - "--output-limit", - "512", - "--idle-flush-seconds", - "0", - ], - ) - - assert input_result.exit_code == 0, input_result.stderr - assert patched_client.calls[0] == ( - "input", - ("job-1", "hello\n"), - {"offset": 5, "timeout": 4.0}, - ) - assert patched_client.init_calls[0]["output_limit"] == 512 - assert patched_client.init_calls[0]["idle_flush_seconds"] == 0.0 - - -def test_shellctl_list_tail_status_terminate_and_delete_map_arguments( - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.results["list_jobs"] = [ - JobInfo( - job_id="job-2", - status=JobStatusName.RUNNING, - created_at="2026-05-21T15:30:12Z", - ) - ] - - list_result = runner.invoke(cli, ["list", "--status", "running", "--limit", "5"]) - - assert list_result.exit_code == 0, list_result.stderr - assert json.loads(list_result.stdout) == [ - { - "job_id": "job-2", - "status": "running", - "created_at": "2026-05-21T15:30:12Z", - } - ] - assert patched_client.calls[0] == ( - "list_jobs", - (), - {"status": JobStatusName.RUNNING, "limit": 5}, - ) - - RecordingShellctlClient.reset() - patched_client.results["tail"] = JobResult( - job_id="job-2", - done=False, - status=JobStatusName.RUNNING, - output_path="/tmp/tail.log", - output="tail", - offset=4, - truncated=False, - ) - tail_result = runner.invoke(cli, ["tail", "job-2", "--output-limit", "16"]) - assert tail_result.exit_code == 0, tail_result.stderr - assert patched_client.calls[0] == ("tail", ("job-2",), {}) - assert patched_client.init_calls[0]["output_limit"] == 16 - - RecordingShellctlClient.reset() - patched_client.results["status"] = JobStatusView( - job_id="job-2", - status=JobStatusName.RUNNING, - done=False, - created_at="2026-05-21T15:30:12Z", - started_at="2026-05-21T15:30:13Z", - offset=4, - ) - status_result = runner.invoke(cli, ["status", "job-2"]) - assert status_result.exit_code == 0, status_result.stderr - assert patched_client.calls[0] == ("status", ("job-2",), {}) - - RecordingShellctlClient.reset() - patched_client.results["terminate"] = JobStatusView( - job_id="job-2", - status=JobStatusName.TERMINATED, - done=True, - created_at="2026-05-21T15:30:12Z", - started_at="2026-05-21T15:30:13Z", - ended_at="2026-05-21T15:30:18Z", - offset=4, - ) - terminate_result = runner.invoke( - cli, - ["terminate", "job-2", "--grace-seconds", "0.25"], - ) - assert terminate_result.exit_code == 0, terminate_result.stderr - assert patched_client.calls[0] == ( - "terminate", - ("job-2",), - {"grace_seconds": 0.25}, - ) - - RecordingShellctlClient.reset() - patched_client.results["delete"] = DeleteJobResponse(job_id="job-2") - delete_result = runner.invoke( - cli, - ["delete", "job-2", "--force", "--grace-seconds", "0.5"], - ) - assert delete_result.exit_code == 0, delete_result.stderr - assert patched_client.calls[0] == ( - "delete", - ("job-2",), - {"force": True, "grace_seconds": 0.5}, - ) - - -def test_shellctl_run_rejects_invalid_env_entry() -> None: - result = runner.invoke(cli, ["run", "printf bad", "--env", "MISSING_EQUALS"]) - - assert result.exit_code == 2 - assert "env entries must use NAME=VALUE format" in result.stderr - - -def test_shellctl_commands_render_sdk_errors_as_json_stderr( - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.error = ShellctlClientError( - 404, - "job_not_found", - "Job missing is not found", - ) - - result = runner.invoke(cli, ["status", "missing"]) - - assert result.exit_code == 1 - assert result.stdout == "" - assert json.loads(result.stderr) == { - "error": { - "code": "job_not_found", - "message": "Job missing is not found", - } - } - - -def test_shellctl_commands_render_transport_errors_as_json_stderr( - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.error = httpx.TransportError("connection refused") - - result = runner.invoke(cli, ["status", "missing"]) - - assert result.exit_code == 1 - assert result.stdout == "" - assert json.loads(result.stderr) == { - "error": { - "code": "connection_error", - "message": "connection refused", - } - } - - -def test_shellctl_commands_render_timeouts_as_json_stderr( - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.error = httpx.TimeoutException("slow server") - - result = runner.invoke(cli, ["status", "missing"]) - - assert result.exit_code == 1 - assert result.stdout == "" - assert json.loads(result.stderr) == { - "error": { - "code": "request_timeout", - "message": "request timed out", - } - } - - -def test_shellctl_base_url_and_auth_token_flags_override_environment( - monkeypatch: pytest.MonkeyPatch, - patched_client: type[RecordingShellctlClient], -) -> None: - patched_client.results["status"] = JobStatusView( - job_id="job-2", - status=JobStatusName.RUNNING, - done=False, - created_at="2026-05-21T15:30:12Z", - started_at="2026-05-21T15:30:13Z", - offset=4, - ) - monkeypatch.setenv("SHELLCTL_BASE_URL", "http://from-env:9999") - monkeypatch.setenv("SHELLCTL_AUTH_TOKEN", "from-env-token") - - result = runner.invoke( - cli, - [ - "status", - "job-2", - "--base-url", - "http://override:8765", - "--auth-token", - "flag-token", - ], - ) - - assert result.exit_code == 0, result.stderr - assert patched_client.init_calls == [ - { - "base_url": "http://override:8765", - "output_limit": 8192, - "idle_flush_seconds": 0.5, - "token": "flag-token", - } - ] - - -def test_shellctl_cli_controller_module_is_removed() -> None: - assert importlib.util.find_spec("shellctl.server.cli_controller") is None - - -def test_importing_shellctl_cli_for_run_help_skips_server_stack() -> None: - package_root = Path(__file__).resolve().parents[3] - command = """ -import json -import sys -from typer.testing import CliRunner -from shellctl.cli import cli - -result = CliRunner().invoke(cli, ["run", "--help"]) -print(json.dumps({ - "exit_code": result.exit_code, - "stdout": result.stdout, - "modules": { - "fastapi": "fastapi" in sys.modules, - "uvicorn": "uvicorn" in sys.modules, - "sqlalchemy": "sqlalchemy" in sys.modules, - "sqlmodel": "sqlmodel" in sys.modules, - "service": "shellctl.server.service" in sys.modules, - "api": "shellctl.server.api" in sys.modules, - "tmux": "shellctl.server.tmux" in sys.modules, - "shared_runtime": "shellctl.shared.runtime" in sys.modules, - "shared_output": "shellctl.shared.output" in sys.modules, - "shared_sanitize": "shellctl.shared.sanitize" in sys.modules, - }, -})) -""" - - result = subprocess.run( - [sys.executable, "-c", command], - capture_output=True, - text=True, - check=False, - cwd=package_root, - env=_package_env(), - ) - - assert result.returncode == 0, result.stderr - payload = json.loads(result.stdout) - assert payload["exit_code"] == 0 - assert "--base-url" in payload["stdout"] - assert payload["modules"] == { - "fastapi": False, - "uvicorn": False, - "sqlalchemy": False, - "sqlmodel": False, - "service": False, - "api": False, - "tmux": False, - "shared_runtime": False, - "shared_output": False, - "shared_sanitize": False, - } - - -def test_importing_server_serve_command_skips_cli_and_sdk_modules() -> None: - package_root = Path(__file__).resolve().parents[3] - command = """ -import json -import sys -from shellctl.server import serve_command - -print(json.dumps({ - "callable": callable(serve_command), - "modules": { - "cli": "shellctl.cli" in sys.modules, - "sdk": "shellctl.client.sdk" in sys.modules, - "client": "shellctl.client" in sys.modules, - }, -})) -""" - - result = subprocess.run( - [sys.executable, "-c", command], - capture_output=True, - text=True, - check=False, - cwd=package_root, - env=_package_env(), - ) - - assert result.returncode == 0, result.stderr - payload = json.loads(result.stdout) - assert payload["callable"] is True - assert payload["modules"] == { - "cli": False, - "sdk": False, - "client": False, - } diff --git a/dify-agent/tests/local/shellctl/test_shellctl_service.py b/dify-agent/tests/local/shellctl/test_shellctl_service.py deleted file mode 100644 index 6cfc30106ea..00000000000 --- a/dify-agent/tests/local/shellctl/test_shellctl_service.py +++ /dev/null @@ -1,1687 +0,0 @@ -from __future__ import annotations - -import importlib -import json -import os -import re -import shutil -import sqlite3 -import subprocess -import sys -from collections.abc import Awaitable, Callable -from pathlib import Path - -import anyio -import httpx2 as httpx -import pytest -from typer.testing import CliRunner - -import shellctl.server.service as server_service_module -from shellctl.server import ( - JobRow, - ShellctlConfig, - ShellctlServerError, - ShellctlService, - cli, - create_app, -) -from shellctl.server.artifacts import ( - JOB_ENV_FILENAME, - PIPE_DRAINED_FILENAME, - PIPE_ERROR_LOG_FILENAME, - PIPE_FAILED_FILENAME, - RUNNER_ENDED_AT_FILENAME, - RUNNER_EXIT_CODE_FILENAME, - job_env_path, - pipe_error_log_path, -) -from shellctl.server.tmux import TmuxController -from shellctl.shared import ( - DEFAULT_AUTH_TOKEN_ENV, - InputJobRequest, - JobStatusName, - JobStatusView, - RunJobRequest, - TerminalSize, - TerminateJobRequest, - WaitJobRequest, - job_pane_target, - job_session_name, -) - -server_serve_module = importlib.import_module("shellctl.server.serve") - - -class FakeTmuxController: - def __init__(self) -> None: - self.sessions: set[str] = set() - self.pipe_active: dict[str, bool | None] = {} - self.cleaned: list[str] = [] - self.touch_ready_on_enable = True - self.pipe_active_on_enable: bool | None = True - self.pipe_error_log_text: str | None = None - self.on_send_interrupt: Callable[[str], Awaitable[None]] | None = None - self.on_send_input: Callable[[str, str], Awaitable[None]] | None = None - - async def start_server(self) -> None: - return None - - async def list_sessions(self) -> set[str]: - return set(self.sessions) - - async def session_exists(self, session_name: str) -> bool: - return session_name in self.sessions - - async def is_output_pipe_active(self, *, job_id: str) -> bool | None: - return self.pipe_active.get(job_id) - - async def create_job_session(self, *, job_id: str, job_dir: Path, cwd: Path, terminal: TerminalSize) -> None: - del job_dir, cwd, terminal - self.sessions.add(job_session_name(job_id)) - self.pipe_active[job_id] = False - - async def enable_output_pipe(self, *, job_id: str, job_dir: Path, ready_file: Path) -> None: - self.pipe_active[job_id] = self.pipe_active_on_enable - if self.pipe_error_log_text is not None: - pipe_error_log_path(job_dir).write_text( - self.pipe_error_log_text, - encoding="utf-8", - ) - if self.touch_ready_on_enable: - ready_file.touch() - - async def send_input(self, *, job_id: str, text: str) -> None: - if self.on_send_input is not None: - await self.on_send_input(job_id, text) - - async def send_interrupt(self, *, job_id: str) -> None: - if self.on_send_interrupt is not None: - await self.on_send_interrupt(job_id) - - async def cleanup_session(self, *, job_id: str) -> None: - self.cleaned.append(job_id) - self.sessions.discard(job_session_name(job_id)) - self.pipe_active.pop(job_id, None) - - -async def _create_service( - tmp_path: Path, *, auth_token: str | None = None -) -> tuple[ShellctlService, FakeTmuxController]: - fake_tmux = FakeTmuxController() - service = ShellctlService( - ShellctlConfig( - auth_token=auth_token, - state_dir=tmp_path / "state", - runtime_dir=tmp_path / "run", - ), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - return service, fake_tmux - - -async def _create_real_service( - tmp_path: Path, - *, - runner_exit_command: tuple[str, ...] | None = None, - sanitize_pty_command: tuple[str, ...] | None = None, -) -> ShellctlService: - defaults = ShellctlConfig( - state_dir=tmp_path / "default-state", - runtime_dir=tmp_path / "default-run", - ) - config = ShellctlConfig( - state_dir=tmp_path / "state", - runtime_dir=tmp_path / "run", - runner_exit_command=runner_exit_command or defaults.runner_exit_command, - sanitize_pty_command=sanitize_pty_command or defaults.sanitize_pty_command, - ) - service = ShellctlService(config) - await service.initialize() - return service - - -def _write_delayed_sanitize_wrapper( - tmp_path: Path, - *, - delay_seconds: float, - sanitize_exit_code: int | None = None, -) -> tuple[str, ...]: - wrapper_path = tmp_path / "delayed-sanitize-wrapper.py" - wrapper_path.write_text( - "\n".join( - [ - "from __future__ import annotations", - "", - "import subprocess", - "import sys", - "import time", - "from pathlib import Path", - "", - (f"REAL = [{sys.executable!r}, '-m', 'shellctl_runtime.sanitize']"), - "args = sys.argv[1:]", - "if '--ready-file' in args:", - " ready_file = Path(args[args.index('--ready-file') + 1])", - " ready_file.touch()", - f" time.sleep({delay_seconds!r})", - (f" raise SystemExit({sanitize_exit_code!r})" if sanitize_exit_code is not None else ""), - "raise SystemExit(subprocess.run(REAL + args, check=False).returncode)", - "", - ] - ), - encoding="utf-8", - ) - return (sys.executable, str(wrapper_path)) - - -def _capture_serve_config( - monkeypatch: pytest.MonkeyPatch, -) -> tuple[dict[str, object], CliRunner]: - captured: dict[str, object] = {} - - def fake_create_app(config: ShellctlConfig) -> object: - captured["config"] = config - return object() - - def fake_run(app: object, *, host: str, port: int, log_level: str) -> None: - captured["app"] = app - captured["host"] = host - captured["port"] = port - captured["log_level"] = log_level - - monkeypatch.setattr(server_serve_module, "_create_app", fake_create_app) - monkeypatch.setattr(server_serve_module, "_uvicorn_run", fake_run) - return captured, CliRunner() - - -async def _seed_job( - service: ShellctlService, - *, - job_id: str, - status: JobStatusName, - exit_code: int | None = None, - ended_at: str | None = None, - created_at: str = "2026-05-21T15:30:12Z", - started_at: str | None = "2026-05-21T15:30:13Z", -) -> JobRow: - job_dir = service.config.jobs_dir / job_id - job_dir.mkdir(mode=0o700) - (job_dir / "script").write_text("printf ready\n", encoding="utf-8") - (job_dir / "output.log").write_text("hello\n", encoding="utf-8") - row = JobRow( - job_id=job_id, - script_path=f"jobs/{job_id}/script", - output_path=f"jobs/{job_id}/output.log", - cwd="/tmp", - terminal_cols=120, - terminal_rows=80, - status=status.value, - session_name=job_session_name(job_id), - pane_target=job_pane_target(job_id), - exit_code=exit_code, - reason=None, - message=None, - created_at=created_at, - started_at=started_at, - ended_at=ended_at, - updated_at=created_at, - ) - inserted = await service._insert_job_row(row) - assert inserted is True - return row - - -class RecordingService(ShellctlService): - def __init__(self, config: ShellctlConfig, *, tmux: FakeTmuxController) -> None: - super().__init__(config, tmux=tmux) - self.deleted_rows: list[str] = [] - - async def _delete_job_row(self, job_id: str) -> None: - self.deleted_rows.append(job_id) - await super()._delete_job_row(job_id) - - -class StartupObservationService(RecordingService): - def __init__(self, config: ShellctlConfig, *, tmux: FakeTmuxController) -> None: - super().__init__(config, tmux=tmux) - self.observed_startup_statuses: list[JobStatusName] = [] - - async def _insert_job_row(self, row: JobRow) -> bool: - inserted = await super()._insert_job_row(row) - if inserted: - startup_view = await self.get_job_status(row.job_id) - self.observed_startup_statuses.append(startup_view.status) - return inserted - - -class ConcurrentDeleteRaceService(RecordingService): - def __init__(self, config: ShellctlConfig, *, tmux: FakeTmuxController) -> None: - super().__init__(config, tmux=tmux) - self.disappear_on_status: set[str] = set() - self.disappear_on_delete: set[str] = set() - - async def get_job_status(self, job_id: str): # type: ignore[override] - if job_id in self.disappear_on_status: - self.disappear_on_status.remove(job_id) - await super()._delete_job_row(job_id) - return await super().get_job_status(job_id) - - async def _delete_job_row(self, job_id: str) -> None: - if job_id in self.disappear_on_delete: - self.disappear_on_delete.remove(job_id) - raise ShellctlServerError(404, "job_not_found", f"Unknown job id: {job_id}") - await super()._delete_job_row(job_id) - - -class InsertConflictService(RecordingService): - def __init__(self, config: ShellctlConfig, *, tmux: FakeTmuxController) -> None: - super().__init__(config, tmux=tmux) - self.failed_insert_job_ids: list[str] = [] - - async def _insert_job_row(self, row: JobRow) -> bool: - if not self.failed_insert_job_ids: - self.failed_insert_job_ids.append(row.job_id) - return False - return await super()._insert_job_row(row) - - -@pytest.mark.anyio -async def test_sqlite_persistence_survives_service_restart(tmp_path: Path) -> None: - service, _fake_tmux = await _create_service(tmp_path) - await _seed_job( - service, - job_id="05211530-k7p", - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2026-05-21T15:30:20Z", - ) - await service.shutdown() - - restarted, _fake_tmux = await _create_service(tmp_path) - row = await restarted._get_job_row("05211530-k7p") - - assert row.status == JobStatusName.EXITED.value - assert row.exit_code == 0 - assert restarted.config.db_path.exists() - await restarted.shutdown() - - -@pytest.mark.anyio -async def test_run_job_does_not_materialize_fresh_row_to_lost_during_startup_window( - tmp_path: Path, -) -> None: - fake_tmux = FakeTmuxController() - service = StartupObservationService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - fake_tmux.touch_ready_on_enable = False - service.config = ShellctlConfig( - state_dir=service.config.state_dir, - runtime_dir=service.config.runtime_dir, - poll_interval_seconds=0.001, - pipe_ready_timeout_seconds=0.01, - ) - - await service.run_job( - RunJobRequest( - script="printf ready\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=0.01, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - assert len(service.observed_startup_statuses) == 1 - assert service.observed_startup_statuses[0] in { - JobStatusName.CREATED, - JobStatusName.STARTING, - JobStatusName.RUNNING, - } - await service.shutdown() - - -@pytest.mark.anyio -async def test_run_job_happy_path_persists_running_row_and_artifacts( - tmp_path: Path, -) -> None: - fake_tmux = FakeTmuxController() - service = RecordingService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - - result = await service.run_job( - RunJobRequest( - script="printf ready\n", - cwd=str(tmp_path), - env={"HELLO": "world", "UNICODE": "盐粒"}, - terminal=TerminalSize(cols=100, rows=40), - timeout=0.01, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - row = await service._get_job_row(result.job_id) - job_dir = service.config.jobs_dir / result.job_id - - assert result.status is JobStatusName.RUNNING - assert result.done is False - assert result.output == "" - assert result.offset == 0 - assert row.status == JobStatusName.RUNNING.value - assert row.cwd == str(tmp_path.resolve()) - assert row.script_path == f"jobs/{result.job_id}/script" - assert row.output_path == f"jobs/{result.job_id}/output.log" - assert row.started_at is not None - assert (job_dir / "script").exists() - assert json.loads(job_env_path(job_dir).read_text(encoding="utf-8")) == { - "HELLO": "world", - "UNICODE": "盐粒", - } - assert (job_dir / "output.log").exists() - assert (job_dir / "start-gate").exists() - assert job_session_name(result.job_id) in fake_tmux.sessions - assert fake_tmux.pipe_active[result.job_id] is True - assert await service.get_job_status(result.job_id) == JobStatusView( - job_id=result.job_id, - status=JobStatusName.RUNNING, - done=False, - exit_code=None, - created_at=row.created_at, - started_at=row.started_at, - ended_at=None, - offset=0, - ) - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_run_job_env_overlay_is_visible_without_replacing_inherited_env( - tmp_path: Path, -) -> None: - service = await _create_real_service(tmp_path) - - try: - result = await service.run_job( - RunJobRequest( - script=( - "python3 - <<'PY'\n" - "import os\n" - "print(os.environ['SHELLCTL_PRESET'])\n" - "print('PATH' in os.environ)\n" - "PY\n" - ), - cwd=str(tmp_path), - env={"SHELLCTL_PRESET": "from-client"}, - terminal=TerminalSize(), - timeout=5, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == "from-client\nTrue\n" - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_run_job_env_overlay_overrides_inherited_value( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("SHELLCTL_PRESET", "parent") - service = await _create_real_service(tmp_path) - - try: - result = await service.run_job( - RunJobRequest( - script=("python3 - <<'PY'\nimport os\nprint(os.environ['SHELLCTL_PRESET'])\nPY\n"), - cwd=str(tmp_path), - env={"SHELLCTL_PRESET": "from-client"}, - terminal=TerminalSize(), - timeout=5, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == "from-client\n" - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_send_input_reaches_real_job_stdin_after_env_bootstrap( - tmp_path: Path, -) -> None: - service = await _create_real_service(tmp_path) - - try: - initial = await service.run_job( - RunJobRequest( - script=("printf 'ready\\n'\nIFS= read -r line\nprintf 'got:%s\\n' \"$line\"\n"), - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=5, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - ready_window = initial - if ready_window.output == "": - ready_window = await service.wait_job( - initial.job_id, - WaitJobRequest( - offset=initial.offset, - timeout=1, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - result = await service.send_input( - initial.job_id, - InputJobRequest( - text="hello from stdin\n", - offset=ready_window.offset, - timeout=1, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - output_after_input = result.output - final = result - if not final.done: - # `send_input()` shares wait semantics with `wait_job()`: it may - # return after the post-input output flushes but before tmux exit - # artifacts have materialized into a terminal DB status on slower CI. - final = await service.wait_job( - initial.job_id, - WaitJobRequest( - offset=result.offset, - timeout=5, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - output_after_input += final.output - - assert ready_window.done is False - assert final.done is True - assert final.status is JobStatusName.EXITED - assert final.exit_code == 0 - assert output_after_input.endswith("got:hello from stdin\n") - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_terminated_job_does_not_emit_python_wrapper_traceback( - tmp_path: Path, -) -> None: - service = await _create_real_service(tmp_path) - - try: - initial = await service.run_job( - RunJobRequest( - script=("trap 'exit 130' INT\nprintf 'ready\\n'\nwhile :; do sleep 1; done\n"), - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=5, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - ready_window = initial - if ready_window.output == "": - ready_window = await service.wait_job( - initial.job_id, - WaitJobRequest( - offset=initial.offset, - timeout=1, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - assert ready_window.done is False - - terminated = await service.terminate_job( - initial.job_id, - TerminateJobRequest(grace_seconds=0.2), - ) - final = await service.wait_job( - initial.job_id, - WaitJobRequest( - offset=ready_window.offset, - timeout=1, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - assert terminated.status is JobStatusName.TERMINATED - assert final.done is True - assert "KeyboardInterrupt" not in final.output - assert "Traceback (most recent call last)" not in final.output - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_run_job_returns_flushed_output_on_first_terminal_result( - tmp_path: Path, -) -> None: - sanitize_pty_command = _write_delayed_sanitize_wrapper(tmp_path, delay_seconds=0.2) - service = await _create_real_service( - tmp_path, - sanitize_pty_command=sanitize_pty_command, - ) - - try: - result = await service.run_job( - RunJobRequest( - script="printf 'delayed-flush\\n'\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=3, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - reread = await service.wait_job( - result.job_id, - WaitJobRequest( - offset=0, - timeout=0.001, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == "delayed-flush\n" - assert reread.output == "delayed-flush\n" - assert reread.offset == result.offset - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -async def test_sanitize_failure_does_not_commit_normal_exit( - tmp_path: Path, -) -> None: - sanitize_pty_command = _write_delayed_sanitize_wrapper( - tmp_path, - delay_seconds=0.2, - sanitize_exit_code=7, - ) - service = await _create_real_service( - tmp_path, - sanitize_pty_command=sanitize_pty_command, - ) - - try: - result = await service.run_job( - RunJobRequest( - script="printf 'missing-output\\n'\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=3, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - reread = await service.wait_job( - result.job_id, - WaitJobRequest( - offset=0, - timeout=0.001, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - job_dir = service.config.jobs_dir / result.job_id - - assert result.done is True - assert result.status is not JobStatusName.EXITED - assert result.exit_code is None - assert result.output == "" - assert reread.output == "" - assert not (job_dir / PIPE_DRAINED_FILENAME).exists() - assert (job_dir / PIPE_FAILED_FILENAME).exists() - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -@pytest.mark.skipif(shutil.which("uv") is None, reason="uv is required") -async def test_uv_quiet_shebang_returns_output_in_first_terminal_result( - tmp_path: Path, -) -> None: - service = await _create_real_service(tmp_path) - - try: - result = await service.run_job( - RunJobRequest( - script=( - "#!/usr/bin/env -S uv run --script --quiet\n" - "# /// script\n" - '# requires-python = ">=3.12"\n' - "# dependencies = []\n" - "# ///\n" - 'print("hello")\n' - ), - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=10, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == "hello\n" - finally: - await service.shutdown() - - -@pytest.mark.anyio -@pytest.mark.skipif(shutil.which("tmux") is None, reason="tmux is required") -@pytest.mark.skipif(shutil.which("bash") is None, reason="bash is required") -@pytest.mark.parametrize( - ("script", "expected_output"), - [ - ("printf 'no-shebang\\n'\n", "no-shebang\n"), - ("#!/bin/sh\nprintf 'sh-shebang\\n'\n", "sh-shebang\n"), - ( - "#!/usr/bin/env bash\nprintf 'bash-shebang\\n'\n", - "bash-shebang\n", - ), - ], -) -async def test_existing_script_modes_still_return_output( - tmp_path: Path, - script: str, - expected_output: str, -) -> None: - service = await _create_real_service(tmp_path) - - try: - result = await service.run_job( - RunJobRequest( - script=script, - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=5, - output_limit=8192, - idle_flush_seconds=1, - ) - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == expected_output - finally: - await service.shutdown() - - -@pytest.mark.anyio -async def test_run_job_retries_after_sqlite_insert_conflict_and_cleans_artifacts( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - fake_tmux = FakeTmuxController() - service = InsertConflictService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - generated_ids = iter(["05211530-k7p", "05211530-abc"]) - - monkeypatch.setattr(server_service_module, "generate_job_id", lambda now=None: next(generated_ids)) - - result = await service.run_job( - RunJobRequest( - script="printf ready\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=0.01, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - failed_job_id = service.failed_insert_job_ids[0] - assert failed_job_id == "05211530-k7p" - assert result.job_id == "05211530-abc" - assert not (service.config.jobs_dir / failed_job_id).exists() - assert (service.config.jobs_dir / result.job_id).exists() - with pytest.raises(ShellctlServerError, match="Unknown job id"): - await service._get_job_row(failed_job_id) - await service.shutdown() - - -@pytest.mark.anyio -async def test_transition_status_does_not_overwrite_terminal_state( - tmp_path: Path, -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - await _seed_job( - service, - job_id="05211530-k7p", - status=JobStatusName.TERMINATED, - ended_at="2026-05-21T15:30:20Z", - ) - - row = await service._transition_status( - "05211530-k7p", - allowed_from={JobStatusName.RUNNING}, - target=JobStatusName.RUNNING, - ) - - assert row.status == JobStatusName.TERMINATED.value - await service.shutdown() - - -@pytest.mark.anyio -async def test_wait_job_flushes_preexisting_unread_output_after_idle( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = True - - async def fail_sleep(_delay: float) -> None: - raise AssertionError("wait_job should not sleep when idle_flush_seconds=0") - - monkeypatch.setattr(server_service_module.anyio, "sleep", fail_sleep) - - result = await service.wait_job( - "05211530-k7p", - WaitJobRequest(offset=0, timeout=0.2, output_limit=8192, idle_flush_seconds=0), - ) - - assert result.output == "hello\n" - assert result.done is False - await service.shutdown() - - -@pytest.mark.anyio -async def test_wait_job_returns_immediately_for_terminal_row_and_reads_remaining_output( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - job_id = "05211530-k7p" - await _seed_job( - service, - job_id=job_id, - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2026-05-21T15:30:20Z", - ) - output_path = service.config.jobs_dir / job_id / "output.log" - output_path.write_text("hello\nworld\n", encoding="utf-8") - - async def fail_sleep(_delay: float) -> None: - raise AssertionError("terminal wait should not sleep") - - monkeypatch.setattr(server_service_module.anyio, "sleep", fail_sleep) - - result = await service.wait_job( - job_id, - WaitJobRequest( - offset=len(b"hello\n"), - timeout=0.5, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - assert result.done is True - assert result.status is JobStatusName.EXITED - assert result.exit_code == 0 - assert result.output == "world\n" - assert result.offset == len(b"hello\nworld\n") - await service.shutdown() - - -@pytest.mark.anyio -async def test_starting_job_is_not_marked_lost_during_local_startup_window( - tmp_path: Path, -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.STARTING) - service._starting_jobs.add("05211530-k7p") - - view = await service.get_job_status("05211530-k7p") - - assert view.status is JobStatusName.STARTING - assert view.done is False - await service.shutdown() - - -@pytest.mark.anyio -async def test_stale_starting_job_with_dead_pipe_materializes_failed( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.STARTING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = False - - view = await service.get_job_status("05211530-k7p") - row = await service._get_job_row("05211530-k7p") - - assert view.status is JobStatusName.FAILED - assert row.reason == "pipe_failed" - await service.shutdown() - - -@pytest.mark.anyio -async def test_terminate_beats_concurrent_runner_exit(tmp_path: Path) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = True - - async def record_exit(job_id: str) -> None: - await service.record_runner_exit(job_id, 130, "2026-05-21T15:30:20Z") - - fake_tmux.on_send_interrupt = record_exit - view = await service.terminate_job("05211530-k7p", TerminateJobRequest(grace_seconds=0)) - - assert view.status is JobStatusName.TERMINATED - assert view.exit_code == 130 - await service.shutdown() - - -@pytest.mark.anyio -async def test_record_runner_exit_materializes_nonterminal_row_to_exited( - tmp_path: Path, -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - - await service.record_runner_exit("05211530-k7p", 7, "2026-05-21T15:30:20Z") - - row = await service._get_job_row("05211530-k7p") - view = await service.get_job_status("05211530-k7p") - - assert row.status == JobStatusName.EXITED.value - assert row.exit_code == 7 - assert row.ended_at == "2026-05-21T15:30:20Z" - assert view.status is JobStatusName.EXITED - assert view.done is True - assert view.exit_code == 7 - assert view.ended_at == "2026-05-21T15:30:20Z" - await service.shutdown() - - -@pytest.mark.anyio -async def test_pipe_monitor_marks_running_job_failed_when_pipe_disappears( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = False - - await service.check_running_jobs_pipe_health() - row = await service._get_job_row("05211530-k7p") - - assert row.status == JobStatusName.FAILED.value - assert row.reason == "pipe_failed" - await service.shutdown() - - -@pytest.mark.anyio -async def test_session_disappearing_between_probes_materializes_lost_not_pipe_failed( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = None - - view = await service.get_job_status("05211530-k7p") - row = await service._get_job_row("05211530-k7p") - - assert view.status is JobStatusName.LOST - assert row.reason == "tmux_session_missing" - await service.shutdown() - - -@pytest.mark.anyio -async def test_drained_uncommitted_normal_exit_self_commits_to_exited( - tmp_path: Path, -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - job_id = "05211530-k7p" - await _seed_job(service, job_id=job_id, status=JobStatusName.RUNNING) - job_dir = service.config.jobs_dir / job_id - (job_dir / RUNNER_EXIT_CODE_FILENAME).write_text("0\n", encoding="utf-8") - (job_dir / RUNNER_ENDED_AT_FILENAME).write_text("2026-05-21T15:30:20Z\n", encoding="utf-8") - (job_dir / PIPE_DRAINED_FILENAME).touch() - - view = await service.get_job_status(job_id) - row = await service._get_job_row(job_id) - - assert view.status is JobStatusName.EXITED - assert view.done is True - assert view.exit_code == 0 - assert row.status == JobStatusName.EXITED.value - assert row.exit_code == 0 - assert row.ended_at == "2026-05-21T15:30:20Z" - await service.shutdown() - - -@pytest.mark.anyio -async def test_list_jobs_ignores_half_created_directory_and_uses_db_rows( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - (service.config.jobs_dir / "half-created").mkdir(mode=0o700) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = True - - jobs = await service.list_jobs(limit=10) - - assert [job.job_id for job in jobs.jobs] == ["05211530-k7p"] - assert fake_tmux.cleaned == [] - await service.shutdown() - - -@pytest.mark.anyio -async def test_reconcile_and_gc_query_sqlite_rows(tmp_path: Path) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job( - service, - job_id="05211530-old", - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2020-05-21T15:30:20Z", - created_at="2020-05-21T15:30:12Z", - ) - await _seed_job(service, job_id="05211530-run", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-run")) - fake_tmux.pipe_active["05211530-run"] = True - - await service.reconcile() - await service.gc_once() - - with pytest.raises(ShellctlServerError, match="Unknown job id"): - await service._get_job_row("05211530-old") - row = await service._get_job_row("05211530-run") - assert row.status == JobStatusName.RUNNING.value - assert not (service.config.jobs_dir / "05211530-old").exists() - await service.shutdown() - - -@pytest.mark.anyio -async def test_delete_job_removes_db_row_before_artifact_cleanup( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - fake_tmux = FakeTmuxController() - service = RecordingService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - await _seed_job( - service, - job_id="05211530-k7p", - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2026-05-21T15:30:20Z", - ) - - observed_db_missing_before_cleanup = False - original_rmtree = server_service_module.shutil.rmtree - - def assert_db_deleted_then_remove(path: Path, *, ignore_errors: bool = False) -> None: - nonlocal observed_db_missing_before_cleanup - with sqlite3.connect(service.config.db_path) as connection: - count = connection.execute( - "SELECT COUNT(*) FROM jobs WHERE job_id = ?", - ("05211530-k7p",), - ).fetchone()[0] - assert count == 0 - observed_db_missing_before_cleanup = True - original_rmtree(path, ignore_errors=ignore_errors) - - monkeypatch.setattr(server_service_module.shutil, "rmtree", assert_db_deleted_then_remove) - - await service.delete_job("05211530-k7p", force=False, grace_seconds=0) - - assert service.deleted_rows == ["05211530-k7p"] - assert observed_db_missing_before_cleanup is True - with pytest.raises(ShellctlServerError, match="Unknown job id"): - await service._get_job_row("05211530-k7p") - await service.shutdown() - - -@pytest.mark.anyio -async def test_list_jobs_skips_row_deleted_mid_iteration(tmp_path: Path) -> None: - fake_tmux = FakeTmuxController() - service = ConcurrentDeleteRaceService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - await _seed_job(service, job_id="05211530-gone", status=JobStatusName.RUNNING) - await _seed_job(service, job_id="05211530-keep", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-gone")) - fake_tmux.sessions.add(job_session_name("05211530-keep")) - fake_tmux.pipe_active["05211530-gone"] = True - fake_tmux.pipe_active["05211530-keep"] = True - service.disappear_on_status.add("05211530-gone") - - jobs = await service.list_jobs(limit=10) - - assert [job.job_id for job in jobs.jobs] == ["05211530-keep"] - await service.shutdown() - - -@pytest.mark.anyio -async def test_reconcile_skips_row_deleted_mid_iteration(tmp_path: Path) -> None: - fake_tmux = FakeTmuxController() - service = ConcurrentDeleteRaceService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - await _seed_job(service, job_id="05211530-gone", status=JobStatusName.RUNNING) - await _seed_job(service, job_id="05211530-keep", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-gone")) - fake_tmux.sessions.add(job_session_name("05211530-keep")) - fake_tmux.pipe_active["05211530-gone"] = True - fake_tmux.pipe_active["05211530-keep"] = True - service.disappear_on_status.add("05211530-gone") - - await service.reconcile() - - row = await service._get_job_row("05211530-keep") - assert row.status == JobStatusName.RUNNING.value - await service.shutdown() - - -@pytest.mark.anyio -async def test_gc_skips_row_deleted_mid_pass(tmp_path: Path) -> None: - fake_tmux = FakeTmuxController() - service = ConcurrentDeleteRaceService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=fake_tmux, - ) - await service.initialize_database() - service._ensure_dir(service.config.jobs_dir) - await _seed_job( - service, - job_id="05211530-gone", - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2020-05-21T15:30:20Z", - created_at="2020-05-21T15:30:12Z", - ) - await _seed_job( - service, - job_id="05211530-keep", - status=JobStatusName.EXITED, - exit_code=0, - ended_at="2020-05-21T15:30:20Z", - created_at="2020-05-21T15:30:13Z", - ) - service.disappear_on_delete.add("05211530-gone") - - await service.gc_once() - - with pytest.raises(ShellctlServerError, match="Unknown job id"): - await service._get_job_row("05211530-keep") - await service.shutdown() - - -@pytest.mark.anyio -async def test_run_job_keeps_start_gate_closed_when_pipe_never_becomes_ready( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - fake_tmux.touch_ready_on_enable = False - service.config = ShellctlConfig( - state_dir=service.config.state_dir, - runtime_dir=service.config.runtime_dir, - poll_interval_seconds=0.001, - pipe_monitor_interval_seconds=0.01, - pipe_ready_timeout_seconds=0.01, - ) - - result = await service.run_job( - RunJobRequest( - script="printf ready\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=0.05, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - assert result.status is JobStatusName.FAILED - assert result.done is True - assert not (service.config.jobs_dir / result.job_id / "start-gate").exists() - await service.shutdown() - - -@pytest.mark.anyio -async def test_run_job_pipe_ready_timeout_records_diagnostics(tmp_path: Path) -> None: - service, fake_tmux = await _create_service(tmp_path) - fake_tmux.touch_ready_on_enable = False - fake_tmux.pipe_active_on_enable = False - fake_tmux.pipe_error_log_text = "Traceback: sanitizer crashed\nsecond line\n" - service.config = ShellctlConfig( - state_dir=service.config.state_dir, - runtime_dir=service.config.runtime_dir, - poll_interval_seconds=0.001, - pipe_ready_timeout_seconds=0.01, - ) - - result = await service.run_job( - RunJobRequest( - script="printf ready\n", - cwd=str(tmp_path), - terminal=TerminalSize(), - timeout=0.05, - output_limit=8192, - idle_flush_seconds=0.01, - ) - ) - - row = await service._get_job_row(result.job_id) - ready_file = service.config.jobs_dir / result.job_id / ".pipe-ready" - - assert result.status is JobStatusName.FAILED - assert row.message is not None - assert "waited " in row.message - assert str(ready_file) in row.message - assert "tmux #{pane_pipe}=0" in row.message - assert PIPE_ERROR_LOG_FILENAME in row.message - assert "Traceback: sanitizer crashed | second line" in row.message - await service.shutdown() - - -@pytest.mark.anyio -async def test_send_input_terminal_race_returns_conflict_not_500( - tmp_path: Path, -) -> None: - service, fake_tmux = await _create_service(tmp_path) - await _seed_job(service, job_id="05211530-k7p", status=JobStatusName.RUNNING) - fake_tmux.sessions.add(job_session_name("05211530-k7p")) - fake_tmux.pipe_active["05211530-k7p"] = True - - async def finish_before_send(job_id: str, _text: str) -> None: - fake_tmux.sessions.clear() - fake_tmux.pipe_active.pop(job_id, None) - await service.record_runner_exit(job_id, 0, "2026-05-21T15:30:20Z") - raise ShellctlServerError( - 409, - "tmux_target_missing", - "no server running on /tmp/tmux.sock", - ) - - fake_tmux.on_send_input = finish_before_send - - with pytest.raises(ShellctlServerError) as exc_info: - await service.send_input( - "05211530-k7p", - InputJobRequest( - text="pwd\n", - offset=0, - timeout=0.05, - output_limit=8192, - idle_flush_seconds=0.01, - ), - ) - - assert exc_info.value.status_code == 409 - assert exc_info.value.code == "job_not_running" - await service.shutdown() - - -@pytest.mark.anyio -async def test_allocate_job_dir_retries_on_atomic_mkdir_collision( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - service, _fake_tmux = await _create_service(tmp_path) - (service.config.jobs_dir / "05211530-k7p").mkdir(mode=0o700) - generated_ids = iter(["05211530-k7p", "05211530-abc"]) - - monkeypatch.setattr(server_service_module, "generate_job_id", lambda now=None: next(generated_ids)) - - job_id, job_dir = service._allocate_job_dir() - - assert job_id == "05211530-abc" - assert job_dir.exists() - await service.shutdown() - - -@pytest.mark.anyio -async def test_service_initialize_allows_missing_auth_token(tmp_path: Path) -> None: - service = ShellctlService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=FakeTmuxController(), - ) - - await service.initialize() - - assert service.config.runner_path.exists() - await service.shutdown() - - -@pytest.mark.anyio -async def test_http_routes_inject_shellctl_service_dependency( - tmp_path: Path, -) -> None: - service, _fake_tmux = await _create_service(tmp_path, auth_token="route-token") - app = create_app(service.config, service=service) - transport = httpx.ASGITransport(app=app) - - async with httpx.AsyncClient(transport=transport, base_url="http://shellctl.test") as client: - health = await client.get("/healthz") - assert health.status_code == 200 - - unauthenticated = await client.get("/v1/jobs") - assert unauthenticated.status_code == 401 - - authenticated = await client.get("/v1/jobs", headers={"Authorization": "Bearer route-token"}) - - assert authenticated.status_code == 200 - assert authenticated.json() == {"jobs": []} - await service.shutdown() - - -@pytest.mark.anyio -async def test_http_routes_enforce_auth_from_environment_fallback( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "route-token") - service, _fake_tmux = await _create_service(tmp_path) - app = create_app(service.config, service=service) - transport = httpx.ASGITransport(app=app) - - async with httpx.AsyncClient(transport=transport, base_url="http://shellctl.test") as client: - unauthenticated = await client.get("/v1/jobs") - authenticated = await client.get("/v1/jobs", headers={"Authorization": "Bearer route-token"}) - - assert unauthenticated.status_code == 401 - assert authenticated.status_code == 200 - await service.shutdown() - - -@pytest.mark.anyio -async def test_create_app_without_explicit_config_reads_auth_token_from_environment( - monkeypatch: pytest.MonkeyPatch, -) -> None: - async def noop_initialize(self: ShellctlService) -> None: - return None - - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "route-token") - monkeypatch.setattr(ShellctlService, "initialize", noop_initialize) - monkeypatch.setattr(ShellctlService, "start_background_gc", lambda self: None) - monkeypatch.setattr(ShellctlService, "start_background_pipe_monitor", lambda self: None) - - app = create_app() - transport = httpx.ASGITransport(app=app) - - async with httpx.AsyncClient(transport=transport, base_url="http://shellctl.test") as client: - response = await client.get("/v1/jobs") - - assert app.state.shellctl_service.config.auth_token == "route-token" - assert response.status_code == 401 - await app.state.shellctl_service.shutdown() - - -@pytest.mark.anyio -async def test_http_routes_skip_auth_when_token_missing(tmp_path: Path) -> None: - service, _fake_tmux = await _create_service(tmp_path) - app = create_app(service.config, service=service) - transport = httpx.ASGITransport(app=app) - - async with httpx.AsyncClient(transport=transport, base_url="http://shellctl.test") as client: - response = await client.get("/v1/jobs") - - assert response.status_code == 200 - assert response.json() == {"jobs": []} - await service.shutdown() - - -@pytest.mark.anyio -async def test_http_routes_skip_auth_when_token_is_empty(tmp_path: Path) -> None: - service, _fake_tmux = await _create_service(tmp_path, auth_token="") - app = create_app(service.config, service=service) - transport = httpx.ASGITransport(app=app) - - async with httpx.AsyncClient(transport=transport, base_url="http://shellctl.test") as client: - response = await client.get("/v1/jobs") - - assert response.status_code == 200 - assert response.json() == {"jobs": []} - await service.shutdown() - - -def test_shellctl_config_reads_auth_token_from_environment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "env-token") - - config = ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") - - assert config.auth_token == "env-token" - - -def test_shellctl_config_treats_empty_environment_auth_token_as_disabled( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "") - - config = ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") - - assert config.auth_token is None - - -def test_serve_cli_passes_direct_auth_token_to_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - captured, runner = _capture_serve_config(monkeypatch) - - result = runner.invoke( - cli, - [ - "serve", - "--listen", - "0.0.0.0:9999", - "--auth-token", - "direct-token", - "--state-dir", - str(tmp_path / "state"), - ], - ) - - assert result.exit_code == 0, result.output - assert captured["host"] == "0.0.0.0" - assert captured["port"] == 9999 - assert captured["log_level"] == "info" - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.auth_token == "direct-token" - - -def test_serve_cli_prefers_explicit_auth_token_over_environment( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "env-token") - captured, runner = _capture_serve_config(monkeypatch) - - result = runner.invoke( - cli, - [ - "serve", - "--auth-token", - "direct-token", - "--state-dir", - str(tmp_path / "state"), - ], - ) - - assert result.exit_code == 0, result.output - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.auth_token == "direct-token" - - -def test_serve_cli_treats_empty_auth_token_as_disabled(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - captured, runner = _capture_serve_config(monkeypatch) - - result = runner.invoke( - cli, - [ - "serve", - "--auth-token", - "", - "--state-dir", - str(tmp_path / "state"), - ], - ) - - assert result.exit_code == 0, result.output - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.auth_token is None - - -def test_serve_cli_explicit_empty_auth_token_beats_environment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "env-token") - captured, runner = _capture_serve_config(monkeypatch) - - result = runner.invoke( - cli, - [ - "serve", - "--auth-token", - "", - "--state-dir", - str(tmp_path / "state"), - ], - ) - - assert result.exit_code == 0, result.output - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.auth_token is None - - -def test_serve_cli_reads_auth_token_from_environment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv(DEFAULT_AUTH_TOKEN_ENV, "env-token") - captured, runner = _capture_serve_config(monkeypatch) - - result = runner.invoke( - cli, - [ - "serve", - "--state-dir", - str(tmp_path / "state"), - ], - ) - - assert result.exit_code == 0, result.output - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.auth_token == "env-token" - - -def test_serve_cli_forwards_state_runtime_and_gc_flags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - captured, runner = _capture_serve_config(monkeypatch) - state_dir = tmp_path / "state" - runtime_dir = tmp_path / "runtime" - - result = runner.invoke( - cli, - [ - "serve", - "--state-dir", - str(state_dir), - "--runtime-dir", - str(runtime_dir), - "--gc-interval-seconds", - "12.5", - "--gc-finished-job-retention-seconds", - "345.0", - ], - ) - - assert result.exit_code == 0, result.output - config = captured["config"] - assert isinstance(config, ShellctlConfig) - assert config.state_dir == state_dir - assert config.runtime_dir == runtime_dir - assert config.gc_interval_seconds == 12.5 - assert config.gc_finished_job_retention_seconds == 345.0 - - -def test_runner_script_records_completion_metadata_without_direct_runner_exit( - tmp_path: Path, -) -> None: - service = ShellctlService( - ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run"), - tmux=FakeTmuxController(), - ) - source = service._runner_script_source() - - assert JOB_ENV_FILENAME in source - assert RUNNER_EXIT_CODE_FILENAME in source - assert RUNNER_ENDED_AT_FILENAME in source - assert "write_atomic" in source - assert 'mv "$tmp" "$dest"' in source - assert re.search(r"\brunner-exit\s+--state-dir\b", source) is None - anyio.run(service.shutdown) - - -def test_shellctl_config_defaults_to_lightweight_sanitize_entrypoint( - tmp_path: Path, -) -> None: - config = ShellctlConfig(state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") - - assert config.pipe_ready_timeout_seconds == 10.0 - assert config.sanitize_pty_command == ("shellctl-sanitize-pty",) - assert config.runner_exit_command == ("shellctl-runner-exit",) - - -def test_pipe_command_finalizer_commits_runner_exit_after_drain(tmp_path: Path) -> None: - controller = TmuxController( - ShellctlConfig( - state_dir=tmp_path / "state", - runtime_dir=tmp_path / "run", - sqlite_busy_timeout_ms=6789, - ) - ) - source = controller._pipe_command_source( - job_id="05211530-k7p", - job_dir=tmp_path / "state" / "jobs" / "05211530-k7p", - ready_file=tmp_path / "state" / "jobs" / "05211530-k7p" / ".pipe-ready", - ) - - assert "shellctl-sanitize-pty" in source - assert PIPE_DRAINED_FILENAME in source - assert PIPE_ERROR_LOG_FILENAME in source - assert PIPE_FAILED_FILENAME in source - assert RUNNER_EXIT_CODE_FILENAME in source - assert RUNNER_ENDED_AT_FILENAME in source - assert "shellctl-runner-exit" in source - assert "--sqlite-busy-timeout-ms 6789" in source - assert "2>>" in source - assert 'exit "$sanitize_status"' in source - assert 'if [ "$sanitize_status" -eq 0 ]' in source - assert "2>" in source - assert source.index("shellctl-sanitize-pty") < source.index(PIPE_DRAINED_FILENAME) - assert source.index(PIPE_DRAINED_FILENAME) < source.index("shellctl-runner-exit") - - -def test_server_cli_no_longer_exposes_sanitize_pty_command() -> None: - result = CliRunner().invoke(cli, ["sanitize-pty"]) - - assert result.exit_code != 0 - - -def test_python_m_shellctl_server_module_entrypoint_shows_cli_help() -> None: - package_root = Path(__file__).resolve().parents[3] - src_path = package_root / "src" - env = dict(os.environ) - current_pythonpath = env.get("PYTHONPATH") - env["PYTHONPATH"] = f"{src_path}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(src_path) - - result = subprocess.run( - [sys.executable, "-m", "shellctl.server", "--help"], - capture_output=True, - text=True, - check=False, - cwd=package_root, - env=env, - ) - - assert result.returncode == 0, result.stderr - assert "sanitize-pty" not in result.stdout - assert "runner-exit" not in result.stdout - - -def test_server_package_keeps_config_access_lazy() -> None: - package_root = Path(__file__).resolve().parents[3] - src_path = package_root / "src" - env = dict(os.environ) - current_pythonpath = env.get("PYTHONPATH") - env["PYTHONPATH"] = f"{src_path}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(src_path) - - script = """ -import json -import sys - -import shellctl.server as server - -server.ShellctlConfig - -print(json.dumps({ - "deny": sorted( - name for name in sys.modules - if name.split(".", 1)[0] in {"fastapi", "typer", "uvicorn"} - ), - "server_modules": sorted( - name for name in sys.modules - if name.startswith("shellctl.server.") - ), -})) -""" - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - check=False, - cwd=package_root, - env=env, - ) - - assert result.returncode == 0, result.stderr - payload = json.loads(result.stdout) - assert payload["deny"] == [] - assert "shellctl.server.api" not in payload["server_modules"] - assert "shellctl.server.cli" not in payload["server_modules"] - assert "shellctl.server.service" not in payload["server_modules"] diff --git a/dify-agent/tests/local/shellctl/test_shellctl_shared.py b/dify-agent/tests/local/shellctl/test_shellctl_shared.py index abd5167bc7d..8034ea3d9ef 100644 --- a/dify-agent/tests/local/shellctl/test_shellctl_shared.py +++ b/dify-agent/tests/local/shellctl/test_shellctl_shared.py @@ -1,9 +1,6 @@ -import json import re from datetime import UTC, datetime -from io import BytesIO from pathlib import Path -from typing import BinaryIO, cast import pytest from pydantic import ValidationError @@ -15,11 +12,6 @@ from shellctl.shared import ( read_output_window, tail_output_window, ) -from shellctl_runtime.sanitize import ( - PtySanitizer, - sanitize_pty_output, - sanitize_pty_stream, -) def test_generate_job_id_matches_proposal_format() -> None: @@ -29,20 +21,6 @@ def test_generate_job_id_matches_proposal_format() -> None: assert all(char in f"{JOB_ID_ALPHABET}-" for char in job_id[9:]) -def test_sanitize_pty_golden_cases() -> None: - fixture_path = Path(__file__).with_name("golden_shellctl_sanitize.json") - cases = json.loads(fixture_path.read_text(encoding="utf-8")) - - for case in cases: - chunks = [bytes.fromhex(chunk) for chunk in case["chunks_hex"]] - sanitizer = PtySanitizer() - streamed = "".join(sanitizer.feed(chunk) for chunk in chunks) + sanitizer.flush() - batch = sanitize_pty_output(b"".join(chunks)) - - assert streamed == case["expected"], case["name"] - assert batch == case["expected"], case["name"] - - def test_read_output_window_preserves_utf8_boundaries(tmp_path: Path) -> None: output_path = tmp_path / "output.log" output_path.write_text("A🙂B", encoding="utf-8") @@ -89,36 +67,6 @@ def test_tail_output_window_skips_partial_utf8_prefix(tmp_path: Path) -> None: assert tail.truncated is False -def test_sanitize_pty_stream_flushes_incrementally() -> None: - class ChunkedInput: - def __init__(self, chunks: list[bytes]) -> None: - self._chunks = chunks - - def read(self, _size: int) -> bytes: - if not self._chunks: - return b"" - return self._chunks.pop(0) - - class FlushCountingOutput(BytesIO): - def __init__(self) -> None: - super().__init__() - self.flush_count = 0 - - def flush(self) -> None: - self.flush_count += 1 - super().flush() - - stdout = FlushCountingOutput() - sanitize_pty_stream( - cast(BinaryIO, cast(object, ChunkedInput([b"ready\n", b"next\n"]))), - stdout, - chunk_size=5, - ) - - assert stdout.getvalue() == b"ready\nnext\n" - assert stdout.flush_count >= 2 - - @pytest.mark.parametrize( ("env", "message"), [ diff --git a/dify-agent/tests/local/shellctl_runtime/golden_shellctl_sanitize.json b/dify-agent/tests/local/shellctl_runtime/golden_shellctl_sanitize.json deleted file mode 100644 index 10eba3678ef..00000000000 --- a/dify-agent/tests/local/shellctl_runtime/golden_shellctl_sanitize.json +++ /dev/null @@ -1,32 +0,0 @@ -[ - { - "name": "utf8_passthrough", - "chunks_hex": ["68656c6c6f20e4b896e7958cf09f99820a"], - "expected": "hello 世界🙂\n" - }, - { - "name": "split_utf8_chunks", - "chunks_hex": ["e4bda0e5a5bd20f09f", "99820a"], - "expected": "你好 🙂\n" - }, - { - "name": "invalid_utf8_replacement", - "chunks_hex": ["666f80ff6f0a"], - "expected": "fo��o\n" - }, - { - "name": "ansi_sequences_removed", - "chunks_hex": ["1b5b33316d7265641b5b306d1b5b4b0a"], - "expected": "red\n" - }, - { - "name": "carriage_return_progress", - "chunks_hex": ["30250d3530250d313030250a"], - "expected": "100%\n" - }, - { - "name": "crlf_normalized", - "chunks_hex": ["6f6b0d0a6e6578740d0a"], - "expected": "ok\nnext\n" - } -] diff --git a/dify-agent/tests/local/shellctl_runtime/test_shellctl_runtime.py b/dify-agent/tests/local/shellctl_runtime/test_shellctl_runtime.py deleted file mode 100644 index 1314c629c3d..00000000000 --- a/dify-agent/tests/local/shellctl_runtime/test_shellctl_runtime.py +++ /dev/null @@ -1,531 +0,0 @@ -from __future__ import annotations - -import ast -import importlib.metadata -import importlib.util -import io -import json -import os -import shutil -import sqlite3 -import subprocess -import sys -import textwrap -from pathlib import Path -from types import TracebackType -from typing import BinaryIO, TypedDict, cast - -import anyio -import pytest - -from shellctl.server import ShellctlConfig, ShellctlService -from shellctl_runtime.runner_exit import record_runner_exit -from shellctl_runtime.sanitize import run_sanitize_pty - -PACKAGE_ROOT = Path(__file__).resolve().parents[3] -SRC_PATH = PACKAGE_ROOT / "src" -HEAVY_IMPORT_ROOTS = { - "fastapi", - "httpx", - "pydantic", - "pygments", - "rich", - "sqlalchemy", - "shellctl", - "typer", - "uvicorn", -} - - -class ModuleHelpStats(TypedDict): - code: int - elapsed: float - deny: list[str] - - -def _pythonpath_env() -> dict[str, str]: - env = dict(os.environ) - current_pythonpath = env.get("PYTHONPATH") - env["PYTHONPATH"] = f"{SRC_PATH}{os.pathsep}{current_pythonpath}" if current_pythonpath else str(SRC_PATH) - return env - - -async def _initialize_database(state_dir: Path) -> None: - service = ShellctlService(ShellctlConfig(state_dir=state_dir, runtime_dir=state_dir / "run")) - await service.initialize_database() - await service.shutdown() - - -def _insert_job_row( - *, - state_dir: Path, - job_id: str, - status: str, - exit_code: int | None = None, - ended_at: str | None = None, - reason: str | None = None, - message: str | None = None, - updated_at: str = "2026-05-21T15:30:13Z", -) -> None: - db_path = state_dir / "shellctl.db" - with sqlite3.connect(db_path) as connection: - connection.execute( - """ - INSERT INTO jobs ( - 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 - ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - job_id, - f"jobs/{job_id}/script", - f"jobs/{job_id}/output.log", - "/tmp", - 80, - 24, - status, - f"shellctl-{job_id}", - f"shellctl-{job_id}:0.0", - exit_code, - reason, - message, - "2026-05-21T15:30:12Z", - "2026-05-21T15:30:13Z", - ended_at, - updated_at, - ), - ) - connection.commit() - - -def _fetch_job_row(state_dir: Path, job_id: str) -> tuple[object, ...]: - with sqlite3.connect(state_dir / "shellctl.db") as connection: - row = connection.execute( - """ - SELECT status, exit_code, ended_at, updated_at, reason, message - FROM jobs - WHERE job_id = ? - """, - (job_id,), - ).fetchone() - assert row is not None - return row - - -def _measure_module_help(module_name: str) -> ModuleHelpStats: - script = textwrap.dedent( - f""" - import contextlib - import io - import json - import runpy - import sys - import time - - sys.path.insert(0, {str(SRC_PATH)!r}) - sys.argv = [{module_name!r}, "--help"] - start = time.perf_counter() - stdout = io.StringIO() - stderr = io.StringIO() - code = 0 - with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): - try: - runpy.run_module({module_name!r}, run_name="__main__") - except SystemExit as exc: - if isinstance(exc.code, int): - code = exc.code - elapsed = time.perf_counter() - start - deny = sorted( - name - for name in sys.modules - if name.split(".", 1)[0] in {sorted(HEAVY_IMPORT_ROOTS)!r} - ) - print(json.dumps({{"code": code, "elapsed": elapsed, "deny": deny}})) - """ - ) - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - check=False, - cwd=PACKAGE_ROOT, - ) - assert result.returncode == 0, result.stderr - return cast(ModuleHelpStats, json.loads(result.stdout)) - - -def test_run_sanitize_pty_touches_ready_file_before_reading(tmp_path: Path) -> None: - ready_file = tmp_path / "ready" - output = io.BytesIO() - - class ReadyCheckingInput: - def __init__(self) -> None: - self._reads = 0 - - def read(self, _size: int) -> bytes: - self._reads += 1 - assert ready_file.exists() - return b"before\rafter\n" if self._reads == 1 else b"" - - run_sanitize_pty( - ready_file, - stdin=cast(BinaryIO, cast(object, ReadyCheckingInput())), - stdout=output, - ) - - assert output.getvalue() == b"after\n" - - -def test_removed_sanitizer_modules_are_not_importable() -> None: - assert importlib.util.find_spec("shellctl.sanitize_pty") is None - assert importlib.util.find_spec("shellctl.shared.sanitize") is None - - -def test_removed_sanitizer_helpers_are_not_shellctl_exports() -> None: - import shellctl as shellctl - import shellctl.shared as shared - - for module in (shellctl, shared): - for name in ("PtySanitizer", "sanitize_pty_output", "sanitize_pty_stream"): - assert not hasattr(module, name) - - -def test_runner_exit_matches_service_state_for_running_job(tmp_path: Path) -> None: - runtime_state_dir = tmp_path / "runtime-state" - service_state_dir = tmp_path / "service-state" - anyio.run(_initialize_database, runtime_state_dir) - anyio.run(_initialize_database, service_state_dir) - _insert_job_row(state_dir=runtime_state_dir, job_id="job-1", status="running") - _insert_job_row(state_dir=service_state_dir, job_id="job-1", status="running") - - record_runner_exit( - state_dir=runtime_state_dir, - job_id="job-1", - exit_code=7, - ended_at="2026-05-21T15:30:20Z", - ) - - async def _record_with_service() -> None: - service = ShellctlService( - ShellctlConfig( - state_dir=service_state_dir, - runtime_dir=service_state_dir / "run", - ) - ) - try: - await service.record_runner_exit("job-1", 7, "2026-05-21T15:30:20Z") - finally: - await service.shutdown() - - anyio.run(_record_with_service) - - assert _fetch_job_row(runtime_state_dir, "job-1") == _fetch_job_row( - service_state_dir, - "job-1", - ) - - -def test_runner_exit_is_idempotent_for_terminal_job(tmp_path: Path) -> None: - state_dir = tmp_path / "state" - anyio.run(_initialize_database, state_dir) - _insert_job_row( - state_dir=state_dir, - job_id="job-2", - status="failed", - exit_code=99, - ended_at="2026-05-21T15:30:18Z", - reason="runner_failed", - message="boom", - updated_at="2026-05-21T15:30:18Z", - ) - - before = _fetch_job_row(state_dir, "job-2") - record_runner_exit( - state_dir=state_dir, - job_id="job-2", - exit_code=7, - ended_at="2026-05-21T15:30:20Z", - ) - - assert _fetch_job_row(state_dir, "job-2") == before - - -def test_runner_exit_uses_explicit_busy_timeout_override(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - state_dir = tmp_path / "state" - anyio.run(_initialize_database, state_dir) - _insert_job_row(state_dir=state_dir, job_id="job-3", status="running") - - captured: dict[str, object] = {} - real_connect = sqlite3.connect - - class RecordingConnection: - def __init__(self, connection: sqlite3.Connection) -> None: - self._connection = connection - - def execute(self, sql: str, parameters: tuple[object, ...] = ()) -> sqlite3.Cursor: - if sql.startswith("PRAGMA busy_timeout="): - captured["pragma"] = sql - return self._connection.execute(sql, parameters) - - def commit(self) -> None: - self._connection.commit() - - def close(self) -> None: - self._connection.close() - - def fake_connect(path: str | Path, timeout: float) -> RecordingConnection: - captured["path"] = str(path) - captured["timeout"] = timeout - return RecordingConnection(real_connect(path, timeout=timeout)) - - monkeypatch.setattr("shellctl_runtime.runner_exit.sqlite3.connect", fake_connect) - - record_runner_exit( - state_dir=state_dir, - job_id="job-3", - exit_code=5, - ended_at="2026-05-21T15:30:20Z", - busy_timeout_ms=6789, - ) - - assert captured["path"] == str(state_dir / "shellctl.db") - assert captured["timeout"] == pytest.approx(6.789) - assert captured["pragma"] == "PRAGMA busy_timeout=6789" - - -def test_runner_exit_does_not_clobber_newer_terminal_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - state_dir = tmp_path / "state" - anyio.run(_initialize_database, state_dir) - _insert_job_row(state_dir=state_dir, job_id="job-4", status="running") - - real_connect = sqlite3.connect - - class RacingConnection: - def __init__(self, connection: sqlite3.Connection) -> None: - self._connection = connection - self._raced = False - - def __enter__(self) -> RacingConnection: - self._connection.__enter__() - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> bool | None: - return self._connection.__exit__(exc_type, exc, tb) - - def execute(self, sql: str, parameters: tuple[object, ...] = ()) -> sqlite3.Cursor: - if sql.lstrip().startswith("UPDATE jobs") and not self._raced: - self._raced = True - with real_connect(state_dir / "shellctl.db") as racing_connection: - racing_connection.execute( - """ - UPDATE jobs - SET status = ?, exit_code = ?, ended_at = ?, updated_at = ?, - reason = ?, message = ? - WHERE job_id = ? - """, - ( - "failed", - 91, - "2026-05-21T15:30:19Z", - "2026-05-21T15:30:19Z", - "pipe_failed", - "concurrent winner", - "job-4", - ), - ) - racing_connection.commit() - return self._connection.execute(sql, parameters) - - def commit(self) -> None: - self._connection.commit() - - def close(self) -> None: - self._connection.close() - - def fake_connect(path: str | Path, timeout: float = 5.0, *args: object, **kwargs: object) -> RacingConnection: - return RacingConnection(real_connect(path, timeout=timeout)) - - monkeypatch.setattr("shellctl_runtime.runner_exit.sqlite3.connect", fake_connect) - - record_runner_exit( - state_dir=state_dir, - job_id="job-4", - exit_code=7, - ended_at="2026-05-21T15:30:20Z", - ) - - assert _fetch_job_row(state_dir, "job-4") == ( - "failed", - 91, - "2026-05-21T15:30:19Z", - "2026-05-21T15:30:19Z", - "pipe_failed", - "concurrent winner", - ) - - -def test_python_m_shellctl_runtime_runner_exit_updates_row(tmp_path: Path) -> None: - state_dir = tmp_path / "state" - anyio.run(_initialize_database, state_dir) - _insert_job_row(state_dir=state_dir, job_id="job-cli", status="running") - - result = subprocess.run( - [ - sys.executable, - "-m", - "shellctl_runtime.runner_exit", - "--state-dir", - str(state_dir), - "--job-id", - "job-cli", - "--exit-code", - "7", - "--ended-at", - "2026-05-21T15:30:20Z", - ], - capture_output=True, - text=True, - check=False, - cwd=PACKAGE_ROOT, - env=_pythonpath_env(), - ) - - assert result.returncode == 0, result.stderr - assert _fetch_job_row(state_dir, "job-cli") == ( - "exited", - 7, - "2026-05-21T15:30:20Z", - "2026-05-21T15:30:20Z", - None, - None, - ) - - -def test_python_m_shellctl_runtime_runner_exit_reports_missing_job( - tmp_path: Path, -) -> None: - state_dir = tmp_path / "state" - anyio.run(_initialize_database, state_dir) - - result = subprocess.run( - [ - sys.executable, - "-m", - "shellctl_runtime.runner_exit", - "--state-dir", - str(state_dir), - "--job-id", - "missing-job", - "--exit-code", - "7", - "--ended-at", - "2026-05-21T15:30:20Z", - ], - capture_output=True, - text=True, - check=False, - cwd=PACKAGE_ROOT, - env=_pythonpath_env(), - ) - - assert result.returncode != 0 - assert "Unknown job id: missing-job" in result.stderr or "Unknown job id: missing-job" in result.stdout - - -def test_runtime_console_scripts_are_installed() -> None: - entry_points = { - entry_point.name: entry_point.value for entry_point in importlib.metadata.entry_points(group="console_scripts") - } - - assert entry_points["shellctl-sanitize-pty"] == "shellctl_runtime.sanitize:main" - assert entry_points["shellctl-runner-exit"] == "shellctl_runtime.runner_exit:main" - assert shutil.which("shellctl-sanitize-pty") is not None - assert shutil.which("shellctl-runner-exit") is not None - - -def test_runtime_modules_only_import_stdlib_or_runtime_helpers() -> None: - forbidden_imports = { - "anyio", - "fastapi", - "httpx", - "pydantic", - "pygments", - "rich", - "shellctl", - "sqlalchemy", - "sqlmodel", - "typer", - "uvicorn", - } - - for filename in ("sanitize.py", "runner_exit.py"): - tree = ast.parse((SRC_PATH / "shellctl_runtime" / filename).read_text("utf-8")) - imports: set[str] = set() - for node in ast.walk(tree): - if isinstance(node, ast.Import): - imports.update(alias.name for alias in node.names) - elif isinstance(node, ast.ImportFrom) and node.module is not None: - imports.add(node.module) - assert not any( - imported == forbidden or imported.startswith(f"{forbidden}.") - for imported in imports - for forbidden in forbidden_imports - ), (filename, sorted(imports)) - - -def test_python_m_shellctl_runtime_sanitize_module_stays_lightweight() -> None: - stats = _measure_module_help("shellctl_runtime.sanitize") - - assert stats["code"] == 0 - assert stats["elapsed"] < 0.35 - assert stats["deny"] == [] - - -def test_python_m_shellctl_runtime_runner_exit_module_stays_lightweight() -> None: - stats = _measure_module_help("shellctl_runtime.runner_exit") - - assert stats["code"] == 0 - assert stats["elapsed"] < 0.35 - assert stats["deny"] == [] - - -def test_python_m_shellctl_runtime_sanitize_sanitizes_output(tmp_path: Path) -> None: - ready_file = tmp_path / "ready" - result = subprocess.run( - [ - sys.executable, - "-m", - "shellctl_runtime.sanitize", - "--ready-file", - str(ready_file), - ], - input=b"before\rafter\n", - capture_output=True, - check=False, - cwd=PACKAGE_ROOT, - env=_pythonpath_env(), - ) - - assert result.returncode == 0, result.stderr.decode("utf-8", errors="replace") - assert ready_file.exists() - assert result.stdout.decode("utf-8") == "after\n" diff --git a/dify-agent/tests/local/test_packaging.py b/dify-agent/tests/local/test_packaging.py index 32e688b72bb..b11bcb650a8 100644 --- a/dify-agent/tests/local/test_packaging.py +++ b/dify-agent/tests/local/test_packaging.py @@ -7,12 +7,10 @@ from pathlib import Path PROJECT_ROOT = Path(__file__).resolve().parents[2] CLIENT_SHARED_DTO_DEPENDENCIES = { - "anyio>=4.12.1,<5.0.0", "httpx==0.28.1", "httpx2>=2.5.0,<3.0.0", "pydantic>=2.12.5,<2.13", "pydantic-ai-slim>=1.102.0,<2.0.0", - "typer>=0.16.1,<0.17", "typing-extensions>=4.12.2,<5.0.0", } @@ -28,13 +26,6 @@ SERVER_RUNTIME_DEPENDENCIES = { "uvicorn[standard]==0.46.0", } -SHELLCTL_SERVER_DEPENDENCIES = { - "aiosqlite>=0.21.0,<1.0.0", - "fastapi==0.136.0", - "sqlmodel>=0.0.24,<0.1.0", - "uvicorn[standard]==0.46.0", -} - GRPC_RUNTIME_DEPENDENCIES = { "grpclib[protobuf]>=0.4.9,<0.5.0", "protobuf>=6.33.5,<7.0.0", @@ -62,7 +53,6 @@ def test_project_dependencies_split_client_and_server_requirements() -> None: assert set(project["dependencies"]) == CLIENT_SHARED_DTO_DEPENDENCIES assert set(project["optional-dependencies"]["grpc"]) == GRPC_RUNTIME_DEPENDENCIES assert set(project["optional-dependencies"]["server"]) == SERVER_RUNTIME_DEPENDENCIES - assert set(project["optional-dependencies"]["shellctl-server"]) == SHELLCTL_SERVER_DEPENDENCIES assert set(pyproject["dependency-groups"]["dev"]) == DEV_DEPENDENCIES @@ -72,7 +62,7 @@ def test_default_package_discovery_excludes_example_packages() -> None: assert find_config["where"] == ["src"] assert "shellctl*" in find_config["include"] - assert "shellctl_runtime*" in find_config["include"] + assert "shellctl_runtime*" not in find_config["include"] assert "agenton_examples*" not in find_config["include"] assert "dify_agent_examples*" not in find_config["include"] @@ -81,42 +71,8 @@ def test_project_declares_console_scripts() -> None: pyproject = _read_pyproject() scripts = pyproject["project"]["scripts"] - assert scripts["dify-agent"] == "dify_agent.agent_stub.cli.main:main" assert scripts["dify-agent-stub-server"] == "dify_agent.agent_stub.server.cli:main" - assert scripts["shellctl"] == "shellctl.cli:main" - assert scripts["shellctl-sanitize-pty"] == "shellctl_runtime.sanitize:main" - assert scripts["shellctl-runner-exit"] == "shellctl_runtime.runner_exit:main" - - -def test_local_sandbox_dockerfile_installs_stub_client_and_shellctl() -> None: - dockerfile = (PROJECT_ROOT / "docker" / "local-sandbox" / "Dockerfile").read_text(encoding="utf-8") - - assert "ARG NODE_VERSION=22.22.1" in dockerfile - assert "ARG PNPM_VERSION=11.9.0" in dockerfile - assert "ARG UV_VERSION=0.8.9" in dockerfile - assert "ARG DIFY_AGENT_TOOL_SPEC=.[grpc,shellctl-server]" in dockerfile - assert "SHELL_SESSION_MANAGER_TOOL_SPEC" not in dockerfile - assert "DIFY_AGENT_STUB_DRIVE_BASE=/mnt/drive" not in dockerfile - assert 'ENV PATH="${UV_TOOL_BIN_DIR}:${PATH}"' not in dockerfile - assert "UV_TOOL_DIR=/opt/dify-agent-tools/envs" in dockerfile - assert "UV_TOOL_BIN_DIR=/opt/dify-agent-tools/bin" in dockerfile - assert 'ENV PATH="/opt/dify-agent-tools/bin:${PATH}"' in dockerfile - assert "bash" in dockerfile - assert "git" in dockerfile - assert "jq" in dockerfile - assert "openssh-client" in dockerfile - assert "ripgrep" in dockerfile - assert 'node_dist="node-v${NODE_VERSION}-linux-${node_arch}"' in dockerfile - assert 'curl -fsSLO "https://nodejs.org/dist/v${NODE_VERSION}/SHASUMS256.txt"' in dockerfile - assert 'tar -xJf "${node_dist}.tar.xz" -C /usr/local --strip-components=1' in dockerfile - assert 'npm install --global "pnpm@${PNPM_VERSION}"' in dockerfile - assert "uv export --frozen --no-dev --all-extras --no-emit-project --no-hashes" in dockerfile - assert "> /tmp/dify-agent-constraints.txt" in dockerfile - assert '--constraints /tmp/dify-agent-constraints.txt --link-mode=copy "${DIFY_AGENT_TOOL_SPEC}"' in dockerfile - assert "COPY --from=tools /opt/dify-agent-tools/envs /opt/dify-agent-tools/envs" in dockerfile - assert "COPY --from=tools /opt/dify-agent-tools/bin /opt/dify-agent-tools/bin" in dockerfile - assert "VIRTUAL_ENV" not in dockerfile - assert "uv sync" not in dockerfile - assert "mkdir -p /mnt/drive" in dockerfile - assert "chown dify:dify /home" in dockerfile - assert '["shellctl", "serve", "--listen", "0.0.0.0:5004"]' in dockerfile + assert "dify-agent" not in scripts + assert "shellctl" not in scripts + assert "shellctl-sanitize-pty" not in scripts + assert "shellctl-runner-exit" not in scripts diff --git a/dify-agent/uv.lock b/dify-agent/uv.lock index a7c321f5137..2531ff07db7 100644 --- a/dify-agent/uv.lock +++ b/dify-agent/uv.lock @@ -19,15 +19,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, ] -[[package]] -name = "aiosqlite" -version = "0.22.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4e/8a/64761f4005f17809769d23e518d915db74e6310474e733e3593cfc854ef1/aiosqlite-0.22.1.tar.gz", hash = "sha256:043e0bd78d32888c0a9ca90fc788b38796843360c855a7262a532813133a0650", size = 14821, upload-time = "2025-12-23T19:25:43.997Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/00/b7/e3bf5133d697a08128598c8d0abc5e16377b51465a33756de24fa7dee953/aiosqlite-0.22.1-py3-none-any.whl", hash = "sha256:21c002eb13823fad740196c5a2e9d8e62f6243bd9e7e4a1f87fb5e44ecb4fceb", size = 17405, upload-time = "2025-12-23T19:25:42.139Z" }, -] - [[package]] name = "annotated-doc" version = "0.0.4" @@ -593,12 +584,10 @@ name = "dify-agent" version = "1.16.0rc1" source = { editable = "." } dependencies = [ - { name = "anyio" }, { name = "httpx" }, { name = "httpx2" }, { name = "pydantic" }, { name = "pydantic-ai-slim" }, - { name = "typer" }, { name = "typing-extensions" }, ] @@ -618,12 +607,6 @@ server = [ { name = "redis" }, { name = "uvicorn", extra = ["standard"] }, ] -shellctl-server = [ - { name = "aiosqlite" }, - { name = "fastapi" }, - { name = "sqlmodel" }, - { name = "uvicorn", extra = ["standard"] }, -] [package.dev-dependencies] dev = [ @@ -644,10 +627,7 @@ docs = [ [package.metadata] requires-dist = [ - { name = "aiosqlite", marker = "extra == 'shellctl-server'", specifier = ">=0.21.0,<1.0.0" }, - { name = "anyio", specifier = ">=4.12.1,<5.0.0" }, { name = "fastapi", marker = "extra == 'server'", specifier = "==0.136.0" }, - { name = "fastapi", marker = "extra == 'shellctl-server'", specifier = "==0.136.0" }, { name = "graphon", marker = "extra == 'server'", specifier = "==0.5.2" }, { name = "grpclib", extras = ["protobuf"], marker = "extra == 'grpc'", specifier = ">=0.4.9,<0.5.0" }, { name = "httpx", specifier = "==0.28.1" }, @@ -661,13 +641,10 @@ requires-dist = [ { name = "pydantic-ai-slim", extras = ["anthropic", "google", "openai"], marker = "extra == 'server'", specifier = ">=1.85.1,<2.0.0" }, { name = "pydantic-settings", marker = "extra == 'server'", specifier = ">=2.12.0,<3.0.0" }, { name = "redis", marker = "extra == 'server'", specifier = ">=7.4.0,<8.0.0" }, - { name = "sqlmodel", marker = "extra == 'shellctl-server'", specifier = ">=0.0.24,<0.1.0" }, - { name = "typer", specifier = ">=0.16.1,<0.17" }, { name = "typing-extensions", specifier = ">=4.12.2,<5.0.0" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = "==0.46.0" }, - { name = "uvicorn", extras = ["standard"], marker = "extra == 'shellctl-server'", specifier = "==0.46.0" }, ] -provides-extras = ["grpc", "server", "shellctl-server"] +provides-extras = ["grpc", "server"] [package.metadata.requires-dev] dev = [ @@ -878,61 +855,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2f/e6/36a3981cd44e7a40a7cd7d374e26f01e02dd49410c5fbbd7df248750d5fb/graphon-0.5.2-py3-none-any.whl", hash = "sha256:11f89399e67ed1ddd2ce1c336accd9c4ad5b8fe2741f9167e6085af0b325cd14", size = 381908, upload-time = "2026-06-16T04:06:20.453Z" }, ] -[[package]] -name = "greenlet" -version = "3.5.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6d/6e/802acd792aebb2256fbbee8cacf2727faaeb6f240ac11008f09eae4414bc/greenlet-3.5.1.tar.gz", hash = "sha256:5a56aeb7d5d9cc4b3a735efb5095bd4b4f6f0e4f93e5ca876d0e2315137b7829", size = 197356, upload-time = "2026-05-20T15:05:03.917Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c4/37/4549f149c9797c21b32c2683c33522af22522099de128b2406672526d005/greenlet-3.5.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:fa4f98af3a528f0c3fd592a26df7f376f93329c8f4d987f6bb979057af8bf5e2", size = 286220, upload-time = "2026-05-20T13:07:28.463Z" }, - { url = "https://files.pythonhosted.org/packages/38/ff/a4f436709716965eaab9f36ea7b906c8a927fbe32fb1372a2071d964f6b1/greenlet-3.5.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ffea73584b216150eab159b6d12348fb253e68757974de1e2c40d8a318ac89ed", size = 601585, upload-time = "2026-05-20T14:00:06.141Z" }, - { url = "https://files.pythonhosted.org/packages/65/ad/54bc3fcee3ad368a61b19b67d88117f7a8c29727bf71fffdeda81fbd946e/greenlet-3.5.1-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1072b4f9edcc1e192d9283a66a3e68d6b84c561de33a83d7858beb9ba1effe10", size = 614215, upload-time = "2026-05-20T14:05:42.675Z" }, - { url = "https://files.pythonhosted.org/packages/40/69/b91cda0647df839483201545913514c2827ebea5e5ccdf931842763bc127/greenlet-3.5.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:add5217d68b31130f0beca584d7fef4878327d2e31642b66618a14eef312b63b", size = 611358, upload-time = "2026-05-20T13:14:26.37Z" }, - { url = "https://files.pythonhosted.org/packages/59/90/3cf77e080350cd02fa307bb2abf05df48f4482c240275bbd2c203ba8bb1c/greenlet-3.5.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a5ea42a752d47a145eae922b605cd1634665ac3d5ec1e72402d5048e8d60d207", size = 1570475, upload-time = "2026-05-20T14:02:25.29Z" }, - { url = "https://files.pythonhosted.org/packages/65/2c/18cece62045e74598c3c393f70dce4a63f56222015ba29a5d4eeb04f764c/greenlet-3.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c5551170cf4f5ff5623e9af81323751979fee2c731e2287b61f73cd27257b823", size = 1635625, upload-time = "2026-05-20T13:14:34.027Z" }, - { url = "https://files.pythonhosted.org/packages/30/f5/310d104ddf41eb5a70f4c268d22508dfb0c3c8e86fec152be34d0d2ed819/greenlet-3.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:3c8bb982ad117d29478ef8f5533e97df21f1e2befd17a299257b0c96d1371c0b", size = 238791, upload-time = "2026-05-20T13:10:39.018Z" }, - { url = "https://files.pythonhosted.org/packages/62/90/ceca11f504cd23a8047a3dea31919adc48df9b626dd0c13f0d858734fdfd/greenlet-3.5.1-cp312-cp312-win_arm64.whl", hash = "sha256:80eb4b04dadc4e67df3fae179a32c4706a3f495bc7f22fc8a81115d5f5512188", size = 235580, upload-time = "2026-05-20T13:08:45.056Z" }, - { url = "https://files.pythonhosted.org/packages/27/69/7f7e5372d998b81001899b1c0823c957aa413ba0f2662e65821611cc31e4/greenlet-3.5.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:51518ff74664078fc51bffcc6fc529b0df5ae58da192691cee765d45ce944a2b", size = 285060, upload-time = "2026-05-20T13:08:51.899Z" }, - { url = "https://files.pythonhosted.org/packages/b1/bf/387f9b6b865fd2ae0d0be09e0004827295a01b71be76ed350dd1e28a91a4/greenlet-3.5.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ffdb3c0bb002c99cd8f298957e046c3dbf6006b5b7cdf11a4e19194624a0a0a", size = 604370, upload-time = "2026-05-20T14:00:07.492Z" }, - { url = "https://files.pythonhosted.org/packages/32/f5/169ce3d4e4c67291bd18f8cbe0299c9f3e45102c7f1fb3c14780c93e4532/greenlet-3.5.1-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7715a5a2c3378ba602c3a440558261e13a820bb53a82693aacd7b7f6d964e283", size = 616987, upload-time = "2026-05-20T14:05:44.237Z" }, - { url = "https://files.pythonhosted.org/packages/ee/e5/7f2e41d5273be07e77560d61ea4e56485b4d6c316d2a84518c62d1364061/greenlet-3.5.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc71ff466927a201b08305acac451ebe1aedfcea002f62f1f2f2ac2ac1e6a135", size = 613911, upload-time = "2026-05-20T13:14:27.539Z" }, - { url = "https://files.pythonhosted.org/packages/c5/a4/fbdc67579b73615a1f91615e814303cc71e06128f7baaba87be79b8fb90c/greenlet-3.5.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cd443683db272ebaaca03af98c0b063ab30db70ea8a31a1559f35e3f7b744ccd", size = 1570689, upload-time = "2026-05-20T14:02:27.225Z" }, - { url = "https://files.pythonhosted.org/packages/e6/b4/77abbe35078be39718a46cd49caf16bceb35662f97a34101dca28aa98e47/greenlet-3.5.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:089fff7a6ce8d9316d1f65ebc00273a56be258c1725b32b94de90a3a979557e1", size = 1635602, upload-time = "2026-05-20T13:14:36.344Z" }, - { url = "https://files.pythonhosted.org/packages/37/f7/129f27ca700845b8ee8ca88ce7f43435a1239c2eddb7677fc938822762cf/greenlet-3.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:110a1ca7b49b014b097f6078272c3f4ed31af45b254de5228b79adba879f6af9", size = 238683, upload-time = "2026-05-20T13:11:50.57Z" }, - { url = "https://files.pythonhosted.org/packages/6d/5c/a485a36e87df8d8fd0632ee01511244f5156a20ed3746cc6599340326395/greenlet-3.5.1-cp313-cp313-win_arm64.whl", hash = "sha256:f16ba1efc0715b680a18b8123d90dad887c6112ae3555b4b5c32c149540c6b4e", size = 235499, upload-time = "2026-05-20T13:12:42.028Z" }, - { url = "https://files.pythonhosted.org/packages/8a/cb/c62454606daf5640369c94d8a9dd540599b1bfc090e2d2180cb77f4038d2/greenlet-3.5.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:d8ab31c9de8651a2facdd5c5bb0011f2380dd1a7af78ce2adf4b56095294fc07", size = 285579, upload-time = "2026-05-20T13:08:56.396Z" }, - { url = "https://files.pythonhosted.org/packages/ec/71/c4270398c2eba968a6071af1dfbdcaeee6ec1c24bc8b435b8cc452700da6/greenlet-3.5.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e300185139abc337ade480c327183adf42a875ac7181bfe66d7d4efea31fbea", size = 651106, upload-time = "2026-05-20T14:00:09.448Z" }, - { url = "https://files.pythonhosted.org/packages/1a/ab/71e34b78a44ec271fb5f550c17bc46d301ddc5953890d935f270b0dcdb5a/greenlet-3.5.1-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7ffdb990dcaa0234cf9845aead5df2e3c3a8b6507d409274dd87e0d5ab05ffc2", size = 663478, upload-time = "2026-05-20T14:05:45.88Z" }, - { url = "https://files.pythonhosted.org/packages/77/96/4efd6fa5c62c85426a0c19077a586258ebc3a2a146ff2493e4312a697a22/greenlet-3.5.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2f82b3597e9d83b63408affed0b48fd0f54935edac4302237b9a837be0dae33c", size = 660800, upload-time = "2026-05-20T13:14:29.129Z" }, - { url = "https://files.pythonhosted.org/packages/7a/e0/6c71401a25cac7000261304e866a2f2cc04dc74810d40e2f118aa4799495/greenlet-3.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c0141e37414c10164e702b8fb1473304221ad98f71600850c6ef7ff4880feba0", size = 1617518, upload-time = "2026-05-20T14:02:28.662Z" }, - { url = "https://files.pythonhosted.org/packages/41/26/c5c06643e8c0af9e7bf18e16cb51d0ab7625155f0392e1c9015d66d556cd/greenlet-3.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:50ae25a67bea74ea41fb14b960bc532df73eb713417b2d61892dced82fe8d3bc", size = 1681593, upload-time = "2026-05-20T13:14:39.417Z" }, - { url = "https://files.pythonhosted.org/packages/8a/bd/e11a108317485075e68af9d23039619b86b28130c3b50d227d42edece64b/greenlet-3.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:8a17c42330e261299766b75ac1ea32caa437a9453c8f65d16a13140db378ecd3", size = 239800, upload-time = "2026-05-20T13:09:30.128Z" }, - { url = "https://files.pythonhosted.org/packages/47/f8/8e8e8417b7bf28639a5a56356ef934d0375e1d0c70a57e04d7701e870ffe/greenlet-3.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:7b5f5fae05b8ac6d176a61b60c394a8cbdc2b5b91b81793066e68745cf165e54", size = 236862, upload-time = "2026-05-20T13:09:10.498Z" }, - { url = "https://files.pythonhosted.org/packages/90/12/41bf27fde4d3605d3773ae57751eda182b8be2f5398011c041173b1d9534/greenlet-3.5.1-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:ea8da1e900d758d078810d4255d8c6aa572181896a31ec79d779eb79c3adc9ad", size = 293637, upload-time = "2026-05-20T13:12:35.529Z" }, - { url = "https://files.pythonhosted.org/packages/44/44/ba14b23e9757707050c2f397d305bbcae62e5d7cad122f8b6baec5ae4a1f/greenlet-3.5.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a19570c52a21420dcbc94e661994bc325c0b5b11304540fed514586da5dc8f2e", size = 650840, upload-time = "2026-05-20T14:00:11.079Z" }, - { url = "https://files.pythonhosted.org/packages/a8/37/5ddc2b686a6844f91abecef43411842426da2e1573f60b49ecf2547f4ae1/greenlet-3.5.1-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3d955c89b75eeca4723d7cc14135f393cd47c32e2a6cb4a8e4c6e760a26b0986", size = 656416, upload-time = "2026-05-20T14:05:47.118Z" }, - { url = "https://files.pythonhosted.org/packages/e1/f0/d17510297c35a2992712f0bf84de3779749999f7d3d63aa1f09db7c62dbe/greenlet-3.5.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de2daaaebd1a5aa88c49045b6baf9310b3263796bd88db713edf37cf53e7bb4e", size = 654397, upload-time = "2026-05-20T13:14:30.696Z" }, - { url = "https://files.pythonhosted.org/packages/37/eb/147387705bb89092645b012586e7273cb5ed3c90ef7eaf3a69173eaf0209/greenlet-3.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3bfbd69cc349e43bf3a8ae1c85548ff0718efc887615c2db16c3833d7b0b072d", size = 1614469, upload-time = "2026-05-20T14:02:30.192Z" }, - { url = "https://files.pythonhosted.org/packages/a6/4e/37ee0da7732b7aa9896f17e15579a9df34b9fcb9dd494f0adfa749af6623/greenlet-3.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:4378720dd888136c27215a0214d32a4d37c3852765d45bc37aad0623423cfd78", size = 1675115, upload-time = "2026-05-20T13:14:40.972Z" }, - { url = "https://files.pythonhosted.org/packages/57/f3/97dfcf4a6eb5077f8a672234216fb5923eb89f2cab7081cb10b2cf75b605/greenlet-3.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:45718441607f9325d948db98cbc691276059316d0358c188c246da4e1d4d23d2", size = 245246, upload-time = "2026-05-20T13:12:22.646Z" }, - { url = "https://files.pythonhosted.org/packages/5d/73/d7f72e34b582f694f4a9b248162db7b09cc458a259ba8f0c0bfa1a34ea7d/greenlet-3.5.1-cp315-cp315-macosx_11_0_universal2.whl", hash = "sha256:2baee5ca02031757ffe8cc3d69f0cc0aec7065ce362622da74f32d3bcab1c541", size = 285575, upload-time = "2026-05-20T13:12:07.043Z" }, - { url = "https://files.pythonhosted.org/packages/df/59/fa9c6e87dc8ad27a95dabe2f29f372b733d05a8a67470f6c901ed9975655/greenlet-3.5.1-cp315-cp315-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b1ec3274918a81d3ea778b9e75b56b72b33f300edb6cf7f3a7fe1dae56683de", size = 656428, upload-time = "2026-05-20T14:00:12.556Z" }, - { url = "https://files.pythonhosted.org/packages/f6/f9/e753408871eaa61dfe35e619cfc67512b036fde99893685d50eea9e07146/greenlet-3.5.1-cp315-cp315-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:111e2390ffffc47d5840b01711dd7fac07d4c09283d0283e7f3264b14e284c64", size = 667064, upload-time = "2026-05-20T14:05:48.662Z" }, - { url = "https://files.pythonhosted.org/packages/96/27/5565b5b40389f1c7753003a07e21892fda8660926787036d5bc0308b8113/greenlet-3.5.1-cp315-cp315-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e630136e905fe5ff43e86945ae41220b6d1470956a39220e708110ac48d01ea5", size = 665697, upload-time = "2026-05-20T13:14:32.943Z" }, - { url = "https://files.pythonhosted.org/packages/cf/82/e7de4178c0c2d1c9a5a3be3cc0b33e46a85b3ee4a77c071bf7ad8600e079/greenlet-3.5.1-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:975eac34b44a7077ca4d421348455b94f0f518246a7f14bc6d2fdcfe5b584368", size = 1621256, upload-time = "2026-05-20T14:02:31.91Z" }, - { url = "https://files.pythonhosted.org/packages/00/10/f2dddcf7dacac17dfc68691809589adad06135eb28930429cf58a6467a2f/greenlet-3.5.1-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:9ab3c3a0b2ae6198e67c898dad5215a49f9ae0d0081b3c3ec59f333e39eeca26", size = 1685956, upload-time = "2026-05-20T13:14:42.55Z" }, - { url = "https://files.pythonhosted.org/packages/22/17/4a232b32133230ada52f70e9d7f5b65b0caef8772f01849bd8d149e7e4ca/greenlet-3.5.1-cp315-cp315-win_amd64.whl", hash = "sha256:cbfc69be86e10dcfef5b1e6269d1d6926552aa89ee39e1de3353360c1b6989ab", size = 239802, upload-time = "2026-05-20T13:13:15.481Z" }, - { url = "https://files.pythonhosted.org/packages/c2/ae/4e623a7e6d4d2a5f4cb8e4c82de4169fc637942caae68d6e676b8a128ac5/greenlet-3.5.1-cp315-cp315-win_arm64.whl", hash = "sha256:92fd6d44ac5e5a887c8a5dc4a8ba0ba908527c31c12f78c6bc7dcfe8aab279f6", size = 236853, upload-time = "2026-05-20T13:15:37.301Z" }, - { url = "https://files.pythonhosted.org/packages/7a/57/816d9cff29119da3505b3d6a5e14a8af89006ac36f47f891ff293ee05af1/greenlet-3.5.1-cp315-cp315t-macosx_11_0_universal2.whl", hash = "sha256:a6fdf2433a5441ef9a95464f7c3e674775da1c8c1177fff311cee1acad4626ed", size = 293877, upload-time = "2026-05-20T13:10:19.078Z" }, - { url = "https://files.pythonhosted.org/packages/23/a1/59b0a7c7d140ff1a75626680b9a9899b79a9176cab298b394968fb023295/greenlet-3.5.1-cp315-cp315t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7546556f0d649f99f6a361098a55f761181bb2ea12ff150bb16d26092ad88244", size = 655333, upload-time = "2026-05-20T14:00:14.758Z" }, - { url = "https://files.pythonhosted.org/packages/72/1b/5efe127597625042218939d01855109f352779050768b670b52edcc16a6c/greenlet-3.5.1-cp315-cp315t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d5ee3ea898009fa898f85f9982255d35278c477bebe185beca249cab42d4526c", size = 659443, upload-time = "2026-05-20T14:05:50.159Z" }, - { url = "https://files.pythonhosted.org/packages/6c/6d/c404246ea4d22d097a7426d0efb5b781bd7eb67715f09e79001bd552ab18/greenlet-3.5.1-cp315-cp315t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a5c81f74d204d3edd136ebfd50dce53acbb776995d721a0fe801626cfc93b8cd", size = 658356, upload-time = "2026-05-20T13:14:35.091Z" }, - { url = "https://files.pythonhosted.org/packages/51/02/f8ee37fb6d2219329f350af241c27fcf12df57e723d11f6fc6d3bacdadaa/greenlet-3.5.1-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:2c18ef16bf6d4dd410e4dd52996888ea1497be26892fe5bbc73580aba4287b8e", size = 1619216, upload-time = "2026-05-20T14:02:33.403Z" }, - { url = "https://files.pythonhosted.org/packages/93/c5/3dc9475ace2c7a3680da12372cddd7f1ac874eb410a1ac48d3e9dab83782/greenlet-3.5.1-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:17d86354f0ae6b61bf9be5148d0dd34e06c3cb7c602c671f79f29ac3b150e659", size = 1678427, upload-time = "2026-05-20T13:14:43.71Z" }, - { url = "https://files.pythonhosted.org/packages/df/4e/750c15c317a41ffb36f0bf40b933e3d744a7dede61889f74443ea69690cf/greenlet-3.5.1-cp315-cp315t-win_amd64.whl", hash = "sha256:e7516cf6ae6b8a582c2770a0caed47b8a48373ed732c33d69a72913ae6ac923e", size = 245225, upload-time = "2026-05-20T13:13:59.366Z" }, - { url = "https://files.pythonhosted.org/packages/4f/fd/d3baea2eeb7b617efd47e87ca06e2ec2c6118d303aa9e918e0ce16eadc10/greenlet-3.5.1-cp315-cp315t-win_arm64.whl", hash = "sha256:5028648bf2253ec4745add746129d3904121fa7fe871a76bed23c5720573ce0a", size = 239590, upload-time = "2026-05-20T13:13:37.382Z" }, -] - [[package]] name = "griffelib" version = "2.0.2" @@ -3651,61 +3573,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/78/d1a1a026ef3af911159398c939b1509d5c36fe524c7b644f34a5146c4e16/spacy_loggers-1.0.5-py3-none-any.whl", hash = "sha256:196284c9c446cc0cdb944005384270d775fdeaf4f494d8e269466cfa497ef645", size = 22343, upload-time = "2023-09-11T12:26:50.586Z" }, ] -[[package]] -name = "sqlalchemy" -version = "2.0.50" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "greenlet", marker = "platform_machine == 'AMD64' or platform_machine == 'WIN32' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'ppc64le' or platform_machine == 'win32' or platform_machine == 'x86_64'" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/57/da/6fbf010c8ebb347679d0d100b22fe9ba5e13fd04046c5df7280d2f0bf706/sqlalchemy-2.0.50.tar.gz", hash = "sha256:af5607d11ef90fd6a5c0549fe0045dce1663d427426bcfb506dcb5346a85a3b9", size = 9907424, upload-time = "2026-05-24T19:20:04.018Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/be/b0/a9d19b43f38f878b1278bca5b00b909f7540d41494396dd2561f9ad0956d/sqlalchemy-2.0.50-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:23ae23d8b9d344d30d0a92f06d45825024a5790f1c1dd4cf452636a50d3e58cb", size = 2159807, upload-time = "2026-05-24T19:27:53.086Z" }, - { url = "https://files.pythonhosted.org/packages/f5/2c/191dd58a248fd2cfd4780fa82c375c505e4ad98c8b522fa69ec492130d77/sqlalchemy-2.0.50-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:47b71b933e7b4ebad407c8fdfd70d2c4f08b78b3238bb30eebdd6eb32ca51b89", size = 3343358, upload-time = "2026-05-24T20:09:29.279Z" }, - { url = "https://files.pythonhosted.org/packages/8a/2b/514fce8a7df81cf5bad7ff7865de7ac0c5776a38cc043475c4703eb7fe8b/sqlalchemy-2.0.50-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:110fdac56ace278949f00de805edacbd6141e382d992f9ba28238b3a0827a600", size = 3357994, upload-time = "2026-05-24T20:17:13.495Z" }, - { url = "https://files.pythonhosted.org/packages/35/a6/a0e283f5494f92b0d77e319ff77e437b1ffe4a051ba67c81d53234825475/sqlalchemy-2.0.50-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0f5e4ac70e9e757f6b3e87c0491ff034442ecd8dfd36d041a50564c322dafc0e", size = 3289399, upload-time = "2026-05-24T20:09:32.239Z" }, - { url = "https://files.pythonhosted.org/packages/b7/96/1b07325ba71752d6a028b77d07bed1483ad545f794e8b1dc89b3ba3b3c68/sqlalchemy-2.0.50-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:724f3dcbe53dd0151e3cb5e7ec4ba4c620bede579caacd16275dc35ce06e8615", size = 3321216, upload-time = "2026-05-24T20:17:15.581Z" }, - { url = "https://files.pythonhosted.org/packages/ed/8e/bad6ed253e8a99edfc99af02f7173ec48a1d3ed1b9b35a1b8bc1700900cc/sqlalchemy-2.0.50-cp312-cp312-win32.whl", hash = "sha256:1208050441471d003b7c8cb4054fb084f185cf35ac3f0ea270803865bca9939a", size = 2119194, upload-time = "2026-05-24T19:50:04.943Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2d/314a6690dda4b9cfc571eab1a63cf6fe6e1470aa3759ccda6aa016ee0f5a/sqlalchemy-2.0.50-cp312-cp312-win_amd64.whl", hash = "sha256:9d1af51558029a156a70986b7df88f042b3d158d7c8d8fb5072912d4b32d89c7", size = 2146186, upload-time = "2026-05-24T19:50:06.74Z" }, - { url = "https://files.pythonhosted.org/packages/0b/c4/c42356b527296e9862f67990efce31ef78b4cf69cd3f80873a528a060320/sqlalchemy-2.0.50-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:06a9210bdc5f4298cff0781087e2ff45683922252dacc452846373a58761f093", size = 2156697, upload-time = "2026-05-24T19:27:54.764Z" }, - { url = "https://files.pythonhosted.org/packages/60/a1/b1a70e3c4365ac7fe9e347f3710f19b562c866fb96d45e3c891588789a7b/sqlalchemy-2.0.50-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b53784972ade4f8174b9aa661f31a06f8a936d2cfdd602913ff3c6dd40ae873", size = 3284260, upload-time = "2026-05-24T20:09:34.195Z" }, - { url = "https://files.pythonhosted.org/packages/3f/4a/f3ac3caa19f263d57b0a47f8c91bbf56583dc2d3fc63acfbf644abb24fe0/sqlalchemy-2.0.50-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31648fa14460537e768a7303b078e4344d208e0d23e06867c1f376a227ed82db", size = 3302280, upload-time = "2026-05-24T20:17:17.825Z" }, - { url = "https://files.pythonhosted.org/packages/66/55/ccada3e3d62254587819749a0bc69f41173eb48a6e385d10e66d32a9c88e/sqlalchemy-2.0.50-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:03f4323c980ad0e918cc9e5369b015f759f4e534db5bbaf4dc36832c10d05064", size = 3231580, upload-time = "2026-05-24T20:09:36.406Z" }, - { url = "https://files.pythonhosted.org/packages/05/f6/6809349130a2de0e109e7f00fd7d431da9565b9b2868b32ee684754f672b/sqlalchemy-2.0.50-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2b9dcc43afef8ac157cd92fce96985d6b8b0cfbd3df4d666f66b4d55a75d202f", size = 3269375, upload-time = "2026-05-24T20:17:20.34Z" }, - { url = "https://files.pythonhosted.org/packages/48/84/278a811ef4e07be9c89dc5cdd7be833268509a66a68c4897cf585e67428f/sqlalchemy-2.0.50-cp313-cp313-win32.whl", hash = "sha256:60922d6599065ddca2c6f376b9aa2f41a6b85a271725e0909490bbc50b1998a5", size = 2117229, upload-time = "2026-05-24T19:50:08.215Z" }, - { url = "https://files.pythonhosted.org/packages/f6/1c/067cc6187ed32d2ec222fe6d2643acc1659a6d0659f8a7cbc5ad3ae83280/sqlalchemy-2.0.50-cp313-cp313-win_amd64.whl", hash = "sha256:287086e67275a212c4582d166a6fb03a65ccc5551d80866270ce0dd9f34eccd3", size = 2143126, upload-time = "2026-05-24T19:50:09.691Z" }, - { url = "https://files.pythonhosted.org/packages/df/32/10ac51b4be7cdecd7e93d069251c86dfbf70b7adbd7c67b48ccea6c49e1c/sqlalchemy-2.0.50-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c966932507a4d7d0a37314927dbfcd89720e3f37d2a1e3352e7ae7939fa8e8a0", size = 2158519, upload-time = "2026-05-24T19:27:56.472Z" }, - { url = "https://files.pythonhosted.org/packages/5a/76/e703d2f7681d7d66c4c891af3f07c7ccf4c76ad7f18351de035b5eda007a/sqlalchemy-2.0.50-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:faffef4bcc20a1892e65e155293d99d60855bbbc79250ab712819cfd56a8e6bb", size = 3282063, upload-time = "2026-05-24T20:09:38.57Z" }, - { url = "https://files.pythonhosted.org/packages/31/26/ef168b184a25701f9995e8fb7e503fafd7a99c1c77cda1bc1a26ea2ed486/sqlalchemy-2.0.50-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c206aec519a2e7bd08abbfb33436e325fd22c632d9c21a9047e376ce241646e", size = 3287069, upload-time = "2026-05-24T20:17:21.942Z" }, - { url = "https://files.pythonhosted.org/packages/c2/15/765acc2bc693bccc43ca4a95d5b69750da8aaf6db1b5c616536e087f8920/sqlalchemy-2.0.50-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:bef4ac756363227ef6402a75fee025a4bc690f92328e825868939b3b3a446a6d", size = 3230453, upload-time = "2026-05-24T20:09:40.398Z" }, - { url = "https://files.pythonhosted.org/packages/63/61/08e03c3adbf5db0087a0b6816746fec8f3032fb2f7fc899a9bb9b2a48ce4/sqlalchemy-2.0.50-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:96fbee6b19c19cd1556c8bf9419447cf2ec149ffcab7ab64348c23e54ef8547f", size = 3252413, upload-time = "2026-05-24T20:17:24.067Z" }, - { url = "https://files.pythonhosted.org/packages/03/0c/370a1f2db38436c615e10134c8a37de3688e74084792380695f3f5083860/sqlalchemy-2.0.50-cp314-cp314-win32.whl", hash = "sha256:8f00e3eb43ba30eb1b238ee03a8a62309486d1321eda3328bb611e0340033ad8", size = 2120063, upload-time = "2026-05-24T19:50:11.08Z" }, - { url = "https://files.pythonhosted.org/packages/7f/a0/fe92bb9817863bc13ba093bda931979a26cc2ca69f8e8f26d07add3d7c6f/sqlalchemy-2.0.50-cp314-cp314-win_amd64.whl", hash = "sha256:15708c613cd5005b7dffe1f66ee6a63ee8f5e46799f71c70ebad74178c676a39", size = 2145830, upload-time = "2026-05-24T19:50:12.452Z" }, - { url = "https://files.pythonhosted.org/packages/cc/ff/e5640a98a0b2f491eb8fde10fb6c773621a2e44340de231fafcc9370f4a9/sqlalchemy-2.0.50-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3699dac4be410e97049a1658e9480da9cde956594aa0f3aebc60b88f21c5ba70", size = 2178435, upload-time = "2026-05-24T19:42:58.889Z" }, - { url = "https://files.pythonhosted.org/packages/b7/85/337116e186f1236375b5fb70c21cfac98e8e8ab0d3a47be838dc47a59e08/sqlalchemy-2.0.50-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f96233858e3df43932ac11589e22520da6e8aeb624b03fedfeebb0e8ea213086", size = 3566059, upload-time = "2026-05-24T20:01:20.848Z" }, - { url = "https://files.pythonhosted.org/packages/96/34/bb0e190e161c3c2c24314a65add57218be14a4a9486886b7f5047c1ff7c8/sqlalchemy-2.0.50-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c4e70c46fad30c3bcc6a4708bc0130a3173e11a5b25f0ea4a9d8911b450f1f52", size = 3535366, upload-time = "2026-05-24T20:03:56.768Z" }, - { url = "https://files.pythonhosted.org/packages/df/5a/a7f759f97e4fd499c5d4e4488c760d5a7fbecf3028b465a04274fcd52384/sqlalchemy-2.0.50-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1918a3cf564d16d95bca7301005f41ab2ad50b07cd3b9da50d3ed986db148d6a", size = 3474879, upload-time = "2026-05-24T20:01:23.058Z" }, - { url = "https://files.pythonhosted.org/packages/9d/d9/2907ea38eb60687d297bf9c39e5ee58053c87b57fe8a9cae97090cecbf10/sqlalchemy-2.0.50-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b00098cdbdbd38c7be3d568b0c9c3122b8c0ec62b911b57cd5e6e0254d60a76d", size = 3486117, upload-time = "2026-05-24T20:03:59.052Z" }, - { url = "https://files.pythonhosted.org/packages/f2/e3/5aa06f167559f8c0bdae487e297d23ba548150ab016a3418265d617a4985/sqlalchemy-2.0.50-cp314-cp314t-win32.whl", hash = "sha256:1fbd55a969d7ac44a98e3dec75016074f809fa08f871585ace58dde110d1bf3e", size = 2150823, upload-time = "2026-05-24T20:08:58.644Z" }, - { url = "https://files.pythonhosted.org/packages/65/9b/112fb8f977582d7489d036e409e3723948bcf5320b3ac465f3c481bbe8f9/sqlalchemy-2.0.50-cp314-cp314t-win_amd64.whl", hash = "sha256:c5c3cdb753a9004183e1ccb634b41611654c989e61bc68617ce878e46d6f1e51", size = 2185794, upload-time = "2026-05-24T20:09:00.319Z" }, - { url = "https://files.pythonhosted.org/packages/d0/10/f7220e9b784d295d241c86ed99aeb537f92afcd469a64861f2717e9bb077/sqlalchemy-2.0.50-py3-none-any.whl", hash = "sha256:92064363517a3ff8212b5a93b8c62876579d8dfd1ca5b561335f30152d884fa9", size = 1943861, upload-time = "2026-05-24T19:59:01.119Z" }, -] - -[[package]] -name = "sqlmodel" -version = "0.0.38" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, - { name = "sqlalchemy" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/64/0d/26ec1329960ea9430131fe63f63a95ea4cb8971d49c891ff7e1f3255421c/sqlmodel-0.0.38.tar.gz", hash = "sha256:d583ec237b14103809f74e8630032bc40ab68cd6b754a610f0813c56911a547b", size = 86710, upload-time = "2026-04-02T21:03:55.571Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/72/c7/10c60af0607ab6fa136264f7f39d205932218516226d38585324ffda705d/sqlmodel-0.0.38-py3-none-any.whl", hash = "sha256:84e3fa990a77395461ded72a6c73173438ce8449d5c1c4d97fbff1b1df692649", size = 27294, upload-time = "2026-04-02T21:03:56.406Z" }, -] - [[package]] name = "srsly" version = "2.5.3"