cetdev/dify
0
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

feat(agent): shellctl rewritten in go (#38841)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Yunlu Wen
2026-07-15 01:58:36 +00:00
committed by GitHub
co-authored by autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
parent 40df83de66
commit 302d7b1e1b
106 changed files with 9224 additions and 12858 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build-push.yml
+8 -6
View File
@@ -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:
.github/workflows/docker-build.yml
+20
View File
@@ -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
.github/workflows/main-ci.yml
+64
View File
@@ -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
.github/workflows/sandbox-runtime-tests.yml
+98
View File
@@ -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
.gitignore
+4
View File
@@ -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
Makefile
+18 -3
View File
@@ -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
api/uv.lock
Generated
+1 -9
View File
@@ -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 = [
dify-agent-runtime/Makefile
+108
View File
@@ -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
dify-agent-runtime/README.md
+79
View File
@@ -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.
dify-agent-runtime/cmd/dify-agent-cli/dumphelp.go
+68
View File
@@ -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 <path> --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")
}
dify-agent-runtime/cmd/dify-agent-cli/dumphelp_test.go
+68
View File
@@ -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)
}
}
}
dify-agent-runtime/cmd/dify-agent-cli/main.go
+438
View File
@@ -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)
}
}
dify-agent-runtime/cmd/dify-agent-cli/main_test.go
+228
View File
@@ -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)
}
}
}
dify-agent-runtime/cmd/runner-exit/main.go
+30
View File
@@ -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)
}
}
dify-agent-runtime/cmd/sanitize-pty/main.go
+22
View File
@@ -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)
}
}
dify-agent-runtime/cmd/shellctl/main.go
+89
View File
@@ -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
}
dify-agent/docker/local-sandbox/Dockerfile → dify-agent-runtime/docker/Dockerfile
+23 -34
View File
@@ -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 \
dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub.pb.go
+515
View File
@@ -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
}
dify-agent-runtime/gen/dify/agent/stub/v1/agent_stub_grpc.pb.go
+197
View File
@@ -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",
}
dify-agent-runtime/go.mod
+28
View File
@@ -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
)
dify-agent-runtime/go.sum
+91
View File
@@ -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=
dify-agent-runtime/internal/agentcli/archive.go
+137
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/client.go
+45
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/client_grpc.go
+69
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/client_http.go
+218
View File
@@ -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)
}
dify-agent-runtime/internal/agentcli/config.go
+498
View File
@@ -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)
}
dify-agent-runtime/internal/agentcli/connect.go
+37
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/drive.go
+353
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/env.go
+150
View File
@@ -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
}
dify-agent-runtime/internal/agentcli/env_test.go
+173
View File
@@ -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")
}
}
dify-agent-runtime/internal/agentcli/file.go
+181
View File
@@ -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
}
}
}
dify-agent-runtime/internal/agentcli/httpclient.go
+230
View File
@@ -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))
}
dify-agent-runtime/internal/runner_exit/runner_exit.go
+109
View File
@@ -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
}
dify-agent-runtime/internal/runner_exit/runner_exit_test.go
+182
View File
@@ -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)
}
}
}
dify-agent-runtime/internal/sanitize/sanitize.go
+194
View File
@@ -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()
}
dify-agent-runtime/internal/sanitize/sanitize_test.go
+83
View File
@@ -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)
}
}
dify-agent-runtime/internal/server/api.go
+285
View File
@@ -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())
}
dify-agent-runtime/internal/server/api_test.go
+170
View File
@@ -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")
}
}
dify-agent-runtime/internal/server/config.go
+128
View File
@@ -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")
}
dify-agent-runtime/internal/server/config_test.go
+56
View File
@@ -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)
}
}
dify-agent-runtime/internal/server/db.go
+368
View File
@@ -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
}
dify-agent-runtime/internal/server/db_test.go
+309
View File
@@ -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())
}
dify-agent-runtime/internal/server/errors.go
+22
View File
@@ -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}
}
dify-agent-runtime/internal/server/output.go
+167
View File
@@ -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
}
dify-agent-runtime/internal/server/output_test.go
+275
View File
@@ -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)
}
}
dify-agent-runtime/internal/server/runtime.go
+35
View File
@@ -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))
}
dify-agent-runtime/internal/server/runtime_test.go
+71
View File
@@ -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)
}
}
dify-agent-runtime/internal/server/server.go
+6
View File
@@ -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
dify-agent-runtime/internal/server/service.go
+943
View File
@@ -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
}
dify-agent-runtime/internal/server/tmux.go
+288
View File
@@ -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, "'", `'\''`) + "'"
}
dify-agent-runtime/internal/server/tmux_test.go
+59
View File
@@ -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)
}
}
}
dify-agent-runtime/internal/server/types.go
+100
View File
@@ -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"`
}
dify-agent-runtime/internal/stubclient/client.go
+143
View File
@@ -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
}
dify-agent-runtime/internal/stubclient/client_test.go
+114
View File
@@ -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)
}
}
dify-agent-runtime/packaging_test.go
+161
View File
@@ -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/<name> directories that must exist.
// Each one produces a binary via `go build ./cmd/<name>`.
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)")
}
}
dify-agent-runtime/tests/acceptance_test.go
+587
View File
@@ -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
}
dify-agent/proto/dify/agent/stub/v1/agent_stub.proto
+2
View File
@@ -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);
dify-agent/pyproject.toml
+5 -13
View File
@@ -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"]
dify-agent/src/dify_agent/agent_stub/cli/__init__.py
-3
View File
@@ -1,3 +0,0 @@
"""Client-safe CLI package for the ``dify-agent`` sandbox command."""
__all__: list[str] = []
dify-agent/src/dify_agent/agent_stub/cli/_agent_stub.py
-20
View File
@@ -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"]
dify-agent/src/dify_agent/agent_stub/cli/_config.py
-300
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/cli/_drive.py
-369
View File
@@ -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
``<drive_path>/SKILL.md`` plus ``<drive_path>/.DIFY-SKILL-FULL.zip``.
* If ``local_path`` is a directory and ``kind`` is ``"dir"``, upload
each regular file under ``drive_path/<relative_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",
]
dify-agent/src/dify_agent/agent_stub/cli/_env.py
-72
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/cli/_files.py
-247
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/cli/main.py
-667
View File
@@ -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: <non-empty string>
description: <string>
---
- 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"]
dify-agent/src/dify_agent/agent_stub/client/__init__.py
-49
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/client/_agent_stub.py
-352
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/client/_agent_stub_grpc.py
-243
View File
@@ -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 <token>`` 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",
]
dify-agent/src/dify_agent/agent_stub/client/_agent_stub_http.py
-597
View File
@@ -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",
]
dify-agent/src/dify_agent/agent_stub/client/_errors.py
-53
View File
@@ -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",
]
dify-agent/src/dify_agent/layers/_agent_cli_help.json
+25
View File
@@ -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"
}
dify-agent/src/dify_agent/layers/_agent_cli_help.py
+28
View File
@@ -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 <args> --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"]
dify-agent/src/dify_agent/layers/config/layer.py
+1 -1
View File
@@ -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,
dify-agent/src/shellctl/cli.py
-587
View File
@@ -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",
]
dify-agent/src/shellctl/server/__init__.py
-58
View File
@@ -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__))
dify-agent/src/shellctl/server/__main__.py
-6
View File
@@ -1,6 +0,0 @@
"""Support `python -m shellctl.server`."""
from shellctl.cli import main
if __name__ == "__main__":
main()
dify-agent/src/shellctl/server/api.py
-198
View File
@@ -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"]
dify-agent/src/shellctl/server/artifacts.py
-62
View File
@@ -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/<job_id>/` 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",
]
dify-agent/src/shellctl/server/cli.py
-17
View File
@@ -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",
]
dify-agent/src/shellctl/server/config.py
-105
View File
@@ -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"]
dify-agent/src/shellctl/server/db.py
-52
View File
@@ -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"]
dify-agent/src/shellctl/server/errors.py
-14
View File
@@ -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"]
dify-agent/src/shellctl/server/serve.py
-103
View File
@@ -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"]
dify-agent/src/shellctl/server/service.py
-1184
View File
File diff suppressed because it is too large Load Diff
dify-agent/src/shellctl/server/tmux.py
-343
View File
@@ -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 <socket>` 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",
]
dify-agent/src/shellctl_runtime/__init__.py
-7
View File
@@ -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.
"""
dify-agent/src/shellctl_runtime/paths.py
-16
View File
@@ -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"]
dify-agent/src/shellctl_runtime/py.typed
-1
View File
@@ -1 +0,0 @@
dify-agent/src/shellctl_runtime/runner_exit.py
-158
View File
@@ -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",
]
dify-agent/src/shellctl_runtime/sanitize.py
-186
View File
@@ -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",
]
dify-agent/tests/local/dify_agent/agent_stub/cli/test_config.py
-337
View File
@@ -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))
dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py
-870
View File
@@ -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",
}
dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py
-365
View File
@@ -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))
dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py
-924
View File
@@ -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: <non-empty string> description: <string> ---" 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
dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_client.py
-667
View File
@@ -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"
dify-agent/tests/local/dify_agent/agent_stub/client/test_agent_stub_config_client.py
-310
View File
@@ -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()
dify-agent/tests/local/dify_agent/test_client_safe_exports.py
+4 -18
View File
@@ -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:
dify-agent/tests/local/dify_agent/test_import_boundaries.py
+10 -27
View File
@@ -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')",
],
)
dify-agent/tests/local/shellctl/golden_shellctl_sanitize.json
-32
View File
@@ -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"
}
]
dify-agent/tests/local/shellctl/test_shellctl_cli.py
-688
View File
@@ -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,
}

Some files were not shown because too many files have changed in this diff Show More