From 248a6eea1215aca749b337cfa2f0b8c860526dbc Mon Sep 17 00:00:00 2001 From: "zhangning.2025" Date: Sat, 10 Oct 2026 14:07:05 +0800 Subject: [PATCH] feat: add managed agents runtime with official Anthropic SDK --- .gitignore | 5 + config.yaml.full | 7 + docker/managed-agents/Dockerfile.worker | 47 + .../Dockerfile.worker.dockerignore | 21 + docker/managed-agents/README.md | 128 ++ docker/managed-agents/README.zh.md | 109 ++ docker/managed-agents/build.sh | 51 + docker/managed-agents/k8s/worker.yaml | 60 + docker/managed-agents/run-claimed.sh | 6 + docker/managed-agents/run-runtime.sh | 36 + docker/managed-agents/verify_anthropic_sdk.py | 65 + .../framework/tools/feishu-channel.en.mdx | 20 + .../docs/framework/tools/feishu-channel.mdx | 20 + examples/16_self_host_sandbox/.env.example | 54 +- examples/16_self_host_sandbox/README.md | 182 +- examples/16_self_host_sandbox/README.zh.md | 169 +- .../16_self_host_sandbox/agentkit_runtime.py | 6 + .../self_host_sandbox_agent/__init__.py | 16 - .../agents/self_host_sandbox_agent/agent.py | 179 +- .../anthropic_gateway_e2e.py | 723 ++++++++ .../16_self_host_sandbox/ark_session_e2e.py | 181 ++ .../build-managed-agents-images.sh | 4 + .../distributed_agent_loop_test.py | 450 +++++ examples/16_self_host_sandbox/event_debug.py | 6 + .../k8s_agent_loop_e2e.py | 333 ++++ .../k8s_capabilities_e2e.py | 201 +++ .../local_agent_loop_test.py | 204 +++ examples/16_self_host_sandbox/main.py | 39 +- .../managed_agent_loop.py | 6 + .../managed_session_resources.py | 6 + .../16_self_host_sandbox/preflight_probe.py | 266 +++ .../16_self_host_sandbox/requirements.txt | 4 +- examples/16_self_host_sandbox/run-runtime.sh | 9 + .../run.managed-agents-sandbox.sh | 6 + examples/16_self_host_sandbox/run.sh | 25 +- .../16_self_host_sandbox/runtime_identity.py | 6 + .../runtime_type_dispatcher_test.py | 199 +++ .../16_self_host_sandbox/sandbox_client.py | 481 +----- .../soak_fault_preflight.py | 629 +++++++ .../16_self_host_sandbox/soak_load_test.py | 1259 ++++++++++++++ .../2026-10-10-non-ui-migration.md | 192 +++ .../2026-10-10-non-ui-migration.zh.md | 165 ++ .../2026-10-10-official-sdk.md | 103 ++ .../2026-10-10-official-sdk.zh.md | 90 + pyproject.toml | 2 +- specs/feishu-channel/README.md | 47 + specs/feishu-channel/README.zh.md | 47 + specs/managed-agent-resources/README.md | 114 ++ specs/managed-agent-resources/README.zh.md | 89 + specs/managed-agent-runtime/README.md | 57 + specs/managed-agent-runtime/README.zh.md | 53 + tests/integrations/mpa/conftest.py | 19 + .../mpa/test_session_client.py} | 152 +- .../mpa/test_session_resources.py | 693 ++++++++ .../models/test_ark_llm_context_management.py | 67 + tests/models/test_ark_llm_expire_at.py | 78 + tests/models/test_ark_llm_fallbacks.py | 302 ++++ tests/runtime/managed_agents/__init__.py | 0 tests/runtime/managed_agents/conftest.py | 19 + .../cross_worker_http_scenario.py | 216 +++ .../test_agentkit_runtime_identity.py | 703 ++++++++ .../test_cross_worker_completion.py | 221 +++ .../runtime/managed_agents/test_dispatcher.py | 524 ++++++ tests/runtime/managed_agents/test_health.py | 30 + .../managed_agents/test_managed_agent_loop.py | 945 +++++++++++ .../test_managed_event_debug.py | 138 ++ .../test_official_tool_config.py | 143 ++ tests/runtime/managed_agents/test_package.py | 26 + .../managed_agents/test_sdk_integrity.py | 55 + .../test_self_host_sandbox_agent.py | 1083 ++++++++++++ .../managed_agents/test_worker_supervision.py | 102 ++ tests/runtime/test_anthropic_gateway_e2e.py | 79 + tests/runtime/test_completion_gate.py | 206 +++ tests/runtime/test_fault_recovery_boundary.py | 349 ++++ tests/runtime/test_managed_agents_delivery.py | 108 ++ tests/runtime/test_preflight_probe.py | 83 + tests/runtime/test_self_host_sandbox_agent.py | 107 -- tests/runtime/test_soak_load_test.py | 1074 ++++++++++++ tests/skills/test_ma_infra_skills_client.py | 233 +++ tests/test_feishu_channel_extension.py | 382 ++++- tests/test_ve_identity_api_key.py | 42 + uv.lock | 2 +- veadk/consts.py | 2 - veadk/extensions/feishu_channel.py | 411 ++++- veadk/integrations/mpa/session_client.py | 508 ++++++ veadk/integrations/mpa/session_resources.py | 305 ++++ .../ve_identity/identity_client.py | 10 +- veadk/models/ark_llm.py | 126 +- veadk/runtime/managed_agents/__init__.py | 15 + veadk/runtime/managed_agents/dispatcher.py | 413 +++++ veadk/runtime/managed_agents/events.py | 134 ++ veadk/runtime/managed_agents/health.py | 53 + veadk/runtime/managed_agents/identity.py | 351 ++++ veadk/runtime/managed_agents/loop.py | 739 ++++++++ veadk/runtime/managed_agents/sandbox.py | 206 +++ veadk/runtime/managed_agents/worker.py | 1511 +++++++++++++++++ veadk/skills/ma_infra.py | 312 ++++ 97 files changed, 19300 insertions(+), 1154 deletions(-) create mode 100644 docker/managed-agents/Dockerfile.worker create mode 100644 docker/managed-agents/Dockerfile.worker.dockerignore create mode 100644 docker/managed-agents/README.md create mode 100644 docker/managed-agents/README.zh.md create mode 100644 docker/managed-agents/build.sh create mode 100644 docker/managed-agents/k8s/worker.yaml create mode 100644 docker/managed-agents/run-claimed.sh create mode 100644 docker/managed-agents/run-runtime.sh create mode 100644 docker/managed-agents/verify_anthropic_sdk.py create mode 100644 examples/16_self_host_sandbox/agentkit_runtime.py create mode 100644 examples/16_self_host_sandbox/anthropic_gateway_e2e.py create mode 100644 examples/16_self_host_sandbox/ark_session_e2e.py create mode 100644 examples/16_self_host_sandbox/build-managed-agents-images.sh create mode 100644 examples/16_self_host_sandbox/distributed_agent_loop_test.py create mode 100644 examples/16_self_host_sandbox/event_debug.py create mode 100644 examples/16_self_host_sandbox/k8s_agent_loop_e2e.py create mode 100644 examples/16_self_host_sandbox/k8s_capabilities_e2e.py create mode 100644 examples/16_self_host_sandbox/local_agent_loop_test.py create mode 100644 examples/16_self_host_sandbox/managed_agent_loop.py create mode 100644 examples/16_self_host_sandbox/managed_session_resources.py create mode 100644 examples/16_self_host_sandbox/preflight_probe.py create mode 100644 examples/16_self_host_sandbox/run-runtime.sh create mode 100644 examples/16_self_host_sandbox/run.managed-agents-sandbox.sh create mode 100644 examples/16_self_host_sandbox/runtime_identity.py create mode 100644 examples/16_self_host_sandbox/runtime_type_dispatcher_test.py create mode 100644 examples/16_self_host_sandbox/soak_fault_preflight.py create mode 100644 examples/16_self_host_sandbox/soak_load_test.py create mode 100644 prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md create mode 100644 prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md create mode 100644 prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.md create mode 100644 prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.zh.md create mode 100644 specs/feishu-channel/README.md create mode 100644 specs/feishu-channel/README.zh.md create mode 100644 specs/managed-agent-resources/README.md create mode 100644 specs/managed-agent-resources/README.zh.md create mode 100644 specs/managed-agent-runtime/README.md create mode 100644 specs/managed-agent-runtime/README.zh.md create mode 100644 tests/integrations/mpa/conftest.py rename tests/{runtime/test_self_host_sandbox_client.py => integrations/mpa/test_session_client.py} (54%) create mode 100644 tests/integrations/mpa/test_session_resources.py create mode 100644 tests/models/test_ark_llm_expire_at.py create mode 100644 tests/runtime/managed_agents/__init__.py create mode 100644 tests/runtime/managed_agents/conftest.py create mode 100644 tests/runtime/managed_agents/cross_worker_http_scenario.py create mode 100644 tests/runtime/managed_agents/test_agentkit_runtime_identity.py create mode 100644 tests/runtime/managed_agents/test_cross_worker_completion.py create mode 100644 tests/runtime/managed_agents/test_dispatcher.py create mode 100644 tests/runtime/managed_agents/test_health.py create mode 100644 tests/runtime/managed_agents/test_managed_agent_loop.py create mode 100644 tests/runtime/managed_agents/test_managed_event_debug.py create mode 100644 tests/runtime/managed_agents/test_official_tool_config.py create mode 100644 tests/runtime/managed_agents/test_package.py create mode 100644 tests/runtime/managed_agents/test_sdk_integrity.py create mode 100644 tests/runtime/managed_agents/test_self_host_sandbox_agent.py create mode 100644 tests/runtime/managed_agents/test_worker_supervision.py create mode 100644 tests/runtime/test_anthropic_gateway_e2e.py create mode 100644 tests/runtime/test_completion_gate.py create mode 100644 tests/runtime/test_fault_recovery_boundary.py create mode 100644 tests/runtime/test_managed_agents_delivery.py create mode 100644 tests/runtime/test_preflight_probe.py delete mode 100644 tests/runtime/test_self_host_sandbox_agent.py create mode 100644 tests/runtime/test_soak_load_test.py create mode 100644 tests/skills/test_ma_infra_skills_client.py create mode 100644 tests/test_ve_identity_api_key.py create mode 100644 veadk/integrations/mpa/session_client.py create mode 100644 veadk/integrations/mpa/session_resources.py create mode 100644 veadk/runtime/managed_agents/__init__.py create mode 100644 veadk/runtime/managed_agents/dispatcher.py create mode 100644 veadk/runtime/managed_agents/events.py create mode 100644 veadk/runtime/managed_agents/health.py create mode 100644 veadk/runtime/managed_agents/identity.py create mode 100644 veadk/runtime/managed_agents/loop.py create mode 100644 veadk/runtime/managed_agents/sandbox.py create mode 100644 veadk/runtime/managed_agents/worker.py create mode 100644 veadk/skills/ma_infra.py diff --git a/.gitignore b/.gitignore index d13eb1d30..53f16ace6 100644 --- a/.gitignore +++ b/.gitignore @@ -234,3 +234,8 @@ mpa-create.config*.yaml # PiAgent release archives are downloaded and verified during image builds. examples/piagent_with_mcp/vendor/pi-*.tar.gz .gstack/ + +# Local Managed Agents verification state; never part of an upstream change. +examples/16_self_host_sandbox/.local-managed-agents/ +examples/16_self_host_sandbox/deploy_records/ +examples/16_self_host_sandbox/soak_results_*/ diff --git a/config.yaml.full b/config.yaml.full index 196cc143e..5c8f44931 100644 --- a/config.yaml.full +++ b/config.yaml.full @@ -77,6 +77,13 @@ tool: app_id: app_secret: transport: ws # `ws` | `webhook` + streaming: false + show_thinking: false + show_tool_calls: false + show_tool_results: false + separate_tool_call_cards: false + separate_thinking_card: false + create_topic: false # [optional] for Volcengine Lake AI Service https://www.volcengine.com/product/las mobile_use: tool_id: #https://console.volcengine.com/ACEP diff --git a/docker/managed-agents/Dockerfile.worker b/docker/managed-agents/Dockerfile.worker new file mode 100644 index 000000000..af49bcc3f --- /dev/null +++ b/docker/managed-agents/Dockerfile.worker @@ -0,0 +1,47 @@ +FROM python:3.12-slim-bookworm + +# Build-only indexes: never export package credentials into the runtime image. +ARG PIP_INDEX_URL=https://pypi.org/simple +ARG UV_DEFAULT_INDEX=https://pypi.org/simple + +ENV PYTHONUNBUFFERED=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + UV_LINK_MODE=copy \ + MANAGED_AGENT_WORK_CONCURRENCY=500 \ + MANAGED_AGENT_WORKDIR=/workspace \ + MANAGED_AGENT_READY_FILE=/tmp/managed-agent-worker-ready \ + PATH=/app/.venv/bin:$PATH + +WORKDIR /app +COPY --from=ghcr.io/astral-sh/uv:0.11.28 /uv /uvx /bin/ +COPY pyproject.toml uv.lock README.md LICENSE ./ +COPY veadk ./veadk +COPY docker/managed-agents/verify_anthropic_sdk.py /tmp/verify_anthropic_sdk.py + +# Export pinned versions and hashes; an index override changes downloads, +# never the locked dependency versions. No adjacent SDK wheel is needed. +RUN --mount=type=cache,id=veadk-managed-agents-uv-python312,target=/root/.cache/uv,sharing=locked \ + uv export --quiet --frozen --no-dev --extra sandbox \ + --no-emit-project --output-file /tmp/requirements.txt \ + && uv venv /app/.venv \ + && PIP_INDEX_URL="$PIP_INDEX_URL" UV_DEFAULT_INDEX="$UV_DEFAULT_INDEX" \ + UV_INDEX_URL="$UV_DEFAULT_INDEX" \ + uv pip install --python /app/.venv/bin/python --no-deps --require-hashes \ + -r /tmp/requirements.txt \ + && PIP_INDEX_URL="$PIP_INDEX_URL" UV_DEFAULT_INDEX="$UV_DEFAULT_INDEX" \ + UV_INDEX_URL="$UV_DEFAULT_INDEX" \ + uv pip install --python /app/.venv/bin/python --no-deps . \ + && python /tmp/verify_anthropic_sdk.py \ + && rm /tmp/requirements.txt /tmp/verify_anthropic_sdk.py \ + && python -c "from veadk.runtime.managed_agents.dispatcher import EnvironmentWorkDispatcher; from anthropic.lib.tools.agent_toolset import beta_agent_toolset_20260401" + +COPY docker/managed-agents/run-runtime.sh /app/run-runtime.sh +COPY docker/managed-agents/run-claimed.sh /opt/gem/run.sh +RUN useradd --create-home --uid 10001 veadk \ + && mkdir -p /workspace \ + && chmod 0755 /app/run-runtime.sh /opt/gem/run.sh \ + && chown veadk:veadk /workspace + +USER 10001:10001 +EXPOSE 8080 +CMD ["bash", "/app/run-runtime.sh"] diff --git a/docker/managed-agents/Dockerfile.worker.dockerignore b/docker/managed-agents/Dockerfile.worker.dockerignore new file mode 100644 index 000000000..a3fa3bc63 --- /dev/null +++ b/docker/managed-agents/Dockerfile.worker.dockerignore @@ -0,0 +1,21 @@ +** +!pyproject.toml +!uv.lock +!README.md +!LICENSE +!veadk/ +!veadk/** +!docker/ +!docker/managed-agents/ +!docker/managed-agents/run-runtime.sh +!docker/managed-agents/run-claimed.sh +!docker/managed-agents/verify_anthropic_sdk.py +**/.env +**/.env.* +**/*.env +**/__pycache__ +**/*.pyc +frontend +veadk/webui +**/deploy_records +**/soak_results_* diff --git a/docker/managed-agents/README.md b/docker/managed-agents/README.md new file mode 100644 index 000000000..5df856155 --- /dev/null +++ b/docker/managed-agents/README.md @@ -0,0 +1,128 @@ +# Managed Agents Worker build + +[中文版](README.zh.md) + +This image runs the installed `veadk.runtime.managed_agents.worker` package and +the AgentKit health listener. The dispatcher and Session adapters ship inside +VeADK; no neighboring repository, private SDK wheel, frontend build, or cloud +deployment is required to build the image. + +## Build and verify locally + +From the repository root: + +```bash +uv sync --extra dev --extra sandbox --extra extensions +uv run --extra dev --extra sandbox pytest tests/runtime/managed_agents tests/integrations/mpa +bash docker/managed-agents/build.sh veadk-managed-agents-worker:local +docker run --rm --entrypoint python veadk-managed-agents-worker:local \ + -m veadk.runtime.managed_agents.worker --help +``` + +The build script defaults to `--load`; publication requires an explicit +`--push` argument and a registry/tag selected by the operator. Building an image +does not update an AgentKit Runtime, Tool, Session, or Kubernetes Deployment. +`MANAGED_AGENTS_PLATFORM` overrides the default `linux/amd64` build platform. + +## Dependency and SDK pitfalls + +Older source images copied a wheel from +`examples/16_self_host_sandbox/.docker-dist/anthropic-*.whl`, built from an adjacent +SDK checkout, and imported `anthropic.lib.environments._dispatcher`. A normal +Anthropic 1.3.0 installation provides Managed Sessions and native tools but does +not provide that private dispatcher. An installed package could import successfully +yet fail to start a Worker. The new dispatcher uses generated Work APIs and the +existing account model-work polling protocol instead; the Docker build validates +its import and native-tool availability before the image is considered built. + +Images default to `https://pypi.org/simple`. If an approved package mirror does +not contain every version in `uv.lock`, use a complete mirror through build args; +do not loosen pins, remove hashes, or regenerate the lock to hide missing packages: + +```bash +PIP_INDEX_URL=https://packages.example.com/simple \ +UV_DEFAULT_INDEX=https://packages.example.com/simple \ + bash docker/managed-agents/build.sh veadk-managed-agents-worker:local +``` + +Indexes are build inputs, not runtime environment variables. This helper accepts +unauthenticated mirror URLs and rejects user info, query strings and fragments +before invoking Docker. Authenticated mirrors require a separately configured +BuildKit secret integration; never put credentials in build arguments, image tags, +source, or logs. Frozen export preserves dependency hashes. The component uses +its own named, locked BuildKit cache rather than sharing an unrelated image +build's cache. This dedicated Worker image installs the sandbox extra. Feishu and +optional data integrations use a separate SDK installation with the extensions +extra; development and Codex packages are not included in the Worker image. + +## Runtime configuration + +Configure gateway/credential/identity settings through the platform's secret and +environment mechanisms. A minimal worker needs `ANTHROPIC_BASE_URL`, a Runtime +credential (`ANTHROPIC_ENVIRONMENT_KEY` or configured Runtime Identity), and its +Work scope. Account workers default to `MANAGED_AGENT_WORK_SCOPE=account` and +`MA_AGENT_LOOP_RUNTIME_TYPE=agentkit_runtime`; environment workers must set +`MANAGED_AGENT_WORK_SCOPE=environment` and `ANTHROPIC_ENVIRONMENT_ID` explicitly. +Session creation additionally requires a configured Agent ID. Session snapshots +remain authoritative for model, tools, Skills, MCP and credentials. + +`MANAGED_AGENT_WORK_CONCURRENCY` defaults to 500 in the image. It counts Work +handlers, including waits, and does not enable simultaneous turns in one Session. +The worker ready file appears only after a valid successful Work poll, and is +removed while draining. HTTP health endpoints alone do not prove worker readiness. +The entry point supervises worker/listener together; failure of either stops both. + +For an AgentKit Tool executing an already claimed Work, the image also retains +`/opt/gem/run.sh`. It selects `MANAGED_AGENT_ENTRYPOINT_MODE=claimed`, uses injected +Work/Session/lease configuration, and defaults tool execution to local inside the +Tool's isolation. The regular image command selects continuous worker mode. + +The process runs as UID/GID 10001. Mount a writable `/workspace` and `/tmp` when +using a read-only filesystem. Native bash must execute within the deployment's OS +isolation; path-confined file tools do not provide a shell security sandbox. + +The [Kubernetes example](k8s/worker.yaml) references an operator-created Secret; +replace the image placeholder before applying. No secret values or real deployment +identifiers are supplied. Kubernetes, AgentKit and live model acceptance require +separate deployment authorization and are not proven by local tests/builds. + +## Official SDK extension strategy + +The current image installs the unchanged official PyPI `anthropic==1.3.0` selected +by `uv.lock`; it does **not** install the historical adjacent custom wheel. +The locked SHA256 matches the [official release](https://pypi.org/project/anthropic/1.3.0/). +The installer validates artifact hashes, and `verify_anthropic_sdk.py` checks every +installed SDK Python file against wheel RECORD, rejecting changed, missing and +unrecorded files. RECORD checking is an integrity check, not an independent origin +proof if the wheel and its RECORD are both replaced; the lock hash supplies that boundary. + +VeADK uses composition rather than subclassing or globally patching the SDK: + +| Behavior | Implementation | +| --- | --- | +| environment Work poll | official `beta.environments.work.with_raw_response.poll`; inspect HTTP 204 before validating Work | +| account model Work poll | public `AsyncAnthropic.get` for `/v1/model-work/poll`, a VeADK gateway extension, not a Claude API endpoint | +| ack / heartbeat / stop and Session events | official generated SDK APIs | +| account/runtime routing headers | public `default_headers` / `extra_headers` | +| Runtime Identity rotation | public HTTP client request/response hooks; preserve Work-scoped Bearer alongside gateway identity | +| native tool execution | official `agent_toolset` factory and public `ToolError` | + +The official SDK documents [raw responses, custom requests and HTTP clients](https://platform.claude.com/docs/en/cli-sdks-libraries/sdks/python#advanced-usage). +Its `EnvironmentWorker` has no arbitrary VeADK handler callback; subclassing would +require overriding private execution methods. No SDK source files or methods are +replaced. The per-Work client is a copy with inherited auth cleared; setting its +public `api_key=None` is necessary because SDK 1.3.0 `copy(api_key=None)` inherits +the parent key. + +Two SDK internals remain version-sensitive: `_skills._download_and_extract` is +shipped by the official package and preserves fail-fast Skill download semantics +(the public helper skips failures); the tool decorator's `_context_manager` is +read for additive cleanup. VeADK performs cleanup locally to avoid the SDK helper +logging raw exception text. Real SDK archive, tool, auth and dispatcher tests cover +these boundaries. A locked official SDK does not require a custom SDK fork. + +For a local VKE Worker with no Session Identity reference, explicitly set +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE=environment` and supply `MODEL_AGENT_API_KEY` +through a dedicated Kubernetes Secret. The default is `session`; AgentKit and any +Session with a provider reference always require Identity and never fall back. +This choice verifies configured model credentials, not cloud Identity delivery. diff --git a/docker/managed-agents/README.zh.md b/docker/managed-agents/README.zh.md new file mode 100644 index 000000000..db97eb756 --- /dev/null +++ b/docker/managed-agents/README.zh.md @@ -0,0 +1,109 @@ +# Managed Agents Worker 构建 + +[English](README.md) + +镜像运行安装后的 `veadk.runtime.managed_agents.worker` 包与 AgentKit 健康监听器。 +dispatcher 和 Session 适配器包含在 VeADK 内，构建不依赖相邻仓库、私有 SDK wheel、 +前端构建或线上部署。 + +## 本地构建与验证 + +在仓库根目录运行： + +```bash +uv sync --extra dev --extra sandbox --extra extensions +uv run --extra dev --extra sandbox pytest tests/runtime/managed_agents tests/integrations/mpa +bash docker/managed-agents/build.sh veadk-managed-agents-worker:local +docker run --rm --entrypoint python veadk-managed-agents-worker:local \ + -m veadk.runtime.managed_agents.worker --help +``` + +构建脚本默认 `--load`。发布必须显式增加 `--push`，并由操作者选择镜像仓库/tag。 +镜像构建不更新 AgentKit Runtime、Tool、Session 或 Kubernetes Deployment。 +`MANAGED_AGENTS_PLATFORM` 可覆盖默认 `linux/amd64` 平台。 + +## 依赖与 SDK 注意事项 + +旧来源镜像复制相邻 SDK checkout 构建的 +`examples/16_self_host_sandbox/.docker-dist/anthropic-*.whl`，并导入 +`anthropic.lib.environments._dispatcher`。普通 Anthropic 1.3.0 安装包含 Managed +Sessions 与原生工具，但不包含该私有 dispatcher，导致包导入成功而 Worker 启动失败。 +新 dispatcher 使用生成的 Work API 与既有账号级模型 Work 轮询协议；Docker 构建结束 +前验证其导入与原生工具可用性。 + +镜像默认使用 `https://pypi.org/simple`。如果批准的镜像源没有 `uv.lock` 中全部版本， +通过 build args 使用完整镜像源；不能放宽版本、删除哈希或重建锁文件掩盖缺失： + +```bash +PIP_INDEX_URL=https://packages.example.com/simple \ +UV_DEFAULT_INDEX=https://packages.example.com/simple \ + bash docker/managed-agents/build.sh veadk-managed-agents-worker:local +``` + +包源只作为构建输入，不写入运行时环境。本脚本只接受无认证镜像源，调用 Docker +前拒绝用户名密码、查询字符串与 fragment。认证镜像源必须单独配置 BuildKit secret +集成；凭据不能写入 build args、镜像 tag、源码或日志。冻结导出保留依赖哈希。组件使用 +独立命名且锁定的 BuildKit 缓存，避免与无关镜像构建共用缓存。专用 Worker 镜像只 +安装 sandbox extra。飞书和可选数据集成使用安装 extensions extra 的独立 SDK 环境； +Worker 镜像不包含开发或 Codex 包。 + +## 运行时配置 + +通过平台 Secret 和环境配置提供网关、凭据及 Identity 信息。最小 Worker 需要 +`ANTHROPIC_BASE_URL`、Runtime 凭据（`ANTHROPIC_ENVIRONMENT_KEY` 或配置后的 +Runtime Identity）与 Work 作用域。账号 Worker 默认 +`MANAGED_AGENT_WORK_SCOPE=account` 和 `MA_AGENT_LOOP_RUNTIME_TYPE=agentkit_runtime`； +Environment Worker 必须显式设置 `MANAGED_AGENT_WORK_SCOPE=environment` 与 +`ANTHROPIC_ENVIRONMENT_ID`。创建 Session 还需要配置 Agent ID。Session 快照是模型、 +工具、Skills、MCP 与凭据的权威来源。 + +镜像默认 `MANAGED_AGENT_WORK_CONCURRENCY=500`，它限制包含等待时间的 Work handler +数量，不允许同一 Session 同时执行多个轮次。只有成功完成合法 Work 轮询后才创建 +就绪文件，排空时删除。HTTP 健康端点不能单独证明 Worker 就绪。启动入口同时监督 +Worker 与监听器，任一进程停止会终止另一个。 + +对于执行已领取 Work 的 AgentKit Tool，镜像保留 `/opt/gem/run.sh`，选择 +`MANAGED_AGENT_ENTRYPOINT_MODE=claimed`，使用注入的 Work/Session/租约配置，并默认 +在 Tool 隔离中本地执行工具。镜像常规命令选择持续 Worker 模式。 + +进程以 UID/GID 10001 运行。使用只读文件系统时，挂载可写 `/workspace` 与 `/tmp`。 +原生 bash 必须在部署提供的 OS 隔离中执行；文件路径限制不能替代 shell 安全沙箱。 + +[Kubernetes 示例](k8s/worker.yaml) 引用操作者创建的 Secret；应用前替换镜像占位符。 +模板不提供密钥值或真实部署标识。Kubernetes、AgentKit 和真实模型验收需要独立部署 +授权，本地测试或构建不能证明线上验收。 + +## 官方 SDK 扩展方案 + +当前镜像安装 `uv.lock` 选定的 **未修改的官方 PyPI `anthropic==1.3.0`**，不会安装 +历史相邻仓库的定制 wheel。锁定 SHA256 与[官方发行版](https://pypi.org/project/anthropic/1.3.0/) +一致。安装器验证发行物哈希，`verify_anthropic_sdk.py` 再将所有已安装 SDK Python +文件与 wheel RECORD 对比，拒绝文件被修改、缺失或额外覆盖。RECORD 是完整性校验， +若 wheel 与 RECORD 同时被替换则不能独立证明来源；此边界由锁定发行物哈希保障。 + +VeADK 采用组合，无需继承或全局修改 SDK 方法： + +| 行为 | 实现 | +| --- | --- | +| 环境级 Work 轮询 | 官方 `beta.environments.work.with_raw_response.poll`，先检查 HTTP 204 再校验 Work | +| 账号级模型 Work 轮询 | 公开 `AsyncAnthropic.get` 请求 `/v1/model-work/poll`；这是 VeADK 网关扩展，非 Claude API 端点 | +| ack / heartbeat / stop 与 Session 事件 | 官方 SDK 生成接口 | +| 账号/Runtime 路由头 | 公开 `default_headers` / `extra_headers` | +| Runtime Identity 轮转 | 公开 HTTP client 请求/响应 hooks，保留 Work Bearer 与网关 Identity 双重认证 | +| 原生工具执行 | 官方 `agent_toolset` 工厂与公开 `ToolError` | + +官方文档说明了[原始响应、自定义请求及 HTTP client](https://platform.claude.com/docs/en/cli-sdks-libraries/sdks/python#advanced-usage)。 +官方 `EnvironmentWorker` 没有任意 VeADK handler 回调；继承需要覆盖私有执行方法。 +本实现不覆盖 SDK 源码或方法。每个 Work 使用清除父认证的客户端副本；必须设置 +副本公开属性 `api_key=None`，因为 SDK 1.3.0 的 `copy(api_key=None)` 会继承父 Key。 + +仍有两个版本敏感的 SDK 内部接口：官方发行包内的 `_skills._download_and_extract` +保留 Skill 下载失败即失败的语义，公开 helper 会跳过失败；工具 decorator 的 +`_context_manager` 用于额外的退出清理。VeADK 在本地执行清理，避免 SDK helper +直接输出异常文本。真实 SDK 的归档、工具、认证与 dispatcher 测试覆盖这些边界。 +锁定官方 SDK 版本不等于需要定制 SDK fork。 + +本地 VKE Worker 没有 Session Identity 引用时，可显式设置 +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE=environment`，通过独立 Kubernetes Secret +提供 `MODEL_AGENT_API_KEY`。默认 `session`；AgentKit 及带 provider 引用的 Session +始终使用 Identity，不会兜底。此路径证明配置的模型凭据，不证明云 Identity 下发。 diff --git a/docker/managed-agents/build.sh b/docker/managed-agents/build.sh new file mode 100644 index 000000000..14bbee5e6 --- /dev/null +++ b/docker/managed-agents/build.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "usage: $0 IMAGE[:TAG] [--push]" >&2 + exit 2 +fi +image="$1" +output="--load" +if [[ $# == 2 ]]; then + if [[ "$2" != "--push" ]]; then + echo "unknown option: $2" >&2 + exit 2 + fi + output="--push" +fi + +PIP_INDEX_URL="${PIP_INDEX_URL:-https://pypi.org/simple}" +UV_DEFAULT_INDEX="${UV_DEFAULT_INDEX:-${UV_INDEX_URL:-$PIP_INDEX_URL}}" + +# Build arguments can survive in image/build metadata. Do not pass credentials. +PIP_INDEX_URL="$PIP_INDEX_URL" UV_DEFAULT_INDEX="$UV_DEFAULT_INDEX" python3 - <<'PY' +import os +import sys +from urllib.parse import urlsplit + +for name in ("PIP_INDEX_URL", "UV_DEFAULT_INDEX"): + try: + value = urlsplit(os.environ[name]) + valid = ( + value.scheme in {"http", "https"} + and value.hostname + and not value.username + and not value.password + and not value.query + and not value.fragment + ) + except ValueError: + valid = False + if not valid: + sys.exit("Package indexes must be HTTP(S) URLs without credentials, query strings or fragments") +PY + +exec docker buildx build --platform "${MANAGED_AGENTS_PLATFORM:-linux/amd64}" \ + "$output" --file "$SCRIPT_DIR/Dockerfile.worker" --tag "$image" \ + --build-arg PIP_INDEX_URL="$PIP_INDEX_URL" \ + --build-arg UV_DEFAULT_INDEX="$UV_DEFAULT_INDEX" \ + "$REPO_ROOT" diff --git a/docker/managed-agents/k8s/worker.yaml b/docker/managed-agents/k8s/worker.yaml new file mode 100644 index 000000000..525add154 --- /dev/null +++ b/docker/managed-agents/k8s/worker.yaml @@ -0,0 +1,60 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: managed-agents-worker +spec: + replicas: 2 + selector: + matchLabels: + app: managed-agents-worker + template: + metadata: + labels: + app: managed-agents-worker + spec: + automountServiceAccountToken: false + terminationGracePeriodSeconds: 60 + securityContext: + runAsNonRoot: true + runAsUser: 10001 + runAsGroup: 10001 + fsGroup: 10001 + seccompProfile: + type: RuntimeDefault + containers: + - name: worker + image: example.invalid/managed-agents-worker:replace-me + envFrom: + - secretRef: + name: managed-agents-worker-config + ports: + - containerPort: 8080 + readinessProbe: + exec: + command: ["test", "-s", "/tmp/managed-agent-worker-ready"] + periodSeconds: 5 + livenessProbe: + httpGet: + path: /ping + port: 8080 + initialDelaySeconds: 30 + periodSeconds: 10 + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: ["ALL"] + resources: + requests: + cpu: 250m + memory: 512Mi + volumeMounts: + - name: workspace + mountPath: /workspace + - name: temporary + mountPath: /tmp + volumes: + - name: workspace + emptyDir: {} + - name: temporary + emptyDir: {} diff --git a/docker/managed-agents/run-claimed.sh b/docker/managed-agents/run-claimed.sh new file mode 100644 index 000000000..2b9ad3455 --- /dev/null +++ b/docker/managed-agents/run-claimed.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +set -euo pipefail +# AgentKit Tool dispatchers start this path for one already claimed Work. +export MANAGED_AGENT_ENTRYPOINT_MODE=claimed +export MANAGED_AGENT_TOOL_EXECUTION="${MANAGED_AGENT_TOOL_EXECUTION:-local}" +exec bash /app/run-runtime.sh "$@" diff --git a/docker/managed-agents/run-runtime.sh b/docker/managed-agents/run-runtime.sh new file mode 100644 index 000000000..3f8f9f550 --- /dev/null +++ b/docker/managed-agents/run-runtime.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +set -euo pipefail + +PYTHON_BIN="${MANAGED_AGENTS_PYTHON:-python}" +listener_pid="" +worker_pid="" +cleanup() { + trap - EXIT INT TERM + for child_pid in "$listener_pid" "$worker_pid"; do + if [[ -n "$child_pid" ]]; then + kill "$child_pid" 2>/dev/null || true + fi + done + wait 2>/dev/null || true +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +export MANAGED_AGENT_WORK_SCOPE="${MANAGED_AGENT_WORK_SCOPE:-account}" +export MA_AGENT_LOOP_RUNTIME_TYPE="${MA_AGENT_LOOP_RUNTIME_TYPE:-agentkit_runtime}" +export MANAGED_AGENT_TOOL_EXECUTION="${MANAGED_AGENT_TOOL_EXECUTION:-remote}" + +case "${MANAGED_AGENT_ENTRYPOINT_MODE:-worker}" in + worker) mode="--managed-agent-worker" ;; + claimed) mode="--managed-agent-work-item" ;; + *) echo "MANAGED_AGENT_ENTRYPOINT_MODE must be worker or claimed" >&2; exit 2 ;; +esac + +"$PYTHON_BIN" -m veadk.runtime.managed_agents.health & +listener_pid=$! +"$PYTHON_BIN" -m veadk.runtime.managed_agents.worker "$mode" "$@" & +worker_pid=$! + +# Failure or completion of either process must stop its sibling. +wait -n "$listener_pid" "$worker_pid" diff --git a/docker/managed-agents/verify_anthropic_sdk.py b/docker/managed-agents/verify_anthropic_sdk.py new file mode 100644 index 000000000..40a2cb6b1 --- /dev/null +++ b/docker/managed-agents/verify_anthropic_sdk.py @@ -0,0 +1,65 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Check the installed SDK is unchanged from its hash-locked wheel RECORD. + +The installer verifies the official artifact hash from uv.lock. This separate +check rejects changed, missing or overlaid Python files after installation; it +is not an independent trust root if RECORD itself has been replaced. +""" + +import base64 +import csv +import hashlib +import importlib.metadata +import io +from pathlib import PurePosixPath + + +def verify_distribution(distribution: importlib.metadata.Distribution) -> int: + record = distribution.read_text("RECORD") + if not record: + raise RuntimeError("Anthropic SDK integrity: missing wheel RECORD") + expected = set() + for name, digest, _size in csv.reader(io.StringIO(record)): + path = PurePosixPath(name) + if not name.startswith("anthropic/") or path.suffix != ".py": + continue + if path.is_absolute() or ".." in path.parts or not digest.startswith("sha256="): + raise RuntimeError("Anthropic SDK integrity: invalid Python file record") + installed = distribution.locate_file(name) + if not installed.is_file() or installed.is_symlink(): + raise RuntimeError(f"Anthropic SDK integrity: missing or linked {name}") + actual = ( + base64.urlsafe_b64encode(hashlib.sha256(installed.read_bytes()).digest()) + .decode() + .rstrip("=") + ) + if actual != digest.removeprefix("sha256="): + raise RuntimeError(f"Anthropic SDK integrity: modified {name}") + expected.add(installed) + if not expected: + raise RuntimeError("Anthropic SDK integrity: no recorded Python files") + actual_files = set(distribution.locate_file("anthropic").rglob("*.py")) + if actual_files != expected: + raise RuntimeError("Anthropic SDK integrity: unrecorded Python files") + return len(expected) + + +if __name__ == "__main__": + sdk = importlib.metadata.distribution("anthropic") + count = verify_distribution(sdk) + print( + f"Official Anthropic SDK {sdk.version}: {count} Python files match wheel RECORD" + ) diff --git a/docs/content/docs/framework/tools/feishu-channel.en.mdx b/docs/content/docs/framework/tools/feishu-channel.en.mdx index f8faeae17..e574c9fd7 100644 --- a/docs/content/docs/framework/tools/feishu-channel.en.mdx +++ b/docs/content/docs/framework/tools/feishu-channel.en.mdx @@ -37,6 +37,12 @@ Environment variables: - `TOOL_FEISHU_CHANNEL_TRANSPORT`: defaults to `ws` - `TOOL_FEISHU_CHANNEL_STREAMING`: whether to enable streaming output, defaults to `false` - `TOOL_FEISHU_CHANNEL_REACTIONS`: whether to reply with a "received" reaction on incoming messages, defaults to `false` +- `TOOL_FEISHU_CHANNEL_SHOW_THINKING`: whether to show model-provided thinking summaries, defaults to `false` +- `TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS`: whether to show tool names and arguments, defaults to `false` +- `TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS`: whether to show tool results, defaults to `false` +- `TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS`: whether each tool call and its result use a dedicated card, defaults to `false` +- `TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD`: whether thinking summaries use a dedicated card, defaults to `false` +- `TOOL_FEISHU_CHANNEL_CREATE_TOPIC`: whether to create a topic from the user message and send every card for that response into it, defaults to `false` Or configure in `config.yaml`: @@ -48,6 +54,12 @@ tool: transport: ws streaming: true reactions: true + show_thinking: true + show_tool_calls: true + show_tool_results: true + separate_tool_call_cards: true + separate_thinking_card: true + create_topic: true ``` ## Minimal example @@ -75,6 +87,10 @@ runner = Runner( channel = FeishuChannelExtension( runner=runner, channel_kwargs={"transport": "ws"}, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, ) @@ -91,3 +107,7 @@ if __name__ == "__main__": - The Feishu `Channel` WebSocket mode is used by default, so as long as the bot has subscribed to message events, you can start the connection directly. - By default replies use `reply_to=`, keeping VeADK output attached to the current Feishu message thread. - Override the default identity mapping via `session_id_factory` and `user_id_factory`. +- Thinking includes only summaries emitted by the model or provider; hidden internal chain-of-thought is not available. +- Common credential fields in tool arguments and results are redacted, and each detail block is limited to 4,000 characters by default. +- Redaction matches structured field names; it does not remove credentials embedded in arbitrary text. Enable tool detail output only for payloads intended for the conversation participants. Set `tool_detail_max_length` on the constructor to change the limit (minimum 100 characters). +- These display controls apply to streaming replies. Separate cards need the corresponding `show_*` option. `create_topic` requires `reply_in_thread=True` and an inbound message ID. diff --git a/docs/content/docs/framework/tools/feishu-channel.mdx b/docs/content/docs/framework/tools/feishu-channel.mdx index 67408c2b2..f119ce546 100644 --- a/docs/content/docs/framework/tools/feishu-channel.mdx +++ b/docs/content/docs/framework/tools/feishu-channel.mdx @@ -37,6 +37,12 @@ pip install lark-oapi - `TOOL_FEISHU_CHANNEL_TRANSPORT`：默认 `ws` - `TOOL_FEISHU_CHANNEL_STREAMING`：是否开启流式输出，默认 `false` - `TOOL_FEISHU_CHANNEL_REACTIONS`：是否在收到消息时回复“收到”表情，默认 `false` +- `TOOL_FEISHU_CHANNEL_SHOW_THINKING`：是否展示模型提供的思考摘要，默认 `false` +- `TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS`：是否展示工具名称和参数，默认 `false` +- `TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS`：是否展示工具执行结果，默认 `false` +- `TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS`：每次工具调用及其结果是否使用独立卡片，默认 `false` +- `TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD`：是否将思考摘要放入独立卡片，默认 `false` +- `TOOL_FEISHU_CHANNEL_CREATE_TOPIC`：是否围绕用户消息创建话题并将本次回复的所有卡片发送到该话题，默认 `false` 或在 `config.yaml` 中配置： @@ -48,6 +54,12 @@ tool: transport: ws streaming: true reactions: true + show_thinking: true + show_tool_calls: true + show_tool_results: true + separate_tool_call_cards: true + separate_thinking_card: true + create_topic: true ``` ## 最小示例 @@ -75,6 +87,10 @@ runner = Runner( channel = FeishuChannelExtension( runner=runner, channel_kwargs={"transport": "ws"}, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, ) @@ -91,3 +107,7 @@ if __name__ == "__main__": - 默认使用飞书 `Channel` 的 WebSocket 模式，因此只要机器人已订阅消息事件，即可直接启动连接。 - 默认回复使用 `reply_to=原消息 message_id`，让 VeADK 的输出挂在当前飞书消息线程下。 - 你可以通过 `session_id_factory` 与 `user_id_factory` 覆盖默认的身份映射逻辑。 +- Thinking 只展示模型或 Provider 实际返回的思考摘要，不会也无法展示模型未返回的内部思维链。 +- 工具参数和结果中的常见凭证字段会被脱敏，单段详情默认最多展示 4000 个字符。 +- 脱敏匹配结构化字段名称，不移除任意文本内嵌的凭据；仅在载荷适合会话参与者查看时启用工具详情。构造参数 `tool_detail_max_length` 可调整限制，最小 100 个字符。 +- 这些显示控制适用于流式回复；独立卡片需要同时启用对应的 `show_*` 参数。`create_topic` 要求 `reply_in_thread=True` 且有入站消息 ID。 diff --git a/examples/16_self_host_sandbox/.env.example b/examples/16_self_host_sandbox/.env.example index 43d9e06fe..05edcc3cd 100644 --- a/examples/16_self_host_sandbox/.env.example +++ b/examples/16_self_host_sandbox/.env.example @@ -1,18 +1,42 @@ -# Copy this file to `.env` and fill in your values: -# cp examples/16_self_host_sandbox/.env.example examples/16_self_host_sandbox/.env +# Copy to .env; configure privately. Explicit process variables take precedence. +ANTHROPIC_BASE_URL=https://gateway.example.com +ANTHROPIC_ENVIRONMENT_KEY=replace-me +# Required for environment polling and for creating a new remote Session. +ANTHROPIC_ENVIRONMENT_ID=replace-me +SANDBOX_AGENT_ID=replace-me +# Set only if your gateway requires an account header. +# X_TOP_ACCOUNT_ID=replace-me -# --- Managed Runtime connection --- -# BOE Runtime 7hw8g3yr exposes its Managed Sessions gateway through this URL. -ANTHROPIC_BASE_URL=https://.apigateway-cn-beijing.volceapi.com -ANTHROPIC_ENVIRONMENT_ID=env_your_runtime_environment_id -ANTHROPIC_ENVIRONMENT_KEY=your_runtime_environment_key -SANDBOX_AGENT_ID=agent_your_managed_agent_id -X_TOP_ACCOUNT_ID=your_volcengine_account_id +MANAGED_AGENT_WORK_SCOPE=account +MA_AGENT_LOOP_RUNTIME_TYPE=agentkit_runtime +MANAGED_AGENT_TOOL_EXECUTION=remote +MANAGED_AGENT_WORK_CONCURRENCY=500 +MANAGED_AGENT_WORKDIR=/workspace +# Runtime Identity may replace the explicit environment key. +# AGENTKIT_RUNTIME_ID=replace-me +# VOLCENGINE_REGION=cn-beijing -# --- Web Server Configuration --- -PORT=8067 -# Comma-separated origins when the Dev UI is accessed through another origin. -# ALLOW_ORIGINS=http://localhost:3000,https://your-dev-proxy.example.com +# Session event history and Ark response IDs provide cross-worker recovery. +# Worker bookkeeping is local; legacy PostgreSQL environment variables are ignored. -# Maximum time to wait for a remote Runtime turn, including tool execution. -# SANDBOX_TIMEOUT_SECONDS=300 +# Diagnostics are disabled by default. Model/poll traces contain safe metadata. +# Full event tracing includes user/tool content; keep outputs private. +# MANAGED_AGENT_MODEL_TRACE=true +# MANAGED_AGENT_POLL_TRACE=true +# MA_DEBUG_EVENTS=1 + +# Long-running Feishu bot channel (install extensions extra). +# TOOL_FEISHU_CHANNEL_APP_ID=replace-me +# TOOL_FEISHU_CHANNEL_APP_SECRET=replace-me +# TOOL_FEISHU_CHANNEL_TRANSPORT=ws +# TOOL_FEISHU_CHANNEL_STREAMING=true +# TOOL_FEISHU_CHANNEL_SHOW_THINKING=true +# TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS=true +# TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS=true +# TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS=true +# TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD=true +# TOOL_FEISHU_CHANNEL_CREATE_TOPIC=true + +# Optional external inspector adapter for cloud acceptance scripts. +# AGENTKIT_TOOL_DEPLOY_SCRIPT=/absolute/path/to/your/agentkit_tool_deploy.py +# AGENTKIT_TOOL_ID=replace-me diff --git a/examples/16_self_host_sandbox/README.md b/examples/16_self_host_sandbox/README.md index 70dbe8adb..09a420dbc 100644 --- a/examples/16_self_host_sandbox/README.md +++ b/examples/16_self_host_sandbox/README.md @@ -1,141 +1,93 @@ -# 16. Self-Hosted Sandbox Agent Demo - -This example runs the conversation and model loop in `veadk web`. Each new VeADK Web session creates exactly one remote Managed Session. Only VeADK-generated `agent.tool_use` events are sent to the Self-Hosted Runtime dispatcher, which creates a TAE Tool session for the configured Sandbox. - ---- - -## 🎯 Architecture - -```text - [ User Prompt ] - │ - ▼ - ┌───────────────────────┐ - │ VeADK Web │ - │ model + agent loop │ - └──────────┬────────────┘ - │ POST agent.tool_use - ▼ - ┌───────────────────────────┐ - │ Runtime 7hw8g3yr │ - │ dispatcher only │ - └──────────┬────────────────┘ - │ creates one tool session - ▼ - ┌────────────────────────────────────────┐ - │ TAE Sandbox Tool m3m24zxs │ - │ executes the Runtime's tool calls │ - └────────────────────────────────────────┘ -``` - ---- +# 16. Self-hosted Managed Agents -## 🚀 Quick Start +[中文版](README.zh.md) -### 1. Configure Environment Variables +VeADK runs the model loop; native tool execution may run locally in an isolated +worker or dispatch through a Managed Agents gateway. The reusable implementation +ships in `veadk.runtime.managed_agents` and `veadk.integrations.mpa`. This directory +contains launchers and acceptance scripts, not production library code. -Copy the example environment configuration file and update it with your credentials: +## Install and run ```bash +uv sync --extra sandbox --extra extensions cp examples/16_self_host_sandbox/.env.example examples/16_self_host_sandbox/.env +# Configure the private .env before starting a worker. +bash examples/16_self_host_sandbox/run.sh --managed-agent-worker ``` -Alternatively, export variables directly: +The launcher loads `.env` as defaults; exported variables take precedence. The +library never loads example configuration or creates a remote Session at import. +Run an installed package directly without this example directory: ```bash -export ANTHROPIC_BASE_URL="https://" -export ANTHROPIC_ENVIRONMENT_ID="env_01SLqXHseguCmohifEqeUAYu" -export ANTHROPIC_ENVIRONMENT_KEY="your-runtime-token" -export SANDBOX_AGENT_ID="agent_your_managed_agent_id" -export X_TOP_ACCOUNT_ID="your-account-id" +python -m veadk.runtime.managed_agents.worker --managed-agent-worker +python -m veadk.runtime.managed_agents.worker --managed-agent-work-item ``` -The local VeADK process needs its normal model configuration. Model reasoning happens in VeADK; only tool execution is delegated to the remote Runtime and Sandbox. - +Worker mode continuously polls Work. Claimed mode requires injected Work/Session +IDs and lease state from the external dispatcher. The Session snapshot owns model, +permissions, tools, Skills, MCP and Identity bindings. Account scope selects the +configured execution pool; environment scope additionally needs an Environment ID. +Local conversation mode creates one Managed Session for each VeADK Session; the +Runner releases it to idle when overlapping local turns finish. -### 2. Run the Demo +## Execution and state -#### Option A: Web UI Mode (`veadk web`) +Native Anthropic tools run through their tool context, retaining file-edit/read +semantics. MCP tools keep their actual MCP transport. Tools requiring confirmation +wait for correlated permission events; custom tools wait for correlated results. +The gateway event ledger reconciles completed inputs and Ark response IDs across +worker replacement. The worker always uses process-local ADK bookkeeping and ignores legacy database +environment variables; durable recovery comes from the event ledger and Ark response IDs. -Run the visual web debugging interface in your browser (defaults to port **8067**, configurable via `PORT` in `.env` or `--port`): +A worker drains on SIGTERM, retains per-Session ordering, and bounds concurrent +Work using `MANAGED_AGENT_WORK_CONCURRENCY`. File tools confine paths to their +workspace. Bash requires OS isolation supplied by the worker deployment. +Cancellation/failure closes runners, MCP/native tool contexts and temporary Skills. +Runtime and Session credentials remain in process memory. -```bash -# Via run.sh helper (automatically switches to agents dir and applies port 8067) -bash examples/16_self_host_sandbox/run.sh --web - -# Or directly in the agents directory -cd examples/16_self_host_sandbox/agents -veadk web --port 8067 -``` +## Feishu channel - - - -#### Option B: CLI Mode +Configure the bot credentials privately and enable the desired optional streaming, +thinking, tool detail, card and topic settings from `.env.example`, then run: ```bash -# Run with run.sh helper -bash examples/16_self_host_sandbox/run.sh - -# Or using uv -uv run --extra sandbox python examples/16_self_host_sandbox/main.py - -# Or select an existing local-to-remote session mapping -python examples/16_self_host_sandbox/main.py \ - --session-id "my-local-session" \ - --prompt "Create a Python script in /workspace and run it with pytest." +bash examples/16_self_host_sandbox/run.sh --feishu ``` +This starts a long-running channel, reconnects WebSocket transport, and drains +in-flight replies on shutdown. Channel settings are opt-in outside this demo. -`ShortTermMemory.after_create_session_callback` creates one remote Managed Session -for each newly created VeADK session. VeADK handles the user message and model -loop locally. `DispatchRuntimeProvider` intercepts each model tool call, posts an -`agent.tool_use`, waits for its matching tool result, and returns that result to -the local VeADK model loop. - -## Docker and Kubernetes deployment - -Use the repository root as the Docker build context: - -```bash -docker build \ - -f examples/16_self_host_sandbox/Dockerfile \ - -t /veadk-self-host-sandbox: \ - . -docker push /veadk-self-host-sandbox: -``` - -Replace the image and Ingress host in `k8s.yaml` for the target environment. Then -create an allowlisted Secret from the local `.env` and deploy the workload: +## Tests and acceptance scripts ```bash -set -a -source examples/16_self_host_sandbox/.env -set +a - -kubectl create secret generic veadk-self-host-sandbox-env \ - --from-literal=ANTHROPIC_BASE_URL="$ANTHROPIC_BASE_URL" \ - --from-literal=ANTHROPIC_ENVIRONMENT_ID="$ANTHROPIC_ENVIRONMENT_ID" \ - --from-literal=ANTHROPIC_ENVIRONMENT_KEY="$ANTHROPIC_ENVIRONMENT_KEY" \ - --from-literal=SANDBOX_AGENT_ID="$SANDBOX_AGENT_ID" \ - --from-literal=X_TOP_ACCOUNT_ID="$X_TOP_ACCOUNT_ID" \ - --dry-run=client -o yaml | kubectl apply -f - - -kubectl apply -f examples/16_self_host_sandbox/k8s.yaml -kubectl rollout status deployment/veadk-self-host-sandbox +uv run --extra dev --extra sandbox pytest tests/runtime/managed_agents tests/integrations/mpa +python examples/16_self_host_sandbox/local_agent_loop_test.py +python examples/16_self_host_sandbox/anthropic_gateway_e2e.py --help +python examples/16_self_host_sandbox/soak_load_test.py --help ``` -Do not create the Secret from the entire `.env`. The Deployment only injects -the allowlisted Runtime connection values above. - -## Development verification - -The repository virtual environment may omit test dependencies when it was -created for runtime use only. Install the development dependency group before -running the focused dispatch tests: - -```bash -uv sync --group dev --extra sandbox -uv run pytest -q tests/runtime/test_self_host_sandbox_client.py \ - tests/runtime/test_self_host_sandbox_agent.py -``` +Unit tests use synthetic services. `anthropic_gateway_e2e.py` offers an official +SDK conversation check and sandbox lifecycle modes; conversation mode does not +require an AgentKit inspector. Optional cloud inspection requires explicit +`AGENTKIT_TOOL_DEPLOY_SCRIPT` and `AGENTKIT_TOOL_ID` inputs. No personal helper path +or deployment ID is supplied. Soak, fault, distributed-worker and Kubernetes +scripts are operator-triggered live checks; examine their required configuration +before running them. Keep local result files and full event logs private. + +The distributed helper requires `ANTHROPIC_BASE_URL`, `ANTHROPIC_ENVIRONMENT_ID`, +`ANTHROPIC_ENVIRONMENT_KEY` and `SANDBOX_AGENT_ID` for existing isolated test +resources. Its Agent must enable Ark Session credentials and native bash. It +checks canonical Session history/continuation and archives only its own Session. +`ark_session_e2e.py` likewise requires explicit gateway and Agent configuration. + +Image build, dependency pitfalls, readiness, writable mounts and the Kubernetes +worker template are covered in the [build guide](../../docker/managed-agents/README.md). +Building, local tests, publication, deployment and cloud E2E are separate outcomes. +The migration deliberately excludes the source's standalone Managed Agents UI. + +Legacy imports `sandbox_client`, `managed_agent_loop`, `runtime_identity`, +`managed_session_resources` and `event_debug` remain example compatibility adapters. +Use packaged imports for new code. See the [runtime contract](../../specs/managed-agent-runtime/README.md) +and [migration design](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md). diff --git a/examples/16_self_host_sandbox/README.zh.md b/examples/16_self_host_sandbox/README.zh.md index 5b4de003a..675bb75f1 100644 --- a/examples/16_self_host_sandbox/README.zh.md +++ b/examples/16_self_host_sandbox/README.zh.md @@ -1,138 +1,83 @@ -# 16. 自建沙箱 Agent 示例 (Self-Hosted Sandbox) - -本示例由 `veadk web` 处理用户消息和模型循环。每个新建的 VeADK Web Session 只创建一个远端 Managed Session；仅将 VeADK 模型产生的 `agent.tool_use` 发送给 Self-Hosted Runtime dispatcher，再由它为配置的 TAE Sandbox 创建 Tool Session。 - ---- - -## 🎯 架构拓扑 - -```text - 【用户指令】: "在沙箱中写个快速排序并运行测试" - │ - ▼ - ┌───────────────────────┐ - │ VeADK Web │ - │ 模型与 Agent 循环 │ - └──────────┬────────────┘ - │ POST agent.tool_use - ▼ - ┌───────────────────────────┐ - │ Runtime 7hw8g3yr │ - │ 仅负责 dispatcher │ - └──────────┬────────────────┘ - │ 为每个 Session 创建 Tool Session - ▼ - ┌────────────────────────────────────────┐ - │ TAE Sandbox Tool m3m24zxs │ - │ 执行 Runtime 产生的工具调用 │ - └────────────────────────────────────────┘ -``` +# 16. 自托管 Managed Agents ---- +[English](README.md) -## 🚀 运行方法 - -### 1. 配置环境变量 +VeADK 执行模型循环；原生工具可在隔离 Worker 内本地执行，或通过 Managed Agents +网关远程派发。复用实现位于 `veadk.runtime.managed_agents` 和 +`veadk.integrations.mpa`，本目录保留启动器与验收脚本，不承载生产库代码。 -复制环境配置文件示例并填入对应的配置与密钥： +## 安装与运行 ```bash +uv sync --extra sandbox --extra extensions cp examples/16_self_host_sandbox/.env.example examples/16_self_host_sandbox/.env +# 启动前填写私有 .env 配置。 +bash examples/16_self_host_sandbox/run.sh --managed-agent-worker ``` -或直接在终端中导出环境变量： +启动器将 `.env` 作为默认值，显式环境变量优先。库导入不加载示例配置，也不创建 +远程 Session。安装后可脱离示例目录直接运行： ```bash -export ANTHROPIC_BASE_URL="https://" -export ANTHROPIC_ENVIRONMENT_ID="env_01SLqXHseguCmohifEqeUAYu" -export ANTHROPIC_ENVIRONMENT_KEY="your-runtime-token" -export SANDBOX_AGENT_ID="agent_your_managed_agent_id" -export X_TOP_ACCOUNT_ID="your-account-id" +python -m veadk.runtime.managed_agents.worker --managed-agent-worker +python -m veadk.runtime.managed_agents.worker --managed-agent-work-item ``` -本地 VeADK 进程需要正常的模型配置。模型推理由 VeADK 完成，只有 Tool 执行委托给远端 Runtime 与 Sandbox。 - - -### 2. 启动示例 +Worker 模式持续轮询 Work。claimed 模式需要外部 dispatcher 注入 Work/Session ID +和租约状态。Session 快照拥有模型、权限、工具、Skills、MCP 和 Identity 绑定。 +账号作用域选择配置的执行池；Environment 作用域还需要 Environment ID。 +本地会话模式将每个 VeADK Session 映射到一个 Managed Session，Runner 在所有 +重叠本地轮次结束后发送 idle。 -#### 方式一：Web 界面交互模式（`veadk web`） +## 执行与状态 -在浏览器中启动可视化 Web 调试交互界面（默认端口为 **8067**，支持在 `.env` 中通过 `PORT` 或 `--port` 自定义）： +Anthropic 原生工具通过自身上下文执行，保留文件编辑与读取语义。MCP 工具使用 +实际 MCP transport。需确认的工具等待关联的权限事件，自定义工具等待关联结果。 +网关事件账本负责跨 Worker 恢复已完成输入与 Ark response ID。Worker 始终使用进程本地 ADK 记账，忽略旧数据库环境变量； +持久恢复依赖事件账本与 Ark response ID。 -```bash -# 方式 A：通过 run.sh 快捷启动（自动切换到 agents 目录并使用 8067 端口） -bash examples/16_self_host_sandbox/run.sh --web - -# 方式 B：进入 agents 目录直接运行 -cd examples/16_self_host_sandbox/agents -veadk web --port 8067 -``` +Worker 在 SIGTERM 时排空，保持单 Session 顺序，并以 +`MANAGED_AGENT_WORK_CONCURRENCY` 限制并发 Work。文件工具限制工作区路径，bash +依赖部署提供 OS 隔离。取消或失败关闭 Runner、MCP/原生工具上下文和临时 Skills。 +Runtime 与 Session 凭据保持在进程内存。 +## 飞书渠道 - - -#### 方式二：CLI 命令行模式 +私下配置机器人凭据，并按需启用 `.env.example` 中的流式、思考、工具详情、卡片 +和话题选项，然后运行： ```bash -# 使用 run.sh 运行 -bash examples/16_self_host_sandbox/run.sh - -# 或使用 uv -uv run --extra sandbox python examples/16_self_host_sandbox/main.py - -# 或指定一个本地到远端的 Session 映射 ID -python examples/16_self_host_sandbox/main.py \ - --session-id "my-local-session" \ - --prompt "请在 /workspace 下创建一个 quick_sort.py 并运行验证" +bash examples/16_self_host_sandbox/run.sh --feishu ``` +该命令启动常驻渠道、重连 WebSocket，并在关闭时排空回复。示例之外的渠道设置 +默认关闭。 -每个新建的 VeADK Session 都通过 `POST /v1/sessions` 创建一个远端 Managed -Session。用户消息和模型循环留在 VeADK；`DispatchRuntimeProvider` 拦截模型生成的 -Tool call，以 `agent.tool_use` 发送给 Runtime，等待匹配的 Tool result 后交还给本地 -VeADK 模型继续生成最终回复。 - -## Docker 与 Kubernetes 部署 - -Docker 构建上下文必须使用仓库根目录： - -```bash -docker build \ - -f examples/16_self_host_sandbox/Dockerfile \ - -t /veadk-self-host-sandbox: \ - . -docker push /veadk-self-host-sandbox: -``` - -将 `k8s.yaml` 中的镜像和 Ingress 域名替换成目标环境的值。然后从本地 `.env` -显式创建白名单 Secret，再部署工作负载： +## 测试与验收脚本 ```bash -set -a -source examples/16_self_host_sandbox/.env -set +a - -kubectl create secret generic veadk-self-host-sandbox-env \ - --from-literal=ANTHROPIC_BASE_URL="$ANTHROPIC_BASE_URL" \ - --from-literal=ANTHROPIC_ENVIRONMENT_ID="$ANTHROPIC_ENVIRONMENT_ID" \ - --from-literal=ANTHROPIC_ENVIRONMENT_KEY="$ANTHROPIC_ENVIRONMENT_KEY" \ - --from-literal=SANDBOX_AGENT_ID="$SANDBOX_AGENT_ID" \ - --from-literal=X_TOP_ACCOUNT_ID="$X_TOP_ACCOUNT_ID" \ - --dry-run=client -o yaml | kubectl apply -f - - -kubectl apply -f examples/16_self_host_sandbox/k8s.yaml -kubectl rollout status deployment/veadk-self-host-sandbox +uv run --extra dev --extra sandbox pytest tests/runtime/managed_agents tests/integrations/mpa +python examples/16_self_host_sandbox/local_agent_loop_test.py +python examples/16_self_host_sandbox/anthropic_gateway_e2e.py --help +python examples/16_self_host_sandbox/soak_load_test.py --help ``` -不要使用整个 `.env` 创建 Secret；Deployment 只注入上述白名单内的 Runtime 连接配置。 - -## 开发验证 - -如果仓库虚拟环境仅安装了运行时依赖，其中可能没有 `pytest`。请先同步开发依赖组， -再运行 Tool dispatch 的定向测试： - -```bash -uv sync --group dev --extra sandbox -uv run pytest -q tests/runtime/test_self_host_sandbox_client.py \ - tests/runtime/test_self_host_sandbox_agent.py -``` +单元测试使用模拟服务。`anthropic_gateway_e2e.py` 提供官方 SDK 对话检查及沙箱生命 +周期模式；对话模式不需要 AgentKit inspector。可选云资源检查必须显式配置 +`AGENTKIT_TOOL_DEPLOY_SCRIPT` 和 `AGENTKIT_TOOL_ID`，没有个人辅助脚本路径或部署 ID +默认值。压测、故障、分布式 Worker 和 Kubernetes 脚本属于操作者主动执行的线上 +检查，运行前检查必填配置。本地结果文件与完整事件日志保持私有。 + +分布式脚本必须通过 `ANTHROPIC_BASE_URL`、`ANTHROPIC_ENVIRONMENT_ID`、 +`ANTHROPIC_ENVIRONMENT_KEY` 和 `SANDBOX_AGENT_ID` 指定已有的隔离测试资源；测试 Agent +必须启用 Ark Session 凭据与原生 bash。脚本检查权威 Session 历史/续接，只归档自己 +创建的 Session。`ark_session_e2e.py` 同样必须显式配置网关与 Agent。 + +镜像构建、依赖注意事项、就绪、可写挂载和 Kubernetes Worker 模板见 +[构建指南](../../docker/managed-agents/README.zh.md)。镜像构建、本地测试、发布、部署 +与云 E2E 是不同结果。本次迁移排除来源的独立 Managed Agents UI。 + +旧 `sandbox_client`、`managed_agent_loop`、`runtime_identity`、 +`managed_session_resources` 和 `event_debug` 导入保留为示例兼容适配器，新代码使用 +包导入。参见[运行时契约](../../specs/managed-agent-runtime/README.zh.md)与 +[迁移设计](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md)。 diff --git a/examples/16_self_host_sandbox/agentkit_runtime.py b/examples/16_self_host_sandbox/agentkit_runtime.py new file mode 100644 index 000000000..465c951f8 --- /dev/null +++ b/examples/16_self_host_sandbox/agentkit_runtime.py @@ -0,0 +1,6 @@ +"""Run the packaged AgentKit health listener.""" + +import runpy + +if __name__ == "__main__": + runpy.run_module("veadk.runtime.managed_agents.health", run_name="__main__") diff --git a/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/__init__.py b/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/__init__.py index d9e3bef2f..aa8e7ed21 100644 --- a/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/__init__.py +++ b/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/__init__.py @@ -1,19 +1,3 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Self-Hosted Sandbox Agent package.""" - from .agent import agent, root_agent __all__ = ["agent", "root_agent"] diff --git a/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/agent.py b/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/agent.py index 1e6a7c309..0e35e4030 100644 --- a/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/agent.py +++ b/examples/16_self_host_sandbox/agents/self_host_sandbox_agent/agent.py @@ -1,182 +1,15 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +"""Optional remote-backed demo Agent for VeADK discovery.""" -"""VeADK agent whose tool calls run through a Self-Hosted Sandbox Runtime.""" - -import asyncio -import logging -import sys -import threading from pathlib import Path -from typing import Any from dotenv import load_dotenv -EXAMPLE_DIR = Path(__file__).resolve().parents[2] -if str(EXAMPLE_DIR) not in sys.path: - sys.path.insert(0, str(EXAMPLE_DIR)) -load_dotenv(EXAMPLE_DIR / ".env", override=False) - -from sandbox_client import SelfHostSandboxClient # noqa: E402 -from veadk import Agent # noqa: E402 -from veadk.memory.short_term_memory import ShortTermMemory # noqa: E402 -from veadk.runtime import DispatchRuntimeProvider, ToolCall # noqa: E402 - -logger = logging.getLogger("veadk.sandbox_agent") - - -class SandboxSessionManager: - """Bind exactly one remote Managed Session to each VeADK session.""" - - def __init__(self) -> None: - self._clients: dict[str, SelfHostSandboxClient] = {} - self._lock = threading.Lock() - - def _new_client(self) -> SelfHostSandboxClient: - client = SelfHostSandboxClient() - client.session_id = None - return client - - def get(self, veadk_session_id: str) -> SelfHostSandboxClient: - key = veadk_session_id or "__default__" - with self._lock: - client = self._clients.get(key) - if client is None: - client = self._new_client() - self._clients[key] = client - return client - - def bind(self, veadk_session_id: str, client: SelfHostSandboxClient) -> None: - key = veadk_session_id or "__default__" - with self._lock: - self._clients[key] = client - - def create_remote_session(self, veadk_session_id: str) -> str: - client = self.get(veadk_session_id) - if not client.session_id: - client.create_session( - title=f"VeADK Self-Hosted Sandbox Session {veadk_session_id}" - ) - return client.session_id or "" - - def release(self, veadk_session_id: str) -> None: - key = veadk_session_id or "__default__" - with self._lock: - client = self._clients.pop(key, None) - if client: - client.post_status_idle() - - -sandbox_sessions = SandboxSessionManager() -sandbox_client = sandbox_sessions.get("__default__") - - -def bash(command: str, timeout: float = 120) -> dict: - """Execute a shell command in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def read_file( - file_path: str, offset: int | None = None, limit: int | None = None -) -> dict: - """Read a text file in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def write_file(file_path: str, content: str) -> dict: - """Create or overwrite a text file in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def edit_file( - file_path: str, old_string: str, new_string: str, replace_all: bool = False -) -> dict: - """Replace an exact string in a remote sandbox file.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def list_files(path: str = "/workspace", max_depth: int = 4) -> dict: - """List files below a directory in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def search_files( - pattern: str, - path: str = "/workspace", - glob: str | None = None, - case_insensitive: bool = False, -) -> dict: - """Search files in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -def python(code: str, workdir: str = "/workspace", timeout: float = 120) -> dict: - """Run Python code in the remote sandbox.""" - raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") - - -async def after_create_session(session: Any) -> None: - veadk_session_id = str(getattr(session, "id", "") or "") - remote_id = await asyncio.to_thread( - sandbox_sessions.create_remote_session, veadk_session_id - ) - logger.info( - "Remote Managed Session created: veadk=%s remote=%s", - veadk_session_id, - remote_id, - ) - - -async def dispatch_task(tool_call: ToolCall) -> dict: - client = sandbox_sessions.get(tool_call.session_id) - return await asyncio.to_thread( - client.dispatch_tool, - tool_call.name, - tool_call.arguments, - dispatch_id=tool_call.id, - ) - - -dispatch_runtime = DispatchRuntimeProvider(dispatch_task, dispatchable_tools=None) -short_term_memory = ShortTermMemory(after_create_session_callback=after_create_session) - -# veadk web calls the service directly, so preserve the per-session callback -# when the UI creates a session outside Runner.create_session(). -_create_local_session = short_term_memory.session_service.create_session - - -async def _create_session_with_remote(*args: Any, **kwargs: Any) -> Any: - session = await _create_local_session(*args, **kwargs) - if session is not None: - await after_create_session(session) - return session +def _demo_agent(): + load_dotenv(Path(__file__).resolve().parents[2] / ".env", override=False) + from veadk.runtime.managed_agents.sandbox import get_default_agent -short_term_memory.session_service.create_session = _create_session_with_remote + return get_default_agent() -agent = Agent( - name="self_host_sandbox_agent", - description="An engineering agent dispatching tool calls to a Self-Hosted Worker.", - instruction=( - f"You are an autonomous engineering assistant connected to Runtime 7hw8g3yr " - f"(Environment: {sandbox_client.environment_id}). Handle user conversation and " - "model reasoning in VeADK. Use the provided tools for every sandbox operation; " - "tool calls are dispatched through Runtime 7hw8g3yr to TAE Sandbox m3m24zxs." - ), - tools=[bash, read_file, write_file, edit_file, list_files, search_files, python], - short_term_memory=short_term_memory, - before_tool_callback=dispatch_runtime.before_tool_callback, -) -root_agent = agent +agent = root_agent = _demo_agent() diff --git a/examples/16_self_host_sandbox/anthropic_gateway_e2e.py b/examples/16_self_host_sandbox/anthropic_gateway_e2e.py new file mode 100644 index 000000000..a55cfe706 --- /dev/null +++ b/examples/16_self_host_sandbox/anthropic_gateway_e2e.py @@ -0,0 +1,723 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Official Anthropic SDK E2E for the Managed Agents public gateway. + +Conversation operations use only ``anthropic.AsyncAnthropic``. Optional +AgentKit control-plane inspection is used solely to prove physical sandbox +creation, reclamation, and replacement; signed endpoints and credentials are +never printed or written to the result file. +""" + +from __future__ import annotations + +import argparse +import asyncio +import importlib.util +import json +import os +import re +import sys +import time +import uuid +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import anthropic +from dotenv import load_dotenv + +DEFAULT_TOOL_ID = os.getenv("AGENTKIT_TOOL_ID", "") +DEFAULT_AGENTKIT_HELPER = "" + + +def now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +def field_value(value: Any, name: str, default: Any = None) -> Any: + return ( + value.get(name, default) + if isinstance(value, dict) + else getattr(value, name, default) + ) + + +def content_text(content: Any) -> str: + if isinstance(content, str): + return content + return "".join( + str(field_value(block, "text", "") or "") + for block in content or [] + if field_value(block, "type") == "text" + ) + + +@dataclass +class ObservedEvent: + type: str + elapsed_seconds: float + event_id: str | None = None + link_id: str | None = None + text: str = "" + + +@dataclass +class TurnResult: + session_id: str + marker: str + started_at: str + elapsed_seconds: float + first_delta_seconds: float | None + final_message_seconds: float | None + event_types: list[str] + tool_use_ids: list[str] + tool_result_ids: list[str] + final_text: str + observations: list[ObservedEvent] = field(default_factory=list) + + def assert_complete(self, *, require_stream: bool, require_tool: bool) -> None: + if "session.status_running" not in self.event_types: + raise AssertionError(f"{self.session_id}: missing session.status_running") + if "session.status_idle" not in self.event_types: + raise AssertionError(f"{self.session_id}: missing session.status_idle") + if self.marker not in self.final_text: + raise AssertionError( + f"{self.session_id}: final response did not contain {self.marker!r}" + ) + if require_stream: + if self.first_delta_seconds is None: + raise AssertionError(f"{self.session_id}: no event_delta was streamed") + if self.final_message_seconds is None: + raise AssertionError(f"{self.session_id}: no final agent.message") + if self.first_delta_seconds >= self.final_message_seconds: + raise AssertionError( + f"{self.session_id}: first delta did not precede the final message" + ) + if require_tool: + if not self.tool_use_ids: + raise AssertionError(f"{self.session_id}: no agent.tool_use") + if not self.tool_result_ids: + raise AssertionError(f"{self.session_id}: no agent.tool_result") + missing = set(self.tool_use_ids) - set(self.tool_result_ids) + if missing: + raise AssertionError( + f"{self.session_id}: tool calls without results: {sorted(missing)}" + ) + + +class AgentKitInspector: + def __init__(self, helper_path: str, tool_id: str) -> None: + path = Path(helper_path) + if not path.is_file(): + raise FileNotFoundError(f"AgentKit helper not found: {path}") + spec = importlib.util.spec_from_file_location("agentkit_tool_e2e_helper", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Could not load AgentKit helper: {path}") + helper_dir = str(path.parent) + if helper_dir not in sys.path: + sys.path.insert(0, helper_dir) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + self.module: ModuleType = module + dotenv = os.getenv("AGENTKIT_DOTENV", str(Path(__file__).with_name(".env"))) + module.load_dotenv(dotenv) + args = SimpleNamespace( + bp=False, + access_key=None, + secret_key=None, + region=None, + service=None, + ) + auth = module.require_credentials(args) + self.client = module.AgentkitClient( + os.getenv("AGENTKIT_ENDPOINT", module.DEFAULT_ENDPOINT), + os.getenv("AGENTKIT_VERSION", module.DEFAULT_VERSION), + auth, + timeout=60, + verbose=False, + ) + self.tool_id = tool_id + + async def sessions(self) -> list[dict[str, str]]: + infos = await asyncio.to_thread( + self.module.list_all_sessions, self.client, self.tool_id, 100 + ) + return [ + { + "session_id": str(item.get("SessionId") or ""), + "user_session_id": str(item.get("UserSessionId") or ""), + "status": str(item.get("Status") or ""), + "created_at": str(item.get("CreatedAt") or ""), + } + for item in infos + ] + + async def for_managed_session(self, session_id: str) -> list[dict[str, str]]: + return [ + item + for item in await self.sessions() + if item["user_session_id"] == session_id + ] + + async def wait_present(self, session_id: str, *, timeout: float) -> dict[str, str]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + matches = await self.for_managed_session(session_id) + ready = next((item for item in matches if item["status"] == "Ready"), None) + if ready: + return ready + await asyncio.sleep(2) + raise TimeoutError(f"AgentKit sandbox did not become Ready for {session_id}") + + async def wait_absent(self, session_id: str, *, timeout: float) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not await self.for_managed_session(session_id): + return + await asyncio.sleep(3) + raise TimeoutError(f"AgentKit sandbox was not reclaimed for {session_id}") + + +async def create_session( + client: anthropic.AsyncAnthropic, *, agent_id: str, environment_id: str, title: str +) -> str: + session = await client.beta.sessions.create( + agent=agent_id, environment_id=environment_id, title=title + ) + return str(session.id) + + +async def run_turn( + client: anthropic.AsyncAnthropic, + session_id: str, + *, + marker: str, + prompt: str, + timeout: float, + ready_gate: asyncio.Event | None = None, + ready_counter: list[int] | None = None, + ready_total: int | None = None, +) -> TurnResult: + started = time.monotonic() + observations: list[ObservedEvent] = [] + event_types: list[str] = [] + tool_uses: list[str] = [] + tool_results: list[str] = [] + final_text_parts: list[str] = [] + first_delta: float | None = None + final_message: float | None = None + + stream = await client.beta.sessions.events.stream( + session_id, event_deltas=["agent.message"] + ) + if ready_gate is not None and ready_counter is not None and ready_total is not None: + ready_counter[0] += 1 + if ready_counter[0] == ready_total: + ready_gate.set() + await ready_gate.wait() + + await client.beta.sessions.events.send( + session_id, + events=[ + {"type": "user.message", "content": [{"type": "text", "text": prompt}]} + ], + ) + + async with stream: + async with asyncio.timeout(timeout): + async for event in stream: + elapsed = time.monotonic() - started + event_type = str(field_value(event, "type", "")) + event_types.append(event_type) + event_id = str(field_value(event, "id", "") or "") or None + link_id: str | None = None + text = "" + if event_type == "event_delta": + link_id = str(field_value(event, "event_id", "") or "") or None + delta = field_value(event, "delta") + text = content_text([field_value(delta, "content")]) + if text and first_delta is None: + first_delta = elapsed + elif event_type == "event_start": + preview = field_value(event, "event") + link_id = str(field_value(preview, "id", "") or "") or None + elif event_type == "agent.message": + text = content_text(field_value(event, "content")) + final_text_parts.append(text) + final_message = elapsed + elif event_type == "agent.tool_use": + call_id = str(field_value(event, "id", "") or "") + if call_id: + tool_uses.append(call_id) + elif event_type == "agent.tool_result": + link_id = str(field_value(event, "tool_use_id", "") or "") or None + if link_id: + tool_results.append(link_id) + text = content_text(field_value(event, "content")) + observations.append( + ObservedEvent(event_type, elapsed, event_id, link_id, text[:500]) + ) + if event_type in {"session.error", "session.status_terminated"}: + raise RuntimeError( + f"{session_id}: terminal failure event {event_type}" + ) + if event_type == "session.status_idle": + break + + result = TurnResult( + session_id=session_id, + marker=marker, + started_at=now_iso(), + elapsed_seconds=time.monotonic() - started, + first_delta_seconds=first_delta, + final_message_seconds=final_message, + event_types=event_types, + tool_use_ids=tool_uses, + tool_result_ids=tool_results, + final_text="\n".join(part for part in final_text_parts if part), + observations=observations, + ) + return result + + +async def smoke( + client: anthropic.AsyncAnthropic, + args: argparse.Namespace, + inspector: AgentKitInspector | None, +) -> dict[str, Any]: + session_id = await create_session( + client, + agent_id=args.agent_id, + environment_id=args.environment_id, + title="SDK streaming smoke", + ) + plain_marker = f"plain-{uuid.uuid4().hex[:12]}" + plain = await run_turn( + client, + session_id, + marker=plain_marker, + prompt=f"Do not use any tool. Reply with exactly: {plain_marker}", + timeout=args.turn_timeout, + ) + plain.assert_complete(require_stream=True, require_tool=False) + if inspector is not None and await inspector.for_managed_session(session_id): + raise AssertionError( + "plain user message unexpectedly started an AgentKit sandbox" + ) + + tool_marker = f"tool-{uuid.uuid4().hex[:12]}" + tool = await run_turn( + client, + session_id, + marker=tool_marker, + prompt=( + f"Use the bash tool exactly once to run: printf {tool_marker}. " + f"Then include exactly {tool_marker} in your final reply." + ), + timeout=args.turn_timeout, + ) + tool.assert_complete(require_stream=True, require_tool=True) + physical = ( + await inspector.wait_present(session_id, timeout=120) if inspector else None + ) + return { + "session_id": session_id, + "plain": asdict(plain), + "tool": asdict(tool), + "physical": physical, + } + + +async def lifecycle( + client: anthropic.AsyncAnthropic, + args: argparse.Namespace, + inspector: AgentKitInspector, +) -> dict[str, Any]: + session_id = await create_session( + client, + agent_id=args.agent_id, + environment_id=args.environment_id, + title="SDK sandbox lifecycle", + ) + first_marker = f"life-a-{uuid.uuid4().hex[:10]}" + first = await run_turn( + client, + session_id, + marker=first_marker, + prompt=f"Use bash exactly once to run: printf {first_marker}. Reply with {first_marker}.", + timeout=args.turn_timeout, + ) + first.assert_complete(require_stream=True, require_tool=True) + first_sandbox = await inspector.wait_present(session_id, timeout=120) + idle_at = time.monotonic() + await inspector.wait_absent(session_id, timeout=args.reclaim_timeout) + wait_remaining = args.restart_after - (time.monotonic() - idle_at) + if wait_remaining > 0: + await asyncio.sleep(wait_remaining) + + second_marker = f"life-b-{uuid.uuid4().hex[:10]}" + second = await run_turn( + client, + session_id, + marker=second_marker, + prompt=f"Use bash exactly once to run: printf {second_marker}. Reply with {second_marker}.", + timeout=args.turn_timeout, + ) + second.assert_complete(require_stream=True, require_tool=True) + second_sandbox = await inspector.wait_present(session_id, timeout=120) + if first_sandbox["session_id"] == second_sandbox["session_id"]: + raise AssertionError("AgentKit reused the reclaimed physical sandbox ID") + return { + "session_id": session_id, + "first": asdict(first), + "second": asdict(second), + "first_sandbox": first_sandbox, + "second_sandbox": second_sandbox, + "restart_after_seconds": args.restart_after, + } + + +async def interrupt( + client: anthropic.AsyncAnthropic, + args: argparse.Namespace, + inspector: AgentKitInspector, +) -> dict[str, Any]: + session_id = await create_session( + client, + agent_id=args.agent_id, + environment_id=args.environment_id, + title="SDK interrupt", + ) + marker = f"interrupt-{uuid.uuid4().hex[:10]}" + started = time.monotonic() + stream = await client.beta.sessions.events.stream( + session_id, event_deltas=["agent.message"] + ) + await client.beta.sessions.events.send( + session_id, + events=[ + { + "type": "user.message", + "content": [ + { + "type": "text", + "text": f"Use bash to run: sleep 120; printf {marker}. Then reply with the output.", + } + ], + } + ], + ) + seen: list[str] = [] + interrupted_at: float | None = None + async with stream: + async with asyncio.timeout(args.turn_timeout): + async for event in stream: + event_type = str(field_value(event, "type", "")) + seen.append(event_type) + if event_type == "agent.tool_use" and interrupted_at is None: + await inspector.wait_present(session_id, timeout=120) + await asyncio.sleep(2) + await client.beta.sessions.events.send( + session_id, events=[{"type": "user.interrupt"}] + ) + interrupted_at = time.monotonic() + if event_type == "agent.tool_result": + text = content_text(field_value(event, "content")) + if marker in text: + raise AssertionError( + "interrupted shell command completed with its marker" + ) + if event_type == "session.status_idle" and interrupted_at is not None: + break + if interrupted_at is None: + raise AssertionError("interrupt test never observed agent.tool_use") + await inspector.wait_absent(session_id, timeout=60) + return { + "session_id": session_id, + "event_types": seen, + "elapsed_seconds": time.monotonic() - started, + "interrupt_to_idle_seconds": time.monotonic() - interrupted_at, + } + + +async def concurrency( + client: anthropic.AsyncAnthropic, + args: argparse.Namespace, + inspector: AgentKitInspector, +) -> dict[str, Any]: + count = args.concurrency + sessions = [] + for index in range(count): + sessions.append( + await create_session( + client, + agent_id=args.agent_id, + environment_id=args.environment_id, + title=f"SDK concurrency {index:02d} {uuid.uuid4().hex[:8]}", + ) + ) + if len(set(sessions)) != count: + raise AssertionError("concurrent Session creation returned duplicate IDs") + markers = [f"c{index:02d}-{uuid.uuid4().hex[:12]}" for index in range(count)] + gate = asyncio.Event() + ready = [0] + + async def one(index: int) -> TurnResult: + marker = markers[index] + result = await run_turn( + client, + sessions[index], + marker=marker, + prompt=( + "Use the bash tool exactly once to run: " + f"started=$(date +%s%3N); sleep {args.concurrency_hold_seconds}; " + "ended=$(date +%s%3N); " + f'printf \'{marker} %s %s\' "$started" "$ended". ' + f"Reply with exactly {marker}." + ), + timeout=args.turn_timeout, + ready_gate=gate, + ready_counter=ready, + ready_total=count, + ) + result.assert_complete(require_stream=True, require_tool=True) + return result + + started = time.monotonic() + turns = asyncio.gather(*(one(index) for index in range(count))) + peak_physical: dict[str, dict[str, str]] = {} + while not turns.done(): + current = { + item["user_session_id"]: item + for item in await inspector.sessions() + if item["user_session_id"] in sessions + } + if len(current) > len(peak_physical): + peak_physical = current + await asyncio.sleep(1) + results = await turns + for index, result in enumerate(results): + all_text = ( + "\n".join(item.text for item in result.observations) + + "\n" + + result.final_text + ) + leaked = [ + marker + for pos, marker in enumerate(markers) + if pos != index and marker in all_text + ] + if leaked: + raise AssertionError(f"{result.session_id}: cross-session markers {leaked}") + physical = { + item["user_session_id"]: item + for item in await inspector.sessions() + if item["user_session_id"] in sessions + } + if len(physical) > len(peak_physical): + peak_physical = physical + if len(peak_physical) != count: + missing = sorted(set(sessions) - set(peak_physical)) + raise AssertionError( + f"expected {count} simultaneously retained physical sandboxes, " + f"observed {len(peak_physical)}; missing={missing}" + ) + physical_ids = [item["session_id"] for item in peak_physical.values()] + if len(set(physical_ids)) != count: + raise AssertionError("concurrent logical Sessions shared a physical sandbox") + + intervals: list[tuple[int, int]] = [] + for index, result in enumerate(results): + tool_text = "\n".join( + item.text + for item in result.observations + if item.type == "agent.tool_result" + ) + match = re.search( + rf"{re.escape(markers[index])}\s+(\d{{10,}})\s+(\d{{10,}})", + tool_text, + ) + if match is None: + raise AssertionError( + f"{result.session_id}: missing tool execution timestamps" + ) + interval = (int(match.group(1)), int(match.group(2))) + if interval[1] <= interval[0]: + raise AssertionError( + f"{result.session_id}: invalid tool execution interval {interval}" + ) + intervals.append(interval) + edges = sorted( + [(start, 1) for start, _ in intervals] + [(end, -1) for _, end in intervals], + key=lambda item: (item[0], -item[1]), + ) + active = 0 + max_tool_overlap = 0 + for _, delta in edges: + active += delta + max_tool_overlap = max(max_tool_overlap, active) + if max_tool_overlap != count: + raise AssertionError( + f"expected {count} overlapping bash executions, observed {max_tool_overlap}" + ) + return { + "concurrency": count, + "elapsed_seconds": time.monotonic() - started, + "sessions": [result.session_id for result in results], + "physical_sessions": peak_physical, + "peak_physical_sessions": len(peak_physical), + "tool_execution_intervals_ms": intervals, + "max_overlapping_tool_executions": max_tool_overlap, + "results": [asdict(result) for result in results], + "first_delta_seconds": [result.first_delta_seconds for result in results], + } + + +def parser() -> argparse.ArgumentParser: + value = argparse.ArgumentParser() + value.add_argument( + "--mode", + choices=( + "conversation", + "smoke", + "interrupt", + "lifecycle", + "concurrency", + "all", + ), + default="all", + ) + value.add_argument( + "--base-url", + default=os.getenv("MANAGED_AGENTS_GATEWAY") or os.getenv("ANTHROPIC_BASE_URL"), + required=not ( + os.getenv("MANAGED_AGENTS_GATEWAY") or os.getenv("ANTHROPIC_BASE_URL") + ), + ) + value.add_argument( + "--auth-token", + default=os.getenv("MANAGED_AGENTS_GATEWAY_TOKEN") + or os.getenv("ANTHROPIC_ENVIRONMENT_KEY"), + required=not ( + os.getenv("MANAGED_AGENTS_GATEWAY_TOKEN") + or os.getenv("ANTHROPIC_ENVIRONMENT_KEY") + ), + ) + value.add_argument( + "--agent-id", + default=os.getenv("MANAGED_AGENTS_TEST_AGENT_ID"), + required=os.getenv("MANAGED_AGENTS_TEST_AGENT_ID") is None, + ) + value.add_argument( + "--environment-id", + default=os.getenv("ANTHROPIC_ENVIRONMENT_ID"), + required=os.getenv("ANTHROPIC_ENVIRONMENT_ID") is None, + ) + value.add_argument("--concurrency", type=int, default=40) + value.add_argument("--concurrency-hold-seconds", type=int, default=90) + value.add_argument("--turn-timeout", type=float, default=600) + value.add_argument("--reclaim-timeout", type=float, default=300) + value.add_argument("--restart-after", type=float, default=300) + value.add_argument("--tool-id", default=DEFAULT_TOOL_ID) + value.add_argument( + "--agentkit-helper", + default=os.getenv("AGENTKIT_TOOL_DEPLOY_SCRIPT", DEFAULT_AGENTKIT_HELPER), + ) + value.add_argument("--output", type=Path) + return value + + +async def async_main(args: argparse.Namespace) -> dict[str, Any]: + if args.concurrency < 1: + raise ValueError("concurrency must be positive") + if args.concurrency_hold_seconds < 1: + raise ValueError("concurrency hold seconds must be positive") + modes = ( + ["smoke", "interrupt", "lifecycle", "concurrency"] + if args.mode == "all" + else [args.mode] + ) + inspector = ( + AgentKitInspector(args.agentkit_helper, args.tool_id) + if any( + mode in {"smoke", "interrupt", "lifecycle", "concurrency"} for mode in modes + ) + else None + ) + output: dict[str, Any] = { + "started_at": now_iso(), + "base_url": args.base_url, + "agent_id": args.agent_id, + "environment_id": args.environment_id, + "modes": modes, + } + async with anthropic.AsyncAnthropic( + base_url=args.base_url.rstrip("/"), + auth_token=args.auth_token, + timeout=args.turn_timeout, + ) as client: + if "conversation" in modes: + session_id = await create_session( + client, + agent_id=args.agent_id, + environment_id=args.environment_id, + title="SDK conversation verification", + ) + print(f"Created test session {session_id}", file=sys.stderr, flush=True) + try: + marker = f"conversation-{uuid.uuid4().hex[:12]}" + result = await run_turn( + client, + session_id, + marker=marker, + prompt=f"Do not use any tool. Reply with exactly: {marker}", + timeout=args.turn_timeout, + ) + result.assert_complete(require_stream=True, require_tool=False) + output["conversation"] = asdict(result) + finally: + # Idle may precede the Work acknowledgement. Interrupt only + # this test Session so queued/running Work cannot block cleanup. + await client.beta.sessions.events.send( + session_id, events=[{"type": "user.interrupt"}] + ) + await client.beta.sessions.events.send( + session_id, events=[{"type": "session.status_terminated"}] + ) + await client.beta.sessions.archive(session_id) + output["session_archived"] = True + if "smoke" in modes: + output["smoke"] = await smoke(client, args, inspector) + if "interrupt" in modes: + assert inspector is not None + output["interrupt"] = await interrupt(client, args, inspector) + if "lifecycle" in modes: + assert inspector is not None + output["lifecycle"] = await lifecycle(client, args, inspector) + if "concurrency" in modes: + assert inspector is not None + output["concurrency"] = await concurrency(client, args, inspector) + output["completed_at"] = now_iso() + output["status"] = "passed" + return output + + +def main() -> int: + # Match this example's configuration without overriding exported values. + load_dotenv(Path(__file__).with_name(".env"), override=False) + args = parser().parse_args() + result = asyncio.run(async_main(args)) + encoded = json.dumps(result, ensure_ascii=False, indent=2) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(encoded + "\n") + print(encoded) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/16_self_host_sandbox/ark_session_e2e.py b/examples/16_self_host_sandbox/ark_session_e2e.py new file mode 100644 index 000000000..0c596c43c --- /dev/null +++ b/examples/16_self_host_sandbox/ark_session_e2e.py @@ -0,0 +1,181 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Run against the deployed worker; credentials stay inside its environment. + +Usage: kubectl exec -i deploy/managed-agents-agent-loop-agentkit -- python - ACTION [SESSION_ID] +with this file on stdin. Actions: probe, start, poll, next, tool, archive. +""" + +import asyncio +import json +import os +import sys +import urllib.request + + +def request(method, path, body=None): + headers = { + "x-api-key": os.environ["TASK_SERVER_API_KEY"], + "x-account-id": os.environ["TASK_SERVER_ACCOUNT_ID"], + "content-type": "application/json", + } + req = urllib.request.Request( + os.environ["ANTHROPIC_BASE_URL"].rstrip("/") + path, + method=method, + headers=headers, + data=json.dumps(body).encode() if body is not None else None, + ) + with urllib.request.urlopen(req, timeout=30) as response: + data = response.read() + return json.loads(data) if data else {} + + +async def probe(): + from volcenginesdkarkruntime import AsyncArk + + async with AsyncArk(api_key=os.environ["MODEL_AGENT_API_KEY"]) as client: + first = await client.responses.create( + model=os.environ["MODEL_AGENT_NAME"], + input="Remember the number 731946. Reply OK.", + store=True, + ) + second = await client.responses.create( + model=os.environ["MODEL_AGENT_NAME"], + input="Which number did I ask you to remember?", + previous_response_id=first.id, + tools=[ + { + "type": "function", + "name": "bash", + "description": "Execute a bash command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + } + ], + store=True, + ) + text = "".join( + part.text + for output in second.output + if output.type == "message" + for part in output.content + if part.type == "output_text" + ) + assert "731946" in text, text + print( + json.dumps( + { + "responses_api": True, + "cloud_context_recall": True, + "status": second.status, + } + ) + ) + + +required = {"ANTHROPIC_BASE_URL", "TASK_SERVER_API_KEY", "TASK_SERVER_ACCOUNT_ID"} +if len(sys.argv) < 2: + raise SystemExit("Usage: ark_session_e2e.py ACTION [SESSION_ID]") +action = sys.argv[1] +if action == "start": + required.update({"SANDBOX_AGENT_ID", "E2E_ENVIRONMENT_ID"}) +elif action == "probe": + required = {"MODEL_AGENT_API_KEY", "MODEL_AGENT_NAME"} +missing = sorted(name for name in required if not os.getenv(name)) +if missing: + raise SystemExit("Missing required configuration: " + ", ".join(missing)) +if action == "probe": + asyncio.run(probe()) +elif action == "start": + agent_id = os.environ["SANDBOX_AGENT_ID"] + result = request( + "POST", + "/v1/sessions", + { + "agent": agent_id, + "environment_id": os.environ["E2E_ENVIRONMENT_ID"], + "title": "Ark Responses worker restart E2E", + "initial_events": [ + { + "type": "user.message", + "content": [ + { + "type": "text", + "text": "Remember the secret test number 731946. Reply only OK. Do not run any tools.", + } + ], + } + ], + }, + ) + print(json.dumps({"session_id": result["id"]})) +else: + session_id = sys.argv[2] + if action == "poll": + session = request("GET", "/v1/sessions/" + session_id) + events = request( + "GET", "/v1/sessions/" + session_id + "/events?limit=1000&order=asc" + )["data"] + uses = {e["id"] for e in events if e["type"] == "agent.tool_use"} + results = [e for e in events if e["type"] == "agent.tool_result"] + print( + json.dumps( + { + "status": session.get("status"), + "tool_uses": len(uses), + "matching_tool_results": sum( + e.get("tool_use_id") in uses for e in results + ), + "events": [ + { + "type": e["type"], + "content": e.get("content"), + "error": e.get("error"), + "has_response_id": bool( + e.get("metadata", {}).get("ark_response_id") + ), + } + for e in events + if e["type"] + in { + "agent.message", + "session.error", + "session.status_idle", + "agent.tool_use", + "agent.tool_result", + "span.model_request_end", + } + ], + } + ) + ) + elif action in {"next", "tool"}: + prompt = "Which secret test number did I ask you to remember? Reply only that number; do not run tools." + if action == "tool": + prompt = "Use bash to calculate 123 * 456 with python, then report the result. You must actually call the bash tool." + request( + "POST", + "/v1/sessions/" + session_id + "/events", + { + "events": [ + { + "type": "user.message", + "content": [{"type": "text", "text": prompt}], + } + ] + }, + ) + print(json.dumps({"sent": action})) + elif action == "archive": + request( + "POST", + "/v1/sessions/" + session_id + "/events", + { + "events": [{"type": "session.status_terminated"}], + }, + ) + request("POST", "/v1/sessions/" + session_id + "/archive", {}) + print(json.dumps({"terminated_and_archived": True})) diff --git a/examples/16_self_host_sandbox/build-managed-agents-images.sh b/examples/16_self_host_sandbox/build-managed-agents-images.sh new file mode 100644 index 000000000..900d59119 --- /dev/null +++ b/examples/16_self_host_sandbox/build-managed-agents-images.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +exec bash "$SCRIPT_DIR/../../docker/managed-agents/build.sh" "$@" diff --git a/examples/16_self_host_sandbox/distributed_agent_loop_test.py b/examples/16_self_host_sandbox/distributed_agent_loop_test.py new file mode 100644 index 000000000..e70d16215 --- /dev/null +++ b/examples/16_self_host_sandbox/distributed_agent_loop_test.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Three-worker canonical Session ledger and Ark continuation smoke test. + +Requires an isolated empty environment and an existing test Agent configured +with an Ark model, Session credential references and native bash permissions.""" + +from __future__ import annotations + +import os +import re +import signal +import subprocess +import sys +import time +import uuid +from pathlib import Path +from typing import Any + +import anthropic + + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parents[1] +STATE_DIR = SCRIPT_DIR / ".local-managed-agents" + + +def _headers() -> dict[str, str]: + headers = {"anthropic-beta": "managed-agents-2026-04-01"} + if account_id := os.getenv("X_TOP_ACCOUNT_ID"): + headers["X-Top-Account-Id"] = account_id + return headers + + +def _client() -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=os.environ["ANTHROPIC_BASE_URL"], + auth_token=os.environ["ANTHROPIC_ENVIRONMENT_KEY"], + default_headers=_headers(), + timeout=60, + ) + + +def _events(client: anthropic.Anthropic, session_id: str) -> list[Any]: + return list(client.beta.sessions.events.list(session_id, limit=1000, order="asc")) + + +def _text(event: Any) -> str: + return "\n".join( + str(getattr(block, "text", "") or "") + for block in (getattr(event, "content", None) or []) + if getattr(block, "type", None) == "text" + ) + + +def _wait_for_turn( + client: anthropic.Anthropic, + session_id: str, + user_event_id: str, + timeout: float = 180, +) -> tuple[list[Any], str]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + events = _events(client, session_id) + start = next( + ( + index + for index, event in enumerate(events) + if getattr(event, "id", None) == user_event_id + ), + None, + ) + if start is not None: + turn = events[start:] + if any(getattr(event, "type", None) == "session.error" for event in turn): + errors = [ + str(getattr(event, "error", "")) + for event in turn + if event.type == "session.error" + ] + raise RuntimeError(f"Agent Loop failed: {errors}") + if any( + getattr(event, "type", None) == "session.status_idle" for event in turn + ): + answers = [ + _text(event) for event in turn if event.type == "agent.message" + ] + return turn, "\n".join(answer for answer in answers if answer) + time.sleep(0.5) + raise TimeoutError(f"Timed out waiting for session {session_id}") + + +def _send(client: anthropic.Anthropic, session_id: str, message: str) -> str: + response = client.beta.sessions.events.send( + session_id, + events=[ + {"type": "user.message", "content": [{"type": "text", "text": message}]} + ], + ) + if not response.data: + raise RuntimeError("Task Server did not return the persisted user.message") + return response.data[0].id + + +def _start_worker(worker_id: str) -> tuple[subprocess.Popen[str], Path]: + STATE_DIR.mkdir(parents=True, exist_ok=True) + log_path = STATE_DIR / f"{worker_id}.log" + log = log_path.open("w") + env = dict(os.environ) + env["PYTHONUNBUFFERED"] = "1" + worker_workdir = STATE_DIR / "workspaces" / worker_id + worker_workdir.mkdir(parents=True, exist_ok=True) + env["MANAGED_AGENT_WORKDIR"] = str(worker_workdir) + env["MANAGED_AGENT_READY_FILE"] = str(STATE_DIR / f"{worker_id}.ready") + env["PYTHONPATH"] = os.pathsep.join( + [str(REPO_ROOT), str(SCRIPT_DIR), env.get("PYTHONPATH", "")] + ) + process = subprocess.Popen( + [ + sys.executable, + str(SCRIPT_DIR / "main.py"), + "--managed-agent-worker", + "--worker-id", + worker_id, + "--max-work-items", + "1", + "--managed-agent-readiness-probe", + ], + cwd=SCRIPT_DIR, + env=env, + stdout=log, + stderr=subprocess.STDOUT, + text=True, + ) + log.close() + return process, log_path + + +def _wait_for_workers_ready( + workers: dict[str, subprocess.Popen[str]], + logs: dict[str, Path], + timeout: float = 30, +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + failed = { + name: process.returncode + for name, process in workers.items() + if (process.poll() is not None and process.returncode != 0) + } + if failed: + raise RuntimeError(f"Workers exited before becoming ready: {failed}") + ready = { + name + for name, path in logs.items() + if path.exists() + and f"MANAGED_AGENT_WORKER_READY worker_id={name}" in path.read_text() + } + if ready == set(workers): + return + time.sleep(0.2) + raise TimeoutError(f"Workers did not all become ready: {sorted(ready)}") + + +def _wait_for_new_claim( + logs: dict[str, Path], + session_id: str, + seen_work_ids: set[str], + timeout: float = 10, +) -> tuple[str, str]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + claims = [ + (name, work_id) + for name, path in logs.items() + for work_id in _claim_ids(path, session_id) + if work_id not in seen_work_ids + ] + if len(claims) == 1: + return claims[0] + if len(claims) > 1: + raise RuntimeError(f"More than one worker claimed the new turn: {claims}") + time.sleep(0.2) + raise TimeoutError(f"No new worker claim found for {session_id}") + + +def _wait_for_work_stopped( + client: anthropic.Anthropic, environment_id: str, work_id: str, timeout: float = 15 +) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + work = client.beta.environments.work.retrieve( + work_id, environment_id=environment_id + ) + if work.state == "stopped": + if not work.acknowledged_at or not work.stopped_at: + raise AssertionError(f"Stopped work lacks lifecycle timestamps: {work}") + return + time.sleep(0.2) + raise TimeoutError(f"Work item {work_id} did not reach stopped") + + +def _claim_ids(path: Path, session_id: str) -> list[str]: + if not path.exists(): + return [] + return re.findall( + rf"MANAGED_AGENT_WORK_CLAIM .*?work_id=(\S+) session_id={re.escape(session_id)}(?:\s|$)", + path.read_text(), + ) + + +def _wait_for_initial_user_message( + client: anthropic.Anthropic, + session_id: str, + expected_text: str, + timeout: float = 10, +) -> str: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for event in _events(client, session_id): + if event.type == "user.message" and _text(event) == expected_text: + return event.id + time.sleep(0.2) + raise TimeoutError(f"Initial user.message was not persisted for {session_id}") + + +def _wait_for_worker_exit( + worker_id: str, process: subprocess.Popen[str], log_path: Path, timeout: float = 15 +) -> None: + try: + return_code = process.wait(timeout=timeout) + except subprocess.TimeoutExpired as error: + raise TimeoutError(f"{worker_id} did not exit after one work item") from error + if return_code != 0: + raise RuntimeError(f"{worker_id} exited with {return_code}; inspect {log_path}") + + +def _stop(process: subprocess.Popen[str]) -> None: + if process.poll() is not None: + return + process.send_signal(signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def main() -> None: + required = ( + "ANTHROPIC_BASE_URL", + "ANTHROPIC_ENVIRONMENT_ID", + "ANTHROPIC_ENVIRONMENT_KEY", + "SANDBOX_AGENT_ID", + ) + missing = [name for name in required if not os.getenv(name)] + if missing: + raise RuntimeError( + f"Missing required environment variables: {', '.join(missing)}" + ) + + first_codeword = f"cedar-{uuid.uuid4().hex[:8]}" + second_codeword = f"maple-{uuid.uuid4().hex[:8]}" + # The caller supplies an isolated environment with an empty Work queue. + environment_id = os.environ["ANTHROPIC_ENVIRONMENT_ID"] + os.environ["ANTHROPIC_ENVIRONMENT_ID"] = environment_id + workers: dict[str, subprocess.Popen[str]] = {} + logs: dict[str, Path] = {} + with _client() as client: + session = None + try: + for name in ("worker-a", "worker-b", "worker-c"): + workers[name], logs[name] = _start_worker(name) + _wait_for_workers_ready(workers, logs) + + agent = client.beta.agents.retrieve(os.environ["SANDBOX_AGENT_ID"]) + first_prompt = ( + f"Use the bash tool to run exactly: printf {first_codeword}. " + "Remember the codeword and include the exact tool output in your reply." + ) + session = client.beta.sessions.create( + agent=agent.id, + environment_id=environment_id, + title="VeADK three-worker stateless Agent Loop test", + initial_events=[ + { + "type": "user.message", + "content": [{"type": "text", "text": first_prompt}], + } + ], + ) + first_id = _wait_for_initial_user_message(client, session.id, first_prompt) + first_events, first_answer = _wait_for_turn(client, session.id, first_id) + seen_work_ids: set[str] = set() + first_owner, first_work_id = _wait_for_new_claim( + logs, session.id, seen_work_ids + ) + seen_work_ids.add(first_work_id) + _wait_for_work_stopped(client, environment_id, first_work_id) + _wait_for_worker_exit(first_owner, workers[first_owner], logs[first_owner]) + if first_codeword not in first_answer: + raise AssertionError( + f"First answer did not contain {first_codeword!r}: {first_answer!r}" + ) + + second_id = _send( + client, + session.id, + f"Without using tools, recall the first codeword and remember a second codeword " + f"{second_codeword}. Reply with both exact codewords.", + ) + second_events, second_answer = _wait_for_turn(client, session.id, second_id) + second_owner, second_work_id = _wait_for_new_claim( + logs, session.id, seen_work_ids + ) + seen_work_ids.add(second_work_id) + _wait_for_work_stopped(client, environment_id, second_work_id) + _wait_for_worker_exit( + second_owner, workers[second_owner], logs[second_owner] + ) + if ( + first_codeword not in second_answer + or second_codeword not in second_answer + ): + raise AssertionError( + f"Second-turn canonical ledger recovery failed: {second_answer!r}" + ) + + third_id = _send( + client, + session.id, + "Without using tools, recall both codewords from the previous two turns. " + "Reply with both exact codewords.", + ) + third_events, third_answer = _wait_for_turn(client, session.id, third_id) + third_owner, third_work_id = _wait_for_new_claim( + logs, session.id, seen_work_ids + ) + seen_work_ids.add(third_work_id) + _wait_for_work_stopped(client, environment_id, third_work_id) + _wait_for_worker_exit(third_owner, workers[third_owner], logs[third_owner]) + if ( + first_codeword not in third_answer + or second_codeword not in third_answer + ): + raise AssertionError( + f"Third-turn canonical ledger recovery failed: {third_answer!r}" + ) + + for number, turn in enumerate( + (first_events, second_events, third_events), start=1 + ): + types = [event.type for event in turn] + for required_type in ( + "user.message", + "session.status_running", + "span.model_request_start", + "span.model_request_end", + "session.status_idle", + ): + if types.count(required_type) != 1: + raise AssertionError( + f"turn {number} expected one {required_type}, got {types}" + ) + if types.count("agent.message") < 1: + raise AssertionError( + f"turn {number} emitted no agent.message: {types}" + ) + first_tool_uses = [ + event for event in first_events if event.type == "agent.tool_use" + ] + if len(first_tool_uses) != 1 or first_tool_uses[0].name != "bash": + raise AssertionError( + f"Expected one bash tool use in turn 1: {first_tool_uses}" + ) + matching_results = [ + event + for event in first_events + if event.type == "agent.tool_result" + and event.tool_use_id == first_tool_uses[0].id + ] + if len(matching_results) != 1 or first_codeword not in str( + matching_results[0].content + ): + raise AssertionError( + f"Missing matching bash result for {first_codeword}: {matching_results}" + ) + later_tool_events = [ + event + for turn in (second_events, third_events) + for event in turn + if event.type in {"agent.tool_use", "agent.tool_result"} + ] + if later_tool_events: + raise AssertionError( + f"Later turns unexpectedly used tools: {later_tool_events}" + ) + + owners = (first_owner, second_owner, third_owner) + if set(owners) != set(workers): + raise AssertionError(f"Each worker must win exactly one turn: {owners}") + claims = {name: _claim_ids(path, session.id) for name, path in logs.items()} + if any(len(ids) != 1 for ids in claims.values()): + raise AssertionError( + f"Expected exactly one claim per worker for this session: {claims}" + ) + work_ids = (first_work_id, second_work_id, third_work_id) + if len(set(work_ids)) != 3: + raise AssertionError( + f"Each turn must have a distinct work item: {work_ids}" + ) + + history = _events(client, session.id) + successful_spans = [ + event + for event in history + if event.type == "span.model_request_end" and not event.is_error + ] + response_ids = [ + getattr(event, "metadata", {}).get("ark_response_id") + for event in successful_spans + ] + if len(successful_spans) != 3 or not all(response_ids): + raise AssertionError("Expected three durable model response references") + print(f"agent_id={agent.id}") + print(f"session_id={session.id}") + print(f"first_worker={first_owner}") + print(f"second_worker={second_owner}") + print(f"third_worker={third_owner}") + print(f"first_work_id={first_work_id}") + print(f"second_work_id={second_work_id}") + print(f"third_work_id={third_work_id}") + print(f"canonical_events={len(history)}") + print(f"model_response_references={len(response_ids)}") + print("Distributed stateless VeADK Agent Loop test passed.") + finally: + for process in workers.values(): + _stop(process) + try: + if session is not None: + client.beta.sessions.archive(session.id) + except Exception as error: # noqa: BLE001 -- cleanup is best effort + print( + f"cleanup warning: Session archive failed ({type(error).__name__})" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/16_self_host_sandbox/event_debug.py b/examples/16_self_host_sandbox/event_debug.py new file mode 100644 index 000000000..e4a826e3f --- /dev/null +++ b/examples/16_self_host_sandbox/event_debug.py @@ -0,0 +1,6 @@ +"""Compatibility import for veadk.runtime.managed_agents.events.""" + +import sys +from importlib import import_module + +sys.modules[__name__] = import_module("veadk.runtime.managed_agents.events") diff --git a/examples/16_self_host_sandbox/k8s_agent_loop_e2e.py b/examples/16_self_host_sandbox/k8s_agent_loop_e2e.py new file mode 100644 index 000000000..6e9cfc59d --- /dev/null +++ b/examples/16_self_host_sandbox/k8s_agent_loop_e2e.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Drive three sequential turns through the deployed Managed Agents gateway.""" + +from __future__ import annotations + +import json +import os +import time +import urllib.error +import urllib.request +import uuid + +BASE_URL = os.environ.get( + "MANAGED_AGENTS_GATEWAY_URL", "http://127.0.0.1:18080" +).rstrip("/") +CONTROL_URL = os.environ.get( + "MANAGED_AGENTS_CONTROL_URL", "http://127.0.0.1:8000" +).rstrip("/") +CONTROL_TOKEN = os.environ.get("ANTHROPIC_ENVIRONMENT_KEY", "") +MODEL = os.environ["MANAGED_AGENTS_MODEL_CARD_ID"] +ENVIRONMENT_ID = os.environ["MANAGED_AGENTS_ENVIRONMENT_ID"] + + +def request(method: str, path: str, body: dict | None = None) -> dict: + data = None if body is None else json.dumps(body).encode() + headers = {"Accept": "application/json"} + if data is not None: + headers["Content-Type"] = "application/json" + req = urllib.request.Request( + BASE_URL + path, data=data, headers=headers, method=method + ) + try: + with urllib.request.urlopen(req, timeout=60) as response: + raw = response.read() + except urllib.error.HTTPError as error: + raise RuntimeError( + f"{method} {path}: HTTP {error.code}: {error.read()[:500]!r}" + ) from error + return json.loads(raw) if raw else {} + + +def control_request(path: str) -> dict: + headers = {"Accept": "application/json"} + if CONTROL_TOKEN: + headers["Authorization"] = f"Bearer {CONTROL_TOKEN}" + req = urllib.request.Request(CONTROL_URL + path, headers=headers) + with urllib.request.urlopen(req, timeout=60) as response: + return json.loads(response.read()) + + +def upload_skill_files(api_path: str, name: str, files: dict[str, bytes]) -> dict: + boundary = f"----managed-agents-{uuid.uuid4().hex}" + body = ( + f"--{boundary}\r\n" + f'Content-Disposition: form-data; name="display_name"\r\n\r\n{name}\r\n' + ).encode() + for path, content in files.items(): + body += ( + f"--{boundary}\r\n" + f'Content-Disposition: form-data; name="files[]"; filename="{name}/{path}"\r\n' + "Content-Type: application/octet-stream\r\n\r\n" + ).encode() + body += content + b"\r\n" + body += f"--{boundary}--\r\n".encode() + req = urllib.request.Request( + BASE_URL + api_path, + data=body, + headers={ + "Accept": "application/json", + "Content-Type": f"multipart/form-data; boundary={boundary}", + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=300) as response: + return json.loads(response.read()) + + +def upload_skill(name: str, files: dict[str, bytes]) -> dict: + return upload_skill_files("/v1/skills", name, files) + + +def upload_skill_version(skill_id: str, name: str, files: dict[str, bytes]) -> dict: + return upload_skill_files(f"/v1/skills/{skill_id}/versions", name, files) + + +def wait_turn( + session_id: str, event_id: str, timeout: float = 240 +) -> tuple[list[dict], str]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + events = request( + "GET", f"/v1/sessions/{session_id}/events?limit=1000&order=asc" + ).get("data", []) + start = next( + ( + index + for index, event in enumerate(events) + if event.get("id") == event_id + ), + None, + ) + if start is not None: + turn = events[start:] + errors = [ + event.get("error") + for event in turn + if event.get("type") == "session.error" + ] + if errors: + raise RuntimeError(f"Agent Loop failed: {errors}") + if any(event.get("type") == "session.status_idle" for event in turn): + text = "\n".join( + str(block.get("text", "")) + for event in turn + if event.get("type") == "agent.message" + for block in (event.get("content") or []) + if block.get("type") == "text" + ) + return turn, text + time.sleep(1) + raise TimeoutError(f"timed out waiting for {event_id}") + + +def send(session_id: str, text: str) -> str: + response = request( + "POST", + f"/v1/sessions/{session_id}/events", + { + "events": [ + {"type": "user.message", "content": [{"type": "text", "text": text}]} + ] + }, + ) + data = response.get("data") or [] + if data and data[0].get("id"): + return str(data[0]["id"]) + # Some Task Server versions return the accepted input shape without the + # generated event ID. Resolve the ID from the durable event ledger instead + # of retrying a write whose outcome is already known. + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + events = request( + "GET", f"/v1/sessions/{session_id}/events?limit=1000&order=asc" + ).get("data", []) + matches = [ + event + for event in events + if event.get("type") == "user.message" + and any( + block.get("type") == "text" and block.get("text") == text + for block in (event.get("content") or []) + ) + ] + if matches and matches[-1].get("id"): + return str(matches[-1]["id"]) + time.sleep(0.5) + raise TimeoutError("accepted user.message was not found in the event ledger") + + +def main() -> None: + first = f"cedar-{uuid.uuid4().hex[:8]}" + second = f"maple-{uuid.uuid4().hex[:8]}" + skill_name = f"managed-recovery-{uuid.uuid4().hex[:8]}" + marker = f"skill-{uuid.uuid4().hex[:8]}" + newer_marker = f"newer-skill-{uuid.uuid4().hex[:8]}" + skill = upload_skill( + skill_name, + { + "SKILL.md": ( + f"---\nname: {skill_name}\ndescription: Verify Session Skill recovery\n---\n" + f"When asked for the skill marker, use bash to run: " + f"cat skills/{skill_name}/references/marker.txt\n" + ).encode(), + "references/marker.txt": (marker + "\n").encode(), + }, + ) + agent = None + session = None + try: + agent = request( + "POST", + "/api/agents", + { + "name": f"k8s-e2e-{uuid.uuid4().hex[:8]}", + "model": {"id": MODEL}, + "system": "Remember codewords across turns and follow attached Skills.", + "skills": [ + {"type": "custom", "skill_id": skill["id"], "version": "latest"} + ], + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": { + "enabled": False, + "permission_policy": {"type": "always_allow"}, + }, + "configs": [ + { + "type": "bash", + "name": "bash", + "enabled": True, + "permission_policy": {"type": "always_allow"}, + } + ], + } + ], + }, + ) + session = request( + "POST", + "/api/sessions", + { + "agent": agent["id"], + "environment_id": ENVIRONMENT_ID, + "title": "K8s three-worker Skills recovery E2E", + }, + ) + pinned = session["agent"]["skills"][0]["version"] + if pinned in {"", "latest", None}: + raise AssertionError(f"Session Skill version was not pinned: {pinned!r}") + newer = upload_skill_version( + skill["id"], + skill_name, + { + "SKILL.md": ( + f"---\nname: {skill_name}\ndescription: Newer version sentinel\n---\n" + f"When asked for the skill marker, use bash to run: " + f"cat skills/{skill_name}/references/marker.txt\n" + ).encode(), + "references/marker.txt": (newer_marker + "\n").encode(), + }, + ) + if newer["id"] == pinned: + raise AssertionError( + "creating a Skill version did not advance its version ID" + ) + prompts = [ + f"Use the attached skill to read its marker. Also remember {first}. Reply with both exact values.", + f"Use the attached skill again. Recall {first}, remember {second}, and reply with all three values.", + "Use the attached skill once more, then reply with its marker and both prior codewords.", + ] + answers = [] + event_counts = [] + for prompt in prompts: + event_id = send(session["id"], prompt) + events, answer = wait_turn(session["id"], event_id) + answers.append(answer) + event_counts.append(len(events)) + if first not in answers[0] or marker not in answers[0]: + raise AssertionError(f"first Worker did not load the Skill: {answers[0]!r}") + if ( + first not in answers[1] + or second not in answers[1] + or marker not in answers[1] + ): + raise AssertionError(f"conversation/Skill recovery failed: {answers[1]!r}") + if any(value not in answers[2] for value in (first, second, marker)): + raise AssertionError( + f"third Worker did not restore all state: {answers[2]!r}" + ) + if any(newer_marker in answer for answer in answers): + raise AssertionError( + "Worker loaded a newer Skill instead of the Session-pinned version" + ) + all_events = request( + "GET", f"/v1/sessions/{session['id']}/events?limit=1000&order=asc" + )["data"] + work = [ + item + for item in control_request( + f"/v1/environments/{ENVIRONMENT_ID}/work?limit=1000" + )["data"] + if item["data"] == {"type": "session", "id": session["id"]} + ] + worker_ids = { + item.get("metadata", {}).get("managed_agent_worker_id") + for item in work + if item.get("metadata", {}).get("managed_agent_worker_id") + } + if len(work) != 3 or len(worker_ids) != 3: + raise AssertionError( + f"expected three one-work Worker processes, got " + f"work={len(work)} workers={worker_ids}" + ) + if any(item.get("state") != "stopped" for item in work): + raise AssertionError(f"one or more Worker runs failed: {work}") + print( + json.dumps( + { + "agent_id": agent["id"], + "session_id": session["id"], + "skill_id": skill["id"], + "pinned_skill_version": pinned, + "newer_skill_version": newer["id"], + "codewords": [first, second], + "skill_marker": marker, + "turn_event_counts": event_counts, + "event_count": len(all_events), + "tool_uses": sum( + event.get("type") == "agent.tool_use" for event in all_events + ), + "tool_results": sum( + event.get("type") == "agent.tool_result" for event in all_events + ), + "worker_ids": sorted(worker_ids), + }, + indent=2, + ) + ) + finally: + cleanup_errors = [] + if session is not None: + try: + request("POST", f"/api/sessions/{session['id']}/archive") + except Exception as error: # noqa: BLE001 -- cleanup is best effort + cleanup_errors.append(f"session: {error}") + if agent is not None: + try: + request("POST", f"/api/agents/{agent['id']}/archive") + except Exception as error: # noqa: BLE001 -- continue deleting Skill + cleanup_errors.append(f"agent: {error}") + try: + request("DELETE", f"/v1/skills/{skill['id']}") + except Exception as error: # noqa: BLE001 -- report after all cleanup attempts + cleanup_errors.append(f"skill: {error}") + if cleanup_errors: + print(json.dumps({"cleanup_warnings": cleanup_errors}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/16_self_host_sandbox/k8s_capabilities_e2e.py b/examples/16_self_host_sandbox/k8s_capabilities_e2e.py new file mode 100644 index 000000000..0e8981da8 --- /dev/null +++ b/examples/16_self_host_sandbox/k8s_capabilities_e2e.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Verify deployed Managed Session MCP and custom-tool restoration.""" + +from __future__ import annotations + +import json +import os +import time +import urllib.error +import urllib.request +import uuid + +BASE_URL = os.environ.get( + "MANAGED_AGENTS_GATEWAY_URL", "http://127.0.0.1:18080" +).rstrip("/") +MODEL = os.environ["MANAGED_AGENTS_MODEL_CARD_ID"] +ENVIRONMENT_ID = os.environ["MANAGED_AGENTS_ENVIRONMENT_ID"] +MCP_URL = os.environ.get( + "MANAGED_AGENTS_TEST_MCP_URL", "http://managed-agents-test-mcp:9000/mcp" +) + + +def request(method: str, path: str, body: dict | None = None) -> dict: + data = None if body is None else json.dumps(body).encode() + headers = {"Accept": "application/json"} + if data is not None: + headers["Content-Type"] = "application/json" + req = urllib.request.Request( + BASE_URL + path, data=data, headers=headers, method=method + ) + try: + with urllib.request.urlopen(req, timeout=60) as response: + raw = response.read() + except urllib.error.HTTPError as error: + raise RuntimeError( + f"{method} {path}: HTTP {error.code}: {error.read()[:500]!r}" + ) from error + return json.loads(raw) if raw else {} + + +def events(session_id: str) -> list[dict]: + return request("GET", f"/v1/sessions/{session_id}/events?limit=1000&order=asc").get( + "data", [] + ) + + +def send_message(session_id: str, text: str) -> None: + request( + "POST", + f"/v1/sessions/{session_id}/events", + { + "events": [ + {"type": "user.message", "content": [{"type": "text", "text": text}]} + ] + }, + ) + + +def wait(session_id: str, predicate, timeout: float = 180) -> dict: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for event in events(session_id): + if predicate(event): + return event + time.sleep(0.5) + raise TimeoutError(f"timed out waiting for Session event in {session_id}") + + +def create_session(agent: dict, title: str) -> dict: + return request( + "POST", + "/api/sessions", + {"agent": agent["id"], "environment_id": ENVIRONMENT_ID, "title": title}, + ) + + +def archive_agent(agent: dict | None) -> None: + if agent is None: + return + try: + request("POST", f"/api/agents/{agent['id']}/archive") + except Exception as error: # noqa: BLE001 -- test cleanup only + print({"cleanup_warning": str(error)}) + + +def test_mcp() -> dict: + agent = None + try: + agent = request( + "POST", + "/api/agents", + { + "name": f"mcp-e2e-{uuid.uuid4().hex[:8]}", + "model": {"id": MODEL}, + "system": "Always call the MCP add_numbers tool with a=19 and b=23, then state the result.", + "mcp_servers": [{"type": "url", "name": "math", "url": MCP_URL}], + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "math", + "default_config": { + "enabled": True, + "permission_policy": {"type": "always_allow"}, + }, + "configs": [], + } + ], + }, + ) + session = create_session(agent, "MCP restore E2E") + send_message(session["id"], "Use the required MCP tool now.") + wait( + session["id"], + lambda item: item.get("type") == "session.status_idle" + and (item.get("stop_reason") or {}).get("type") == "end_turn", + ) + ledger = events(session["id"]) + uses = [item for item in ledger if item.get("type") == "agent.mcp_tool_use"] + results = [ + item for item in ledger if item.get("type") == "agent.mcp_tool_result" + ] + answers = [item for item in ledger if item.get("type") == "agent.message"] + assert uses and uses[-1]["mcp_server_name"] == "math" + assert results and "42" in str(results[-1].get("content")) + assert answers and "42" in str(answers[-1].get("content")) + return { + "session_id": session["id"], + "use": uses[-1]["name"], + "result": results[-1]["content"], + } + finally: + archive_agent(agent) + + +def test_custom() -> dict: + agent = None + try: + agent = request( + "POST", + "/api/agents", + { + "name": f"custom-e2e-{uuid.uuid4().hex[:8]}", + "model": {"id": MODEL}, + "system": "Always call approve_release with name=demo, then report its returned ticket.", + "tools": [ + { + "type": "custom", + "name": "approve_release", + "description": "Request release approval", + "input_schema": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + }, + } + ], + }, + ) + session = create_session(agent, "Custom tool restore E2E") + send_message(session["id"], "Request the approval now.") + call = wait( + session["id"], lambda item: item.get("type") == "agent.custom_tool_use" + ) + request( + "POST", + f"/v1/sessions/{session['id']}/events", + { + "events": [ + { + "type": "user.custom_tool_result", + "custom_tool_use_id": call["id"], + "content": [{"type": "text", "text": "ticket-custom-e2e-42"}], + } + ] + }, + ) + wait( + session["id"], + lambda item: item.get("type") == "session.status_idle" + and (item.get("stop_reason") or {}).get("type") == "end_turn", + ) + ledger = events(session["id"]) + answers = [item for item in ledger if item.get("type") == "agent.message"] + assert answers and "ticket-custom-e2e-42" in str(answers[-1].get("content")) + return { + "session_id": session["id"], + "custom_tool_use_id": call["id"], + "answer": answers[-1]["content"], + } + finally: + archive_agent(agent) + + +if __name__ == "__main__": + print( + json.dumps( + {"mcp": test_mcp(), "custom": test_custom()}, ensure_ascii=False, indent=2 + ) + ) diff --git a/examples/16_self_host_sandbox/local_agent_loop_test.py b/examples/16_self_host_sandbox/local_agent_loop_test.py new file mode 100644 index 000000000..ddccb0921 --- /dev/null +++ b/examples/16_self_host_sandbox/local_agent_loop_test.py @@ -0,0 +1,204 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Local HTTP/SSE smoke test for the VeADK Managed Agents Loop. + +This script uses the installed Managed Agents package and public Anthropic SDK. +It uses no real credentials, model, or remote service. +""" + +from __future__ import annotations + +import json +import asyncio +import threading +from typing import Any, ClassVar +from datetime import datetime, timezone +from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler +from types import SimpleNamespace +from urllib.parse import urlparse + +from typing_extensions import override + +from managed_agent_loop import ManagedAgentsLoop +from sandbox_client import SelfHostSandboxClient + +SESSION_ID = "session_veadk_local" + + +class LocalSessionEvents(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + events: ClassVar[list[dict[str, Any]]] = [] + changed: ClassVar[threading.Condition] = threading.Condition() + + @override + def log_message(self, format: str, *args: object) -> None: + return + + def do_POST(self) -> None: + if urlparse(self.path).path != f"/v1/sessions/{SESSION_ID}/events": + self.send_error(404) + return + length = int(self.headers.get("content-length", "0")) + payload = json.loads(self.rfile.read(length)) + persisted: list[dict[str, Any]] = [] + with self.changed: + for event in payload["events"]: + item = dict(event) + item.setdefault("id", f"event_{len(self.events) + 1}") + item["processed_at"] = datetime.now(timezone.utc).isoformat() + self.events.append(item) + persisted.append(item) + self.changed.notify_all() + self._json({"data": persisted}) + + def do_GET(self) -> None: + path = urlparse(self.path).path + if path == f"/v1/sessions/{SESSION_ID}/events": + with self.changed: + events = [dict(event) for event in self.events] + self._json({"data": events, "has_more": False}) + return + if path != f"/v1/sessions/{SESSION_ID}/events/stream": + self.send_error(404) + return + + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("cache-control", "no-cache") + self.send_header("connection", "close") + self.end_headers() + cursor = 0 + while True: + with self.changed: + if cursor >= len(self.events): + self.changed.wait(timeout=10) + if cursor >= len(self.events): + return + pending = [dict(event) for event in self.events[cursor:]] + cursor += len(pending) + for event in pending: + frame = ( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + ) + self.wfile.write(frame) + self.wfile.flush() + if event["type"] in {"session.status_idle", "session.error"}: + return + + def _json(self, value: object) -> None: + body = json.dumps(value).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +class LocalSessionService: + def __init__(self) -> None: + self.session: object | None = None + + async def get_session(self, **kwargs: object) -> object | None: + return self.session + + async def create_session(self, **kwargs: object) -> object: + self.session = object() + return self.session + + +class LocalRunner: + app_name = "veadk-managed-agent-local-test" + + def __init__(self) -> None: + self.session_service = LocalSessionService() + + async def run_async(self, **kwargs: object): + message = kwargs["new_message"] + prompt = message.parts[0].text + yield SimpleNamespace( + partial=False, + content=SimpleNamespace( + parts=[SimpleNamespace(text="Thinking locally", thought=True)] + ), + usage_metadata=None, + get_function_calls=lambda: [], + ) + yield SimpleNamespace( + partial=False, + content=SimpleNamespace( + parts=[SimpleNamespace(text=str(prompt).upper(), thought=False)] + ), + usage_metadata=SimpleNamespace( + prompt_token_count=3, + candidates_token_count=3, + cached_content_token_count=0, + ), + get_function_calls=lambda: [], + ) + + +def main() -> None: + LocalSessionEvents.events = [] + server = ThreadingHTTPServer(("127.0.0.1", 0), LocalSessionEvents) + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + client = SelfHostSandboxClient( + base_url=f"http://127.0.0.1:{server.server_port}", + environment_id="env_local", + agent_id="agent_local", + session_id=SESSION_ID, + bearer_token="local-test-token", # gitleaks:allow loopback test credential + trust_env=False, + ) + + loop = ManagedAgentsLoop(runner=LocalRunner(), session_client=client) + loop_thread = threading.Thread( + target=lambda: asyncio.run(loop.run(max_turns=1)), daemon=True + ) + loop_thread.start() + try: + client.client.beta.sessions.events.send( + SESSION_ID, + events=[ + { + "type": "user.message", + "content": [{"type": "text", "text": "hello veadk loop"}], + } + ], + ) + loop_thread.join(timeout=10) + if loop_thread.is_alive(): + raise RuntimeError("VeADK Agent Loop did not finish") + + event_types = [event["type"] for event in client.list_events()] + expected = [ + "user.message", + "session.status_running", + "span.model_request_start", + "agent.thinking", + "agent.message", + "span.model_request_end", + "session.status_idle", + ] + if event_types != expected: + raise RuntimeError(f"Unexpected events: {event_types!r}") + print("VeADK Managed Agents Loop local test passed.") + finally: + server.shutdown() + server.server_close() + + +if __name__ == "__main__": + main() diff --git a/examples/16_self_host_sandbox/main.py b/examples/16_self_host_sandbox/main.py index 9b7cf261b..881fa1cb4 100644 --- a/examples/16_self_host_sandbox/main.py +++ b/examples/16_self_host_sandbox/main.py @@ -1,43 +1,16 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +"""Run the packaged self-hosted Managed Agents worker or demo.""" -"""Run the VeADK agent with remotely dispatched sandbox tools.""" - -import argparse import asyncio -import uuid +from pathlib import Path -from agents.self_host_sandbox_agent.agent import agent -from veadk import Runner +from dotenv import load_dotenv async def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--prompt", - default=( - "Use the bash tool to run: printf 'veadk-self-host-ok'. " - "Then reply with the exact output." - ), - ) - parser.add_argument("--session-id", default=None) - args = parser.parse_args() + load_dotenv(Path(__file__).with_name(".env"), override=False) + from veadk.runtime.managed_agents.worker import main as run - session_id = args.session_id or f"veadk-{uuid.uuid4()}" - runner = Runner(agent=agent, app_name="self_host_sandbox_demo") - output = await runner.run(messages=args.prompt, session_id=session_id) - print(output) + await run() if __name__ == "__main__": diff --git a/examples/16_self_host_sandbox/managed_agent_loop.py b/examples/16_self_host_sandbox/managed_agent_loop.py new file mode 100644 index 000000000..e95ca15bf --- /dev/null +++ b/examples/16_self_host_sandbox/managed_agent_loop.py @@ -0,0 +1,6 @@ +"""Compatibility import for veadk.runtime.managed_agents.loop.""" + +import sys +from importlib import import_module + +sys.modules[__name__] = import_module("veadk.runtime.managed_agents.loop") diff --git a/examples/16_self_host_sandbox/managed_session_resources.py b/examples/16_self_host_sandbox/managed_session_resources.py new file mode 100644 index 000000000..7db07f3de --- /dev/null +++ b/examples/16_self_host_sandbox/managed_session_resources.py @@ -0,0 +1,6 @@ +"""Compatibility import for veadk.integrations.mpa.session_resources.""" + +import sys +from importlib import import_module + +sys.modules[__name__] = import_module("veadk.integrations.mpa.session_resources") diff --git a/examples/16_self_host_sandbox/preflight_probe.py b/examples/16_self_host_sandbox/preflight_probe.py new file mode 100644 index 000000000..a4466a146 --- /dev/null +++ b/examples/16_self_host_sandbox/preflight_probe.py @@ -0,0 +1,266 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Minimal single-session preflight against the ma-infra unified gateway. + +Proves one logical Session runs a plain streamed turn plus one real bash tool +turn end-to-end (user.message -> event_delta -> agent.tool_use/result -> final +-> session.status_idle). Prints the logical Session ID and, when the +AgentKit control-plane helper is available, the mapped physical sandbox +Session ID. No credentials are printed. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import re +import sys +import time +import uuid +from pathlib import Path + +import anthropic + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from anthropic_gateway_e2e import ( # noqa: E402 + AgentKitInspector, + content_text, + field_value, +) + +_SAFE_EVENT_VALUE = re.compile(r"^[A-Za-z0-9_.:-]{1,128}$") + + +def safe_event_value(value): + text = str(value or "").strip() + return text if _SAFE_EVENT_VALUE.fullmatch(text) else "" + + +def append_event_trace(path, *, turn, event=None, error_type=None): + if path is None: + return + payload = {"turn": safe_event_value(turn)} + if event is not None: + payload.update( + { + "type": safe_event_value(field_value(event, "type", "")), + "id": safe_event_value(field_value(event, "id", "")), + "event_id": safe_event_value(field_value(event, "event_id", "")), + "tool_use_id": safe_event_value(field_value(event, "tool_use_id", "")), + "is_error": field_value(event, "is_error", None) + if isinstance(field_value(event, "is_error", None), bool) + else None, + "sequence": int(field_value(event, "sequence", 0) or 0), + } + ) + if error_type is not None: + payload.update( + { + "type": "probe.error", + "error_type": safe_event_value(error_type), + } + ) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as output: + output.write(json.dumps(payload, sort_keys=True, separators=(",", ":"))) + output.write("\n") + + +async def run_turn(client, session_id, *, marker, prompt, timeout, event_output=None): + started = time.monotonic() + types_seen = [] + tool_uses = [] + tool_results = [] + final_parts = [] + first_delta = None + saw_current_turn = False + after_sequence = 0 + async for existing in client.beta.sessions.events.list( + session_id, limit=1, order="desc" + ): + after_sequence = int(field_value(existing, "sequence", 0) or 0) + break + try: + stream = await client.beta.sessions.events.stream( + session_id, + event_deltas=["agent.message"], + extra_query={"page": str(after_sequence)}, + ) + await client.beta.sessions.events.send( + session_id, + events=[ + { + "type": "user.message", + "content": [{"type": "text", "text": prompt}], + } + ], + ) + async with stream: + async with asyncio.timeout(timeout): + async for event in stream: + et = str(field_value(event, "type", "")) + sequence = int(field_value(event, "sequence", 0) or 0) + if sequence and sequence <= after_sequence: + continue + append_event_trace(event_output, turn=marker, event=event) + types_seen.append(et) + if et in { + "event_delta", + "agent.message", + "agent.tool_use", + "agent.tool_result", + "session.error", + "session.status_terminated", + }: + saw_current_turn = True + if et == "event_delta" and first_delta is None: + delta = field_value(event, "delta") + if content_text([field_value(delta, "content")]): + first_delta = time.monotonic() - started + elif et == "agent.message": + final_parts.append(content_text(field_value(event, "content"))) + elif et == "agent.tool_use": + cid = str(field_value(event, "id", "") or "") + if cid: + tool_uses.append(cid) + elif et == "agent.tool_result": + tid = str(field_value(event, "tool_use_id", "") or "") + if tid: + tool_results.append(tid) + elif et in {"session.error", "session.status_terminated"}: + raise RuntimeError(f"terminal event {et}") + if et == "session.status_idle" and saw_current_turn: + break + except Exception as exc: + append_event_trace(event_output, turn=marker, error_type=type(exc).__name__) + raise + return { + "marker": marker, + "elapsed": round(time.monotonic() - started, 3), + "first_delta": first_delta, + "types": types_seen, + "tool_uses": tool_uses, + "tool_results": tool_results, + "final": "\n".join(p for p in final_parts if p), + } + + +async def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--base-url", default=os.getenv("MANAGED_AGENTS_GATEWAY")) + p.add_argument("--agent-id", default=os.getenv("MANAGED_AGENTS_TEST_AGENT_ID")) + p.add_argument("--environment-id", default=os.getenv("ANTHROPIC_ENVIRONMENT_ID")) + p.add_argument( + "--auth-token", + default=os.getenv("MANAGED_AGENTS_GATEWAY_TOKEN", "gateway-client"), + ) + p.add_argument("--turn-timeout", type=float, default=300) + p.add_argument("--tool-id", default=os.getenv("AGENTKIT_TOOL_ID")) + p.add_argument("--with-inspector", action="store_true") + p.add_argument( + "--event-output", + type=Path, + default=Path( + os.getenv( + "MANAGED_AGENTS_PREFLIGHT_EVENTS", + Path(__file__).with_name("deploy_records") / "tool_probe_events.jsonl", + ) + ), + ) + args = p.parse_args() + + inspector = None + session_id = None + success = False + if args.with_inspector and args.tool_id: + try: + inspector = AgentKitInspector( + os.getenv( + "AGENTKIT_TOOL_DEPLOY_SCRIPT", + "", + ), + args.tool_id, + ) + except Exception as exc: # noqa: BLE001 + print(f"[inspector-disabled] {type(exc).__name__}: {exc}") + + async with anthropic.AsyncAnthropic( + base_url=args.base_url.rstrip("/"), + auth_token=args.auth_token, + timeout=args.turn_timeout, + ) as client: + session = await client.beta.sessions.create( + agent=args.agent_id, environment_id=args.environment_id, title="preflight" + ) + sid = str(session.id) + session_id = sid + print(f"logical_session={sid}") + + pm = f"plain-{uuid.uuid4().hex[:8]}" + plain = await run_turn( + client, + sid, + marker=pm, + prompt=f"Do not use any tool. Reply with exactly: {pm}", + timeout=args.turn_timeout, + event_output=args.event_output, + ) + print( + f"[plain] elapsed={plain['elapsed']}s first_delta={plain['first_delta']} " + f"final_ok={pm in plain['final']}" + ) + plain_ok = plain["first_delta"] is not None and pm in plain["final"] + + tm = f"tool-{uuid.uuid4().hex[:8]}" + tool = await run_turn( + client, + sid, + marker=tm, + prompt=f"Use the bash tool exactly once to run: printf {tm}. " + f"Then include exactly {tm} in your final reply.", + timeout=args.turn_timeout, + event_output=args.event_output, + ) + matched = set(tool["tool_uses"]) and set(tool["tool_uses"]) <= set( + tool["tool_results"] + ) + print( + f"[tool] elapsed={tool['elapsed']}s tool_uses={len(tool['tool_uses'])} " + f"tool_results={len(tool['tool_results'])} matched={bool(matched)} " + f"final_ok={tm in tool['final']}" + ) + tool_ok = ( + bool(matched) + and tm in tool["final"] + and "session.status_idle" in tool["types"] + ) + + if inspector is not None: + try: + phys = await inspector.for_managed_session(sid) + print(f"physical_sandboxes={[x['session_id'] for x in phys]}") + for x in phys: + print( + f" physical: session_id={x['session_id']} status={x['status']} " + f"user_session_id={x['user_session_id']}" + ) + except Exception as exc: # noqa: BLE001 + print(f"[inspector-query-failed] {type(exc).__name__}: {exc}") + success = plain_ok and tool_ok + try: + await client.beta.sessions.archive(sid) + print(f"archived_session={sid}") + except Exception as exc: # noqa: BLE001 + print(f"[archive-failed] {type(exc).__name__}: {exc}") + success = False + if not success: + print(f"preflight_failed session={session_id}") + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(main())) diff --git a/examples/16_self_host_sandbox/requirements.txt b/examples/16_self_host_sandbox/requirements.txt index 8726fd273..9e8fb7046 100644 --- a/examples/16_self_host_sandbox/requirements.txt +++ b/examples/16_self_host_sandbox/requirements.txt @@ -1,4 +1,6 @@ # Dependencies for Self-Hosted Sandbox Example -anthropic>=0.40.0 +anthropic>=1.3.0 python-dotenv>=1.0.0 httpx[socks]>=0.27.0 +lark-channel-sdk +lark-oapi diff --git a/examples/16_self_host_sandbox/run-runtime.sh b/examples/16_self_host_sandbox/run-runtime.sh new file mode 100644 index 000000000..739c93642 --- /dev/null +++ b/examples/16_self_host_sandbox/run-runtime.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +export PYTHONPATH="$REPO_ROOT:${PYTHONPATH:-}" +if [[ -x "$REPO_ROOT/.venv/bin/python" ]]; then + export MANAGED_AGENTS_PYTHON="$REPO_ROOT/.venv/bin/python" +fi +exec bash "$REPO_ROOT/docker/managed-agents/run-runtime.sh" "$@" diff --git a/examples/16_self_host_sandbox/run.managed-agents-sandbox.sh b/examples/16_self_host_sandbox/run.managed-agents-sandbox.sh new file mode 100644 index 000000000..5b393b665 --- /dev/null +++ b/examples/16_self_host_sandbox/run.managed-agents-sandbox.sh @@ -0,0 +1,6 @@ +#!/usr/bin/env bash +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export MANAGED_AGENT_ENTRYPOINT_MODE=claimed +export MANAGED_AGENT_TOOL_EXECUTION="${MANAGED_AGENT_TOOL_EXECUTION:-local}" +exec bash "$SCRIPT_DIR/run-runtime.sh" "$@" diff --git a/examples/16_self_host_sandbox/run.sh b/examples/16_self_host_sandbox/run.sh index 9c85527c7..a67dea8a4 100755 --- a/examples/16_self_host_sandbox/run.sh +++ b/examples/16_self_host_sandbox/run.sh @@ -7,6 +7,11 @@ REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" # Ensure veadk and example root are in PYTHONPATH export PYTHONPATH="$REPO_ROOT:$SCRIPT_DIR:${PYTHONPATH:-}" +# Ensure internal gateway domains bypass proxy by default +DEFAULT_NO_PROXY=".volceapi.com,localhost,127.0.0.1,.volces.com" +export no_proxy="${no_proxy:-$DEFAULT_NO_PROXY}" +export NO_PROXY="${NO_PROXY:-$no_proxy}" + # Load .env if present if [ -f "$SCRIPT_DIR/.env" ]; then # Treat .env as defaults. Explicit caller-provided values are useful for @@ -14,7 +19,20 @@ if [ -f "$SCRIPT_DIR/.env" ]; then declare -A EXPLICIT_ENV=() for key in \ ANTHROPIC_BASE_URL ANTHROPIC_ENVIRONMENT_ID ANTHROPIC_ENVIRONMENT_KEY \ - SANDBOX_AGENT_ID X_TOP_ACCOUNT_ID SANDBOX_TIMEOUT_SECONDS PORT; do + ANTHROPIC_SESSION_ID \ + SANDBOX_AGENT_ID X_TOP_ACCOUNT_ID SANDBOX_TIMEOUT_SECONDS PORT \ + VEADK_MANAGED_SESSION_DB_URL DATABASE_POSTGRESQL_HOST \ + DATABASE_POSTGRESQL_PORT DATABASE_POSTGRESQL_USER \ + DATABASE_POSTGRESQL_PASSWORD DATABASE_POSTGRESQL_DATABASE \ + DATABASE_POSTGRESQL_SCHEMA MANAGED_AGENT_WORKDIR \ + TOOL_FEISHU_CHANNEL_APP_ID TOOL_FEISHU_CHANNEL_APP_SECRET \ + TOOL_FEISHU_CHANNEL_TRANSPORT TOOL_FEISHU_CHANNEL_STREAMING \ + TOOL_FEISHU_CHANNEL_REACTIONS TOOL_FEISHU_CHANNEL_SHOW_THINKING \ + TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS \ + TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS \ + TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS \ + TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD \ + TOOL_FEISHU_CHANNEL_CREATE_TOPIC; do if [[ -v $key ]]; then EXPLICIT_ENV[$key]="${!key}" fi @@ -35,6 +53,11 @@ fi PORT="${PORT:-8067}" +if [ "${1:-}" = "--feishu" ]; then + shift + set -- --feishu "$@" +fi + # Check if web mode is requested if [ "${1:-}" = "--web" ]; then shift diff --git a/examples/16_self_host_sandbox/runtime_identity.py b/examples/16_self_host_sandbox/runtime_identity.py new file mode 100644 index 000000000..06fa9d7fb --- /dev/null +++ b/examples/16_self_host_sandbox/runtime_identity.py @@ -0,0 +1,6 @@ +"""Compatibility import for veadk.runtime.managed_agents.identity.""" + +import sys +from importlib import import_module + +sys.modules[__name__] = import_module("veadk.runtime.managed_agents.identity") diff --git a/examples/16_self_host_sandbox/runtime_type_dispatcher_test.py b/examples/16_self_host_sandbox/runtime_type_dispatcher_test.py new file mode 100644 index 000000000..8529185e3 --- /dev/null +++ b/examples/16_self_host_sandbox/runtime_type_dispatcher_test.py @@ -0,0 +1,199 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +import asyncio +import base64 +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import anthropic +import pytest +from veadk.runtime.managed_agents.dispatcher import EnvironmentWorkDispatcher + + +def _secret(token: str) -> str: + payload = json.dumps({"sessions_token": token}).encode() + return base64.urlsafe_b64encode(payload).decode().rstrip("=") + + +def _work(environment_id: str) -> dict[str, object]: + return { + "id": f"work_{environment_id}", + "type": "work", + "environment_id": environment_id, + "data": {"type": "session", "id": f"session_{environment_id}"}, + "metadata": {}, + "runtime_type": "codex", + "state": "starting", + "secret": _secret("session-token"), + "created_at": "2026-09-16T00:00:00Z", + } + + +class _Service(ThreadingHTTPServer): + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _Handler) + self.requests: list[dict[str, object]] = [] + self.drain = False + self.long_poll_started = threading.Event() + self.release_long_poll = threading.Event() + + +class _Handler(BaseHTTPRequestHandler): + server: _Service + + def log_message(self, *_args: object) -> None: + return + + def _record(self, body: object = None) -> None: + self.server.requests.append( + { + "method": self.command, + "path": self.path, + "headers": {key.lower(): value for key, value in self.headers.items()}, + "body": body, + } + ) + + def _send(self, value: object, status: int = 200) -> None: + payload = json.dumps(value).encode() + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + try: + self.wfile.write(payload) + except BrokenPipeError: + pass + + def do_GET(self) -> None: + self._record() + if "/work/poll" not in self.path and not self.path.startswith( + "/v1/model-work/poll" + ): + self._send({}, 404) + return + if self.server.drain: + self.server.long_poll_started.set() + self.server.release_long_poll.wait(timeout=5) + self._send({}, 204) + return + environment_id = ( + "env_account" + if self.path.startswith("/v1/model-work/poll") + else "env_direct" + ) + self._send(_work(environment_id)) + + def do_POST(self) -> None: + length = int(self.headers.get("content-length", "0")) + body = json.loads(self.rfile.read(length)) if length else {} + self._record(body) + path = self.path.split("?", 1)[0] + environment_id = path.split("/")[3] + if path.endswith("/heartbeat"): + self._send( + { + "type": "work_heartbeat", + "state": "active", + "last_heartbeat": "2026-09-16T00:00:01Z", + "lease_extended": True, + "ttl_seconds": 2, + } + ) + return + if path.endswith(("/ack", "/stop")): + value = _work(environment_id) + value["state"] = "stopped" if path.endswith("/stop") else "active" + self._send(value) + return + self._send({}, 404) + + +@pytest.fixture +def service(): + instance = _Service() + thread = threading.Thread(target=instance.serve_forever, daemon=True) + thread.start() + try: + yield instance + finally: + instance.release_long_poll.set() + instance.shutdown() + instance.server_close() + thread.join(timeout=2) + + +def _client(service: _Service, token: str) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=f"http://127.0.0.1:{service.server_port}", + auth_token=token, + http_client=anthropic.DefaultAsyncHttpxClient(trust_env=False), + timeout=5, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("account_work", [False, True], ids=["environment", "account"]) +async def test_runtime_type_dispatcher_lifecycle( + service: _Service, account_work: bool +) -> None: + handled: list[str] = [] + + async def handler(item: object, _client: anthropic.AsyncAnthropic) -> None: + assert getattr(item, "runtime_type") == "codex" + handled.append(getattr(item, "id")) + await asyncio.sleep(0.05) + + async with _client( + service, "account-token" if account_work else "environment-token" + ) as client: + dispatcher = EnvironmentWorkDispatcher( + client, + handler=handler, + environment_id="" if account_work else "env_direct", + environment_key="account-token" if account_work else "environment-token", + account_work=account_work, + worker_id="runtime-type-test", + extra_headers={"X-Runtime-Type": "codex"}, + ) + assert await dispatcher.run(max_items=1) == 1 + + environment_id = "env_account" if account_work else "env_direct" + paths = [str(request["path"]) for request in service.requests] + assert handled == [f"work_{environment_id}"] + assert any( + "/work/poll" in path or path.startswith("/v1/model-work/poll") for path in paths + ) + assert any("/ack?" in path for path in paths) + assert any("/heartbeat?" in path for path in paths) + assert any("/stop?" in path for path in paths) + assert all( + request["headers"].get("x-runtime-type") == "codex" + for request in service.requests + ) + + +@pytest.mark.asyncio +async def test_runtime_type_dispatcher_drain_cancels_empty_poll( + service: _Service, +) -> None: + service.drain = True + + async def handler(_item: object, _client: anthropic.AsyncAnthropic) -> None: + raise AssertionError("empty queue produced work") + + async with _client(service, "environment-token") as client: + dispatcher = EnvironmentWorkDispatcher( + client, + handler=handler, + environment_id="env_direct", + environment_key="environment-token", + extra_headers={"X-Runtime-Type": "codex"}, + ) + task = asyncio.create_task(dispatcher.run()) + assert await asyncio.to_thread(service.long_poll_started.wait, 2) + dispatcher.drain() + assert await asyncio.wait_for(task, timeout=2) == 0 + + assert service.requests[0]["headers"].get("x-runtime-type") == "codex" diff --git a/examples/16_self_host_sandbox/sandbox_client.py b/examples/16_self_host_sandbox/sandbox_client.py index 4b4f1c2f6..5ac98ff6f 100644 --- a/examples/16_self_host_sandbox/sandbox_client.py +++ b/examples/16_self_host_sandbox/sandbox_client.py @@ -1,479 +1,6 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +"""Compatibility import for veadk.integrations.mpa.session_client.""" -"""Self-Hosted Sandbox client using official Anthropic Python SDK (client.beta.sessions & events). +import sys +from importlib import import_module -Endpoints managed via Anthropic SDK: -1. client.beta.sessions.create: - Creates a session and enqueues work into the WorkQueue for the worker. -2. client.beta.sessions.events.send: - Publishes ``agent.tool_use`` and ``session.status_idle`` events. -3. client.beta.sessions.events.list: - Retrieves session events to wait for ``user.tool_result`` from workers. -""" - -import base64 -import html -import json -import logging -import os -import re -import shlex -import time -import uuid -from typing import Any, Dict, List, Optional - -import anthropic - -logger = logging.getLogger("veadk.sandbox") - - -class SelfHostSandboxClient: - """Client for Managed Agents Self-Hosted Sandbox backed by official Anthropic SDK.""" - - def __init__( - self, - base_url: Optional[str] = None, - environment_id: Optional[str] = None, - agent_id: Optional[str] = None, - session_id: Optional[str] = None, - bearer_token: Optional[str] = None, - account_id: Optional[str] = None, - remote_bash_tool_name: Optional[str] = None, - timeout: int = 120, - ): - self.base_url = ( - base_url - or os.getenv("ANTHROPIC_BASE_URL") - or os.getenv("SANDBOX_BASE_URL", "") - ).rstrip("/") - self.environment_id = ( - environment_id - or os.getenv("ANTHROPIC_ENVIRONMENT_ID") - or os.getenv("SANDBOX_ENVIRONMENT_ID", "env_01SLqXHseguCmohifEqeUAYu") - ) - self.agent_id = agent_id or os.getenv( - "SANDBOX_AGENT_ID", "agent_011CSd8hFhXGpz33bM1pBw7y" - ) - self.session_id = session_id or os.getenv("SANDBOX_SESSION_ID") - - raw_token = ( - bearer_token - or os.getenv("ANTHROPIC_ENVIRONMENT_KEY") - or os.getenv("ANTHROPIC_API_KEY") - or os.getenv("SANDBOX_BEARER_TOKEN") - or os.getenv("SANDBOX_API_KEY", "") - ) - if raw_token.startswith("Bearer "): - raw_token = raw_token[7:].strip() - self.bearer_token = raw_token - self.account_id = account_id or os.getenv("X_TOP_ACCOUNT_ID", "") - self.remote_bash_tool_name = remote_bash_tool_name or os.getenv( - "SANDBOX_BASH_TOOL_NAME", "bash" - ) - self.timeout = int(os.getenv("SANDBOX_TIMEOUT_SECONDS", str(timeout))) - - if not self.base_url: - raise ValueError( - "ANTHROPIC_BASE_URL (or SANDBOX_BASE_URL) must be configured in .env or arguments." - ) - if not self.bearer_token: - raise ValueError( - "ANTHROPIC_ENVIRONMENT_KEY (or SANDBOX_BEARER_TOKEN) must be configured in .env or arguments." - ) - - # The Runtime gateway needs the Managed Agents beta flag and, when - # supplied, the Volcengine top-account routing header. - default_headers = { - "anthropic-beta": "managed-agents-2026-04-01", - } - if self.account_id: - default_headers["X-Top-Account-Id"] = str(self.account_id) - - self.client = anthropic.Anthropic( - base_url=self.base_url, - auth_token=self.bearer_token, - default_headers=default_headers, - timeout=float(self.timeout), - ) - - def create_session( - self, - title: str = "VeADK Self-Hosted Sandbox Session", - ) -> Dict[str, Any]: - """POST /v1/sessions via Anthropic SDK (client.beta.sessions.create). - - Creates a session and enqueues a work item into managed_selfhost_work_items. - """ - sess = self.client.beta.sessions.create( - agent=self.agent_id, - environment_id=self.environment_id, - title=title, - ) - self.session_id = sess.id - logger.info( - "Session %s created on remote Runtime via Anthropic SDK.", self.session_id - ) - if hasattr(sess, "model_dump"): - return sess.model_dump(warnings=False) - return { - "id": sess.id, - "environment_id": self.environment_id, - "agent": self.agent_id, - } - - def post_events(self, events: List[Dict[str, Any]]) -> Dict[str, Any]: - """POST /v1/sessions/{session_id}/events via Anthropic SDK (client.beta.sessions.events.send). - - Dispatches tool calls and tool results into the Session event stream. - """ - if not self.session_id: - self.create_session() - - response = self.client.beta.sessions.events.send( - session_id=self.session_id, - events=events, - ) - if hasattr(response, "model_dump"): - return response.model_dump(warnings=False) - return {"status": "ok"} - - def send_tool_result( - self, tool_use_id: str, content: str, is_error: bool = False - ) -> Dict[str, Any]: - """Send tool execution result to POST /v1/sessions/{session_id}/events.""" - event = { - "type": "user.tool_result", - "tool_use_id": tool_use_id, - "content": [{"type": "text", "text": content}], - "is_error": is_error, - } - return self.post_events([event]) - - def list_events( - self, - *, - limit: int = 100, - order: str = "asc", - ) -> List[Dict[str, Any]]: - """Read and normalize events from GET /v1/sessions/{id}/events via Anthropic SDK.""" - if not self.session_id: - raise RuntimeError("Cannot read events before a Session is created.") - - page = self.client.beta.sessions.events.list( - session_id=self.session_id, - limit=limit, - order=order, - ) - events: List[Dict[str, Any]] = [] - for item in page.data: - if hasattr(item, "model_dump"): - d = item.model_dump(warnings=False) - elif isinstance(item, dict): - d = dict(item) - else: - d = { - "id": getattr(item, "id", None), - "type": getattr(item, "type", None), - "name": getattr(item, "name", None), - "input": getattr(item, "input", None), - "content": getattr(item, "content", None), - "is_error": getattr(item, "is_error", False), - "tool_use_id": getattr(item, "tool_use_id", None), - } - events.append(d) - return events - - def post_status_idle(self, stop_reason: str = "end_turn") -> None: - """Publish a ``session.status_idle`` event so workers know the turn ended and can release their lease.""" - if not self.session_id: - return - try: - self.post_events( - [ - { - "type": "session.status_idle", - "stop_reason": { - "type": stop_reason, - }, - } - ] - ) - logging.info( - "Session %s marked status_idle (%s).", self.session_id, stop_reason - ) - except Exception as e: - logging.warning( - "Failed to post session.status_idle for %s: %s", self.session_id, e - ) - - def dispatch_tool( - self, - tool_name: str, - arguments: Dict[str, Any], - *, - dispatch_id: str = "", - ) -> Dict[str, Any]: - """Convert a VeADK tool to bash and publish an ``agent.tool_use``.""" - command, timeout = self._tool_to_bash(tool_name, arguments) - return self.execute_command( - command, - timeout=timeout, - dispatch_id=dispatch_id, - ) - - def _tool_to_bash( - self, tool_name: str, arguments: Dict[str, Any] - ) -> tuple[str, float]: - """Translate supported non-MCP tools into one remote bash invocation.""" - timeout = self._positive_timeout(arguments.get("timeout", self.timeout)) - if tool_name == "bash": - cmd = self._required_string(arguments, "command", tool_name) - return html.unescape(cmd), timeout - - if tool_name == "read_file": - file_path = self._required_string(arguments, "file_path", tool_name) - offset = arguments.get("offset") - limit = arguments.get("limit") - quoted_path = shlex.quote(file_path) - if offset is None and limit is None: - return f"cat -- {quoted_path}", timeout - start = self._positive_integer( - offset if offset is not None else 1, "offset" - ) - if limit is None: - condition = f"NR >= {start}" - else: - line_limit = self._positive_integer(limit, "limit") - condition = f"NR >= {start} && NR < {start + line_limit}" - awk_program = f'{condition} {{printf "%d\\t%s\\n", NR, $0}}' - return f"awk {shlex.quote(awk_program)} {quoted_path}", timeout - - if tool_name == "write_file": - file_path = self._required_string(arguments, "file_path", tool_name) - content = self._required_string( - arguments, "content", tool_name, allow_empty=True - ) - encoded = base64.b64encode(content.encode("utf-8")).decode("ascii") - quoted_path = shlex.quote(file_path) - command = ( - f"mkdir -p $(dirname -- {quoted_path}) && " - f"printf %s {shlex.quote(encoded)} | base64 -d > {quoted_path}" - ) - return command, timeout - - if tool_name == "edit_file": - payload = { - "file_path": self._required_string(arguments, "file_path", tool_name), - "old_string": self._required_string(arguments, "old_string", tool_name), - "new_string": self._required_string( - arguments, "new_string", tool_name, allow_empty=True - ), - "replace_all": bool(arguments.get("replace_all", False)), - } - encoded = base64.b64encode( - json.dumps(payload, ensure_ascii=False).encode("utf-8") - ).decode("ascii") - script = ( - "import base64,json,pathlib,sys;" - "d=json.loads(base64.b64decode(sys.argv[1]));" - "p=pathlib.Path(d['file_path']);s=p.read_text();o=d['old_string'];" - "n=s.count(o);" - "(_ for _ in ()).throw(ValueError('old_string not found')) if n==0 else None;" - "(_ for _ in ()).throw(ValueError(f'old_string appears {n} times')) " - "if n>1 and not d['replace_all'] else None;" - "p.write_text(s.replace(o,d['new_string'],-1 if d['replace_all'] else 1))" - ) - return ( - f"python3 -c {shlex.quote(script)} {shlex.quote(encoded)}", - timeout, - ) - - if tool_name == "list_files": - path = str(arguments.get("path") or "/workspace") - depth = self._positive_integer(arguments.get("max_depth", 4), "max_depth") - return ( - f"find {shlex.quote(path)} -maxdepth {depth} -type f | sort | head -n 500", - timeout, - ) - - if tool_name == "search_files": - pattern = self._required_string(arguments, "pattern", tool_name) - path = str(arguments.get("path") or "/workspace") - glob = arguments.get("glob") - rg_flags = "-n --no-heading" - grep_flags = "-RIn" - if arguments.get("case_insensitive"): - rg_flags += " -i" - grep_flags += " -i" - if glob: - rg_flags += f" --glob {shlex.quote(str(glob))}" - grep_flags += f" --include {shlex.quote(str(glob))}" - quoted_pattern = shlex.quote(pattern) - quoted_path = shlex.quote(path) - return ( - "if command -v rg >/dev/null 2>&1; then " - f"rg {rg_flags} -- {quoted_pattern} {quoted_path}; " - "else " - f"grep {grep_flags} -- {quoted_pattern} {quoted_path}; fi", - timeout, - ) - - if tool_name == "python": - code = self._required_string(arguments, "code", tool_name, allow_empty=True) - workdir = str(arguments.get("workdir") or "/workspace") - return ( - f"cd {shlex.quote(workdir)} && python3 -c {shlex.quote(code)}", - timeout, - ) - - raise ValueError( - f"Unsupported non-MCP tool {tool_name!r}; add an explicit bash adapter" - ) - - def execute_command( - self, - command: str, - *, - timeout: float | None = None, - dispatch_id: str = "", - ) -> Dict[str, Any]: - """Publish ``agent.tool_use`` and wait for the Worker's tool result.""" - if not self.session_id: - self.create_session() - - timeout_seconds = timeout or float(self.timeout) - command = html.unescape(command) - tool_use_id = dispatch_id or f"toolu_{uuid.uuid4().hex}" - self.post_events( - [ - { - "type": "agent.tool_use", - "id": tool_use_id, - "name": self.remote_bash_tool_name, - "input": { - "command": command, - "timeout_ms": int(timeout_seconds * 1000), - "timeout": int(timeout_seconds * 1000), - }, - } - ] - ) - return self._wait_for_tool_result( - timeout=timeout_seconds, - tool_use_id=tool_use_id, - dispatch_id=dispatch_id or tool_use_id, - ) - - def _wait_for_tool_result( - self, - *, - timeout: float, - tool_use_id: str, - dispatch_id: str, - ) -> Dict[str, Any]: - deadline = time.monotonic() + timeout - - while time.monotonic() < deadline: - events = self.list_events(limit=100, order="asc") - for event in events: - event_type = str(event.get("type") or "") - if event_type not in {"agent.tool_result", "user.tool_result"}: - continue - result_tool_use_id = str(event.get("tool_use_id") or "") - if result_tool_use_id == tool_use_id: - return self._normalize_bash_result( - event, - tool_use_id=tool_use_id, - dispatch_id=dispatch_id, - ) - time.sleep(0.5) - - raise TimeoutError( - f"Timed out after {timeout:g}s waiting for tool result {tool_use_id}." - ) - - @staticmethod - def _required_string( - arguments: Dict[str, Any], - key: str, - tool_name: str, - *, - allow_empty: bool = False, - ) -> str: - value = arguments.get(key) - if not isinstance(value, str) or (not allow_empty and not value.strip()): - raise ValueError(f"{tool_name} requires a valid {key}") - return value - - @staticmethod - def _positive_integer(value: Any, name: str) -> int: - if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise ValueError(f"{name} must be a positive integer") - return value - - @staticmethod - def _positive_timeout(value: Any) -> float: - if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: - raise ValueError("timeout must be greater than zero") - return float(value) - - def _normalize_bash_result( - self, - event: Dict[str, Any], - *, - tool_use_id: str, - dispatch_id: str, - ) -> Dict[str, Any]: - content = self._content_text(event.get("content")) - match = re.match(r"^exit=(-?\d+)\n([\s\S]*)$", content) - if match: - exit_code = int(match.group(1)) - stdout = match.group(2) - else: - exit_code = 1 if event.get("is_error") else 0 - stdout = content - return { - "dispatch_id": dispatch_id, - "tool_use_id": tool_use_id, - "environment_id": self.environment_id, - "session_id": self.session_id, - "status": "failed" if exit_code else "completed", - "exit_code": exit_code, - "stdout": stdout, - "stderr": "", - } - - @staticmethod - def _event_seq(event: Dict[str, Any]) -> int: - try: - return int(event.get("_seq", 0) or 0) - except (TypeError, ValueError): - return 0 - - @staticmethod - def _content_text(content: Any) -> str: - if isinstance(content, str): - return content - if isinstance(content, list): - parts = [] - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - parts.append(str(block.get("text") or "")) - elif isinstance(block, str): - parts.append(block) - return "\n".join(part for part in parts if part) - if content is None: - return "" - return json.dumps(content, ensure_ascii=False) +sys.modules[__name__] = import_module("veadk.integrations.mpa.session_client") diff --git a/examples/16_self_host_sandbox/soak_fault_preflight.py b/examples/16_self_host_sandbox/soak_fault_preflight.py new file mode 100644 index 000000000..5cd8da528 --- /dev/null +++ b/examples/16_self_host_sandbox/soak_fault_preflight.py @@ -0,0 +1,629 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Dedicated sandbox-loss preflight; never include it in soak statistics. + +This script is deliberately separate from ``soak_load_test.py``. It creates a +single logical Session, waits until a long tool call has a bound physical +Session, deletes exactly that physical ID while the Work lease is active, then +measures the next tool turn's lease-based recovery. Run only after deployment +authorization. +""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import os +import sys +import time +import uuid +from datetime import datetime, timezone +from pathlib import Path + +import anthropic +import httpx + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from anthropic_gateway_e2e import ( # noqa: E402 + AgentKitInspector, + content_text, + field_value, +) +from soak_load_test import sanitize_error_detail, verify_turn_settled # noqa: E402 + + +def assert_tool_turn(result: dict, marker: str) -> None: + if result.get("error"): + raise AssertionError(f"turn collection failed: {result['error']}") + uses = result["tool_uses"] + results = result["tool_results"] + if ( + not uses + or len(uses) != len(set(uses)) + or len(results) != len(set(results)) + or set(uses) != set(results) + ): + raise AssertionError(f"tool IDs do not match: uses={uses} results={results}") + if result.get("tool_error_ids"): + raise AssertionError( + f"tool execution failed: IDs={result['tool_error_ids']} " + f"details={result.get('tool_error_details', [])}" + ) + if marker not in result["final"]: + raise AssertionError("final response is missing marker") + if "session.status_idle" not in result["types"]: + raise AssertionError("turn did not reach idle") + + +def event_sequence(event) -> int: + value = field_value(event, "seq", field_value(event, "sequence")) + if isinstance(value, bool) or not str(value).isdigit() or int(value) < 1: + raise RuntimeError("canonical event is missing a positive sequence") + return int(value) + + +async def send_turn(client, session_id: str, prompt: str) -> tuple[str, int]: + """Use the persisted send acknowledgement as the authoritative turn boundary.""" + response = await client.beta.sessions.events.send( + session_id, + events=[ + {"type": "user.message", "content": [{"type": "text", "text": prompt}]} + ], + ) + accepted = field_value(response, "data", []) + if len(accepted) != 1 or field_value(accepted[0], "type") != "user.message": + raise RuntimeError("send did not acknowledge exactly one user.message") + event_id = field_value(accepted[0], "id") + if not event_id: + raise RuntimeError("acknowledged user.message has no ID") + return str(event_id), event_sequence(accepted[0]) + + +async def run_verified_tool_turn( + client, + session_id: str, + *, + marker: str, + prompt: str, + timeout: float, + environment_id: str | None = None, +) -> dict: + started = time.monotonic() + types_seen = [] + tool_uses = [] + tool_results = [] + tool_error_ids = [] + tool_error_details = [] + final_parts = [] + summaries = [] + input_id = None + input_seq = None + boundary_seen = False + error = None + seen: dict[int, str] = {} + last_seq = 0 + final_seq = 0 + last_tool_seq = 0 + idle_seq = 0 + completion = {} + try: + async with asyncio.timeout(timeout): + input_id, input_seq = await send_turn(client, session_id, prompt) + # Include the acknowledged input itself, even if the worker finished + # before the stream opens. Durable SSE replay cannot lose that turn. + stream = await client.beta.sessions.events.stream( + session_id, extra_headers={"Last-Event-ID": str(input_seq - 1)} + ) + async with stream: + async for event in stream: + seq = event_sequence(event) + if seq < input_seq: + continue + event_type = str(field_value(event, "type", "")) + event_id = str(field_value(event, "id", "") or "") + raw = ( + event.model_dump(mode="json") + if hasattr(event, "model_dump") + else event + ) + signature = hashlib.sha256( + json.dumps(raw, sort_keys=True, default=str).encode() + ).hexdigest() + if seq in seen: + if seen[seq] != signature: + raise RuntimeError("conflicting duplicate event sequence") + continue + if seq <= last_seq: + raise RuntimeError("event sequence moved backwards") + seen[seq] = signature + last_seq = seq + summaries.append( + { + "sequence": seq, + "id": event_id, + "type": event_type, + "tool_use_id": field_value(event, "tool_use_id"), + "is_error": field_value(event, "is_error"), + "observed_at": datetime.now(timezone.utc).isoformat(), + } + ) + if not boundary_seen: + if ( + seq != input_seq + or event_id != input_id + or event_type != "user.message" + ): + raise RuntimeError( + "stream did not replay acknowledged user.message" + ) + boundary_seen = True + elif event_type in {"user.message", "user.interrupt"}: + raise RuntimeError(f"unexpected concurrent {event_type}") + types_seen.append(event_type) + if event_type == "agent.tool_use": + if not event_id: + raise RuntimeError("tool use has no ID") + tool_uses.append(event_id) + last_tool_seq = seq + elif event_type == "agent.tool_result": + tool_id = str(field_value(event, "tool_use_id", "") or "") + if not tool_id or tool_id not in tool_uses: + raise RuntimeError( + "tool result has no matching preceding use" + ) + tool_results.append(tool_id) + last_tool_seq = seq + if field_value(event, "is_error", False) is True: + tool_error_ids.append(tool_id) + detail = content_text(field_value(event, "content")) + tool_error_details.append( + sanitize_error_detail(detail or "tool execution failed") + ) + elif event_type == "agent.message": + # A model may speak before invoking its tool. Only a + # message after all results is evidence of the final. + if tool_results and set(tool_uses) == set(tool_results): + final_parts.append( + content_text(field_value(event, "content")) + ) + final_seq = seq + elif event_type in {"session.error", "session.status_terminated"}: + raise RuntimeError(f"terminal event {event_type}") + if event_type == "session.status_idle": + reason = field_value( + field_value(event, "stop_reason", {}), "type" + ) + if reason != "end_turn": + raise RuntimeError("idle is not end_turn") + if final_seq <= last_tool_seq or not final_parts: + raise RuntimeError( + "idle arrived before final after tool results" + ) + idle_seq = seq + break + if not boundary_seen: + raise RuntimeError("stream ended before acknowledged user.message") + if environment_id: + if not idle_seq: + raise RuntimeError("stream ended before end_turn idle") + await verify_turn_settled( + client, + session_id, + environment_id, + input_id=input_id, + input_sequence=input_seq, + idle_sequence=idle_seq, + evidence=completion, + timeout=max(0.001, timeout - (time.monotonic() - started)), + ) + except Exception as exc: # noqa: BLE001 + # Return partial evidence so assertions, timeouts and transport errors + # cannot erase the failed recovery turn from the final JSON artifact. + error = sanitize_error_detail(f"{type(exc).__name__}: {exc}", limit=300) + return { + "input_event_id": input_id, + "input_sequence": input_seq, + "events": summaries, + "error": error, + "completion": completion, + "elapsed": round(time.monotonic() - started, 3), + "types": types_seen, + "tool_uses": tool_uses, + "tool_results": tool_results, + "tool_error_ids": tool_error_ids, + "tool_error_details": tool_error_details, + "final": sanitize_error_detail( + "\n".join(part for part in final_parts if part), limit=2000 + ), + } + + +async def run_loss_during_active_tool( + client, inspector, session_id: str, *, marker: str, timeout: float +) -> dict: + started = time.monotonic() + stream = await client.beta.sessions.events.stream(session_id) + await client.beta.sessions.events.send( + session_id, + events=[ + { + "type": "user.message", + "content": [ + { + "type": "text", + "text": ( + f"Use bash exactly once to run: sleep 120; printf {marker}. " + f"Then reply with {marker}." + ), + } + ], + } + ], + ) + event_types = [] + physical = None + deleted_at = None + error = None + try: + async with stream: + async with asyncio.timeout(timeout): + async for event in stream: + event_type = str(field_value(event, "type", "")) + event_types.append(event_type) + if event_type == "agent.tool_use" and physical is None: + physical = await inspector.wait_present(session_id, timeout=120) + await asyncio.to_thread( + inspector.client.call, + "DeleteSession", + { + "ToolId": inspector.tool_id, + "SessionId": physical["session_id"], + }, + ) + deleted_at = time.monotonic() + if event_type == "session.status_idle": + break + except Exception as exc: # noqa: BLE001 + error = sanitize_error_detail(f"{type(exc).__name__}: {exc}", limit=300) + if physical is None and error is not None: + raise RuntimeError(f"physical lookup failed: {error}") + if physical is None or deleted_at is None: + raise AssertionError("fault preflight never observed a bound physical sandbox") + return { + "elapsed_seconds": round(time.monotonic() - started, 3), + "deleted_at_offset_seconds": round(deleted_at - started, 3), + "event_types": event_types, + "physical": physical, + "error": error, + } + + +async def session_work_snapshot( + base_url: str, + auth_token: str, + environment_id: str, + session_id: str, + *, + expected_work_ids: set[str] | None = None, +) -> dict: + rows: list[dict] = [] + query_started_at = datetime.now(timezone.utc).isoformat() + page: str | None = None + seen_pages: set[str] = set() + headers = {"Authorization": f"Bearer {auth_token}"} + async with httpx.AsyncClient(timeout=30) as http: + for page_index in range(100): + params = {"limit": 100} + if page: + params["page"] = page + response = await http.get( + f"{base_url.rstrip('/')}/v1/environments/{environment_id}/work", + headers=headers, + params=params, + ) + response.raise_for_status() + payload = response.json() + if not isinstance(payload, dict) or not isinstance( + payload.get("data"), list + ): + raise RuntimeError("Work response must be an object with a data list") + if "next_page" not in payload: + raise RuntimeError("Work response is missing next_page") + for item in payload["data"]: + if not isinstance(item, dict): + raise RuntimeError("Work response contains a non-object item") + work_id = item.get("work_id") or item.get("id") + owner = item.get("data") + if not isinstance(work_id, str) or not work_id: + raise RuntimeError("Work response item is missing work_id/id") + if not isinstance(owner, dict) or not isinstance(owner.get("id"), str): + raise RuntimeError(f"Work {work_id} has invalid data.id ownership") + if item.get("type") != "work" or item.get("work_type") != "session": + raise RuntimeError(f"Work {work_id} has an invalid type/work_type") + if owner.get("type") != "session": + raise RuntimeError( + f"Work {work_id} has invalid data.type ownership" + ) + if item.get("environment_id") != environment_id: + raise RuntimeError( + f"Work {work_id} belongs to an unexpected environment" + ) + if item.get("state") not in {"queued", "starting", "active", "stopped"}: + raise RuntimeError(f"Work {work_id} has an unknown state") + if owner["id"] == session_id: + rows.append(item) + next_page = payload["next_page"] + if next_page is None: + found = {str(item.get("work_id") or item.get("id")) for item in rows} + missing = set(expected_work_ids or ()) - found + if missing: + raise RuntimeError( + f"Work snapshot lost known IDs: {','.join(sorted(missing))}" + ) + return { + "complete": True, + "query_started_at": query_started_at, + "query_completed_at": datetime.now(timezone.utc).isoformat(), + "page_count": page_index + 1, + "rows": rows, + } + if not isinstance(next_page, str) or not next_page: + raise RuntimeError("Work next_page must be a non-empty string or null") + if next_page in seen_pages: + raise RuntimeError("Work pagination repeated next_page") + seen_pages.add(next_page) + page = next_page + raise RuntimeError("Work pagination exceeded 100 pages") + + +def summarize_work(rows: list[dict]) -> list[dict]: + return [ + { + "work_id": item.get("work_id") or item.get("id"), + "state": item.get("state"), + "created_at": item.get("created_at"), + "started_at": item.get("started_at"), + "latest_heartbeat_at": item.get("latest_heartbeat_at"), + "stop_requested_at": item.get("stop_requested_at"), + "stopped_at": item.get("stopped_at"), + } + for item in rows + ] + + +async def main(args: argparse.Namespace) -> dict: + inspector = AgentKitInspector(args.inspector_script, args.tool_id) + async with anthropic.AsyncAnthropic( + base_url=args.base_url.rstrip("/"), + auth_token=args.auth_token, + timeout=args.turn_timeout, + max_retries=0, + ) as client: + session = await client.beta.sessions.create( + agent=args.agent_id, + environment_id=args.environment_id, + title=f"soak fault preflight {uuid.uuid4().hex[:8]}", + ) + logical_id = str(session.id) + cleanup = "pending" + result = { + "logical_session_id": logical_id, + "note": "Never merge this fault result into soak statistics.", + "recovery_ok": False, + } + known_work_ids: set[str] = set() + try: + fault_marker = f"fault-a-{uuid.uuid4().hex[:8]}" + fault = await run_loss_during_active_tool( + client, + inspector, + logical_id, + marker=fault_marker, + timeout=args.turn_timeout, + ) + result["fault_turn"] = fault + if fault.get("error"): + raise RuntimeError(f"fault turn did not finish: {fault['error']}") + work_after_delete_snapshot = await session_work_snapshot( + args.base_url, args.auth_token, args.environment_id, logical_id + ) + work_after_delete_rows = work_after_delete_snapshot["rows"] + if not work_after_delete_rows: + raise RuntimeError( + "Work snapshot contains no Work for the fault Session" + ) + known_work_ids = { + str(item.get("work_id") or item.get("id")) + for item in work_after_delete_rows + } + work_after_delete = summarize_work(work_after_delete_rows) + absence_query_started_at = datetime.now(timezone.utc).isoformat() + await inspector.wait_absent(logical_id, timeout=120) + absence_query_completed_at = datetime.now(timezone.utc).isoformat() + + recovery_started = time.monotonic() + second_marker = f"fault-b-{uuid.uuid4().hex[:8]}" + recovery_error = None + second = None + second_physical = None + try: + second = await run_verified_tool_turn( + client, + logical_id, + marker=second_marker, + prompt=f"Use bash exactly once to run: printf {second_marker}. Reply with {second_marker}.", + timeout=args.recovery_timeout, + environment_id=args.environment_id, + ) + assert_tool_turn(second, second_marker) + second_physical = await inspector.wait_present( + logical_id, timeout=args.recovery_timeout + ) + if second_physical["session_id"] == fault["physical"]["session_id"]: + raise AssertionError("recovery reused the deleted physical Session") + except Exception as exc: # noqa: BLE001 + recovery_error = sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ) + result.update( + { + "second_turn": second, + "second_physical": second_physical, + "recovery_error": recovery_error, + } + ) + work_after_recovery_snapshot = await session_work_snapshot( + args.base_url, + args.auth_token, + args.environment_id, + logical_id, + expected_work_ids=known_work_ids, + ) + work_after_recovery = summarize_work(work_after_recovery_snapshot["rows"]) + known_work_ids.update( + item["work_id"] for item in work_after_recovery if item["work_id"] + ) + result.update( + { + "fault_turn": fault, + "second_turn": second, + "second_physical": second_physical, + "recovery_ok": recovery_error is None, + "recovery_error": recovery_error, + "recovery_turn_seconds": round( + time.monotonic() - recovery_started, 3 + ), + "work_after_delete": work_after_delete, + "work_after_delete_query": { + key: value + for key, value in work_after_delete_snapshot.items() + if key != "rows" + }, + "work_after_recovery": work_after_recovery, + "work_after_recovery_query": { + key: value + for key, value in work_after_recovery_snapshot.items() + if key != "rows" + }, + "physical_absence_query_started_at": absence_query_started_at, + "physical_absence_query_completed_at": absence_query_completed_at, + } + ) + except Exception as exc: # noqa: BLE001 + result["recovery_ok"] = False + result["error"] = sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ) + finally: + result["cleanup_started_at"] = datetime.now(timezone.utc).isoformat() + physical_cleanup_query_started_at = None + physical_cleanup_query_completed_at = None + physical_cleanup_reclaimed = False + physical_cleanup_error = None + try: + await client.beta.sessions.events.send( + logical_id, events=[{"type": "user.interrupt"}] + ) + await client.beta.sessions.archive(logical_id) + cleanup = "archived" + except Exception as exc: # noqa: BLE001 + cleanup = "archive_failed" + physical_cleanup_error = sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ) + physical_cleanup_query_started_at = datetime.now(timezone.utc).isoformat() + try: + await inspector.wait_absent(logical_id, timeout=args.recovery_timeout) + physical_cleanup_reclaimed = True + except Exception as exc: # noqa: BLE001 + detail = sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ) + physical_cleanup_error = ( + f"{physical_cleanup_error}; {detail}" + if physical_cleanup_error + else detail + ) + finally: + physical_cleanup_query_completed_at = datetime.now( + timezone.utc + ).isoformat() + print(f"cleanup={cleanup}", file=sys.stderr) + try: + work_after_cleanup_snapshot = await session_work_snapshot( + args.base_url, + args.auth_token, + args.environment_id, + logical_id, + expected_work_ids=known_work_ids, + ) + except Exception as exc: # noqa: BLE001 + result["recovery_ok"] = False + work_after_cleanup_snapshot = { + "complete": False, + "rows": [], + "error": sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ), + } + work_after_cleanup = summarize_work(work_after_cleanup_snapshot["rows"]) + active_states = {"queued", "starting", "active"} + result["cleanup"] = cleanup + result["physical_cleanup_query_started_at"] = physical_cleanup_query_started_at + result["physical_cleanup_query_completed_at"] = ( + physical_cleanup_query_completed_at + ) + result["physical_cleanup_reclaimed"] = physical_cleanup_reclaimed + result["physical_cleanup_error"] = physical_cleanup_error + result["work_after_cleanup"] = work_after_cleanup + result["work_after_cleanup_query"] = { + key: value + for key, value in work_after_cleanup_snapshot.items() + if key != "rows" + } + result["orphan_work_after_cleanup"] = [ + item for item in work_after_cleanup if item["state"] in active_states + ] + return result + + +def parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--base-url", default=os.getenv("MANAGED_AGENTS_GATEWAY")) + p.add_argument("--agent-id", default=os.getenv("MANAGED_AGENTS_TEST_AGENT_ID")) + p.add_argument("--environment-id", default=os.getenv("ANTHROPIC_ENVIRONMENT_ID")) + p.add_argument( + "--auth-token", + default=os.getenv("MANAGED_AGENTS_GATEWAY_TOKEN", "gateway-client"), + ) + p.add_argument("--tool-id", default=os.getenv("AGENTKIT_TOOL_ID"), required=False) + p.add_argument("--turn-timeout", type=float, default=300) + p.add_argument("--recovery-timeout", type=float, default=600) + p.add_argument("--output", type=Path, required=True) + p.add_argument( + "--inspector-script", + default=os.getenv( + "AGENTKIT_TOOL_DEPLOY_SCRIPT", + "", + ), + ) + return p + + +if __name__ == "__main__": + args = parser().parse_args() + if not all((args.base_url, args.agent_id, args.environment_id, args.tool_id)): + raise SystemExit("base-url, agent-id, environment-id and tool-id are required") + result = asyncio.run(main(args)) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n") + if ( + not result.get("recovery_ok") + or result.get("cleanup") != "archived" + or not result.get("physical_cleanup_reclaimed") + or result.get("orphan_work_after_cleanup") + ): + raise SystemExit(1) diff --git a/examples/16_self_host_sandbox/soak_load_test.py b/examples/16_self_host_sandbox/soak_load_test.py new file mode 100644 index 000000000..252549b2a --- /dev/null +++ b/examples/16_self_host_sandbox/soak_load_test.py @@ -0,0 +1,1259 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Sustained soak / load test for the Managed Agents unified gateway. + +Unlike ``anthropic_gateway_e2e.py --mode concurrency`` (a one-shot isolation +check with no time axis), this harness holds a target number of concurrent +logical Sessions busy for a fixed duration by continuously *replenishing* +requests: whenever a worker slot finishes a turn it immediately starts the +next one, so the system stays at target concurrency for the whole window. + +It records every request incrementally (JSONL) and in aggregate: + * total requests, successes, failures, success rate + * throughput (requests/min) + * end-to-end latency percentiles (p50/p90/p95/p99, min/max/mean) + * time-to-first-delta (streaming) percentiles + * error classification (by type/status/short reason) + * actual in-flight logical requests and low-frequency, run-scoped physical + sandbox samples (via AgentKit helper) + +It never prints or writes credentials. Signed endpoints stay out of the +result file. Sessions created for the run are best-effort archived on exit. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import re +import statistics +import sys +import time +import uuid +from collections import Counter +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import anthropic + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from anthropic_gateway_e2e import ( # noqa: E402 + AgentKitInspector, + content_text, + field_value, +) + + +def now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +def percentile(values: list[float], pct: float) -> float | None: + if not values: + return None + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + rank = (len(ordered) - 1) * (pct / 100.0) + low = int(rank) + high = min(low + 1, len(ordered) - 1) + frac = rank - low + return ordered[low] + (ordered[high] - ordered[low]) * frac + + +@dataclass +class RequestRecord: + request_id: str + worker: int + session_id: str + started_at: str + completed_at: str + started_offset_seconds: float + completed_offset_seconds: float + completion_phase: str + elapsed_seconds: float + first_delta_seconds: float | None + ok: bool + require_tool: bool + tool_uses: int + tool_results: int + tool_use_ids: list[str] = field(default_factory=list) + tool_result_ids: list[str] = field(default_factory=list) + tool_error_ids: list[str] = field(default_factory=list) + tool_error_details: list[str] = field(default_factory=list) + cleanup_status: str | None = None + cleanup_error: str | None = None + physical_query_started_at: str | None = None + physical_query_completed_at: str | None = None + physical_query_count: int = 0 + physical_reclaimed: bool | None = None + error_type: str | None = None + error_category: str | None = None + error_reason: str | None = None + completion: dict[str, Any] | None = None + + +@dataclass +class Summary: + started_at: str + completed_at: str + duration_seconds_target: float + duration_seconds_actual: float + concurrency: int + require_tool: bool + total_requests: int + successes: int + failures: int + success_rate: float + request_throughput_per_min: float + success_throughput_per_min: float + formal_window_seconds: float + drain_seconds: float + formal_completions: int + drain_completions: int + peak_in_flight_requests: int + success_latency: dict[str, float | None] + latency_including_failures: dict[str, float | None] + first_text_event: dict[str, float | None] + error_breakdown: dict[str, int] + observed_peak_physical_sessions: int + physical_session_ids: list[str] = field(default_factory=list) + physical_samples: int = 0 + physical_sample_errors: int = 0 + physical_reclaim_confirmed: int = 0 + physical_reclaim_unconfirmed: int = 0 + physical_reclaim_unverified: int = 0 + total_sessions_created: int = 0 + sample_sessions_created: list[str] = field(default_factory=list) + + +class DurableJsonlWriter: + """Durably append evidence so an interrupted run remains auditable.""" + + def __init__(self, path: Path | None) -> None: + self.path = path + self._file = None + + def __enter__(self): + if self.path is not None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self._file = self.path.open("w", encoding="utf-8") + return self + + def append(self, value: Any) -> None: + if self._file is None: + return + payload = asdict(value) if hasattr(value, "__dataclass_fields__") else value + self._file.write(json.dumps(payload, ensure_ascii=False) + "\n") + self._file.flush() + os.fsync(self._file.fileno()) + + def __exit__(self, *_args) -> None: + if self._file is not None: + self._file.close() + + +def sanitize_error_detail(value: str, limit: int = 500) -> str: + """Keep useful failure detail while removing common credential forms.""" + value = re.sub(r"(?i)bearer\s+[a-z0-9._~+/-]+", "Bearer ", value) + credential_key = ( + r"(?:authorization|x[-_]api[-_]key|api[-_]?key|access[-_]?key|" + r"secret[-_]?access[-_]?key|secret[-_]?key|client[-_]?secret|" + r"session[-_]?token|security[-_]?token|auth[-_]?token|token|password|" + r"private[-_]?key|secret)" + ) + value = re.sub( + rf"(?i)(?P['\"]?)(?P{credential_key})(?P=quote)" + r"\s*(?P[:=])\s*" + r"(?:(?P['\"])(?P.*?)(?P=value_quote)|" + r"(?P[^\s,;&}]+))", + lambda match: ( + f"{match.group('quote')}{match.group('key')}{match.group('quote')}" + f"{match.group('sep')}" + ), + value, + ) + return value[:limit] + + +def error_category(reason: str) -> str: + for needle, category in ( + ("expire_at is invalid", "expire_at_invalid"), + ("InvalidParameter", "invalid_parameter"), + ("rate limit", "rate_limit"), + ("QuotaExceeded", "quota_exceeded"), + ("timed out", "timeout"), + ("Timeout", "timeout"), + ("PreviousResponseNotFound", "previous_response_not_found"), + ): + if needle.lower() in reason.lower(): + return category + return "unclassified" + + +def _error_record( + request_id: str, + worker: int, + session_id: str, + started_iso: str, + started: float, + exc: Exception, + *, + require_tool: bool, + run_start: float, + deadline: float, + first_delta: float | None = None, + tool_uses: int = 0, + tool_results: int = 0, + tool_use_ids: list[str] | None = None, + tool_result_ids: list[str] | None = None, + tool_error_ids: list[str] | None = None, + tool_error_details: list[str] | None = None, +) -> RequestRecord: + reason = str(exc) + etype = type(exc).__name__ + status = getattr(exc, "status_code", None) + if status is not None: + etype = f"{etype}[{status}]" + return RequestRecord( + request_id=request_id, + worker=worker, + session_id=session_id, + started_at=started_iso, + completed_at=now_iso(), + started_offset_seconds=round(started - run_start, 3), + completed_offset_seconds=round(time.monotonic() - run_start, 3), + completion_phase=("formal" if time.monotonic() <= deadline else "drain"), + elapsed_seconds=time.monotonic() - started, + first_delta_seconds=first_delta, + ok=False, + require_tool=require_tool, + tool_uses=tool_uses, + tool_results=tool_results, + tool_use_ids=list(tool_use_ids or []), + tool_result_ids=list(tool_result_ids or []), + tool_error_ids=list(tool_error_ids or []), + tool_error_details=list(tool_error_details or []), + error_type=etype, + error_category=error_category(reason), + error_reason=sanitize_error_detail(reason), + ) + + +async def _bounded_pages(api, *args, **options) -> list[Any]: + """Read explicit public pages; never accept a partial/repeating snapshot.""" + rows = [] + seen = set() + for _ in range(100): + response = await api(*args, limit=100, **options) + data = field_value(response, "data") + missing = object() + page = field_value(response, "next_page", missing) + if not isinstance(data, list) or page is missing: + raise RuntimeError("incomplete public snapshot response") + rows.extend(data) + if page is None: + if field_value(response, "has_more", False): + raise RuntimeError("snapshot has_more without next_page") + return rows + if not isinstance(page, str) or not page or page in seen: + raise RuntimeError("invalid/repeated snapshot page") + seen.add(page) + options["page"] = page + raise RuntimeError("snapshot exceeded 100 pages") + + +async def verify_turn_settled( + client, + session_id: str, + environment_id: str, + *, + input_id: str, + input_sequence: int, + idle_sequence: int, + evidence: dict, + timeout: float, +) -> None: + """Require a closed model-Work graph, then recheck it and canonical events. + + ma-infra stops a Work and creates its successor in one Session-locked + transaction. Two complete matching Work/event snapshots avoid accepting a + pagination race. An active/queued successor keeps the gate open until the + deadline; elapsed quiet time alone is never success. Tool Work can retain + its sandbox lease after the model finishes (runtime_type is model-only). + """ + evidence.update( + input_event_id=input_id, + input_sequence=input_sequence, + idle_sequence=idle_sequence, + settled=False, + work=[], + events=[], + ) + previous = None + known = set() + async with asyncio.timeout(timeout): + while True: + rows = await _bounded_pages( + client.beta.environments.work.list, environment_id + ) + owned = [ + row + for row in rows + if field_value(field_value(row, "data", {}), "id") == session_id + ] + work = [ + { + **{ + key: field_value(row, key) + for key in ( + "id", + "state", + "runtime_type", + "created_at", + "stopped_at", + ) + }, + "worker_id": field_value( + field_value(row, "metadata", {}), "managed_agent_worker_id" + ), + } + for row in owned + ] + evidence["work"] = work + ids = {row["id"] for row in work} + if None in ids or not known.issubset(ids): + raise RuntimeError("completion snapshot lost known Work IDs") + known.update(ids) + events = await _bounded_pages( + client.beta.sessions.events.list, + session_id, + order="asc", + page=str(input_sequence - 1), + ) + summaries = [ + { + "id": field_value(e, "id"), + "type": field_value(e, "type"), + "sequence": field_value(e, "seq", field_value(e, "sequence")), + "tool_use_id": field_value(e, "tool_use_id"), + "stop_reason": field_value( + field_value(e, "stop_reason", {}), "type" + ), + } + for e in events + ] + evidence["events"] = summaries + seen_input = seen_idle = False + signatures = {} + last = input_sequence - 1 + for e in summaries: + seq = e["sequence"] + if isinstance(seq, bool) or not str(seq).isdigit(): + raise RuntimeError("completion event has no canonical sequence") + seq = int(seq) + signature = json.dumps(e, sort_keys=True, default=str) + if seq in signatures: + if signatures[seq] != signature: + raise RuntimeError("conflicting completion event sequence") + continue + if seq <= last: + raise RuntimeError("completion event sequence moved backwards") + signatures[seq] = signature + last = seq + if ( + seq == input_sequence + and e["id"] == input_id + and e["type"] == "user.message" + ): + seen_input = True + elif e["type"] in {"user.message", "user.interrupt"}: + raise RuntimeError( + "unexpected concurrent input during completion gate" + ) + if seq == idle_sequence and e["type"] == "session.status_idle": + if e["stop_reason"] != "end_turn": + raise RuntimeError("completion idle is not end_turn") + seen_idle = True + if seq > idle_sequence and ( + e["type"].startswith( + ( + "agent.", + "span.model_request", + "user.tool", + "user.custom_tool", + ) + ) + or e["type"] + in { + "session.status_running", + "session.status_idle", + "session.error", + "session.status_terminated", + } + ): + raise RuntimeError( + "new execution after end_turn idle without a new input" + ) + if not seen_input or not seen_idle: + raise RuntimeError("completion history lost input/idle boundary") + model = [row for row in work if row["runtime_type"]] + closed = bool(model) and all(row["state"] == "stopped" for row in model) + signature = json.dumps([work, summaries], sort_keys=True, default=str) + evidence["observed_at"] = now_iso() + if closed and signature == previous: + evidence["settled"] = True + return + previous = signature if closed else None + # Polling waits for an actual Work transition, never a quiet timer. + if not closed: + await asyncio.sleep(0.1) + + +async def run_one_turn( + client: anthropic.AsyncAnthropic, + session_id: str, + *, + marker: str, + prompt: str, + timeout: float, + require_tool: bool, + worker: int, + request_id: str, + run_start: float, + deadline: float, + environment_id: str | None = None, +) -> RequestRecord: + started = time.monotonic() + started_iso = now_iso() + first_delta: float | None = None + tool_use_ids: list[str] = [] + tool_result_ids: list[str] = [] + tool_error_ids: list[str] = [] + tool_error_details: list[str] = [] + final_parts: list[str] = [] + saw_idle = False + completion: dict[str, Any] = {} + input_id = None + input_sequence = idle_sequence = None + try: + stream = await client.beta.sessions.events.stream( + session_id, event_deltas=["agent.message"] + ) + accepted = await client.beta.sessions.events.send( + session_id, + events=[ + {"type": "user.message", "content": [{"type": "text", "text": prompt}]} + ], + ) + if environment_id: + inputs = field_value(accepted, "data", []) + if len(inputs) != 1: + raise RuntimeError("missing canonical input acknowledgement") + input_id = field_value(inputs[0], "id") + input_sequence = field_value( + inputs[0], "seq", field_value(inputs[0], "sequence") + ) + if not input_id or not str(input_sequence).isdigit(): + raise RuntimeError("missing canonical input identity/sequence") + async with stream: + async with asyncio.timeout(timeout): + async for event in stream: + et = str(field_value(event, "type", "")) + if et == "event_delta" and first_delta is None: + delta = field_value(event, "delta") + if content_text([field_value(delta, "content")]): + first_delta = time.monotonic() - started + elif et == "agent.message": + final_parts.append(content_text(field_value(event, "content"))) + elif et == "agent.tool_use": + tool_id = str(field_value(event, "id", "") or "") + if tool_id: + tool_use_ids.append(tool_id) + elif et == "agent.tool_result": + tool_id = str(field_value(event, "tool_use_id", "") or "") + if tool_id: + tool_result_ids.append(tool_id) + if field_value(event, "is_error", False) is True: + tool_error_ids.append(tool_id) + detail = content_text(field_value(event, "content")) + tool_error_details.append( + sanitize_error_detail( + detail or "tool execution failed" + ) + ) + elif et in {"session.error", "session.status_terminated"}: + detail = "" + for key in ("error", "message", "content", "reason"): + val = field_value(event, key) + if val: + detail = str(val) + break + raise RuntimeError(f"{et}: {detail[:300]}") + if et == "session.status_idle": + if environment_id: + if ( + field_value( + field_value(event, "stop_reason", {}), "type" + ) + != "end_turn" + ): + raise RuntimeError("idle is not end_turn") + idle_sequence = field_value( + event, "seq", field_value(event, "sequence") + ) + saw_idle = True + break + if environment_id and saw_idle: + await verify_turn_settled( + client, + session_id, + environment_id, + input_id=input_id, + input_sequence=int(input_sequence), + idle_sequence=int(idle_sequence), + evidence=completion, + timeout=max(0.001, timeout - (time.monotonic() - started)), + ) + final_text = "\n".join(p for p in final_parts if p) + ok = saw_idle and (marker in final_text) + if require_tool: + ok = ( + ok + and bool(tool_use_ids) + and len(tool_use_ids) == len(set(tool_use_ids)) + and set(tool_use_ids) == set(tool_result_ids) + and len(tool_result_ids) == len(set(tool_result_ids)) + and not tool_error_ids + ) + if ok: + category = reason = None + elif tool_error_ids: + category = "tool_result_error" + reason = f"tool_result reported is_error=true for IDs: {','.join(tool_error_ids)}" + else: + category = "turn_assertion_failed" + reason = "assertion: incomplete turn or marker/tool missing" + return RequestRecord( + request_id=request_id, + worker=worker, + session_id=session_id, + started_at=started_iso, + completed_at=now_iso(), + started_offset_seconds=round(started - run_start, 3), + completed_offset_seconds=round(time.monotonic() - run_start, 3), + completion_phase=("formal" if time.monotonic() <= deadline else "drain"), + elapsed_seconds=time.monotonic() - started, + first_delta_seconds=first_delta, + ok=ok, + require_tool=require_tool, + tool_uses=len(tool_use_ids), + tool_results=len(tool_result_ids), + tool_use_ids=tool_use_ids, + tool_result_ids=tool_result_ids, + tool_error_ids=tool_error_ids, + tool_error_details=tool_error_details, + error_type=None if ok else "AssertionError", + error_category=category, + error_reason=reason, + completion=completion or None, + ) + except Exception as exc: # noqa: BLE001 + record = _error_record( + request_id, + worker, + session_id, + started_iso, + started, + exc, + require_tool=require_tool, + run_start=run_start, + deadline=deadline, + first_delta=first_delta, + tool_uses=len(tool_use_ids), + tool_results=len(tool_result_ids), + tool_use_ids=tool_use_ids, + tool_result_ids=tool_result_ids, + tool_error_ids=tool_error_ids, + tool_error_details=tool_error_details, + ) + record.completion = completion or None + return record + + +async def worker_loop( + client: anthropic.AsyncAnthropic, + worker: int, + *, + deadline: float, + args: argparse.Namespace, + records: list[RequestRecord], + lock: asyncio.Lock, + active: set[str], + in_flight: dict[str, str], + created: list[str], + writer: DurableJsonlWriter, + cleanup_records: list[dict[str, str | None]], + cleanup_writer: DurableJsonlWriter, + cleanup_tasks: list[asyncio.Task[None]], + peak_in_flight: list[int], + run_start: float, + run_id: str, +) -> None: + """Continuously replenish requests until deadline. + + A fresh logical Session is created for every turn. Reusing a Session and + re-opening its event stream replays prior history (including an earlier + ``session.status_idle``), which would end a turn instantly; a new Session + per turn avoids that and keeps physical-sandbox concurrency honest. + """ + while time.monotonic() < deadline: + request_id = uuid.uuid4().hex + async with lock: + in_flight[request_id] = "" + peak_in_flight[0] = max(peak_in_flight[0], len(in_flight)) + marker = f"w{worker:02d}-{uuid.uuid4().hex[:10]}" + if args.require_tool: + prompt = ( + f"Use the bash tool exactly once to run: printf {marker}. " + f"Then include exactly {marker} in your final reply." + ) + else: + prompt = f"Do not use any tool. Reply with exactly: {marker}" + started_iso = now_iso() + started = time.monotonic() + session_id = "" + try: + session = await client.beta.sessions.create( + agent=args.agent_id, + environment_id=args.environment_id, + title=f"soak {run_id} w{worker:02d} {uuid.uuid4().hex[:8]}", + ) + session_id = str(session.id) + async with lock: + active.add(session_id) + in_flight[request_id] = session_id + created.append(session_id) + except Exception as exc: # noqa: BLE001 + record = _error_record( + request_id, + worker, + session_id, + started_iso, + started, + exc, + require_tool=args.require_tool, + run_start=run_start, + deadline=deadline, + ) + async with lock: + record.cleanup_status = "not_created" + records.append(record) + writer.append(record) + in_flight.pop(request_id, None) + if args.request_gap > 0: + await asyncio.sleep(args.request_gap) + continue + + record = await run_one_turn( + client, + session_id, + marker=marker, + prompt=prompt, + timeout=args.turn_timeout, + require_tool=args.require_tool, + worker=worker, + request_id=request_id, + run_start=run_start, + deadline=deadline, + environment_id=args.environment_id, + ) + # Persist the complete turn before cleanup so a slow or failed archive + # cannot erase successful request evidence. Cleanup has its own ledger. + record.cleanup_status = "pending" + async with lock: + records.append(record) + writer.append(record) + in_flight.pop(request_id, None) + cleanup_tasks.append( + asyncio.create_task( + cleanup_session( + client, + session_id, + record, + lock=lock, + active=active, + cleanup_records=cleanup_records, + cleanup_writer=cleanup_writer, + archive_timeout=args.archive_timeout, + ) + ) + ) + if args.request_gap > 0: + await asyncio.sleep(args.request_gap) + + +async def cleanup_session( + client: anthropic.AsyncAnthropic, + session_id: str, + record: RequestRecord, + *, + lock: asyncio.Lock, + active: set[str], + cleanup_records: list[dict[str, Any]], + cleanup_writer: DurableJsonlWriter, + archive_timeout: float, +) -> None: + try: + async with asyncio.timeout(archive_timeout): + await client.beta.sessions.archive(session_id, timeout=archive_timeout) + record.cleanup_status = "archived" + except Exception as exc: # noqa: BLE001 + record.cleanup_status = "archive_failed" + record.cleanup_error = sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ) + + cleanup_record: dict[str, Any] = { + "session_id": session_id, + "status": record.cleanup_status, + "error": record.cleanup_error, + "archive_completed_at": now_iso(), + "physical_query_started_at": record.physical_query_started_at, + "physical_query_completed_at": record.physical_query_completed_at, + "physical_reclaimed": None, + "physical_query_count": 0, + } + async with lock: + cleanup_records.append(cleanup_record) + cleanup_writer.append(cleanup_record) + + +async def confirm_physical_reclaims( + inspector: AgentKitInspector | None, + active: set[str], + records: list[RequestRecord], + cleanup_records: list[dict[str, Any]], + cleanup_writer: DurableJsonlWriter, + *, + timeout: float, + interval: float = 10.0, +) -> None: + by_session = {record.session_id: record for record in records if record.session_id} + cleanup_by_session = {item["session_id"]: item for item in cleanup_records} + if inspector is None: + for session_id in sorted(active): + record = by_session[session_id] + record.cleanup_status += "_physical_unverified" + cleanup_by_session[session_id]["status"] = record.cleanup_status + cleanup_writer.append(cleanup_by_session[session_id]) + active.discard(session_id) + return + + started = time.monotonic() + pending = set(active) + while pending and time.monotonic() - started < timeout: + query_started_at = now_iso() + rows = await run_scoped_physical_sessions(inspector, pending) + query_completed_at = now_iso() + present = {item["user_session_id"] for item in rows} + for session_id in sorted(pending): + record = by_session[session_id] + record.physical_query_started_at = ( + record.physical_query_started_at or query_started_at + ) + record.physical_query_completed_at = query_completed_at + record.physical_query_count += 1 + reclaimed = pending - present + for session_id in reclaimed: + record = by_session[session_id] + record.physical_reclaimed = True + record.cleanup_status += "_physical_absent" + item = cleanup_by_session[session_id] + item.update( + status=record.cleanup_status, + physical_query_started_at=record.physical_query_started_at, + physical_query_completed_at=record.physical_query_completed_at, + physical_query_count=record.physical_query_count, + physical_reclaimed=True, + ) + cleanup_writer.append(item) + active.discard(session_id) + pending -= reclaimed + if pending: + await asyncio.sleep(interval) + + for session_id in sorted(pending): + record = by_session[session_id] + record.physical_reclaimed = False + record.cleanup_status += "_physical_unconfirmed" + detail = f"physical sandbox still present after {timeout:.1f}s" + record.cleanup_error = ( + f"{record.cleanup_error}; {detail}" if record.cleanup_error else detail + ) + item = cleanup_by_session[session_id] + item.update( + status=record.cleanup_status, + error=record.cleanup_error, + physical_query_started_at=record.physical_query_started_at, + physical_query_completed_at=record.physical_query_completed_at, + physical_query_count=record.physical_query_count, + physical_reclaimed=False, + ) + cleanup_writer.append(item) + + +async def sample_physical( + inspector: AgentKitInspector | None, + active: set[str], + in_flight: dict[str, str], + lock: asyncio.Lock, + stop: asyncio.Event, + samples: list[dict[str, Any]], + writer: DurableJsonlWriter, + interval: float, +) -> None: + if inspector is None: + return + while not stop.is_set(): + watch: set[str] = set() + inflight_snapshot: dict[str, str] = {} + try: + async with lock: + inflight_snapshot = dict(in_flight) + watch = set(active) + snapshot_started_at = now_iso() + rows = await run_scoped_physical_sessions(inspector, watch) + current = { + item["user_session_id"]: item + for item in rows + if item["user_session_id"] in watch + } + sample = { + "snapshot_started_at": snapshot_started_at, + "snapshot_completed_at": now_iso(), + "in_flight_requests": len(inflight_snapshot), + "in_flight": inflight_snapshot, + "logical_session_ids": sorted(watch), + "physical_sessions": sorted( + current.values(), key=lambda item: item["session_id"] + ), + "error": None, + } + samples.append(sample) + writer.append(sample) + except Exception as exc: # noqa: BLE001 + sample = { + "snapshot_started_at": snapshot_started_at, + "snapshot_completed_at": now_iso(), + "in_flight_requests": len(inflight_snapshot), + "in_flight": inflight_snapshot, + "logical_session_ids": sorted(watch), + "physical_sessions": [], + "error": sanitize_error_detail( + f"{type(exc).__name__}: {exc}", limit=300 + ), + } + samples.append(sample) + writer.append(sample) + try: + await asyncio.wait_for(stop.wait(), timeout=interval) + except asyncio.TimeoutError: + pass + + +async def run_scoped_physical_sessions( + inspector: AgentKitInspector, managed_session_ids: set[str] +) -> list[dict[str, str]]: + """List only physical sessions bound to this run's active logical IDs.""" + if not managed_session_ids: + return [] + custom = getattr(inspector, "sessions_for_managed_sessions", None) + if callable(custom): + rows = await custom(managed_session_ids) + return _validated_physical_rows(rows, managed_session_ids) + + sessions: list[dict[str, str]] = [] + for managed_session_chunk in _chunks(sorted(managed_session_ids), 50): + sessions.extend( + await _list_physical_session_chunk(inspector, managed_session_chunk) + ) + return _validated_physical_rows(sessions, managed_session_ids) + + +def _validated_physical_rows( + rows: Any, managed_session_ids: set[str] +) -> list[dict[str, str]]: + if not isinstance(rows, list): + raise RuntimeError("physical session query must return a list") + for item in rows: + if not isinstance(item, dict): + raise RuntimeError("physical session item must be an object") + if not item.get("session_id"): + raise RuntimeError("physical session item is missing session_id") + if item.get("user_session_id") not in managed_session_ids: + raise RuntimeError( + f"physical query returned unexpected UserSessionId: {item.get('user_session_id')}" + ) + return rows + + +def _chunks(values: list[str], size: int) -> list[list[str]]: + return [values[index : index + size] for index in range(0, len(values), size)] + + +async def _list_physical_session_chunk( + inspector: AgentKitInspector, managed_session_ids: list[str] +) -> list[dict[str, str]]: + sessions: list[dict[str, str]] = [] + next_token = None + seen_tokens: set[str] = set() + while True: + body: dict[str, Any] = { + "ToolId": inspector.tool_id, + "MaxResults": 100, + "Filters": [ + { + "Name": "UserSessionId", + "Values": managed_session_ids, + } + ], + } + if next_token: + body["NextToken"] = next_token + payload = await asyncio.to_thread(inspector.client.call, "ListSessions", body) + if not isinstance(payload, dict): + raise RuntimeError("ListSessions response must be an object") + raw_sessions = payload.get("SessionInfos") + if not isinstance(raw_sessions, list): + raise RuntimeError("ListSessions response must contain SessionInfos list") + for item in raw_sessions: + if not isinstance(item, dict): + raise RuntimeError("ListSessions SessionInfos item must be an object") + session_id = item.get("SessionId") + user_session_id = item.get("UserSessionId") + if not isinstance(session_id, str) or not session_id: + raise RuntimeError("ListSessions item is missing SessionId") + if user_session_id not in managed_session_ids: + raise RuntimeError( + f"ListSessions returned unexpected UserSessionId: {user_session_id}" + ) + sessions.append( + { + "session_id": session_id, + "user_session_id": user_session_id, + "status": str(item.get("Status") or ""), + "created_at": str(item.get("CreatedAt") or ""), + } + ) + next_token = payload.get("NextToken") + if next_token is None or next_token == "": + return sessions + if not isinstance(next_token, str): + raise RuntimeError("ListSessions NextToken must be a string or null") + if next_token in seen_tokens: + raise RuntimeError("AgentKit ListSessions repeated NextToken") + seen_tokens.add(next_token) + + +async def async_main(args: argparse.Namespace) -> dict[str, Any]: + run_id = args.run_id or ( + f"run-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}-{uuid.uuid4().hex[:8]}" + ) + inspector = None + if args.with_inspector and args.tool_id: + inspector = AgentKitInspector( + os.getenv( + "AGENTKIT_TOOL_DEPLOY_SCRIPT", + "", + ), + args.tool_id, + ) + + records: list[RequestRecord] = [] + lock = asyncio.Lock() + active: set[str] = set() + in_flight: dict[str, str] = {} + created: list[str] = [] + cleanup_records: list[dict[str, Any]] = [] + cleanup_tasks: list[asyncio.Task[None]] = [] + peak_in_flight = [0] + physical_samples: list[dict[str, Any]] = [] + started_iso = now_iso() + run_start = time.monotonic() + deadline = run_start + args.duration + + records_output = args.records_output + if records_output is None and args.output is not None: + records_output = args.output.with_suffix(".records.jsonl") + samples_output = args.samples_output + if samples_output is None and args.output is not None: + samples_output = args.output.with_suffix(".samples.jsonl") + cleanup_output = args.cleanup_output + if cleanup_output is None and args.output is not None: + cleanup_output = args.output.with_suffix(".cleanup.jsonl") + with ( + DurableJsonlWriter(records_output) as writer, + DurableJsonlWriter(samples_output) as sample_writer, + DurableJsonlWriter(cleanup_output) as cleanup_writer, + ): + async with anthropic.AsyncAnthropic( + base_url=args.base_url.rstrip("/"), + auth_token=args.auth_token, + timeout=args.turn_timeout, + max_retries=0, + ) as client: + stop = asyncio.Event() + sampler = asyncio.create_task( + sample_physical( + inspector, + active, + in_flight, + lock, + stop, + physical_samples, + sample_writer, + args.sample_interval, + ) + ) + + loops = [ + worker_loop( + client, + index, + deadline=deadline, + args=args, + records=records, + lock=lock, + active=active, + in_flight=in_flight, + created=created, + writer=writer, + cleanup_records=cleanup_records, + cleanup_writer=cleanup_writer, + cleanup_tasks=cleanup_tasks, + peak_in_flight=peak_in_flight, + run_start=run_start, + run_id=run_id, + ) + for index in range(args.concurrency) + ] + await asyncio.gather(*loops) + if cleanup_tasks: + await asyncio.gather(*cleanup_tasks) + stop.set() + await sampler + await confirm_physical_reclaims( + inspector, + active, + records, + cleanup_records, + cleanup_writer, + timeout=args.physical_reclaim_timeout, + ) + + actual = time.monotonic() - run_start + successes = sum(1 for r in records if r.ok) + failures = len(records) - successes + latencies = [r.elapsed_seconds for r in records if r.ok] + all_latencies = [r.elapsed_seconds for r in records] + first_deltas = [ + r.first_delta_seconds + for r in records + if r.ok and r.first_delta_seconds is not None + ] + errors = Counter( + f"{r.error_category}: {r.error_type}" + for r in records + if not r.ok and r.error_type + ) + + def stats(values: list[float]) -> dict[str, float | None]: + if not values: + return {k: None for k in ("min", "mean", "p50", "p90", "p95", "p99", "max")} + return { + "min": round(min(values), 3), + "mean": round(statistics.fmean(values), 3), + "p50": round(percentile(values, 50), 3), + "p90": round(percentile(values, 90), 3), + "p95": round(percentile(values, 95), 3), + "p99": round(percentile(values, 99), 3), + "max": round(max(values), 3), + } + + summary = Summary( + started_at=started_iso, + completed_at=now_iso(), + duration_seconds_target=args.duration, + duration_seconds_actual=round(actual, 1), + concurrency=args.concurrency, + require_tool=args.require_tool, + total_requests=len(records), + successes=successes, + failures=failures, + success_rate=round(successes / len(records), 4) if records else 0.0, + request_throughput_per_min=round(len(records) / (actual / 60.0), 2) + if actual > 0 + else 0.0, + success_throughput_per_min=round(successes / (actual / 60.0), 2) + if actual > 0 + else 0.0, + formal_window_seconds=args.duration, + drain_seconds=round(max(0.0, actual - args.duration), 1), + formal_completions=sum(r.completion_phase == "formal" for r in records), + drain_completions=sum(r.completion_phase == "drain" for r in records), + peak_in_flight_requests=peak_in_flight[0], + success_latency=stats(latencies), + latency_including_failures=stats(all_latencies), + first_text_event=stats(first_deltas), + error_breakdown=dict(errors), + observed_peak_physical_sessions=max( + (len(sample["physical_sessions"]) for sample in physical_samples), + default=0, + ), + physical_session_ids=sorted( + { + item["session_id"] + for sample in physical_samples + for item in sample["physical_sessions"] + } + ), + physical_samples=len(physical_samples), + physical_sample_errors=sum( + bool(sample["error"]) for sample in physical_samples + ), + physical_reclaim_confirmed=sum( + item.get("physical_reclaimed") is True for item in cleanup_records + ), + physical_reclaim_unconfirmed=sum( + item.get("physical_reclaimed") is False for item in cleanup_records + ), + physical_reclaim_unverified=sum( + item.get("physical_reclaimed") is None for item in cleanup_records + ), + total_sessions_created=len(created), + sample_sessions_created=created[:10], + ) + + return { + "summary": asdict(summary), + "run_id": run_id, + "base_url": args.base_url, + "agent_id": args.agent_id, + "environment_id": args.environment_id, + "records_output": str(records_output) if records_output else None, + "samples_output": str(samples_output) if samples_output else None, + "cleanup_output": str(cleanup_output) if cleanup_output else None, + "records": [asdict(r) for r in records], + "physical_samples": physical_samples, + "cleanup_records": cleanup_records, + "unreclaimed_logical_session_ids": sorted(active), + } + + +def parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--base-url", default=os.getenv("MANAGED_AGENTS_GATEWAY")) + p.add_argument("--agent-id", default=os.getenv("MANAGED_AGENTS_TEST_AGENT_ID")) + p.add_argument("--environment-id", default=os.getenv("ANTHROPIC_ENVIRONMENT_ID")) + p.add_argument( + "--auth-token", + default=os.getenv("MANAGED_AGENTS_GATEWAY_TOKEN", "gateway-client"), + ) + p.add_argument("--concurrency", type=int, default=50) + p.add_argument( + "--duration", type=float, default=1200, help="seconds (default 20min)" + ) + p.add_argument("--turn-timeout", type=float, default=180) + p.add_argument( + "--request-gap", + type=float, + default=0.0, + help="seconds to wait between a worker's turns (0 = replenish immediately)", + ) + p.add_argument( + "--require-tool", + action="store_true", + help="each turn must run one bash tool (exercises sandbox path)", + ) + p.add_argument( + "--with-inspector", + action="store_true", + help="sample AgentKit physical sandbox Sessions for peak concurrency", + ) + p.add_argument("--tool-id", default=os.getenv("AGENTKIT_TOOL_ID")) + p.add_argument("--sample-interval", type=float, default=30.0) + p.add_argument("--physical-reclaim-timeout", type=float, default=120.0) + p.add_argument( + "--archive-timeout", + type=float, + default=30.0, + help="total wall-clock limit per Session archive, including SDK retries", + ) + p.add_argument( + "--run-id", help="unique non-secret run label (generated by default)" + ) + p.add_argument("--output", type=Path, required=True) + p.add_argument( + "--records-output", + type=Path, + help="durable per-turn JSONL (default: .records.jsonl)", + ) + p.add_argument( + "--samples-output", + type=Path, + help="durable low-frequency sample JSONL (default: .samples.jsonl)", + ) + p.add_argument( + "--cleanup-output", + type=Path, + help="durable cleanup JSONL (default: .cleanup.jsonl)", + ) + p.add_argument("--progress-interval", type=float, default=60.0) + return p + + +async def _with_progress(args: argparse.Namespace) -> dict[str, Any]: + return await async_main(args) + + +def main() -> int: + args = parser().parse_args() + if not args.base_url or not args.agent_id or not args.environment_id: + print("base-url, agent-id and environment-id are required", file=sys.stderr) + return 2 + if args.with_inspector and not args.tool_id: + print("--with-inspector requires --tool-id", file=sys.stderr) + return 2 + result = asyncio.run(_with_progress(args)) + s = result["summary"] + encoded = json.dumps(result, ensure_ascii=False, indent=2) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(encoded + "\n") + # Human-readable summary to stdout; full JSON only to file when requested. + print(json.dumps(s, ensure_ascii=False, indent=2)) + cleanup_failed = any( + not str(item.get("status") or "").startswith("archived") + for item in result["cleanup_records"] + ) + physical_incomplete = bool(result["unreclaimed_logical_session_ids"]) + if args.with_inspector: + physical_incomplete = ( + physical_incomplete + or result["summary"]["physical_sample_errors"] > 0 + or any( + item.get("physical_reclaimed") is not True + for item in result["cleanup_records"] + ) + ) + return 1 if cleanup_failed or physical_incomplete else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md b/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md new file mode 100644 index 000000000..129a188a5 --- /dev/null +++ b/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md @@ -0,0 +1,192 @@ +# Non-UI Managed Agents upstream migration + +[中文版](2026-10-10-non-ui-migration.zh.md) + +Change ID: `managed-agents-upstream`. Created/revised: 2026-10-10. Status: implemented (local migration; unpublished). + +## Background and evidence + +The source repository's `main` at `89a90e02` and the official VeADK `main` at +`171d8d86` diverge from `80e968d7`. Comparing complete trees would overwrite +official Codex, Studio, MPA provisioning, and compatibility improvements. +The source's new runtime implementation lives in `examples/16_self_host_sandbox`; +its library modules import bare example modules and its images require an +adjacent Managed Agents Anthropic SDK checkout. The target already owns +`veadk/runtime/provider.py` and `veadk/integrations/mpa/managed` provisioning. + +## Goals and non-goals + +Deliver all source-owned non-UI changes on the current official base, with a +packaged runtime, maintained example entry points, isolated tests, and reproducible +image inputs. Source is read-only. UI, Studio-only backend routes/repositories, +built web assets, unrelated upstream code, personal deployment records, and live +cloud mutation are excluded. This does not authorize commit, push, or PR creation. + +## Scenarios and requirements + +- `FR-1`: A caller imports the packaged loop/client without an example directory, + credentials, network activity, or a global remote-backed demo Agent. +- `FR-2`: A worker executes native tools, pinned Session Skills and MCP resources, + retaining Work polling, lease/heartbeat, cancellation, cross-worker replay, + credential rotation, event ordering, and cleanup behavior from the source. +- `FR-3`: Ark continuation retains tool declarations, uses server retention by + default, preserves explicit expiration, and emits opt-in safe metadata traces. +- `FR-4`: Runtime Identity supports explicit pool lookup and Session-scoped model + and MCP credentials; Skills operations preserve their bounded account scope. +- `FR-5`: Feishu streaming controls are opt-in and preserve the official branch's + graceful shutdown and application-loop lifecycle. +- `FR-6`: Non-UI build, example, and test scripts point to installed library code. + Images use public package indexes by default, keep locked hashes, never require + a hardcoded registry or local SDK checkout, and fail early when SDK capability + requirements are unmet. Image publishing and cloud E2E remain separate actions. + +## Design and contract impact + +Dependency direction is runtime -> MPA Session integration -> Anthropic/AgentKit +SDK. `veadk/runtime/managed_agents` owns worker, loop, native tools, event tracing, +Identity adapter and health entry point. `veadk/integrations/mpa/session_client.py` +and `session_resources.py` own the protocol client and Session resource snapshots. +Existing MPA provisioning is reused without confusing its `managed/worker.py` +(sandbox provisioning) with model-work execution. Existing RuntimeProvider stays +unchanged. Compatibility wrappers in the example import packaged implementations; +example launchers may load their local environment, library imports may not. +Deployment-specific IDs and defaults are replaced by required configuration. +`veadk/models`, `veadk/skills`, `veadk/integrations/ve_identity`, and +`veadk/extensions/feishu_channel.py` retain their established ownership. + +State/event wire formats remain those of the source; no backend schema or cloud +resource migration is introduced. Cancellation must unwind runners/tools and +temporary Skill materializations; terminal replay must not execute completed +inputs again. Credentials remain runtime inputs, are never persisted in docs, +images, traces, or Git, and error reporting must remain redacted. Supported +Python 3.10-3.13 imports must use compatible datetime and asyncio primitives. + +Component contracts: [runtime](../../../specs/managed-agent-runtime/README.md), +[resources](../../../specs/managed-agent-resources/README.md), +[Feishu](../../../specs/feishu-channel/README.md). + +## Tasks and ownership + +| Task | Requirements | Files / owner | +| --- | --- | --- | +| `T-1` | `FR-1`, `FR-2` | Runtime/Session packages, runtime tests, example Python wrappers: runtime agent | +| `T-2` | `FR-3`, `FR-4` | Ark, IdentityClient, Skills client, model/resource tests: model/skills agent | +| `T-3` | `FR-5` | Feishu extension, channel tests and paired user docs: channel agent | +| `T-4` | `FR-6` | Docker, dependency metadata/lock, non-UI scripts, example docs: root agent | +| `T-5` | all | Integration, regression, packaging, secret scan, documentation reconciliation: root agent | + +Writers have disjoint file ownership. No shared Git index operations or commits +occur while agents work. Root reconciles package metadata and runs final gates. + +## Verification and acceptance + +| Requirements | Acceptance | Verification | Status | +| --- | --- | --- | --- | +| `FR-1`, `FR-2` | `AC-1`: installable library with equivalent lifecycle | targeted runtime/Session and actual SDK transport/native-tool tests | pass; included in final 387 focused tests | +| `FR-3`, `FR-4` | `AC-2`: existing official behavior plus source resource/model delta | model/Identity/Skills tests | pass; 109 initially, included in final focused run | +| `FR-5` | `AC-3`: streaming and shutdown tests pass | channel tests | pass; 18, included in final focused run | +| all | `AC-4`: integrated regression and repository gates | parallel regression plus all-file/changed-file pre-commit | pass for scoped regression; unfiltered gate has 7 reproduced official UI failures (see below) | +| `FR-6` | `AC-5`: wheel/import/CLI/build-input validation | wheel/import/CLI, shell syntax, image build and container checks | pass; local image only | + +Tests use synthetic services and temporary state. Live cloud/provider E2E and +image publication: not_run, outside the authorized local migration. UI checks: +not_applicable because UI and Studio backend changes are excluded. + +## Risks and recovery + +The source's EnvironmentWorkDispatcher currently requires a Managed Agents SDK +capability not present in a generic Anthropic installation. Check available SDK +capabilities and implement a supported explicit integration, rather than silently +claiming worker readiness. The final record must distinguish simulated runtime +tests, installed package imports, image build, and real cloud E2E. Original source +and official base remain intact; abandoning the new branch restores the base. + +## Review and delivery record + +The user approved the architecture and non-UI scope with “ok, just do it, do it +best” and explicitly authorized multiple subagents. Direct design review found +and resolved: use the official base; preserve existing MPA and runtime interfaces; +exclude UI backend wiring; keep SDK requirements explicit; avoid import-time +credential access; preserve Python 3.10 compatibility; validate async cleanup. +No `review-spec` skill is installed, so the repository's equivalent direct review +is used. Component owners review their bilingual contracts before production +edits. All `T-1` through `T-5` tasks are complete within the non-UI scope. + +### Source-to-target reconciliation + +| Source | Delivered location / decision | +| --- | --- | +| `main.py`, `managed_agent_loop.py`, `event_debug.py`, `runtime_identity.py` | `veadk/runtime/managed_agents/{worker,loop,events,identity}.py` | +| demo agent/session lifecycle | packaged `sandbox.py`; example discovery factory remains a thin adapter | +| `sandbox_client.py`, `managed_session_resources.py` | `veadk/integrations/mpa/{session_client,session_resources}.py` | +| missing private SDK dispatcher | packaged dispatcher using the published SDK, with real protocol tests | +| Ark/Identity/Skills/Feishu | existing owning SDK modules plus standalone `veadk/skills/ma_infra.py` | +| Worker/Runtime launch and claimed Tool startup | `docker/managed-agents/`, installed module entry points, `/opt/gem/run.sh`, thin example launchers | +| conversation/distributed/Kubernetes/soak/fault checks | portable example scripts and packaged runtime/Session/completion-gate tests | +| personal rollout/build-candidate scripts, RDS and external ma-server/task-server deployments | replaced by public Worker build and generic Kubernetes template; no deployment-specific state copied | +| UI, Studio-only Skills repository/routes, gateway UI image and built assets | excluded as approved | +| historical deployment proof and one-off analysis documents | not copied as current evidence; this design records new local verification | + +### Verification record (2026-10-10) + +The tested scope is the uncommitted feature diff on official `171d8d86`, including +new package files. Locked package versions/hashes and unrelated platform markers +are unchanged; only the sandbox Anthropic minimum changes from `>=0.40.0` to the +source's required `>=1.3.0`. + +- **pass:** Python 3.12 final focused run: 387 tests, including actual SDK HTTP, + native file/bash, MCP, canonical cross-worker replay, cleanup and supervision. +- **pass:** Python 3.10.20 isolated frozen environment: 356 focused tests plus + actual SDK transient publishing. The discovered generic-alias SDK cast issue is + fixed using `typing.Dict`; no unsupported production import is hidden by mocks. +- **pass:** `MODEL_AGENT_API_KEY=offline-regression-test-key .venv/bin/python -m + pytest -n 4 -m "not codex_smoke and not piagent_smoke" -q --tb=short`, with the + seven explicit baseline deselections below: 6963 passed, 42 skipped, 4 xfailed, + 6 subtests passed. Four workers were selected for this final large regression; + the initial default two-worker run exposed test setup and baseline issues. +- **fail, official baseline:** the unfiltered regression's seven + `tests/frontend/test_migration_delivery_recovery.py` cases also fail on an + untouched official worktree with identical dependencies. They are: + `test_mirror_delivery_writes_the_documents_the_cli_would_have_written`, + `test_mirror_delivery_keeps_the_sequence_the_cli_already_started`, + `test_mirror_delivery_refuses_a_delivery_the_cli_already_settled`, + `test_verification_mirrors_the_cli_finding_projection`, + `test_an_unparseable_findings_file_reports_no_findings_at_all`, + `test_a_startup_fallback_degrades_the_delivery`, and + `test_a_rebuild_that_lands_settles_the_task_on_the_cli_delivery_contract`. + Each was deselected with `--deselect ::` for the scoped regression; + no marker, production UI code or assertion was changed to suppress them. +- The initial unfiltered run also exposed existing examples/Harness tests relying + on an unrelated example test's import-time dummy model key. New tests no longer + mutate global credentials during collection. The final command provides an + explicit offline fixture key instead. Test package names now avoid a + `test_dispatcher` collision with Studio scheduler tests. +- **pass:** all-file and changed-file pre-commit (Ruff check/format, gitleaks, + YAML secret scan); `git diff --check`; bilingual Markdown links and no-UI scope. +- **pass:** `uv build --wheel`; extracted installed wheel imports and CLI help + from a temporary directory with a sanitized environment and no example path. +- **pass:** `npm --prefix docs run types:check` and `npm --prefix docs run build`. + Dependency installation uses the repository's frozen pnpm lock; generated + workspace policy and Corepack's unintended parent configuration change were + removed from the feature scope. +- **pass:** local Docker image `veadk-managed-agents-worker:upstream-local` built + from locked sandbox dependencies with public indexes. Local network/proxy + settings were supplied only to the build. The dedicated Worker excludes + optional integration/development/Codex extras. No neighboring SDK wheel is used. +- **pass:** container SDK imports/CLI and `/opt/gem/run.sh` checked as UID 10001 + with network disabled and read-only filesystem; installed health server returns + HTTP 200 on `/ping` in the same isolation. These checks do not exercise a real + Work gateway, model provider or AgentKit deployment. +- **not_run:** live cloud/model E2E, image publication, Git commit/push and PR + creation. The image is local, the branch tracks `volcengine/main`, and the + original source worktree remains clean. + +### Review fixes + +Independent review resolved import-time credential access, failed-turn completion +and secret-bearing errors, Skill workspace collisions/symlink cleanup, closing all +tool contexts after a cleanup failure, Work lease fencing/Session FIFO/auth refresh, +first-valid-poll readiness, Python 3.10 casts, credential-bearing index arguments, +obsolete database assumptions in distributed tests, and named BuildKit cache +isolation. Claimed Tool execution retains its historical launch path and local +tool default inside the deployment's OS isolation. diff --git a/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md b/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md new file mode 100644 index 000000000..d7f14f836 --- /dev/null +++ b/prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md @@ -0,0 +1,165 @@ +# Managed Agents 非 UI 改动合入 + +[English](2026-10-10-non-ui-migration.md) + +变更 ID：`managed-agents-upstream`。创建/修订：2026-10-10。状态：implemented（本地迁移，未发布）。 + +## 背景与证据 + +来源仓库 `main` 的 `89a90e02` 与官方 VeADK `main` 的 `171d8d86` 从 +`80e968d7` 分叉。直接比较完整代码树会覆盖官方 Codex、Studio、MPA 资源准备和 +兼容性更新。来源的新增运行时位于 `examples/16_self_host_sandbox`，库逻辑通过 +裸模块名导入示例代码，镜像依赖相邻 Managed Agents Anthropic SDK checkout。 +目标已经拥有 `veadk/runtime/provider.py` 与 `veadk/integrations/mpa/managed`。 + +## 目标与非目标 + +在官方当前基线上交付来源全部非 UI 增量，包含包化运行时、兼容示例入口、隔离测试 +和可复现镜像输入。来源只读。排除 UI、仅服务于 Studio 的后端路由和仓储、构建后的 +网页资源、无关上游代码、个人部署记录和线上云资源修改。此范围不授权 commit、push +或创建 PR。 + +## 场景与需求 + +- `FR-1`：调用方不依赖示例目录即可导入循环与客户端；导入不需要凭据，不访问网络， + 不创建全局远程示例 Agent。 +- `FR-2`：Worker 执行原生工具、Session 固定的 Skills/MCP 资源，保留来源的 Work + 轮询、租约/心跳、取消、跨 Worker 重放、凭据轮转、事件顺序和清理行为。 +- `FR-3`：Ark 续接请求保留工具声明；默认使用服务端保留期，保留显式过期参数， + 仅在开启时记录安全的元数据跟踪。 +- `FR-4`：Runtime Identity 支持显式凭据池查询与 Session 级模型/MCP 凭据； + Skills 操作保持受限的账号作用域。 +- `FR-5`：飞书流式配置默认关闭，保留官方优雅关闭与应用事件循环生命周期。 +- `FR-6`：非 UI 构建、示例与测试脚本使用安装的库；镜像默认使用公共包源，保留 + 锁定哈希，不依赖固定镜像仓库或本地 SDK checkout，SDK 缺少所需能力时提前失败。 + 镜像发布与线上云 E2E 属于独立操作。 + +## 设计与契约影响 + +依赖方向为运行时 → MPA Session 集成 → Anthropic/AgentKit SDK。 +`veadk/runtime/managed_agents` 负责 Worker、循环、原生工具、事件跟踪、Identity +适配与健康入口；`veadk/integrations/mpa/session_client.py` 和 +`session_resources.py` 负责协议客户端与 Session 资源快照。复用现有 MPA 资源准备， +其 `managed/worker.py` 是沙箱资源创建，不能混同于模型 Work 执行。 +保持现有 RuntimeProvider。示例兼容包装导入包实现；示例启动器可加载本地环境， +库导入不能加载。部署专属 ID 和默认值改为必填配置。模型、Skills、IdentityClient +和飞书扩展保持原有目录归属。 + +状态和事件格式保持来源契约，不引入后端 schema 或云资源迁移。取消必须关闭 Runner、 +工具及临时 Skill 目录；终态重放不能重复执行已完成输入。凭据只作为运行时输入， +不写入文档、镜像、跟踪或 Git；错误报告保持脱敏。Python 3.10-3.13 的导入必须使用 +兼容的 datetime 与 asyncio 原语。 + +组件契约：[运行时](../../../specs/managed-agent-runtime/README.zh.md)、 +[资源](../../../specs/managed-agent-resources/README.zh.md)、 +[飞书](../../../specs/feishu-channel/README.zh.md)。 + +## 任务与职责 + +| 任务 | 需求 | 文件 / 负责人 | +| --- | --- | --- | +| `T-1` | `FR-1`、`FR-2` | Runtime/Session 包、运行时测试、示例 Python 包装：runtime agent | +| `T-2` | `FR-3`、`FR-4` | Ark、IdentityClient、Skills 客户端、模型/资源测试：model/skills agent | +| `T-3` | `FR-5` | 飞书扩展、渠道测试、双语用户文档：channel agent | +| `T-4` | `FR-6` | Docker、依赖/锁文件、非 UI 脚本、示例文档：root agent | +| `T-5` | 全部 | 集成、回归、打包、密钥扫描、文档同步：root agent | + +文件写入范围互不重叠。并行期间不操作共享 Git index 或提交；root 统一依赖并执行 +最终检查。 + +## 验证与验收 + +| 需求 | 验收 | 验证 | 状态 | +| --- | --- | --- | --- | +| `FR-1`、`FR-2` | `AC-1`：可安装且生命周期等价 | Runtime/Session 定向测试与真实 SDK transport/原生工具 | pass；包含在最终 387 项定向测试中 | +| `FR-3`、`FR-4` | `AC-2`：保留官方行为并加入资源/模型增量 | 模型/Identity/Skills 测试 | pass；首轮 109 项，纳入最终定向回归 | +| `FR-5` | `AC-3`：流式与关闭测试通过 | 渠道测试 | pass；18 项，纳入最终定向回归 | +| 全部 | `AC-4`：集成回归和仓库检查 | 并行回归与全文件/变更文件 pre-commit | 范围内回归 pass；无排除的检查有 7 个已复现官方 UI 失败，见下文 | +| `FR-6` | `AC-5`：wheel/导入/CLI/构建输入检查 | wheel/导入/CLI、shell 语法、镜像构建与容器检查 | pass；仅本地镜像 | + +测试使用模拟服务与临时状态。线上云/模型 E2E 和镜像发布为 not_run，超出本次本地 +迁移授权。UI 检查为 not_applicable，因为 UI 和 Studio 后端修改均排除。 + +## 风险与恢复 + +来源的 EnvironmentWorkDispatcher 需要普通 Anthropic 安装没有的 Managed Agents +SDK 能力。必须检查 SDK 能力并实现明确的受支持集成，不能虚报 Worker 就绪。 +交付记录区分模拟运行时测试、安装包导入、镜像构建与真实云 E2E。原来源和官方基线 +保留完整；放弃新分支即可恢复到基线。 + +## 审查与交付记录 + +用户以“ok, just do it, do it best”批准架构与非 UI 范围，并明确授权多个 subagent。 +直接设计审查已解决：使用官方基线、保留 MPA 与运行时接口、排除 UI 后端接线、明确 +SDK 要求、禁止导入时访问凭据、保持 Python 3.10 兼容、验证异步清理。 +未安装 `review-spec` skill，因此采用仓库要求的等价直接审查。组件负责人在生产代码 +修改前审查双语契约。非 UI 范围内的 `T-1` 至 `T-5` 均完成。 + +### 来源与目标核对 + +| 来源 | 交付位置 / 决策 | +| --- | --- | +| `main.py`、`managed_agent_loop.py`、`event_debug.py`、`runtime_identity.py` | `veadk/runtime/managed_agents/{worker,loop,events,identity}.py` | +| 示例 Agent/Session 生命周期 | 包内 `sandbox.py`；示例发现工厂保留薄包装 | +| `sandbox_client.py`、`managed_session_resources.py` | `veadk/integrations/mpa/{session_client,session_resources}.py` | +| 缺失的私有 SDK dispatcher | 基于公开 SDK 的包内 dispatcher，具备真实协议测试 | +| Ark/Identity/Skills/飞书 | 保持原 SDK 模块归属，新增独立 `veadk/skills/ma_infra.py` | +| Worker/Runtime 启动与 claimed Tool 入口 | `docker/managed-agents/`、安装包入口、`/opt/gem/run.sh` 与示例薄启动器 | +| 对话/分布式/Kubernetes/压测/故障检查 | 可移植示例脚本及包化 Runtime/Session/completion-gate 测试 | +| 个人 rollout/build-candidate、RDS 和外部 ma-server/task-server 部署 | 以公共 Worker 构建及通用 Kubernetes 模板替代，不复制部署专属状态 | +| UI、Studio 专用 Skills 仓储/路由、UI 网关镜像与构建资源 | 按批准范围排除 | +| 历史部署证明与一次性分析文档 | 不作为当前证据复制；本设计记录新执行的本地验证 | + +### 验证记录（2026-10-10） + +验证范围是基于官方 `171d8d86` 的未提交功能 diff，包含新增包文件。锁定版本、哈希 +与无关平台 marker 均保留，只将 sandbox 的 Anthropic 下限从 `>=0.40.0` 改为来源 +实现所需的 `>=1.3.0`。 + +- **pass**：Python 3.12 最终定向回归 387 项，包含真实 SDK HTTP、原生文件/bash、 + MCP、跨 Worker 权威历史重放、清理与进程监督。 +- **pass**：Python 3.10.20 独立冻结环境 356 项定向测试，加真实 SDK transient + 发布测试；已用 `typing.Dict` 修复 SDK generic-alias cast 兼容问题，没有用 mock + 掩盖不受支持的生产导入。 +- **pass**：`MODEL_AGENT_API_KEY=offline-regression-test-key .venv/bin/python -m + pytest -n 4 -m "not codex_smoke and not piagent_smoke" -q --tb=short`，显式排除 + 下述 7 个基线失败后：6963 passed、42 skipped、4 xfailed、6 subtests passed。 + 最终大规模回归选用 4 个 worker；首次默认 2-worker 回归暴露测试环境与基线问题。 +- **fail，官方基线**：无排除回归中的 7 个 + `tests/frontend/test_migration_delivery_recovery.py` 用例，在同依赖的未修改官方 + worktree 也失败，分别是： + `test_mirror_delivery_writes_the_documents_the_cli_would_have_written`、 + `test_mirror_delivery_keeps_the_sequence_the_cli_already_started`、 + `test_mirror_delivery_refuses_a_delivery_the_cli_already_settled`、 + `test_verification_mirrors_the_cli_finding_projection`、 + `test_an_unparseable_findings_file_reports_no_findings_at_all`、 + `test_a_startup_fallback_degrades_the_delivery`、 + `test_a_rebuild_that_lands_settles_the_task_on_the_cli_delivery_contract`。 + 范围内回归逐项使用 `--deselect ::`，未修改 marker、生产 UI 代码或 + 断言来掩盖失败。 +- 首次无排除回归还暴露既有示例/Harness 测试依赖无关示例测试在导入时偷偷设置的 + 模型 Key。新测试不再在收集时修改全局凭据；最终命令明确提供离线 fixture Key。 + 测试包命名也已避免与 Studio scheduler 的 `test_dispatcher` 冲突。 +- **pass**：全文件与变更文件 pre-commit（Ruff check/format、gitleaks、YAML 密钥 + 扫描）、`git diff --check`、双语 Markdown 链接和无 UI 改动范围检查。 +- **pass**：`uv build --wheel`；在临时目录和清理后的环境中导入解压安装包并执行 + CLI help，不依赖示例路径。 +- **pass**：`npm --prefix docs run types:check` 和 `npm --prefix docs run build`。 + 依赖安装使用仓库冻结的 pnpm lock；生成的 workspace policy 与 Corepack 意外修改 + 的父目录配置均已从功能范围中撤销。 +- **pass**：本地 `veadk-managed-agents-worker:upstream-local` 镜像通过公共包源和锁定 + sandbox 依赖构建；网络/proxy 设置只用于本地构建。专用 Worker 不安装可选集成、 + 开发或 Codex extras，不使用相邻 SDK wheel。 +- **pass**：禁用网络、只读文件系统、UID 10001 的容器内验证 SDK 导入/CLI 与 + `/opt/gem/run.sh`；同样隔离的安装包健康服务 `/ping` 返回 HTTP 200。这些检查 + 不执行真实 Work 网关、模型提供方或 AgentKit 部署。 +- **not_run**：真实云/模型 E2E、镜像发布、Git commit/push 和 PR 创建。镜像仅在 + 本地，分支跟踪 `volcengine/main`，原来源 worktree 保持干净。 + +### 审查修复 + +独立审查解决了：导入时访问凭据、失败轮次错误完成及泄漏错误文本、Skill 工作区 +冲突/符号链接清理、清理失败后仍关闭全部工具、Work 租约栅栏/Session FIFO/凭据 +刷新、首次合法轮询就绪、Python 3.10 cast、含凭据的包源参数、分布式脚本过时的 +数据库假设与独立命名 BuildKit 缓存。claimed Tool 保留历史启动路径，并在部署的 +OS 隔离中默认本地执行工具。 diff --git a/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.md b/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.md new file mode 100644 index 000000000..e5b4de337 --- /dev/null +++ b/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.md @@ -0,0 +1,103 @@ +# Official Anthropic SDK adapter verification + +[中文版](2026-10-10-official-sdk.zh.md) + +Change ID: `managed-agents-official-sdk`. Created/revised: 2026-10-10. Status: implemented. + +## Background and scope + +The original example's Dockerfile installed an adjacent custom Anthropic wheel to +provide `anthropic.lib.environments._dispatcher`. The migrated Worker already uses +published `anthropic==1.3.0` and a VeADK dispatcher. Its locked wheel hash matches +[PyPI](https://pypi.org/project/anthropic/1.3.0/). The user's follow-up asks to research +and implement the same function using the official SDK, optionally through Python +extension techniques. The approved scope includes image publication, local Kubernetes deployment and real +SDK acceptance in the specified namespace. Cloud AgentKit deployment remains excluded. + +## Requirements and design + +- `FR-1`: install the unchanged official SDK; the image checks its installed Python + files against distribution RECORD hashes to reject modified package contents. +- `FR-2`: environment polling uses generated `work.with_raw_response.poll`, preserving + HTTP 204 handling, strict Work validation, scoped auth, routing headers, leases, + cancellation and readiness. Account `/v1/model-work/poll` is a VeADK gateway extension + accessed with documented `client.get`; it is not an official Claude endpoint. +- `FR-3`: tests exercise VeADK's real scoped-client adapter rather than importing the + SDK's private auth helper. Parent auth is removed; transport identity injection + preserves Runtime gateway auth alongside Work lease auth. +- `FR-4`: build guides explain the original dependency, official installation, extension + boundaries and remaining version-sensitive helpers. + +Use composition. Official EnvironmentWorker has no arbitrary handler callback; +subclassing it would override private execution and SessionToolRunner behavior. Global +monkeypatching would affect unrelated clients. Public headers/http-client/raw-response +options provide the required transport extension without altering SDK code. +`ToolError` uses its public tools export. Published internal `_download_and_extract` +remains necessary because public download_session_skills silently skips failed pinned +Skills; this dependency is documented and covered by real SDK archive tests. + +## Tasks and acceptance + +| Task | Requirement | Files and verification | +| --- | --- | --- | +| `T-1` | `FR-1` | Docker SDK-integrity checker, image installation gate, tamper regression | +| `T-2` | `FR-2`, `FR-3` | dispatcher and identity protocol tests: public generated poll, empty/malformed responses, all-request headers and credential refresh | +| `T-3` | `FR-4` | paired build guides, runtime spec, this design; link and whitespace checks | + +`AC-1`: unmodified installed SDK passes and modified SDK files fail integrity checks. +`AC-2`: targeted Runtime/MPA tests pass on locked official SDK, including Python 3.10. +`AC-3`: pre-commit and Docker build/isolated imports pass; no private wheel input, +version churn, SDK source overwrite or cloud changes. Cloud E2E remains not_run. + +## Review and risk + +The user explicitly requested research followed by implementation, continuing the +approved official-SDK migration approach. Direct review confirms public raw-response +support and composition feasibility; no review-spec skill is available. Update both +runtime component languages for FR-1/FR-2. Other resource and model contracts are +unchanged. SDK-internal helper names remain version-sensitive; locked versions and +integration tests bound this risk. RECORD verification detects modified installed +files, while the lock's official artifact hash verifies origin during installation. +Neither verifies server behavior. No live credential or endpoint is recorded. + + +Independent review also found that the official internal cleanup helper logs raw exceptions. `T-2` includes a local cleanup adapter using close/aclose and additive context-manager exits; it reports only exception types and propagates cancellation. Regression tests verify cleanup continues without exposing exception text. + +## Delivery evidence + + +- **pass (2026-10-10):** 213 targeted Python 3.12 tests; official-SDK adapter + subset 109 tests on Python 3.10, plus 41 credential tests after the local-key + opt-in addition. All-file pre-commit passed. +- **pass:** published immutable Worker image; official SDK 1.3.0 RECORD matches + all 1389 Python files. Running Kubernetes Worker source hashes match the local + dispatcher/worker/identity/Session implementation. Worker is Ready in the + authorized namespace. The official SDK 1.5.0 acceptance client is installed + separately, preserving AKX's existing version. +- **pass:** the user approved deleting one historical Key. Archived/zero-active state was + revalidated before deletion; a dedicated Environment Key was created. Official SDK 1.5 + passed two real model turns, six native tools, pinned Skill instructions, cross-Worker + continuation and stopped model Work; test Session/Agent/Skill cleanup passed. Official + tool config type discriminators now work with legacy name fallback; real-SDK regression + coverage is included in 213 targeted tests. Public APIG, cloud Identity and remote tools + remain outside this acceptance. + + +## Authorized local Kubernetes deployment + +The user subsequently authorized deployment in `ma-infra-stg-zn-ppe`, deployment +assets in `/home/mofanke/gitcode/akx`, and final acceptance with official Anthropic +SDK. `T-4` adds a dedicated environment-scoped Worker and a new VKE template; +existing Helm revisions/images are saved and existing backend/frontend/cloud +Runtime defaults remain intact. Official SDK creates isolated resources and checks +real Work polling, model/tool events, two-turn continuation and pinned Skills. + +Current workers require Session model Identity even for local VKE. Add opt-in +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE=environment` for non-AgentKit Sessions +without a provider reference, reading only `MODEL_AGENT_API_KEY`. The default +`session` retains Identity behavior; AgentKit or any explicit provider reference +always uses Session Identity without fallback. Unknown sources/missing keys fail +closed. A dedicated Kubernetes Secret provides the local model key; no credentials +enter repository files. This acceptance proves configured model credentials, not +cloud Runtime Identity. Rollback deletes the new Worker and restores only any +explicitly changed config; test Agents/Sessions/Skills are cleaned up. diff --git a/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.zh.md b/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.zh.md new file mode 100644 index 000000000..87a4aa7d9 --- /dev/null +++ b/prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.zh.md @@ -0,0 +1,90 @@ +# 官方 Anthropic SDK 适配验证 + +[English](2026-10-10-official-sdk.md) + +变更 ID：`managed-agents-official-sdk`。创建/修订：2026-10-10。状态：implemented。 + +## 背景与范围 + +原示例 Dockerfile 安装相邻仓库的定制 Anthropic wheel，以提供 +`anthropic.lib.environments._dispatcher`。迁移后的 Worker 已使用发行版 +`anthropic==1.3.0` 与 VeADK dispatcher；锁定 wheel 哈希与 +[PyPI](https://pypi.org/project/anthropic/1.3.0/) 一致。用户此次要求研究并实现官方 +SDK 下的同等功能，可考虑 Python 扩展技术。用户随后批准镜像发布、指定 namespace 的本地 Kubernetes 部署与真实 SDK +验收；云 AgentKit 部署仍不在范围内。 + +## 需求与设计 + +- `FR-1`：安装未修改的官方 SDK；镜像将安装后的 Python 文件与发行物 RECORD 哈希 + 比较，拒绝修改过的包内容。 +- `FR-2`：环境级轮询使用生成的 `work.with_raw_response.poll`，保留 HTTP 204、 + 严格 Work 校验、作用域认证、路由头、租约、取消与就绪语义。账号级 + `/v1/model-work/poll` 是 VeADK 网关扩展，用官方支持的 `client.get` 访问，不能 + 宣称为 Claude 官方端点。 +- `FR-3`：测试调用真实 VeADK scoped-client 适配器，不再导入 SDK 私有认证 helper。 + 清除父客户端认证，传输层 Identity 注入保留 Runtime 网关与 Work 租约的双重认证。 +- `FR-4`：构建指南说明原始依赖、官方安装方式、扩展边界和仍存在的版本敏感 helper。 + +采用组合。官方 EnvironmentWorker 没有任意 handler 回调，继承需要覆盖私有执行 +方法与 SessionToolRunner 行为；全局 monkeypatch 会影响无关客户端。公开请求头、 +HTTP client 和 raw-response 配置足以完成传输扩展，无需修改 SDK 源码。 +`ToolError` 改用 tools 公开导出。发行版内部 `_download_and_extract` 仍需保留， +因为公开 download_session_skills 会跳过下载失败的固定 Skill；此依赖明确记录， +并用真实 SDK 归档测试验证。 + +## 任务与验收 + +| 任务 | 需求 | 文件与验证 | +| --- | --- | --- | +| `T-1` | `FR-1` | Docker SDK 完整性检查器、镜像安装门禁、修改文件反例测试 | +| `T-2` | `FR-2`、`FR-3` | dispatcher 与 Identity 协议测试：公开生成轮询、空/错误响应、全部请求头及凭据刷新 | +| `T-3` | `FR-4` | 双语构建指南、Runtime 契约与本设计；链接及空白检查 | + +`AC-1`：未修改的安装包通过，SDK 文件被修改时完整性检查失败。 +`AC-2`：锁定官方 SDK 的 Runtime/MPA 定向测试通过，包含 Python 3.10。 +`AC-3`：pre-commit、Docker 构建及隔离导入通过；无私有 wheel 输入、版本变动、 +SDK 源码覆盖或云端修改。云 E2E 为 not_run。 + +## 审查与风险 + +用户明确要求研究后实现，延续已批准的官方 SDK 迁移方案。直接审查确认公开 +raw-response 支持与组合可行性；当前无 review-spec skill。同步更新 Runtime +双语契约以覆盖 FR-1/FR-2；其他资源与模型契约保持不变。SDK 内部 helper 名称仍有 +版本敏感性，由锁定版本与集成测试控制风险。RECORD 校验检测已安装文件被修改， +锁文件中的官方发行物哈希在安装时验证来源；两者均不能证明服务端行为。 +不记录真实凭据或端点。 + + +独立审查还发现官方内部清理 helper 会输出原始异常。`T-2` 增加本地清理适配，调用 close/aclose 并保留额外 context-manager 退出，仅报告异常类型，传播取消；回归验证清理继续且不泄漏异常文本。 + +## 交付证据 + + +- **pass（2026-10-10）**：Python 3.12 定向测试 213 项；Python 3.10 官方 SDK 适配 + 子集 109 项，增加本地模型凭据 opt-in 后另跑 41 项凭据测试。全文件 pre-commit 通过。 +- **pass**：发布固定 digest Worker 镜像；官方 SDK 1.3.0 的 1389 个 Python 文件 + 均匹配 RECORD。K8s 内 Worker 的 dispatcher/worker/Identity/Session 源码哈希 + 与本地一致，目标 namespace 内已 Ready。独立安装官方 SDK 1.5.0 验收客户端， + 保留 AKX 原有版本。 +- **pass**：用户明确批准清理单个历史 Key；复核归档与零活跃状态后删除成功， + 专用环境 Key 创建成功。官方 SDK 1.5 完成两轮真实模型对话、六种原生工具、 + 固定版本 Skill、跨 Worker 续接与模型 Work stopped 验证；自建 Session/Agent/Skill + 清理成功。另修复官方 SDK 工具 configs 的 type 判别，兼容 legacy name， + 真实 SDK 对象回归已计入 213 项测试。公网 APIG、云 Identity 与远程工具不在本次证明范围。 + + +## 已授权的本地 Kubernetes 部署 + +用户随后授权部署 `ma-infra-stg-zn-ppe`，部署资产放在 +`/home/mofanke/gitcode/akx`，最终用官方 Anthropic SDK 验收。`T-4` 新增限定到专用 +Environment 的 Worker 与 VKE 模板；保存现有 Helm revision/镜像，保留后端、前端与 +云 Runtime 默认配置。官方 SDK 创建独立资源，验证真实 Work 轮询、模型/工具事件、 +两轮续接与固定版本 Skills。 + +当前 Worker 即使在本地 VKE 也要求 Session 模型 Identity。增加显式选择 +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE=environment`，仅用于非 AgentKit 且无 +provider 引用的 Session，只读取 `MODEL_AGENT_API_KEY`。默认 `session` 保持 +Identity 行为；AgentKit 或任何显式 provider 引用始终使用 Session Identity， +不允许兜底。未知来源或缺少 Key 明确失败。模型 Key 由独立 Kubernetes Secret +提供，不进入仓库文件。本次验收证明配置的模型凭据，不证明云 Runtime Identity。 +回滚删除新 Worker 并恢复明确改过的配置；清理测试 Agent/Session/Skill。 diff --git a/pyproject.toml b/pyproject.toml index ca58dfb44..f6d84f3bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,7 +123,7 @@ harness-sidecar = [ ] cli = [] sandbox = [ - "anthropic>=0.40.0", + "anthropic>=1.3.0", "python-dotenv>=1.0.0", "httpx[socks]>=0.27.0", ] diff --git a/specs/feishu-channel/README.md b/specs/feishu-channel/README.md new file mode 100644 index 000000000..1f7788019 --- /dev/null +++ b/specs/feishu-channel/README.md @@ -0,0 +1,47 @@ +# Feishu Channel Streaming + +[中文版](README.zh.md) + +Component: `feishu-channel`. Revised: 2026-10-10. Status: implemented, locally verified. +Change: [non-UI migration PRD](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md), `FR-5` / `AC-3`. + +## Scope +Port non-UI Feishu channel improvements onto official VeADK while preserving upstream graceful shutdown and app-loop task tracking. +The extension owns presentation of Runner events and channel lifecycle. Runner owns execution; the Lark SDK owns message/card delivery. No Studio UI or backend wiring is included. + +## Contract +- `CON-FEISHU-1`: Default replies remain answer-only. Thinking, tool inputs, and tool results are opt-in constructor/environment options. Provider thought text is whitespace-normalized before display; hidden model reasoning is unavailable. +- `CON-FEISHU-2`: Separate thinking and per-call tool cards can reply in an existing or newly requested topic. Existing call IDs correlate out-of-order results; a result without a known call gets a standalone card. Missing stages produce no content cards. +- `CON-FEISHU-3`: Repeated completed provider events do not duplicate thinking or accumulated answer text. Tool records are deduplicated by nonempty call/result ID; events without IDs cannot provide the same correlation guarantee. +- `CON-FEISHU-4`: Displayed structured tool fields are recursively redacted by sensitive field names, bounded in length, and Markdown-fence escaped. Serialization fallback must use the redacted copy. Redaction does not inspect secrets embedded in arbitrary strings. +- `CON-FEISHU-5`: Streaming tasks settle on successful completion and on failures. Card errors propagate unless a Runner error/cancellation is already propagating; cleanup preserves that original failure. Upstream shutdown closes the socket then drains accepted messages on the application loop, cancelling unfinished messages after its timeout. + +## Interface and configuration + +Public entry point: `veadk.extensions.FeishuChannelExtension(runner=..., streaming=True, show_tool_calls=True)`. Streaming controls affect `channel.stream` only; ordinary handlers and non-streaming Runner replies retain their text formatter. + +| Constructor option | Environment variable | Default | +| --- | --- | --- | +| `show_thinking` | `TOOL_FEISHU_CHANNEL_SHOW_THINKING` | `False` | +| `show_tool_calls` | `TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS` | `False` | +| `show_tool_results` | `TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS` | `False` | +| `separate_tool_call_cards` | `TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS` | `False` | +| `separate_thinking_card` | `TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD` | `False` | +| `create_topic` | `TOOL_FEISHU_CHANNEL_CREATE_TOPIC` | `False` | +| `tool_detail_max_length` | none | `4000`, minimum `100` | + +Boolean options combine constructor true with environment values `1`, `true`, `yes`, `on` (case insensitive); explicit false does not override an enabled environment variable. Topic creation requires `reply_in_thread=True` and an inbound message ID, and sends `reply_to` plus `reply_in_thread=True` to Lark. Existing topic IDs also enable threaded replies. No payload or credentials are persisted by the formatter. + +## Compatibility and failures + +Preserve upstream lifecycle, identity mapping, parent/history retrieval, reconnect behavior, and dependency fallbacks. Tool failure icons recognize a truthy `error`, statuses `error`/`failed`/`failure`, and nonzero `exit_code`; they do not change execution status. Sensitive markers are `api_key`, `apikey`, `access_key`, `authorization`, `cookie`, `credential`, `password`, `secret`, `token` after lowercase and hyphen-to-underscore normalization. Card delivery is not guaranteed by a local test. Shutdown's default timeout remains the upstream `TOOL_FEISHU_CHANNEL_DRAIN_TIMEOUT` value, 300 seconds. + +## Verification +Run existing Feishu extension tests plus source regressions for ordered stream sections, deduplication, separate cards, failed results, topic options, and redacted bounded payloads. Live Feishu delivery requires credentials and is outside local validation. + +Command: `.venv/bin/python -m pytest tests/test_feishu_channel_extension.py -q`. +`CON-FEISHU-1..3` map to detailed streaming, separate-card, topic, and missing-stage tests; `CON-FEISHU-4` maps to truncation and serialization-fallback tests; `CON-FEISHU-5` maps to graceful shutdown, Runner-failure settlement, and card-cleanup failure propagation tests. + +## Change record + +2026-10-10: Ported source channel controls onto official `171d8d86`, preserving its lifecycle changes; fixed redaction fallback and background-card cleanup error propagation. Existing upstream baseline: 8 tests passed. Source regressions fail before the new helper is installed, then pass after migration. Added negative-path regressions expose both formatter fallback leakage and swallowed card errors before their fixes. Local validation is recorded by the parent migration PRD; live delivery is not_run. diff --git a/specs/feishu-channel/README.zh.md b/specs/feishu-channel/README.zh.md new file mode 100644 index 000000000..1b63ff151 --- /dev/null +++ b/specs/feishu-channel/README.zh.md @@ -0,0 +1,47 @@ +# 飞书渠道流式输出 + +[English](README.md) + +组件：`feishu-channel`。修订：2026-10-10。状态：已实现，本地验证通过。 +变更：[非 UI 迁移 PRD](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md)，`FR-5` / `AC-3`。 + +## 范围 +将非 UI 的飞书渠道增强移植到官方 VeADK，保留上游优雅关闭与应用事件循环任务追踪。 +扩展负责呈现 Runner 事件与渠道生命周期；Runner 负责执行，Lark SDK 负责消息和卡片投递。范围不包含 Studio UI 或后台接线。 + +## 行为约定 +- `CON-FEISHU-1`：默认仅输出回答；思考、工具输入、工具结果通过构造参数或环境变量显式启用。服务商返回的思考文本展示前规范化空白；无法获得模型未返回的内部推理。 +- `CON-FEISHU-2`：可将思考与每次工具调用拆成独立卡片，回复到已有或主动创建的话题。调用 ID 关联乱序结果；找不到对应调用的结果使用独立卡片。缺失阶段不产生内容卡片。 +- `CON-FEISHU-3`：服务商重复发送的完成事件不得重复显示思考或累计回答。工具记录按非空调用/结果 ID 去重；无 ID 事件不能提供相同的关联保证。 +- `CON-FEISHU-4`：工具结构化字段按敏感字段名递归脱敏、限制长度并转义 Markdown 代码围栏。序列化回退必须使用已脱敏副本；任意字符串内嵌的秘密不在字段脱敏范围内。 +- `CON-FEISHU-5`：成功与异常路径均回收流式任务。卡片错误向上传播，若 Runner 错误或取消正在传播则保留原始失败。保留上游先关闭套接字、再等待应用循环内已接收消息完成的行为；超时后取消未完成消息。 + +## 接口与配置 + +公开入口：`veadk.extensions.FeishuChannelExtension(runner=..., streaming=True, show_tool_calls=True)`。流式控制仅影响 `channel.stream`；普通处理函数及非流式 Runner 回复保留文本格式器。 + +| 构造参数 | 环境变量 | 默认值 | +| --- | --- | --- | +| `show_thinking` | `TOOL_FEISHU_CHANNEL_SHOW_THINKING` | `False` | +| `show_tool_calls` | `TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS` | `False` | +| `show_tool_results` | `TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS` | `False` | +| `separate_tool_call_cards` | `TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS` | `False` | +| `separate_thinking_card` | `TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD` | `False` | +| `create_topic` | `TOOL_FEISHU_CHANNEL_CREATE_TOPIC` | `False` | +| `tool_detail_max_length` | 无 | `4000`，最小 `100` | + +布尔参数由构造参数 true 或环境值 `1`、`true`、`yes`、`on`（不区分大小写）启用；显式 false 不覆盖已启用的环境变量。创建话题要求 `reply_in_thread=True` 且有入站消息 ID，向 Lark 传递 `reply_to` 与 `reply_in_thread=True`。已有话题 ID 也启用话题内回复。格式器不持久化载荷或凭据。 + +## 兼容与失败 + +保留上游生命周期、身份映射、父消息/历史获取、重连行为与依赖回退。失败图标识别真值 `error`、`error`/`failed`/`failure` 状态与非零 `exit_code`，不改变执行状态。敏感标记为 `api_key`、`apikey`、`access_key`、`authorization`、`cookie`、`credential`、`password`、`secret`、`token`，字段名转小写并将连字符替换为下划线后匹配。本地测试不能保证真实卡片投递。关闭默认超时保留上游 `TOOL_FEISHU_CHANNEL_DRAIN_TIMEOUT` 的 300 秒。 + +## 验证 +执行现有飞书扩展测试及源仓库回归测试，覆盖输出顺序、去重、独立卡片、失败工具结果、话题参数、脱敏与长度限制。真实飞书投递依赖凭据，不属于本地验证。 + +命令：`.venv/bin/python -m pytest tests/test_feishu_channel_extension.py -q`。 +`CON-FEISHU-1..3` 对应详细流式、独立卡片、话题及缺失阶段测试；`CON-FEISHU-4` 对应截断与序列化回退测试；`CON-FEISHU-5` 对应优雅关闭、Runner 失败回收、卡片清理错误传播测试。 + +## 变更记录 + +2026-10-10：将源仓库渠道控制移植到官方 `171d8d86`，保留其生命周期变更；修复脱敏回退与后台卡片清理错误传播。上游基线 8 个测试通过。源仓库回归测试在新辅助函数安装前失败，迁移后通过。新增异常路径测试在修复前揭示格式器回退泄露与卡片错误吞没。父级迁移 PRD 记录本地验证结果；真实投递为 not_run。 diff --git a/specs/managed-agent-resources/README.md b/specs/managed-agent-resources/README.md new file mode 100644 index 000000000..438920fa1 --- /dev/null +++ b/specs/managed-agent-resources/README.md @@ -0,0 +1,114 @@ +# Managed Agent resource adapters + +[中文版](README.zh.md) + +Component ID: `managed-agent-resources`. Status: `active`. Revised: 2026-10-10. +Related approved PRD: [managed-agents-upstream](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md), `FR-3`, `FR-4`, `T-2`, `AC-2`. +Dependent component: [managed-agent-runtime](../managed-agent-runtime/README.md). + +## Scope and structure + +Migrate the source backend changes without copying frontend or Studio APIs. Ark +request behavior remains in `veadk/models/ark_llm.py`; credential API requests stay +in `veadk/integrations/ve_identity/identity_client.py`; the standalone Managed +Agents skills transport lives in `veadk/skills/ma_infra.py`. Runtime credential +resolution belongs to the runtime adapter, which calls IdentityClient rather +than embedding a second SDK client. Existing upstream workload/user-pool helpers +and default environment handling remain intact. + +## Behavioral contract + +- `CON-1`: Ark continuation requests preserve currently authorized tool declarations; + absence of tools does not imply any additional authorization. +- `CON-2`: Default request construction does not inject a one-hour `expire_at`. Explicit + caller retention is preserved; otherwise the remote service chooses retention. +- `CON-3`: `retry_expired_response` defaults to the existing replay behavior. A managed + runtime can disable replay to prevent silently discarding prior conversation. +- `CON-4`: `MANAGED_AGENT_MODEL_TRACE=true` enables a dedicated metadata-only JSON trace: + model/tool names, tool selection, presence of a continuation, and event/output + types. Prompts, responses, API keys, tool arguments and continuation IDs are + excluded. Tracing is disabled by default. +- `CON-5`: IdentityClient accepts optional `pool_name` for API-key resolution; omission + preserves the SDK default. Credentials remain in process memory. +- `CON-6`: The skills client uses the configured `/v1/skills` endpoint, account header and + API-key header. It exposes synchronous operations and thread-backed asynchronous + counterparts. Pagination metadata remains available; `latest` resolves to a + concrete version before downloading content. Archives exceeding 32 MiB fail. +- `CON-7`: HTTP and service errors become typed skills errors with status, code and retry + classification. Multipart publishing lets the HTTP library set its boundary. + +## Entry points, configuration and state + +Public imports are `veadk.models.ark_llm.ArkLlm`, +`veadk.integrations.ve_identity.identity_client.IdentityClient`, and +`veadk.skills.ma_infra.{MaInfraSkillConfig,MaInfraSkillsClient,MaInfraSkillError}`. +Ark calls its existing SDK Responses API; IdentityClient calls the Identity SDK +`get_resource_api_key`. For example, `client.get_api_key(provider_name="model", +agent_identity_token=token, pool_name="customer")` selects that credential pool. +The credentials remain runtime values, never serialized to this document. + +`MaInfraSkillConfig.from_env()` reads `MA_SKILLS_BASE_URL`, `MA_SKILLS_API_KEY`, +and `MA_SKILLS_ACCOUNT_ID`; an absent base URL disables configuration by +returning `None`. Explicit configuration takes precedence by being passed directly +to the client. Default timeout is 30 seconds; redirects are disabled. Calls cover +GET/POST `/v1/skills`, GET/DELETE `/v1/skills/{id}`, and GET/POST versions/content +beneath that skill. The control plane owns authorization and registry state. +The client owns no local persistent cache, version migration, deduplication or +automatic mutation retries. Callers must `close()` the synchronous HTTP client. +Async methods use `asyncio.to_thread`; cancelling the awaiting coroutine does not +interrupt an in-flight synchronous HTTP request, which remains timeout bounded. + +## Security, compatibility and failures + +Skills requests attach only the configured API-key and account headers. External +resource identifiers and archive bytes are passed to the control plane; downloads +are returned as bytes and this component does not extract archives. Runtime Skill +extraction and temporary-directory cleanup belong to the dependent runtime. +Only whitelisted Ark trace metadata is emitted. SDK/service failure behavior and +caller-owned credentials remain unchanged; no credential is logged by the new +Identity pool forwarding. Typed skills errors expose `status_code`, `code`, and +`retryable`: network errors map to retryable 502; 429 and 5xx are retryable, other +HTTP errors are not. Server error text is returned to callers; it is not logged +by the skills client. The 32 MiB download cap is checked after transfer; upload +size enforcement belongs to the service. + +Python 3.10-3.13 remains supported. Existing Identity calls omit `pool_name` and +preserve the service default. Ark replay remains enabled by default; runtime +callers explicitly disable it when full history cannot be replayed. No schema or +SDK pin changes are required by these resource adapters. + +## Acceptance + +Local tests cover continuation tool preservation, retention defaults and explicit +values, expired-context fail-fast behavior, trace field redaction, Identity pool +forwarding, skills headers/pagination/version pinning/error handling and async +helpers. No real cloud mutation or credential-bearing configuration is needed. + +## Local verification record + +On 2026-10-10, `pytest tests/models tests/skills +tests/test_ve_identity_api_key.py tests/integrations/test_mpa_identity.py -q` +passed **109 tests** against the target tree. New regression tests first failed +on the official baseline for tool removal, expiration override and missing pool +scope. Ruff 0.11.12 lint and formatting passed for the changed component files. +This is offline protocol/model coverage; live Identity and skills services were +not invoked. + +Contract-to-test mapping: `CON-1`/`CON-2` use model context/expire-at tests; +`CON-3`/`CON-4` use model fallback/trace tests; `CON-5` uses +`tests/test_ve_identity_api_key.py`; `CON-6`/`CON-7` use +`tests/skills/test_ma_infra_skills_client.py`. Negative cases include absent +configuration/tools, expiration, oversized archives and service/network failures. +The 2026-10-10 migration changes only these contracts; no unresolved resource +contract questions remain. Cloud E2E and actual registry mutation remain outside +local acceptance. + +Minimum-version verification: Python **3.10.20**, an isolated environment created +with `uv sync --frozen --python 3.10 --extra sandbox --extra extensions --extra dev`, +Anthropic **1.3.0**, Google ADK **2.2.0**. The focused runtime, MPA Session, model, +Skills, Identity and Feishu suites passed **356 tests** (17 deprecation warnings, +24.24 seconds). The new real-SDK transient-event transport regression also passed +individually. A Python 3.10 SDK GenericAlias response-cast incompatibility was +resolved with `typing.Dict` in the transient publisher and transport tests; no +dependency pin or target environment change was necessary. This verifies local +protocol behavior, not remote cloud deployment. diff --git a/specs/managed-agent-resources/README.zh.md b/specs/managed-agent-resources/README.zh.md new file mode 100644 index 000000000..5efa98ea4 --- /dev/null +++ b/specs/managed-agent-resources/README.zh.md @@ -0,0 +1,89 @@ +# Managed Agent 资源适配器 + +[English](README.md) + +组件 ID：`managed-agent-resources`。状态：`active`。修订：2026-10-10。 +关联已批准 PRD：[managed-agents-upstream](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md)，`FR-3`、`FR-4`、`T-2`、`AC-2`。 +依赖组件：[managed-agent-runtime](../managed-agent-runtime/README.zh.md)。 + +## 范围与结构 + +迁移来源仓库的后端增量，不复制前端或 Studio 接口。Ark 请求行为放在 +`veadk/models/ark_llm.py`，凭据 SDK 请求放在 +`veadk/integrations/ve_identity/identity_client.py`，独立技能客户端放在 +`veadk/skills/ma_infra.py`。运行时凭据解析由运行时适配器调用 IdentityClient， +不另建 SDK 客户端。保留官方已有的工作负载池、用户池辅助方法和环境配置行为。 + +## 行为约定 + +- `CON-1`：Ark 续接请求保留本次授权的工具声明；未声明工具时不增加工具权限。 +- `CON-2`：默认请求不注入一小时的 `expire_at`。保留调用方显式保留时间；否则由服务端决定。 +- `CON-3`：`retry_expired_response` 默认保留已有重试行为；托管运行时可以关闭重放，避免静默丢弃历史对话。 +- `CON-4`：仅在 `MANAGED_AGENT_MODEL_TRACE=true` 时开启独立 JSON 元数据日志，记录模型与工具名称、 + 工具选择、是否续接、事件与输出类型。排除提示词、回复正文、API Key、工具参数和续接 ID。 +- `CON-5`：IdentityClient 的 API Key 查询接受可选 `pool_name`；省略时保留 SDK 默认行为。 + 凭据只保存在进程内存。 +- `CON-6`：技能客户端访问配置的 `/v1/skills`，携带账号与 API Key 请求头；提供同步和线程包装的异步接口。 + 保留分页元数据；下载前把 `latest` 解析为具体版本；超过 32 MiB 的压缩包报错。 +- `CON-7`：HTTP 与服务错误转换为带状态、错误码和可重试标志的技能异常；上传由 HTTP 库生成 multipart boundary。 + +## 入口、配置与状态 + +公开入口为 `veadk.models.ark_llm.ArkLlm`、 +`veadk.integrations.ve_identity.identity_client.IdentityClient` 和 +`veadk.skills.ma_infra.{MaInfraSkillConfig,MaInfraSkillsClient,MaInfraSkillError}`。 +Ark 继续调用 SDK Responses API，IdentityClient 调用 Identity SDK 的 +`get_resource_api_key`。例如 `client.get_api_key(provider_name="model", +agent_identity_token=token, pool_name="customer")` 选择对应凭据池。 +凭据属于运行时输入，不序列化到本文。 + +`MaInfraSkillConfig.from_env()` 读取 `MA_SKILLS_BASE_URL`、`MA_SKILLS_API_KEY`、 +`MA_SKILLS_ACCOUNT_ID`；缺少 base URL 时返回 `None`，表示未配置。 +直接向客户端传配置时使用显式值。默认超时 30 秒，禁止跟随重定向。 +接口包括 GET/POST `/v1/skills`、GET/DELETE `/v1/skills/{id}`，以及对应技能下的 +版本与内容 GET/POST。控制面负责授权与注册表状态。客户端不持久缓存、不迁移版本、 +不去重，也不自动重试写入。调用方负责 `close()` 同步 HTTP 客户端。 +异步接口通过 `asyncio.to_thread` 调用同步方法；取消等待中的协程不会中断线程中的 +HTTP 请求，该请求仍受超时约束。 + +## 安全、兼容与失败 + +技能请求携带配置的 API Key 与账号请求头，向控制面传递资源 ID 和压缩包内容。 +下载返回 bytes，本组件不解压；运行时组件负责技能解压与临时目录清理。 +Ark 日志只记录白名单元数据。SDK 与服务失败行为保持兼容；新增 Identity 池参数 +转发不记录凭据。技能异常包含 `status_code`、`code`、`retryable`；网络失败映射为 +可重试 502，429 和 5xx 可重试，其他 HTTP 错误不可重试。服务错误文本返回调用方， +技能客户端不记录它。32 MiB 下载上限在传输后检查，上传大小限制由服务端负责。 + +继续支持 Python 3.10-3.13。已有 Identity 调用省略 `pool_name` 时保留服务默认池。 +Ark 默认允许过期重放，无法恢复完整历史的运行时显式关闭重放。 +资源适配器不要求修改 schema 或 SDK 固定版本。 + +## 验收 + +本地测试覆盖续接工具保留、默认与显式保留时间、过期上下文直接失败、日志字段脱敏、 +Identity 池参数转发、技能请求头、分页、版本固定、错误处理和异步辅助方法。 +不需要云端写入或含凭据配置。 + +## 本地验证记录 + +2026-10-10，在目标代码树运行 `pytest tests/models tests/skills +tests/test_ve_identity_api_key.py tests/integrations/test_mpa_identity.py -q`， +**109 项通过**。新增回归测试先在官方基线上复现工具被删除、保留时间被覆盖和缺少池参数。 +修改组件文件通过 Ruff 0.11.12 lint 和格式化检查。这里验证的是离线协议与模型行为， +未调用真实 Identity 或技能服务。 + +契约与测试对应：`CON-1`/`CON-2` 使用模型上下文与 expire-at 测试， +`CON-3`/`CON-4` 使用模型 fallback 与 trace 测试，`CON-5` 使用 +`tests/test_ve_identity_api_key.py`，`CON-6`/`CON-7` 使用 +`tests/skills/test_ma_infra_skills_client.py`。负向场景覆盖缺少配置或工具、 +上下文过期、超限压缩包、服务与网络失败。2026-10-10 迁移只修改这些契约， +没有未解决的资源契约问题。云端 E2E 与真实注册表写入不属于本地验收。 + +最低版本验证：Python **3.10.20**，使用 +`uv sync --frozen --python 3.10 --extra sandbox --extra extensions --extra dev` +建立独立环境；Anthropic **1.3.0**，Google ADK **2.2.0**。运行时、MPA Session、 +模型、Skills、Identity 与飞书聚焦测试共 **356 项通过**，17 条弃用告警，耗时 24.24 秒。 +新增真实 SDK transient-event 传输回归测试单独运行也通过。Python 3.10 的 SDK +GenericAlias 响应转换兼容问题已通过 transient publisher 和传输测试使用 +`typing.Dict` 修复，不需改依赖固定版本或目标环境。这里验证本地协议行为，不代表云端部署。 diff --git a/specs/managed-agent-runtime/README.md b/specs/managed-agent-runtime/README.md new file mode 100644 index 000000000..6765dd285 --- /dev/null +++ b/specs/managed-agent-runtime/README.md @@ -0,0 +1,57 @@ +# Managed Agent Runtime + +[中文版](README.zh.md) + +Component ID: `managed-agent-runtime`. Status: `active`. Revised: 2026-10-10. +Related approved PRD: [managed-agents-upstream](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.md), `FR-1`, `FR-2`, `FR-6`, `T-1`, `T-4`, `AC-1`, `AC-5`. +This component packages the self-hosted Managed Agents execution loop. It supplements existing runtime providers without changing their public contracts. + +## Ownership and interfaces + +`veadk/runtime/managed_agents/` owns worker polling, Work execution, event reconciliation, native tools and credential caches. `dispatcher.py` implements public SDK raw poll requests and generated ack/heartbeat/stop operations; no adjacent SDK checkout or private dispatcher import is required. Readiness is announced only after a valid successful poll. `veadk/integrations/mpa/session_client.py` owns the Anthropic Session transport; `session_resources.py` owns pinned Skill materialization and Session-bound MCP toolsets. Neither package imports example modules. The example delegates execution to the package and loads its own optional dotenv configuration. + +`ManagedAgentsLoop` consumes pending `user.message` events and publishes model/tool events. `serve_managed_agent_worker` polls Work; `serve_claimed_managed_agent_work` executes an externally claimed Work. `SelfHostSandboxClient` preserves environment variable precedence and injected Work authentication. Examples retain old module imports as compatibility adapters. + +## State and isolation + +The frozen Session snapshot controls model, system instruction, enabled tools, permissions, Skills and MCP servers. Each Work gets local ADK bookkeeping and its own native tool context, Skill workspace and MCP toolsets. No shared Agent is mutated to apply a Session's bindings. Event IDs/cursors reconcile prior turns and avoid replaying completed inputs. Ark response IDs restore conversational context across worker replacement; worker bookkeeping always uses the local ADK backend and ignores legacy database settings. + +## Tools, failures and cleanup + +Native Anthropic tools execute directly. File tools preserve bytes/content and reject workspace escapes. Native bash can run commands permitted by its host OS; the file workspace is not an OS sandbox. Mutually untrusted Sessions require separate container/OS isolation; concurrent contexts in a shared worker do not provide shell isolation. Session IDs must use letters, numbers, underscores, dots or hyphens without leading/trailing dots; unsafe IDs are rejected rather than mapped to colliding directories. `always_ask` waits for matching permission responses; custom tools wait for correlated results. The default tool budget is 300 seconds and the outer action timeout must cover it. Cancellation, timeout, terminal Session status and transport failures remain visible. Failed model turns publish a sanitized error and propagate failure, retaining the input for recovery; replay does not treat `session.error` as successful completion. Worker shutdown drains dispatcher tasks and closes toolsets and native tool contexts. Failed or cancelled Skill downloads remove partial workspaces; completed Work removes its tracked download roots without following directory symlinks. Every toolset and native tool context is closed even if another close fails. + +## Credentials and diagnostics + +Runtime, Session-model and MCP credentials remain process-local and isolated by identity context. Session-model cache capacity is 2000 with a 3600-second TTL. Credentials are never written to environment variables or persistent storage. Explicit Work credentials take precedence over Runtime Identity. Identity errors exposed to callers are sanitized; authentication retries refresh credentials and honor shutdown. Optional event and model/poll diagnostics redact credential fields and avoid enabling broad HTTP logging. + +## Verification + +| Contract | Verification | +| --- | --- | +| `CON-1`: packaged imports and CLI require no example paths or credentials | package import and installed wheel/CLI checks | +| `CON-2`: valid poll readiness, lease fencing, Session FIFO and bounded shutdown | dispatcher HTTP and lifecycle tests | +| `CON-3`: canonical history and Ark continuation restore completed turns across workers | fresh-worker replay and actual SDK transport tests | +| `CON-4`: Session-local tools/resources and complete cleanup after failure or cancellation | native tool, MCP, Skill workspace and cancellation tests | +| `CON-5`: scoped authentication and sanitized diagnostics | credential refresh, isolation and error-redaction tests | +| `CON-6`: portable worker and claimed entry points | process supervision, local image build and isolated container checks | + +Package tests cover event reconciliation, cross-worker response restoration, cancellation, authentication retries, real native file/bash tools, correlated custom/MCP results, local archive streaming and cleanup. Live cloud deployment is a separate acceptance step and is not implied by these tests. + +## Official SDK compatibility + +The dispatcher composes the unchanged official SDK. Environment polling uses +`work.with_raw_response.poll`; account model polling uses public `get` for the +VeADK-specific gateway endpoint. HTTP client hooks implement gateway Identity +rotation; generated APIs retain Work authentication. The image validates SDK +Python files against the installed wheel RECORD after hash-locked installation. +Tool cleanup preserves close/aclose and additive context-manager exit while +logging only exception types. Published internal Skill extraction and the tool +context-manager attribute remain version-sensitive, covered by integration tests. +See [official-SDK design](../../prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.md) +and [build guide](../../docker/managed-agents/README.md#official-sdk-extension-strategy). + +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE` defaults to `session`. Explicit +`environment` permits `MODEL_AGENT_API_KEY` only for non-AgentKit Sessions without +an `ark_api_key_provider`; explicit provider references and AgentKit always require +Session Identity. Unknown sources and missing local keys fail closed. Local +Kubernetes acceptance using this option does not verify cloud Runtime Identity. diff --git a/specs/managed-agent-runtime/README.zh.md b/specs/managed-agent-runtime/README.zh.md new file mode 100644 index 000000000..cd10f537b --- /dev/null +++ b/specs/managed-agent-runtime/README.zh.md @@ -0,0 +1,53 @@ +# Managed Agent Runtime + +[English](README.md) + +组件 ID：`managed-agent-runtime`。状态：`active`。修订日期：2026-10-10。 +关联已批准 PRD：[managed-agents-upstream](../../prd-spec/features/managed-agents-upstream/2026-10-10-non-ui-migration.zh.md)，`FR-1`、`FR-2`、`FR-6`、`T-1`、`T-4`、`AC-1`、`AC-5`。 +本组件将自托管 Managed Agents 执行循环移入正式包，扩展现有 runtime provider，不改变其公开契约。 + +## 职责与接口 + +`veadk/runtime/managed_agents/` 负责 Worker 轮询、Work 执行、事件对账、原生工具及凭据缓存。`dispatcher.py` 使用 SDK 公开原始请求执行 poll，并用生成接口执行 ack/heartbeat/stop，无需相邻 SDK checkout 或私有 dispatcher 导入；只有轮询成功且响应有效后才宣布就绪。`veadk/integrations/mpa/session_client.py` 负责 Anthropic Session 传输，`session_resources.py` 负责固定版本 Skill 落盘和 Session 绑定的 MCP 工具集。正式包不导入示例模块；示例调用正式包并加载自己的可选 dotenv 配置。 + +`ManagedAgentsLoop` 消费待处理的 `user.message`，发布模型与工具事件。`serve_managed_agent_worker` 轮询 Work，`serve_claimed_managed_agent_work` 执行外部已认领的 Work。`SelfHostSandboxClient` 保留环境变量优先级及注入的 Work 认证；旧示例模块保留为兼容适配器。 + +## 状态与隔离 + +冻结的 Session 快照决定模型、系统指令、启用工具、权限、Skills 和 MCP 服务。每个 Work 拥有本地 ADK 状态及独立的原生工具上下文、Skill 工作目录和 MCP 工具集；不修改共享 Agent 来应用 Session 绑定。事件 ID 和游标用于对账既有轮次，避免重复执行已完成输入。Ark response ID 支持 Worker 更换后恢复对话上下文；Worker 状态始终使用本地 ADK 后端，忽略旧数据库配置。 + +## 工具、失败与清理 + +直接执行 Anthropic 原生工具。文件工具保留原始内容并拒绝工作目录逃逸。原生 bash 可执行宿主操作系统允许的命令，文件工作目录不构成操作系统沙箱。互不信任的 Session 必须使用独立容器或操作系统隔离，共享 Worker 中并发上下文不提供 shell 隔离。Session ID 必须使用字母、数字、下划线、点或连字符且不以点开头或结尾；拒绝不安全 ID，避免映射到冲突目录。`always_ask` 等待匹配的权限响应，自定义工具等待关联结果。默认工具预算为 300 秒，外层 action timeout 必须覆盖此预算。取消、超时、Session 终态和传输失败保持可见。模型轮次失败时发布脱敏错误并传播异常，保留输入以供恢复；对账不会把 `session.error` 视为成功完成。Worker 停止时排空 dispatcher 任务并关闭 MCP 工具集与原生工具上下文。Skill 下载失败或取消时删除部分工作目录，Work 完成后清理已跟踪的下载根目录，不跟随目录符号链接。即使某次关闭失败，也尝试关闭所有工具集和原生工具上下文。 + +## 凭据与诊断 + +Runtime、Session 模型和 MCP 凭据仅保存在进程内，并按身份上下文隔离。Session 模型缓存容量为 2000、TTL 为 3600 秒。凭据不写入环境变量或持久存储。显式 Work 凭据优先于 Runtime Identity；对外 Identity 错误经过脱敏，认证重试刷新凭据且遵守停止信号。可选事件、模型和轮询诊断对凭据字段脱敏，不开启通用 HTTP 日志。 + +## 验证 + +| 契约 | 验证 | +| --- | --- | +| `CON-1`：正式包导入与 CLI 不需要示例路径或凭据 | 包导入与安装包/CLI 检查 | +| `CON-2`：合法轮询就绪、租约栅栏、Session FIFO 与有界停止 | dispatcher HTTP 与生命周期测试 | +| `CON-3`：权威历史与 Ark 续接在 Worker 更换后恢复已完成轮次 | 新 Worker 重放与真实 SDK 传输测试 | +| `CON-4`：Session 独立工具/资源，失败或取消后完整清理 | 原生工具、MCP、Skill 工作目录与取消测试 | +| `CON-5`：作用域内认证与脱敏诊断 | 凭据刷新、隔离与错误脱敏测试 | +| `CON-6`：可移植 Worker 与 claimed 启动入口 | 进程监督、本地镜像构建与隔离容器检查 | + +正式包测试覆盖事件对账、跨 Worker response 恢复、取消、认证重试、真实原生文件及 bash 工具、自定义与 MCP 结果关联、本地归档流式下载及清理。云端部署验收是独立步骤，不能由这些测试推定已完成。 + +## 官方 SDK 兼容性 + +dispatcher 组合未修改的官方 SDK。环境级轮询使用 `work.with_raw_response.poll`， +账号级模型轮询通过公开 `get` 访问 VeADK 专属网关端点。HTTP client hooks 实现 +网关 Identity 轮转，生成接口保留 Work 认证。镜像在哈希锁定安装后，将 SDK Python +文件与 wheel RECORD 校验。工具清理保留 close/aclose 和额外 context-manager +退出，仅记录异常类型。发行包内部的 Skill 解包与工具 context-manager 属性仍为 +版本敏感依赖，由集成测试覆盖。详见[官方 SDK 设计](../../prd-spec/refactors/managed-agents-official-sdk/2026-10-10-official-sdk.zh.md) +与[构建指南](../../docker/managed-agents/README.zh.md#官方-sdk-扩展方案)。 + +`MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE` 默认 `session`。显式 `environment` 仅允许 +非 AgentKit 且无 `ark_api_key_provider` 的 Session 使用 `MODEL_AGENT_API_KEY`； +显式 provider 引用及 AgentKit 始终要求 Session Identity。未知来源或缺少本地 Key +明确失败。使用此选项的本地 Kubernetes 验收不证明云 Runtime Identity。 diff --git a/tests/integrations/mpa/conftest.py b/tests/integrations/mpa/conftest.py new file mode 100644 index 000000000..bfbeb60df --- /dev/null +++ b/tests/integrations/mpa/conftest.py @@ -0,0 +1,19 @@ +"""Keep Managed Agents tests isolated from user credentials and other suites.""" + +import pytest + + +@pytest.fixture(autouse=True) +def isolated_managed_agent_environment(monkeypatch): + from veadk.runtime.managed_agents.sandbox import get_default_agent + + for key, value in { + "ANTHROPIC_BASE_URL": "https://sandbox.example.com", + "ANTHROPIC_ENVIRONMENT_KEY": "test-token", + "MODEL_AGENT_API_KEY": "test-model-api-key", + "MODEL_AGENT_NAME": "test-model", + }.items(): + monkeypatch.setenv(key, value) + get_default_agent.cache_clear() + yield + get_default_agent.cache_clear() diff --git a/tests/runtime/test_self_host_sandbox_client.py b/tests/integrations/mpa/test_session_client.py similarity index 54% rename from tests/runtime/test_self_host_sandbox_client.py rename to tests/integrations/mpa/test_session_client.py index 6878de267..548579b5c 100644 --- a/tests/runtime/test_self_host_sandbox_client.py +++ b/tests/integrations/mpa/test_session_client.py @@ -12,21 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import importlib.util -import subprocess -from pathlib import Path +import pytest + +from veadk.integrations.mpa import session_client as MODULE -CLIENT_PATH = ( - Path(__file__).parents[2] - / "examples" - / "16_self_host_sandbox" - / "sandbox_client.py" -) -SPEC = importlib.util.spec_from_file_location("self_host_sandbox_client", CLIENT_PATH) -assert SPEC and SPEC.loader -MODULE = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(MODULE) SelfHostSandboxClient = MODULE.SelfHostSandboxClient @@ -41,6 +31,20 @@ def _client(): ) +def test_injected_anthropic_session_id_takes_precedence(monkeypatch): + monkeypatch.setenv("ANTHROPIC_SESSION_ID", "session-from-work-item") + monkeypatch.setenv("SANDBOX_SESSION_ID", "legacy-session") + + client = SelfHostSandboxClient( + base_url="https://sandbox.example.com", + environment_id="env-123", + agent_id="agent-123", + bearer_token="token", + ) + + assert client.session_id == "session-from-work-item" + + def test_execute_command_posts_agent_tool_use_and_correlates_worker_result( monkeypatch, ): @@ -89,9 +93,14 @@ def test_execute_command_posts_agent_tool_use_and_correlates_worker_result( "input": { "command": "echo hello", "timeout_ms": 5000, - "timeout": 5000, }, - } + }, + { + "type": "agent.tool_result", + "tool_use_id": "call-123", + "content": "exit=0\nhello\n", + "is_error": False, + }, ] assert result == { "dispatch_id": "call-123", @@ -102,6 +111,8 @@ def test_execute_command_posts_agent_tool_use_and_correlates_worker_result( "exit_code": 0, "stdout": "hello\n", "stderr": "", + "content": "exit=0\nhello\n", + "is_error": False, } @@ -126,62 +137,65 @@ def test_send_tool_result_uses_payload_shape(monkeypatch): ] -def test_file_tools_are_converted_to_working_bash_commands(tmp_path): +@pytest.mark.parametrize( + "name,input", + [ + ("read", {"file_path": "a 'quoted' $file.txt", "view_range": [1, -1]}), + ( + "write", + { + "file_path": "notes.txt", + "content": "' \" `$HOME` $(touch nope) &\n中文", + }, + ), + ( + "edit", + { + "file_path": "notes.txt", + "old_string": "old", + "new_string": "new", + "replace_all": True, + }, + ), + ("glob", {"pattern": "**/*.py"}), + ("grep", {"pattern": "a.*b", "path": "."}), + ("bash", {"restart": True, "timeout_ms": 5000}), + ("search_internal_db", {"query": "select 'special'"}), + ], +) +def test_dispatch_preserves_registered_name_and_json(monkeypatch, name, input): client = _client() - target = tmp_path / "nested" / "example.txt" - - write_command, _ = client._tool_to_bash( - "write_file", - {"file_path": str(target), "content": "alpha\nbeta\n"}, - ) - subprocess.run(write_command, shell=True, check=True) - assert target.read_text() == "alpha\nbeta\n" - - edit_command, _ = client._tool_to_bash( - "edit_file", - { - "file_path": str(target), - "old_string": "beta", - "new_string": "gamma", - }, - ) - subprocess.run(edit_command, shell=True, check=True) - assert target.read_text() == "alpha\ngamma\n" - - read_command, _ = client._tool_to_bash( - "read_file", - {"file_path": str(target), "offset": 2, "limit": 1}, - ) - result = subprocess.run( - read_command, - shell=True, - check=True, - capture_output=True, - text=True, + posted = [] + result_content = [{"type": "text", "text": "done"}] + monkeypatch.setattr(client, "post_events", lambda events: posted.extend(events)) + monkeypatch.setattr( + client, + "list_events", + lambda **kw: [ + { + "type": "user.tool_result", + "tool_use_id": "native-1", + "content": result_content, + "is_error": False, + } + ], ) - assert result.stdout == "2\tgamma\n" + result = client.dispatch_tool(name, input, dispatch_id="native-1") + assert posted[0] == { + "type": "agent.tool_use", + "id": "native-1", + "name": name, + "input": input, + } + assert result["content"] == result_content + assert posted[1]["content"] == result_content + assert not hasattr(client, "tool_to_bash") -def test_non_mcp_tools_are_mapped_to_bash_and_unknown_tools_fail(): +def test_bash_command_is_not_html_decoded(monkeypatch): client = _client() - - list_command, _ = client._tool_to_bash( - "list_files", {"path": "/workspace", "max_depth": 2} - ) - search_command, _ = client._tool_to_bash( - "search_files", {"pattern": "needle", "glob": "*.py"} - ) - python_command, _ = client._tool_to_bash( - "python", {"code": "print('ok')", "workdir": "/workspace"} - ) - - assert list_command.startswith("find /workspace") - assert "command -v rg" in search_command - assert "python3 -c" in python_command - - try: - client._tool_to_bash("unknown_tool", {}) - except ValueError as error: - assert "add an explicit bash adapter" in str(error) - else: - raise AssertionError("unknown non-MCP tools must fail closed") + posted = [] + monkeypatch.setattr(client, "post_events", lambda events: posted.extend(events)) + monkeypatch.setattr(client, "_wait_for_tool_result", lambda **kw: {}) + client.execute_command("printf '%s' '&'", timeout=2) + assert posted[0]["input"] == {"command": "printf '%s' '&'", "timeout_ms": 2000} diff --git a/tests/integrations/mpa/test_session_resources.py b/tests/integrations/mpa/test_session_resources.py new file mode 100644 index 000000000..06fd5d651 --- /dev/null +++ b/tests/integrations/mpa/test_session_resources.py @@ -0,0 +1,693 @@ +from __future__ import annotations + +import asyncio +import importlib +import io +import json +import logging +import socket +import subprocess +import sys +import textwrap +import time +import zipfile +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import anthropic +import httpx2 +import pytest + + +resources = importlib.import_module("veadk.integrations.mpa.session_resources") +main = importlib.import_module("veadk.runtime.managed_agents.worker") + + +def test_managed_worker_short_term_memory_is_local_without_database_env( + monkeypatch, +) -> None: + for name in ( + "VEADK_MANAGED_SESSION_DB_URL", + "DATABASE_POSTGRESQL_HOST", + "DATABASE_POSTGRESQL_USER", + "DATABASE_POSTGRESQL_PASSWORD", + "DATABASE_POSTGRESQL_DATABASE", + ): + monkeypatch.delenv(name, raising=False) + + memory = main.managed_short_term_memory() + + assert memory.backend == "local" + assert memory.db_url == "" + assert type(memory.session_service).__name__ == "InMemorySessionService" + + +def test_managed_worker_entry_reaches_dispatcher_without_database_env( + monkeypatch, tmp_path +) -> None: + for name in ( + "VEADK_MANAGED_SESSION_DB_URL", + "DATABASE_POSTGRESQL_HOST", + "DATABASE_POSTGRESQL_USER", + "DATABASE_POSTGRESQL_PASSWORD", + "DATABASE_POSTGRESQL_DATABASE", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("MANAGED_AGENT_READY_FILE", str(tmp_path / "ready")) + entered_poll = False + + class ClientContext: + environment_id = "env-test" + bearer_token = "synthetic-token" + runtime_type = "default" + _default_headers = {"X-Runtime-Type": "default"} + + def create_async_client(self): + sdk = SimpleNamespace(base_url="http://controlled.invalid") + + class Context: + async def __aenter__(self): + return sdk + + async def __aexit__(self, *_args): + return None + + return Context() + + class Dispatcher: + def __init__(self, _sdk, **kwargs): + assert kwargs["account_work"] is False + assert kwargs["environment_id"] == "env-test" + + async def run(self, **_kwargs): + nonlocal entered_poll + entered_poll = True + return 0 + + def drain(self): + return None + + monkeypatch.setattr(main, "SelfHostSandboxClient", ClientContext) + dispatcher_module = ModuleType("anthropic.lib.environments._dispatcher") + dispatcher_module.EnvironmentWorkDispatcher = Dispatcher + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.dispatcher", dispatcher_module + ) + + result = asyncio.run(main.serve_managed_agent_worker(max_work_items=1)) + + assert result == 0 + assert entered_poll is True + assert not (tmp_path / "ready").exists() + + +def test_managed_agent_config_uses_durable_responses_contract() -> None: + config = main.managed_agent_config( + { + "name": "managed", + "model": "test-model", + "tools": [], + "skills": [], + } + ) + + assert config["enable_responses"] is True + assert config["enable_responses_cache"] is True + assert config["model_extra_config"] == { + "store": True, + "retry_expired_response": False, + "extra_body": {"caching": {"type": "disabled"}}, + } + assert config["tools"] == [] + + +def test_managed_agent_config_does_not_add_disabled_tools() -> None: + config = main.managed_agent_config( + { + "name": "managed", + "model": "test-model", + "skills": [], + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": { + "enabled": False, + "permission_policy": {"type": "always_allow"}, + }, + "configs": [ + { + "type": "bash", + "name": "bash", + "enabled": False, + } + ], + } + ], + } + ) + + assert config["tools"] == [] + assert config["before_tool_callback"] is None + + +def test_managed_agent_poll_trace_disabled_by_default(monkeypatch, caplog) -> None: + monkeypatch.delenv("MANAGED_AGENT_POLL_TRACE", raising=False) + with caplog.at_level(logging.INFO): + main.configure_managed_agent_poll_trace( + worker_id="worker-a", + base_url="https://user:password@example.test/private?secret_key=value", + ) + assert not [ + record + for record in caplog.records + if record.name == "anthropic.managed_agent_poll" + ] + + +def test_managed_agent_poll_trace_ready_is_safe(monkeypatch, caplog) -> None: + monkeypatch.setenv("MANAGED_AGENT_POLL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_poll") + logger.addHandler(caplog.handler) + try: + main.configure_managed_agent_poll_trace( + worker_id="worker-a", + base_url="https://user:password@example.test/private?secret_key=synthetic-secret", + ) + finally: + logger.removeHandler(caplog.handler) + record = next( + record + for record in caplog.records + if record.name == "anthropic.managed_agent_poll" + ) + payload = json.loads(record.message) + assert payload == { + "account_scope": True, + "event": "poll_trace_ready", + "origin": "https://example.test", + "path": "/v1/model-work/poll", + "worker_id": "worker-a", + } + assert "synthetic-secret" not in record.message + assert "password" not in record.message + assert logging.getLogger("anthropic._base_client").level == logging.WARNING + assert logging.getLogger("httpx").level == logging.WARNING + + +def test_managed_agent_poll_trace_installs_info_handler_when_root_filters_it( + monkeypatch, +) -> None: + monkeypatch.setenv("MANAGED_AGENT_POLL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_poll") + existing = list(logger.handlers) + try: + main.configure_managed_agent_poll_trace( + worker_id="worker-a", base_url="https://example.test" + ) + assert any( + getattr(handler, "_managed_agent_poll_trace", False) + and handler.level <= logging.INFO + for handler in logger.handlers + ) + finally: + logger.handlers[:] = existing + + +def test_managed_agent_model_trace_disabled_by_default(monkeypatch) -> None: + monkeypatch.delenv("MANAGED_AGENT_MODEL_TRACE", raising=False) + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + try: + logger.handlers.clear() + main.configure_managed_agent_model_trace() + assert not any( + getattr(handler, "_managed_agent_model_trace", False) + for handler in logger.handlers + ) + finally: + logger.handlers[:] = existing + + +def test_managed_agent_model_trace_installs_isolated_handler(monkeypatch) -> None: + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + try: + main.configure_managed_agent_model_trace() + assert logger.propagate is False + assert any( + getattr(handler, "_managed_agent_model_trace", False) + and handler.level <= logging.INFO + for handler in logger.handlers + ) + assert logging.getLogger("anthropic._base_client").level == logging.WARNING + assert logging.getLogger("httpx").level == logging.WARNING + finally: + logger.handlers[:] = existing + + +def test_managed_agent_model_trace_ready_uses_owner_module(monkeypatch, caplog) -> None: + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + try: + logger.handlers.clear() + logger.addHandler(caplog.handler) + main.configure_managed_agent_model_trace() + record = next( + record + for record in caplog.records + if record.name == logger.name and "model_trace_ready" in record.message + ) + payload = json.loads(record.message) + assert payload["event"] == "model_trace_ready" + assert payload["module"].endswith("/veadk/models/ark_llm.py") + finally: + logger.handlers[:] = existing + + +def free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_materialize_session_skills_keeps_pinned_package_files( + monkeypatch, tmp_path +) -> None: + calls = [] + session = SimpleNamespace( + agent=SimpleNamespace( + skills=[SimpleNamespace(type="custom", skill_id="skill_a", version="v1")] + ) + ) + client = object() + + async def download(received_client, *, workdir, session): + calls.append((received_client, Path(workdir), session)) + directory = Path(workdir) / "skills" / "skill-a" + directory.mkdir(parents=True) + (directory / "SKILL.md").write_text("# Pinned instructions") + return [directory] + + monkeypatch.setattr(resources, "download_session_skills", download) + + resolved = asyncio.run( + resources.materialize_session_skills(session, tmp_path, client=client) + ) + + assert len(calls) == 1 + assert calls[0][0] is client + assert calls[0][1].parent == tmp_path.resolve() + assert calls[0][2] is session + assert (resolved[0] / "SKILL.md").read_text() == "# Pinned instructions" + assert "# Pinned instructions" in resources.skill_instructions(resolved) + asyncio.run(resources.cleanup_session_skills(resolved)) + + +def skill_archive() -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("demo/SKILL.md", "# Demo\nRead supporting.txt") + archive.writestr("demo/supporting.txt", "supporting file marker") + return buffer.getvalue() + + +def pinned_session(): + return SimpleNamespace( + id="session-test", + agent=SimpleNamespace( + skills=[SimpleNamespace(skill_id="skill_a", version="v1")] + ), + ) + + +@pytest.mark.asyncio +async def test_real_sdk_skill_files_are_readable_isolated_and_cleaned(tmp_path): + paths = [] + archive = skill_archive() + + def respond(request): + paths.append(request.url.path) + if request.url.path.endswith("/content"): + return httpx2.Response(200, content=archive) + return httpx2.Response(200, json={"id": "resolved-v1", "name": "demo"}) + + async with anthropic.AsyncAnthropic( + api_key="test-key", + base_url="https://skills.example.test", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(respond)), + ) as client: + first = await resources.materialize_session_skills( + pinned_session(), tmp_path, client=client + ) + second = await resources.materialize_session_skills( + pinned_session(), tmp_path, client=client + ) + try: + assert first[0] != second[0] + assert first[0].is_relative_to(tmp_path) + assert second[0].is_relative_to(tmp_path) + assert str(first[0] / "SKILL.md") in resources.skill_instructions(first) + result = await main._run_managed_tool( + "read", {"file_path": str(first[0] / "supporting.txt")}, workdir=tmp_path + ) + assert result["status"] == "completed" + assert "supporting file marker" in result["stdout"] + await resources.cleanup_session_skills(first) + assert not first[0].exists() + assert (second[0] / "supporting.txt").is_file() + assert paths.count("/v1/skills/skill_a/versions/resolved-v1/content") == 2 + finally: + await resources.cleanup_session_skills(first + second) + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.asyncio +async def test_real_sdk_interrupted_archive_stream_is_retryable(tmp_path): + class BrokenStream(httpx2.AsyncByteStream): + async def __aiter__(self): + yield b"partial archive" + raise httpx2.ReadError("connection lost while streaming") + + def respond(request): + if request.url.path.endswith("/content"): + return httpx2.Response(200, stream=BrokenStream()) + return httpx2.Response(200, json={"id": "resolved-v1", "name": "demo"}) + + async with anthropic.AsyncAnthropic( + api_key="test-key", + base_url="https://skills.example.test", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(respond)), + ) as client: + with pytest.raises(resources.SessionSkillMaterializationError) as error: + await resources.materialize_session_skills( + pinned_session(), tmp_path, client=client + ) + assert error.value.retryable is True + assert isinstance(error.value.__cause__, httpx2.ReadError) + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["retrieve", "download"]) +@pytest.mark.parametrize("failure", ["connection", "timeout", 429, 500, 503, 403, 404]) +async def test_real_sdk_skill_failure_retains_retryability(tmp_path, phase, failure): + def respond(request): + is_download = request.url.path.endswith("/content") + if phase == "download" and not is_download: + return httpx2.Response(200, json={"id": "resolved-v1", "name": "demo"}) + if failure == "connection": + raise httpx2.ConnectError("connection failed", request=request) + if failure == "timeout": + raise httpx2.ReadTimeout("request timed out", request=request) + return httpx2.Response(failure, json={"error": {"message": "failed"}}) + + async with anthropic.AsyncAnthropic( + api_key="test-key", + base_url="https://skills.example.test", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(respond)), + ) as client: + with pytest.raises(resources.SessionSkillMaterializationError) as error: + await resources.materialize_session_skills( + pinned_session(), tmp_path, client=client + ) + assert error.value.retryable is (failure not in (403, 404)) + assert isinstance(error.value.__cause__, anthropic.APIError) + assert list(tmp_path.iterdir()) == [] + + +def test_materialize_session_skills_tempfile_default_and_cleanup( + monkeypatch, +) -> None: + session = SimpleNamespace( + id="session-xyz-12345678", + agent=SimpleNamespace( + skills=[SimpleNamespace(type="custom", skill_id="skill_b", version="v1")] + ), + ) + + async def download(client, *, workdir, session): + directory = Path(workdir) / "skills" / "skill-b" + directory.mkdir(parents=True) + (directory / "SKILL.md").write_text("# Tempfile instructions") + return [directory] + + monkeypatch.setattr(resources, "download_session_skills", download) + + resolved = asyncio.run( + resources.materialize_session_skills(session, client=object()) + ) + + assert len(resolved) == 1 + temp_root = resolved[0].parents[1] + assert temp_root in resources._ACTIVE_TEMP_DIRS + assert temp_root.exists() + assert (resolved[0] / "SKILL.md").read_text() == "# Tempfile instructions" + + asyncio.run(resources.cleanup_session_skills(resolved)) + + assert not temp_root.exists() + assert temp_root not in resources._ACTIVE_TEMP_DIRS + + +def test_materialize_session_skills_fails_when_pinned_version_cannot_resolve( + monkeypatch, tmp_path +) -> None: + async def download(client, *, workdir, session): + raise RuntimeError("cannot resolve skill_a@v1") + + monkeypatch.setattr(resources, "download_session_skills", download) + session = SimpleNamespace( + agent=SimpleNamespace( + skills=[SimpleNamespace(skill_id="skill_a", version="v1")] + ) + ) + + with pytest.raises( + resources.SessionSkillMaterializationError, + match="failed to materialize pinned Session Skills", + ) as excinfo: + asyncio.run( + resources.materialize_session_skills(session, tmp_path, client=object()) + ) + assert excinfo.value.retryable is False + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert "cannot resolve skill_a@v1" in str(excinfo.value.__cause__) + + +def test_materialize_session_skills_rejects_partial_sdk_result( + monkeypatch, tmp_path +) -> None: + async def download(client, *, workdir, session): + return [] + + monkeypatch.setattr(resources, "download_session_skills", download) + session = SimpleNamespace( + agent=SimpleNamespace( + skills=[SimpleNamespace(skill_id="skill_a", version="v1")] + ) + ) + + with pytest.raises(resources.SessionSkillMaterializationError) as excinfo: + asyncio.run( + resources.materialize_session_skills(session, tmp_path, client=object()) + ) + + assert excinfo.value.retryable is False + assert "materialized 0 of 1" in str(excinfo.value.__cause__) + + +def test_materialize_session_skills_preserves_retryable_transport_failure( + monkeypatch, tmp_path +) -> None: + class RetryableError(RuntimeError): + retryable = True + + async def download(client, *, workdir, session): + raise RetryableError("temporary transport failure") + + monkeypatch.setattr(resources, "download_session_skills", download) + session = SimpleNamespace( + agent=SimpleNamespace( + skills=[SimpleNamespace(skill_id="skill_a", version="v1")] + ) + ) + + with pytest.raises(resources.SessionSkillMaterializationError) as excinfo: + asyncio.run( + resources.materialize_session_skills(session, tmp_path, client=object()) + ) + + assert excinfo.value.retryable is True + assert isinstance(excinfo.value.__cause__, RetryableError) + + +def test_real_mcp_server_is_restored_from_session_snapshot(tmp_path) -> None: + port = free_port() + server = tmp_path / "mcp_server.py" + server.write_text( + textwrap.dedent( + f""" + from fastmcp import FastMCP + mcp = FastMCP("managed-session-test") + @mcp.tool + def echo_marker(value: str) -> str: + return "mcp-restored:" + value + mcp.run(transport="http", host="127.0.0.1", port={port}, path="/mcp", show_banner=False) + """ + ) + ) + process = subprocess.Popen( + [sys.executable, str(server)], stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + try: + for _ in range(100): + with socket.socket() as sock: + if sock.connect_ex(("127.0.0.1", port)) == 0: + break + time.sleep(0.05) + else: + raise AssertionError("MCP server did not start") + + snapshot = { + "mcp_servers": [ + {"type": "url", "name": "echo", "url": f"http://127.0.0.1:{port}/mcp"} + ], + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "echo", + "default_config": { + "enabled": True, + "permission_policy": {"type": "always_allow"}, + }, + "configs": [], + } + ], + } + toolset = resources.mcp_toolsets(snapshot)[0] + + async def use_tool(): + tools = await toolset.get_tools_with_prefix() + assert [tool.name for tool in tools] == ["mcp__echo__echo_marker"] + result = await tools[0].run_async( + args={"value": "ok"}, tool_context=SimpleNamespace(state={}) + ) + await toolset.close() + return result + + assert "mcp-restored:ok" in str(asyncio.run(use_tool())) + finally: + process.terminate() + process.wait(timeout=5) + + +def test_custom_tool_waits_for_matching_result_and_resumes() -> None: + batches = [] + + class Page: + def __aiter__(self): + async def iterate(): + yield SimpleNamespace( + type="user.custom_tool_result", + custom_tool_use_id="call-1", + content=[{"type": "text", "text": "approved-result"}], + is_error=False, + ) + + return iterate() + + async def send(session_id, *, events): + batches.append((session_id, list(events))) + + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace( + events=SimpleNamespace(send=send, list=lambda *args, **kwargs: Page()) + ) + ) + ) + tool = main.ManagedCustomTool( + { + "type": "custom", + "name": "approve", + "description": "Approve", + "input_schema": {"type": "object"}, + }, + sdk, + "session-1", + ) + result = asyncio.run( + tool.run_async( + args={"item": "x"}, + tool_context=SimpleNamespace(function_call_id="call-1"), + ) + ) + assert result == {"result": "approved-result"} + assert batches[0][1][0]["type"] == "agent.custom_tool_use" + assert batches[0][1][1]["stop_reason"]["action_type"] == "custom_tool_result" + assert batches[1][1] == [{"type": "session.status_running"}] + + +def test_mcp_uses_session_binding_without_mutating_shared_agent(): + snapshot = { + "model": "model", + "mcp_servers": [ + {"type": "url", "name": "orders", "url": "https://mcp.example/mcp"} + ], + "tools": [{"type": "mcp_toolset", "mcp_server_name": "orders"}], + } + refs = [ + { + "Name": "user-key", + "PoolName": "team", + "MCPServerName": "orders", + "HeaderName": "x-api-key", + "Prefix": "", + } + ] + config = main.managed_agent_config(snapshot, credential_keys=refs) + headers = config["tools"][0].identity_headers + assert headers.name == "user-key" and headers.pool_name == "team" + assert headers.header == "x-api-key" and headers.prefix == "" + assert "authorization" not in snapshot["mcp_servers"][0] + other = main.managed_agent_config( + snapshot, credential_keys=[{**refs[0], "Name": "other-key"}] + ) + assert other["tools"][0].identity_headers.name == "other-key" + assert headers.name == "user-key" + assert main.managed_agent_config(snapshot)["tools"][0].identity_headers is None + with pytest.raises(ValueError, match="binding"): + main.managed_agent_config( + snapshot, credential_keys=[{**refs[0], "MCPServerName": "unknown"}] + ) + with pytest.raises(ValueError, match="binding"): + main.managed_agent_config(snapshot, credential_keys=refs + refs) + + +def test_session_workdir_rejects_colliding_unsafe_ids(tmp_path): + with pytest.raises(ValueError): + resources.session_workdir(tmp_path, "a/b") + assert resources.session_workdir(tmp_path, "a_b") == tmp_path / "sessions" / "a_b" + + +@pytest.mark.asyncio +async def test_skill_cleanup_does_not_follow_replaced_directory_symlink(tmp_path): + root = tmp_path / "owned" + skill = root / "skills" / "skill" + skill.mkdir(parents=True) + outside = tmp_path / "outside" + outside.mkdir() + marker = outside / "keep" + marker.write_text("preserve") + resources._ACTIVE_TEMP_DIRS.add(root) + skill.rmdir() + skill.symlink_to(outside, target_is_directory=True) + await resources.cleanup_session_skills([skill]) + assert marker.read_text() == "preserve" + assert not root.exists() + assert root not in resources._ACTIVE_TEMP_DIRS diff --git a/tests/models/test_ark_llm_context_management.py b/tests/models/test_ark_llm_context_management.py index 7c85debbd..76dd10e1e 100644 --- a/tests/models/test_ark_llm_context_management.py +++ b/tests/models/test_ark_llm_context_management.py @@ -39,3 +39,70 @@ def test_request_reorganization_preserves_context_management(): ) assert request["context_management"] == context_management + + +def test_request_reorganization_preserves_tools_with_previous_response_id(): + tools = [ + { + "type": "function", + "name": "bash", + "description": "Run a command", + "parameters": {"type": "object"}, + } + ] + + request = request_reorganization_by_ark( + { + "model": "openai/test-model", + "input": [{"role": "user", "content": "next"}], + "previous_response_id": "response-1", + "tools": tools, + } + ) + + assert request["previous_response_id"] == "response-1" + assert request["tools"] == tools + + +def test_request_reorganization_does_not_add_tools_when_none_are_authorized(): + request = request_reorganization_by_ark( + { + "model": "openai/test-model", + "input": [{"role": "user", "content": "next"}], + "previous_response_id": "response-1", + } + ) + + assert "tools" not in request + + +def test_disabling_responses_cache_removes_previous_id_but_preserves_tools(): + tools = [{"type": "function", "name": "bash", "parameters": {}}] + + request = request_reorganization_by_ark( + { + "model": "openai/test-model", + "input": [], + "previous_response_id": "response-1", + "tools": tools, + "store": True, + }, + enable_responses_cache=False, + ) + + assert "previous_response_id" not in request + assert request["tools"] == tools + + +# Regression: cloud conversation storage must not be shortened to one hour. +def test_response_retention_uses_server_default_or_explicit_value(): + from veadk.models.ark_llm import request_reorganization_by_ark + + default = request_reorganization_by_ark( + {"model": "openai/test", "input": [], "store": True} + ) + assert "expire_at" not in default["extra_body"] + explicit = request_reorganization_by_ark( + {"model": "openai/test", "input": [], "extra_body": {"expire_at": 123456789}} + ) + assert explicit["extra_body"]["expire_at"] == 123456789 diff --git a/tests/models/test_ark_llm_expire_at.py b/tests/models/test_ark_llm_expire_at.py new file mode 100644 index 000000000..2465c31e2 --- /dev/null +++ b/tests/models/test_ark_llm_expire_at.py @@ -0,0 +1,78 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from google.adk.models import LlmRequest + +from veadk.consts import DEFAULT_MODEL_EXTRA_CONFIG +from veadk.models import ark_llm +from veadk.models.ark_llm import ArkLlm, ArkLlmClient + + +def _request(extra_body=None): + request = {"model": "openai/test-model", "input": []} + if extra_body is not None: + request["extra_body"] = extra_body + return request + + +def test_no_default_shortens_conversation_retention(): + """Merged origin/main policy: no injected expire_at, Ark's default applies.""" + assert "expire_at" not in DEFAULT_MODEL_EXTRA_CONFIG["extra_body"] + + request = ark_llm.request_reorganization_by_ark(_request()) + + assert "expire_at" not in request["extra_body"] + + +def test_explicit_expire_at_has_priority_over_server_default(): + request = ark_llm.request_reorganization_by_ark( + _request({"expire_at": 9_999_999, "custom_option": "kept"}) + ) + + assert request["extra_body"] == { + "expire_at": 9_999_999, + "custom_option": "kept", + } + + +async def _consume(model): + return [response async for response in model.generate_content_async(LlmRequest())] + + +@pytest.mark.asyncio +async def test_same_model_instance_keeps_explicit_expire_at(monkeypatch): + requests = [] + + class FakeClient(ArkLlmClient): + async def aresponses(self, **kwargs): + requests.append(kwargs) + return object() + + monkeypatch.setattr( + ark_llm, "ark_response_to_generate_content_response", lambda _: None + ) + + explicit_model = ArkLlm( + model="openai/test-model", + llm_client=FakeClient(), + extra_body={"expire_at": 9_999_999}, + ) + await _consume(explicit_model) + + default_model = ArkLlm(model="openai/test-model", llm_client=FakeClient()) + await _consume(default_model) + + assert requests[0]["extra_body"]["expire_at"] == 9_999_999 + assert "expire_at" not in requests[1]["extra_body"] diff --git a/tests/models/test_ark_llm_fallbacks.py b/tests/models/test_ark_llm_fallbacks.py index dfd2f0bd9..32a155b87 100644 --- a/tests/models/test_ark_llm_fallbacks.py +++ b/tests/models/test_ark_llm_fallbacks.py @@ -13,11 +13,18 @@ # limitations under the License. import httpx +import json +import logging +from types import SimpleNamespace import pytest from google.adk.models import LlmResponse +from google.adk.models import LlmRequest +from google.genai import types from volcenginesdkarkruntime._exceptions import ArkBadRequestError +from volcenginesdkarkruntime.types.responses import ResponseFunctionToolCall from veadk.models.ark_llm import ArkLlm +from veadk.models import ark_llm @pytest.mark.asyncio @@ -136,3 +143,298 @@ async def fake_generate(self, responses_args, stream=False): ("openai/fallback-model", None), ] assert responses[0].model_version == "openai/fallback-model" + + +@pytest.mark.asyncio +async def test_ark_llm_does_not_replay_expired_response_when_disabled(monkeypatch): + model = ArkLlm( + model="openai/primary-model", + retry_expired_response=False, + ) + calls = [] + request = httpx.Request("POST", "https://ark.example/v1/responses") + expired_error = ArkBadRequestError( + "previous response expired", + response=httpx.Response(400, request=request), + body={"code": "InvalidParameter.PreviousResponseNotFound"}, + request_id="request-id", + ) + + async def fake_generate(self, responses_args, stream=False): + calls.append(responses_args.get("previous_response_id")) + raise expired_error + yield # pragma: no cover + + monkeypatch.setattr(ArkLlm, "generate_content_via_responses", fake_generate) + + with pytest.raises(RuntimeError, match="context expired; start a new Session"): + async for _ in model._generate_content_with_fallbacks( + {"input": [], "previous_response_id": "expired-id"} + ): + pass + + assert calls == ["expired-id"] + assert "retry_expired_response" not in model._additional_args + + +def test_model_trace_disabled_by_default(monkeypatch, caplog): + monkeypatch.delenv("MANAGED_AGENT_MODEL_TRACE", raising=False) + with caplog.at_level(logging.INFO, logger="anthropic.managed_agent_model"): + ark_llm._trace_model_request( + {"model": "model-a", "tools": [{"name": "bash", "type": "function"}]} + ) + assert not [ + record + for record in caplog.records + if record.name == "anthropic.managed_agent_model" + ] + + +def test_model_trace_owner_configures_ready_without_root_logger(monkeypatch, capsys): + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + try: + logger.handlers.clear() + ark_llm.configure_managed_agent_model_trace() + captured = capsys.readouterr() + payload = json.loads(captured.err.strip()) + assert payload["event"] == "model_trace_ready" + assert payload["module"].endswith("/veadk/models/ark_llm.py") + assert logger.propagate is False + assert len(logger.handlers) == 1 + finally: + logger.handlers[:] = existing + + +def test_model_trace_whitelists_request_and_response_fields(monkeypatch, caplog): + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_model") + existing_handlers = list(logger.handlers) + existing_propagate = logger.propagate + logger.handlers.clear() + logger.propagate = True + try: + with caplog.at_level(logging.INFO, logger=logger.name): + ark_llm._trace_model_request( + { + "model": "model-a", + "tools": [ + { + "name": "bash", + "type": "function", + "parameters": {"secret": "no"}, + }, + {"name": "bad name", "type": "custom"}, + ], + "tool_choice": {"type": "function", "function": {"name": "bash"}}, + "previous_response_id": "secret-response-id", + "input": "secret prompt", + "api_key": "secret-key", + } + ) + ark_llm._trace_model_response( + SimpleNamespace( + output=[ + SimpleNamespace(type="function_call"), + SimpleNamespace(type="message"), + ], + secret="no", + ) + ) + finally: + logger.handlers[:] = existing_handlers + logger.propagate = existing_propagate + payloads = [ + json.loads(record.message) + for record in caplog.records + if record.name == logger.name + ] + assert payloads == [ + { + "event": "model_request", + "has_previous_response_id": True, + "model": "model-a", + "tool_choice": {"name": "bash", "type": "function"}, + "tools": [ + {"name": "bash", "type": "function"}, + {"name": "", "type": "custom"}, + ], + }, + {"event": "model_response", "output_types": ["function_call", "message"]}, + ] + rendered = "\n".join(record.message for record in caplog.records) + assert "secret prompt" not in rendered + assert "secret-key" not in rendered + assert "secret-response-id" not in rendered + + +@pytest.mark.asyncio +async def test_controlled_function_call_response_is_traced_and_converted( + monkeypatch, caplog +): + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + raw_response = SimpleNamespace( + output=[ + ResponseFunctionToolCall( + type="function_call", + call_id="toolu-safe-1", + name="bash", + arguments='{"command":"printf controlled"}', + id="output-safe-1", + status="completed", + ) + ], + status="completed", + incomplete_details=None, + model="model-a", + usage=None, + id="response-secret-not-logged", + error=None, + ) + + class Client(ark_llm.ArkLlmClient): + async def aresponses(self, **kwargs): + assert kwargs["tools"][0]["name"] == "bash" + return raw_response + + model = ArkLlm(model="openai/model-a", llm_client=Client()) + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + logger.handlers.clear() + logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.INFO, logger=logger.name): + responses = [ + item + async for item in model.generate_content_via_responses( + { + "model": "openai/model-a", + "input": [], + "tools": [{"name": "bash", "type": "function"}], + "tool_choice": {"type": "function", "name": "bash"}, + } + ) + ] + finally: + logger.handlers[:] = existing + + calls = responses[0].content.parts[0].function_call + assert calls.name == "bash" + assert calls.id == "toolu-safe-1" + payloads = [ + json.loads(record.message) + for record in caplog.records + if record.name == logger.name + ] + assert [payload["event"] for payload in payloads] == [ + "model_trace_ready", + "model_request", + "model_response", + ] + assert payloads[1]["tools"] == [{"name": "bash", "type": "function"}] + assert payloads[1]["tool_choice"] == {"name": "bash", "type": "function"} + assert payloads[2] == {"event": "model_response", "output_types": ["function_call"]} + assert "printf controlled" not in "\n".join( + record.message for record in caplog.records + ) + assert "response-secret-not-logged" not in "\n".join( + record.message for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_generate_path_self_configures_model_trace(monkeypatch, capsys): + monkeypatch.setenv("MANAGED_AGENT_MODEL_TRACE", "true") + logger = logging.getLogger("anthropic.managed_agent_model") + existing = list(logger.handlers) + + class Client(ark_llm.ArkLlmClient): + async def aresponses(self, **kwargs): + return SimpleNamespace( + output=[ + ResponseFunctionToolCall( + type="function_call", + call_id="toolu-self-configured", + name="bash", + arguments="{}", + id="output-self-configured", + status="completed", + ) + ], + status="completed", + incomplete_details=None, + model="model-a", + usage=None, + id="response-not-logged", + error=None, + ) + + try: + logger.handlers.clear() + model = ArkLlm(model="openai/model-a", llm_client=Client()) + responses = [ + response + async for response in model.generate_content_via_responses( + { + "model": "openai/model-a", + "input": [], + "tools": [{"name": "bash", "type": "function"}], + } + ) + ] + assert responses[0].content.parts[0].function_call.name == "bash" + events = [ + json.loads(line)["event"] for line in capsys.readouterr().err.splitlines() + ] + assert events == ["model_trace_ready", "model_request", "model_response"] + finally: + logger.handlers[:] = existing + + +@pytest.mark.asyncio +async def test_generate_content_preserves_authorized_tools_with_previous_id( + monkeypatch, +): + requests = [] + + class Client(ark_llm.ArkLlmClient): + async def aresponses(self, **kwargs): + requests.append(kwargs) + return SimpleNamespace( + output=[], + status="completed", + incomplete_details=None, + model="model-a", + usage=None, + id="response-2", + error=None, + ) + + model = ArkLlm(model="openai/model-a", llm_client=Client()) + monkeypatch.setattr( + ark_llm, + "ark_response_to_generate_content_response", + lambda response: LlmResponse(model_version=response.model), + ) + request = LlmRequest( + previous_interaction_id="response-1", + config=types.GenerateContentConfig( + tools=[ + types.Tool( + function_declarations=[ + types.FunctionDeclaration( + name="bash", + description="Run a command", + parameters_json_schema={"type": "object"}, + ) + ] + ) + ] + ), + ) + + _ = [item async for item in model.generate_content_async(request)] + + assert requests[0]["previous_response_id"] == "response-1" + assert [tool["name"] for tool in requests[0]["tools"]] == ["bash"] diff --git a/tests/runtime/managed_agents/__init__.py b/tests/runtime/managed_agents/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/runtime/managed_agents/conftest.py b/tests/runtime/managed_agents/conftest.py new file mode 100644 index 000000000..bfbeb60df --- /dev/null +++ b/tests/runtime/managed_agents/conftest.py @@ -0,0 +1,19 @@ +"""Keep Managed Agents tests isolated from user credentials and other suites.""" + +import pytest + + +@pytest.fixture(autouse=True) +def isolated_managed_agent_environment(monkeypatch): + from veadk.runtime.managed_agents.sandbox import get_default_agent + + for key, value in { + "ANTHROPIC_BASE_URL": "https://sandbox.example.com", + "ANTHROPIC_ENVIRONMENT_KEY": "test-token", + "MODEL_AGENT_API_KEY": "test-model-api-key", + "MODEL_AGENT_NAME": "test-model", + }.items(): + monkeypatch.setenv(key, value) + get_default_agent.cache_clear() + yield + get_default_agent.cache_clear() diff --git a/tests/runtime/managed_agents/cross_worker_http_scenario.py b/tests/runtime/managed_agents/cross_worker_http_scenario.py new file mode 100644 index 000000000..e79402a85 --- /dev/null +++ b/tests/runtime/managed_agents/cross_worker_http_scenario.py @@ -0,0 +1,216 @@ +"""Invoked by ma-infra's isolated PostgreSQL/HTTP contract test. + +Two persistent Python processes have independent Worker caches. Their only +shared state is the public API. Model/tool outputs are deterministic fixtures; +this is not a live LLM or sandbox test. +""" + +import asyncio +import json +import os +import subprocess +import sys +from types import SimpleNamespace + +import anthropic +import httpx2 +from veadk.runtime.managed_agents.dispatcher import sessions_token + +from veadk.runtime.managed_agents.loop import ManagedAgentsLoop + + +class SessionService: + async def get_session(self, **_kwargs): + return object() + + +class Runner: + app_name = "contract" + session_service = SessionService() + + def __init__(self, session_id): + self.session_id = session_id + self.calls = [] + + async def run_async(self, *, new_message, **_kwargs): + prompt = new_message.parts[0].text + self.calls.append(prompt) + call_id = f"call-{prompt}" + await self.sdk.beta.sessions.events.send( + self.session_id, + events=[ + {"id": call_id, "type": "agent.tool_use", "name": "bash", "input": {}}, + {"type": "user.tool_result", "tool_use_id": call_id, "content": []}, + {"type": "agent.tool_result", "tool_use_id": call_id, "content": []}, + ], + ) + yield SimpleNamespace( + partial=False, + content=SimpleNamespace( + parts=[SimpleNamespace(text=prompt, thought=False, function_call=None)] + ), + ) + + +async def worker(): + config = json.loads(sys.stdin.readline()) + runner = Runner(config["session"]) + loop = ManagedAgentsLoop(runner=runner, session_id=config["session"]) + for line in sys.stdin: + command = json.loads(line) + async with anthropic.AsyncAnthropic( + base_url=config["url"], auth_token=command["secret"], max_retries=0 + ) as sdk: + runner.sdk = sdk + if command.get("work"): + await sdk.beta.environments.work.ack( + command["work"], environment_id="env" + ) + await sdk.beta.environments.work.heartbeat( + command["work"], + environment_id="env", + expected_last_heartbeat="NO_HEARTBEAT", + desired_ttl_seconds=90, + ) + count = await loop.run_pending(sdk) + if command.get("work"): + await sdk.beta.environments.work.stop( + command["work"], environment_id="env" + ) + print(json.dumps({"turns": count, "calls": runner.calls}), flush=True) + + +async def scenario(): + config = { + "url": os.environ["MANAGED_AGENTS_TEST_URL"], + "session": os.environ["MANAGED_AGENTS_TEST_SESSION"], + } + # Worker children explicitly receive no database/DSN configuration. + child_env = { + k: v + for k, v in os.environ.items() + if not any( + s in k.upper() + for s in ("POSTGRES", "DATABASE", "DSN", "MANAGED_AGENTS_TEST_TOKEN") + ) + } + children = [ + subprocess.Popen( + [sys.executable, __file__, "worker"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + env=child_env, + ) + for _ in range(2) + ] + try: + for child in children: + child.stdin.write(json.dumps(config) + "\n") + child.stdin.flush() + async with anthropic.AsyncAnthropic( + base_url=config["url"], + auth_token=os.environ["MANAGED_AGENTS_TEST_TOKEN"], + max_retries=0, + ) as sdk: + async with httpx2.AsyncClient(base_url=config["url"]) as raw: + for path in ( + "/v1/model-work/poll", + f"/v1/sessions/{config['session']}/events", + ): + assert (await raw.get(path)).status_code == 401 + assert ( + await raw.get(path, headers={"Authorization": "Bearer invalid"}) + ).status_code == 401 + assert ( + await raw.get("/internal/v1/model-work/poll") + ).status_code == 404 + assert ( + await raw.get( + f"/v1/sessions/{config['session']}/events", + headers={ + "Sec-Fetch-Site": "same-origin", + "Sec-Fetch-Mode": "cors", + }, + ) + ).status_code == 200 + for index, pod in enumerate((0, 1, 0)): + prompt = f"input-{index}" + await sdk.beta.sessions.events.send( + config["session"], + events=[ + { + "id": prompt, + "type": "user.message", + "content": [{"type": "text", "text": prompt}], + } + ], + ) + from anthropic.types.beta.environments import BetaSelfHostedWork + + response = await sdk.get("/v1/model-work/poll", cast_to=httpx2.Response) + assert response.status_code == 200 + claimed = BetaSelfHostedWork.model_validate(response.json()) + assert claimed is not None + child = children[pod] + lease = sessions_token(claimed.secret) + assert lease + child.stdin.write( + json.dumps({"work": claimed.id, "secret": lease}) + "\n" + ) + child.stdin.flush() + reply = json.loads( + await asyncio.wait_for( + asyncio.to_thread(child.stdout.readline), timeout=30 + ) + ) + assert reply["turns"] == 1 and reply["calls"][-1] == prompt, reply + empty = await sdk.get("/v1/model-work/poll", cast_to=httpx2.Response) + assert empty is not None and empty.status_code == 204, ( + "completed input created successor" + ) + + # Stopped lease remains fenced through the edge, never becomes app identity. + async with httpx2.AsyncClient(base_url=config["url"]) as raw: + stale = await raw.get( + f"/v1/sessions/{config['session']}/events", + headers={"Authorization": f"Bearer {lease}"}, + ) + assert stale.status_code in (401, 412), stale.status_code + history = [ + e + async for e in sdk.beta.sessions.events.list( + config["session"], order="asc" + ) + ] + idle = history[-1] + assert idle.type == "session.status_idle" + rows = [ + row + async for row in sdk.beta.environments.work.list( + "env", session_id=config["session"] + ) + ] + assert rows and all(row.state == "stopped" for row in rows) + # Fresh Worker reconstructs completion using public history alone. + fresh = ManagedAgentsLoop( + runner=Runner(config["session"]), session_id=config["session"] + ) + assert await fresh.run_pending(sdk) == 0 + print( + "two isolated Worker processes: three inputs, three executions; terminal Work graph verified" + ) + finally: + for child in children: + child.stdin.close() + for child in children: + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + child.kill() + child.wait() + assert child.returncode == 0, child.returncode + + +if __name__ == "__main__": + asyncio.run(worker() if sys.argv[1:] == ["worker"] else scenario()) diff --git a/tests/runtime/managed_agents/test_agentkit_runtime_identity.py b/tests/runtime/managed_agents/test_agentkit_runtime_identity.py new file mode 100644 index 000000000..90a6b4bac --- /dev/null +++ b/tests/runtime/managed_agents/test_agentkit_runtime_identity.py @@ -0,0 +1,703 @@ +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from veadk.runtime.managed_agents import identity as runtime_identity + + +@pytest.fixture(autouse=True) +def isolate_gateway_credentials(monkeypatch): + monkeypatch.setattr(runtime_identity, "_gateway_providers", {}) + + +class FakeIdentity: + def __init__(self, key: Any = None, fail=None, runtime_id="runtime-id"): + self.key = SimpleNamespace(api_key="secret") if key is None else key + self.fail = fail + self.runtime_id = runtime_id + self.calls = 0 + + def get_workload_access_token(self, *, workload_name): + assert workload_name == self.runtime_id + return SimpleNamespace(workload_access_token="token") + + def get_api_key(self, *, provider_name, agent_identity_token): + self.calls += 1 + assert provider_name == f"{self.runtime_id}-ma" + assert agent_identity_token == "token" + if self.fail: + raise self.fail + return self.key + + +@pytest.mark.asyncio +async def test_missing_workload_is_authentication_failure_not_missing_provider(): + class MissingWorkload(FakeIdentity): + def get_workload_access_token(self, *, workload_name): + raise RuntimeError("Workload NotFound raw-secret") + + fake = MissingWorkload() + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", + client_factory=lambda: fake, + ) + with pytest.raises(runtime_identity.RuntimeIdentityError) as info: + await provider.get() + assert info.value.category == "authentication_failed" + assert fake.calls == 0 + assert "raw-secret" not in str(info.value) + + +@pytest.mark.asyncio +async def test_identity_extracts_sdk_response_and_caches_in_memory(): + fake = FakeIdentity() + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", + client_factory=lambda: fake, + ) + assert await provider.get() == "secret" + assert await provider.get() == "secret" + assert fake.calls == 1 + assert "secret" not in repr(provider) + + +def test_gateway_key_survives_elapsed_time_and_refreshes_explicitly(monkeypatch): + now = [100.0] + monkeypatch.setattr(runtime_identity.time, "monotonic", lambda: now[0]) + fake = FakeIdentity(key="first-key") + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", client_factory=lambda: fake + ) + assert provider.get_sync() == "first-key" + for elapsed in (30, 60, 3600, 86400): + now[0] += elapsed + assert provider.get_sync() == "first-key" + assert fake.calls == 1 + fake.key = "rotated-key" + assert provider.get_sync(force_refresh=True) == "rotated-key" + assert fake.calls == 2 + fake.fail = RuntimeError("AccessDenied raw-secret") + with pytest.raises(runtime_identity.RuntimeIdentityError): + provider.get_sync(force_refresh=True) + with pytest.raises(runtime_identity.RuntimeIdentityError): + provider.get_sync() + assert fake.calls == 4, "Failed refresh must not fall back to a stale key" + + +def test_gateway_credentials_are_shared_during_concurrent_first_use(monkeypatch): + from concurrent.futures import ThreadPoolExecutor + from threading import Barrier + + fake = FakeIdentity() + monkeypatch.setattr(runtime_identity, "IdentityClient", lambda **kwargs: fake) + barrier = Barrier(8) + + def resolve(_): + barrier.wait(timeout=10) + provider = runtime_identity.get_runtime_credential_provider( + runtime_id="runtime-id" + ) + return provider.get_sync() + + with ThreadPoolExecutor(max_workers=8) as executor: + assert list(executor.map(resolve, range(8))) == ["secret"] * 8 + assert fake.calls == 1 + + +def test_gateway_credentials_are_isolated_by_runtime_and_region(monkeypatch): + calls = [] + + class Identity: + def __init__(self, *, region): + self.region = region + + def get_workload_access_token(self, *, workload_name): + return SimpleNamespace(workload_access_token=workload_name) + + def get_api_key(self, *, provider_name, agent_identity_token): + calls.append((agent_identity_token, self.region)) + return f"{provider_name}-{self.region}" + + monkeypatch.setattr(runtime_identity, "IdentityClient", Identity) + references = [ + ("runtime-a", "cn-beijing"), + ("runtime-b", "cn-beijing"), + ("runtime-a", "cn-shanghai"), + ] + for _ in range(2): + for runtime_id, region in references: + provider = runtime_identity.get_runtime_credential_provider( + runtime_id=runtime_id, region=region + ) + assert provider.get_sync() == f"{runtime_id}-ma-{region}" + assert calls == references + + +@pytest.mark.asyncio +async def test_identity_error_is_sanitized(): + fake = FakeIdentity(fail=RuntimeError("AccessDenied raw-secret")) + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", + client_factory=lambda: fake, + ) + with pytest.raises(runtime_identity.RuntimeIdentityError) as info: + await provider.get() + assert info.value.category == "authentication_failed" + assert "raw-secret" not in str(info.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "raw", + [ + "value", + {"ApiKey": "value"}, + {"api_key": "value"}, + SimpleNamespace(api_key="value"), + ], +) +async def test_sdk_key_shapes(raw): + fake = FakeIdentity(key=raw) + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", + client_factory=lambda: fake, + ) + assert await provider.get() == "value" + provider.invalidate() + assert await provider.get() == "value" + assert fake.calls == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "message, category, calls", + [ + ("Provider AccessDenied secret", "authentication_failed", 1), + ("Provider NotFound secret", "provider_not_found", 1), + ("ConnectionError secret", "network_error", 3), + ], +) +async def test_identity_classifies_and_bounds_retries(message, category, calls): + fake = FakeIdentity(fail=RuntimeError(message)) + provider = runtime_identity.RuntimeCredentialProvider( + runtime_id="runtime-id", + client_factory=lambda: fake, + ) + with pytest.raises(runtime_identity.RuntimeIdentityError) as info: + await provider.get() + assert info.value.category == category + assert "secret" not in str(info.value) + assert fake.calls == calls + + +@pytest.mark.asyncio +async def test_existing_client_uses_identity_key_for_sync_and_async_sdk(monkeypatch): + import sys + + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.identity", runtime_identity + ) + monkeypatch.setenv("MA_RUNTIME_ID", "runtime-id") + monkeypatch.delenv("MA_IDENTITY_PROVIDER_NAME", raising=False) + for name in ( + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_API_KEY", + "SANDBOX_BEARER_TOKEN", + "SANDBOX_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + fake = FakeIdentity() + monkeypatch.setattr(runtime_identity, "IdentityClient", lambda **kwargs: fake) + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + client = SelfHostSandboxClient(base_url="https://example.test") + try: + assert client.bearer_token == "secret" + assert client.client.auth_token == "secret" + assert fake.calls == 1 + async with client.create_async_client() as sdk: + assert sdk.auth_token == "secret" + peer = SelfHostSandboxClient(base_url="https://example.test") + try: + assert peer.bearer_token == "secret" + assert fake.calls == 1 + finally: + peer.client.close() + assert "ANTHROPIC_ENVIRONMENT_KEY" not in __import__("os").environ + finally: + client.client.close() + + +def test_explicit_work_token_keeps_existing_authentication(monkeypatch): + import sys + + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.identity", runtime_identity + ) + monkeypatch.setenv("MA_RUNTIME_ID", "runtime-id") + + def unexpected_identity(**kwargs): + raise AssertionError("Explicit Work tokens must not call Identity") + + monkeypatch.setattr(runtime_identity, "IdentityClient", unexpected_identity) + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + client = SelfHostSandboxClient( + base_url="https://example.test", bearer_token="work-token" + ) + try: + assert client.bearer_token == "work-token" + finally: + client.client.close() + + +@pytest.mark.parametrize( + "environment, expected", + [ + ( + { + "AGENTKIT_RUNTIME_ID": "explicit-agentkit", + "MA_RUNTIME_ID": "explicit-ma", + "MODEL_AGENT_CLIENT_REQ_ID": "agentkit/platform-runtime", + }, + "explicit-agentkit", + ), + ( + { + "MA_RUNTIME_ID": "explicit-ma", + "MODEL_AGENT_CLIENT_REQ_ID": "agentkit/platform-runtime", + }, + "explicit-ma", + ), + ( + { + "AGENTKIT_RUNTIME_ID": " ", + "MA_RUNTIME_ID": "explicit-ma", + "MODEL_AGENT_CLIENT_REQ_ID": "agentkit/platform-runtime", + }, + "explicit-ma", + ), + ({"MODEL_AGENT_CLIENT_REQ_ID": "agentkit/r-yev123"}, "r-yev123"), + ({"MODEL_AGENT_CLIENT_REQ_ID": "veadk/0.5.0"}, ""), + ({"MODEL_AGENT_CLIENT_REQ_ID": "agentkit/"}, ""), + ({"MODEL_AGENT_CLIENT_REQ_ID": "agentkit/r-one/extra"}, ""), + ({"MODEL_AGENT_CLIENT_REQ_ID": "prefix-agentkit/r-one"}, ""), + ], +) +def test_runtime_id_environment_resolution_is_strict_and_prioritized( + monkeypatch, environment, expected +): + for name in ( + "AGENTKIT_RUNTIME_ID", + "MA_RUNTIME_ID", + "MODEL_AGENT_CLIENT_REQ_ID", + ): + monkeypatch.delenv(name, raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + from veadk.integrations.mpa.session_client import _runtime_id_from_environment + + assert _runtime_id_from_environment() == expected + + +def test_platform_client_request_id_enables_runtime_identity(monkeypatch): + import sys + + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.identity", runtime_identity + ) + monkeypatch.delenv("AGENTKIT_RUNTIME_ID", raising=False) + monkeypatch.delenv("MA_RUNTIME_ID", raising=False) + monkeypatch.setenv("MODEL_AGENT_CLIENT_REQ_ID", "agentkit/r-platform") + for name in ( + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_API_KEY", + "SANDBOX_BEARER_TOKEN", + "SANDBOX_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + fake = FakeIdentity(runtime_id="r-platform") + monkeypatch.setattr(runtime_identity, "IdentityClient", lambda **kwargs: fake) + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + client = SelfHostSandboxClient(base_url="https://example.test") + try: + assert client._identity_credentials.runtime_id == "r-platform" + assert client.bearer_token == "secret" + finally: + client.client.close() + + +@pytest.mark.parametrize( + "key_name", + [ + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_API_KEY", + "SANDBOX_BEARER_TOKEN", + "SANDBOX_API_KEY", + ], +) +def test_environment_key_skips_identity(monkeypatch, key_name): + import sys + + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.identity", runtime_identity + ) + for name in ( + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_API_KEY", + "SANDBOX_BEARER_TOKEN", + "SANDBOX_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv(key_name, "Bearer supplied-environment-key") + monkeypatch.setenv("MA_RUNTIME_ID", "runtime-id") + + def unexpected_identity(**kwargs): + raise AssertionError("Environment keys must not call Identity") + + monkeypatch.setattr(runtime_identity, "IdentityClient", unexpected_identity) + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + client = SelfHostSandboxClient(base_url="https://example.test") + try: + assert client.client.auth_token == "supplied-environment-key" + finally: + client.client.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_status", [401, 403]) +async def test_transport_pairs_identity_header_with_sdk_work_lease( + monkeypatch, auth_status +): + import sys + from typing import Dict + + import anthropic + import httpx2 + from veadk.runtime.managed_agents.dispatcher import scoped_client + + monkeypatch.setitem( + sys.modules, "veadk.runtime.managed_agents.identity", runtime_identity + ) + monkeypatch.setenv("MA_RUNTIME_ID", "runtime-id") + monkeypatch.delenv("MA_IDENTITY_PROVIDER_NAME", raising=False) + for name in ( + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_API_KEY", + "SANDBOX_BEARER_TOKEN", + "SANDBOX_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + fake = FakeIdentity(key="first-key") + monkeypatch.setattr(runtime_identity, "IdentityClient", lambda **kwargs: fake) + sent = [] + response_status = [200] + + def respond(request): + sent.append(dict(request.headers)) + return httpx2.Response(response_status[0], json={"ok": True}) + + sync_http = anthropic.DefaultHttpxClient + async_http = anthropic.DefaultAsyncHttpxClient + monkeypatch.setattr( + anthropic, + "DefaultHttpxClient", + lambda **kwargs: sync_http(transport=httpx2.MockTransport(respond), **kwargs), + ) + monkeypatch.setattr( + anthropic, + "DefaultAsyncHttpxClient", + lambda **kwargs: async_http(transport=httpx2.MockTransport(respond), **kwargs), + ) + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + client = SelfHostSandboxClient(base_url="https://example.test") + try: + client.client.get("/v1/agents", cast_to=Dict[str, bool]) + assert sent[-1]["x-api-key"] == "first-key" + assert "authorization" not in sent[-1] + async with client.create_async_client() as sdk: + poll = scoped_client(sdk, client.bearer_token) + await poll.get("/v1/model-work/poll", cast_to=Dict[str, bool]) + assert sent[-1]["x-api-key"] == "first-key" + assert "authorization" not in sent[-1] + work = scoped_client(sdk, "work-lease") + await work.get("/v1/sessions/session-id", cast_to=Dict[str, bool]) + assert sent[-1]["x-api-key"] == "first-key" + assert sent[-1]["authorization"] == "Bearer work-lease" + fake.key = "rotated-key" + response_status[0] = auth_status + with pytest.raises(anthropic.APIStatusError): + await work.get("/v1/sessions/session-id", cast_to=Dict[str, bool]) + response_status[0] = 200 + await work.get("/v1/sessions/session-id", cast_to=Dict[str, bool]) + assert sent[-1]["x-api-key"] == "rotated-key" + assert sent[-1]["authorization"] == "Bearer work-lease" + assert fake.calls == 2 + client.client.get("/v1/agents", cast_to=Dict[str, bool]) + assert sent[-1]["x-api-key"] == "rotated-key" + assert fake.calls == 2 + finally: + client.client.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "volcengine_region, region, expected", + [ + ("cn-shanghai", "cn-guangzhou", "cn-shanghai"), + (None, "cn-guangzhou", "cn-guangzhou"), + ("", "cn-guangzhou", "cn-guangzhou"), + (None, None, "cn-beijing"), + ], +) +async def test_session_model_key_uses_configured_identity_region( + monkeypatch, volcengine_region, region, expected +): + monkeypatch.setenv("AGENTKIT_RUNTIME_ID", "runtime-id") + for name, value in (("VOLCENGINE_REGION", volcengine_region), ("REGION", region)): + monkeypatch.delenv(name, raising=False) + if value is not None: + monkeypatch.setenv(name, value) + regions = [] + + class ModelIdentity(FakeIdentity): + def __init__(self, *, region): + super().__init__() + regions.append(region) + + def get_api_key(self, **kwargs): + return "test-model-key" + + monkeypatch.setattr(runtime_identity, "IdentityClient", ModelIdentity) + cache = runtime_identity.SessionModelCredentials() + assert ( + await cache.get({"ark_api_key_provider": "ma-ark-" + "a" * 52}) + == "test-model-key" + ) + assert regions == [expected] + + +@pytest.mark.asyncio +async def test_session_model_key_lru_isolation_and_failure(monkeypatch): + monkeypatch.setenv("AGENTKIT_RUNTIME_ID", "runtime-id") + calls = [] + + class ModelIdentity(FakeIdentity): + def get_api_key(self, *, provider_name, agent_identity_token): + calls.append(provider_name) + if provider_name == "ma-ark-" + "d" * 52: + raise RuntimeError("AccessDenied raw-secret") + return "key-for-" + provider_name + + cache = runtime_identity.SessionModelCredentials( + capacity=2, client_factory=ModelIdentity + ) + a, b, c, d = [{"ark_api_key_provider": "ma-ark-" + ch * 52} for ch in "abcd"] + assert await cache.get(a) != await cache.get(b) + assert await cache.get(a) == "key-for-" + a["ark_api_key_provider"] + await cache.get(c) + await cache.get(b) + assert calls == [ + a["ark_api_key_provider"], + b["ark_api_key_provider"], + c["ark_api_key_provider"], + b["ark_api_key_provider"], + ] + assert len(cache._providers) == 2 + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="missing_provider_reference" + ) as error: + await cache.get({}) + assert error.value.category == "missing_provider_reference" + assert not error.value.retryable + for _ in range(2): + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="authentication_failed" + ) as error: + await cache.get(d) + assert "raw-secret" not in str(error.value) + assert calls[-2:] == [d["ark_api_key_provider"]] * 2 + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="invalid_provider_reference" + ): + await cache.get({"ark_api_key_provider": "runtime-id-ma"}) + monkeypatch.delenv("AGENTKIT_RUNTIME_ID") + monkeypatch.delenv("MA_RUNTIME_ID", raising=False) + monkeypatch.delenv("MODEL_AGENT_CLIENT_REQ_ID", raising=False) + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="missing_workload_identity" + ): + await cache.get(a) + + +@pytest.mark.asyncio +async def test_session_model_key_default_capacity_and_one_hour_ttl(monkeypatch): + import base64 + + monkeypatch.setenv("AGENTKIT_RUNTIME_ID", "runtime-id") + now = [100.0] + monkeypatch.setattr(runtime_identity.time, "monotonic", lambda: now[0]) + calls = [] + + class ModelIdentity(FakeIdentity): + def get_api_key(self, *, provider_name, agent_identity_token): + calls.append(provider_name) + return "test-model-key" + + cache = runtime_identity.SessionModelCredentials(client_factory=ModelIdentity) + assert cache.capacity == 2000 + + def session(index): + digest = ( + base64.b32encode(index.to_bytes(32, "big")).decode().rstrip("=").lower() + ) + return {"ark_api_key_provider": "ma-ark-" + digest} + + first = session(0) + await cache.get(first) + now[0] += 3599 + await cache.get(first) + assert len(calls) == 1 + now[0] += 1 + await cache.get(first) + assert len(calls) == 2, "TTL expires at one hour; reads must not extend it" + for index in range(1, 2000): + await cache.get(session(index)) + assert len(cache._providers) == 2000 + await cache.get(first) + await cache.get(session(2000)) + assert len(cache._providers) == 2000 + assert ("runtime-id", first["ark_api_key_provider"]) in cache._providers + assert ("runtime-id", session(1)["ark_api_key_provider"]) not in cache._providers + + +@pytest.mark.asyncio +async def test_mcp_identity_headers_use_workload_and_cache_without_model_state( + monkeypatch, +): + monkeypatch.setenv("MA_RUNTIME_ID", "mcp-runtime") + name = "user-orders-key" + calls = [] + + class Identity: + def get_workload_access_token(self, *, workload_name): + assert workload_name == "mcp-runtime" + return SimpleNamespace(workload_access_token="workload-token") + + def get_api_key(self, *, provider_name, agent_identity_token, pool_name): + assert pool_name == "team-pool" + calls.append(provider_name) + assert provider_name == name and agent_identity_token == "workload-token" + return "mcp-private-key" + + headers = runtime_identity.MCPIdentityHeaders( + {"Name": name, "PoolName": "team-pool", "MCPServerName": "orders"}, + client_factory=Identity, + ) + await headers.prepare() + assert headers(None) == {"Authorization": "Bearer mcp-private-key"} + await headers.prepare() + assert calls == [name] + assert "mcp-private-key" not in repr(headers) + with pytest.raises(runtime_identity.RuntimeIdentityError): + runtime_identity.RuntimeCredentialProvider( + runtime_id="mcp-runtime", provider_name=name + ) + + +@pytest.mark.asyncio +async def test_mcp_identity_failure_does_not_fall_back(monkeypatch): + monkeypatch.setenv("MA_RUNTIME_ID", "mcp-runtime") + + class Denied: + def get_workload_access_token(self, **kwargs): + raise RuntimeError("403 sensitive-key") + + headers = runtime_identity.MCPIdentityHeaders( + { + "Name": "user-orders-key", + "PoolName": "default", + "MCPServerName": "orders", + }, + client_factory=Denied, + ) + with pytest.raises(runtime_identity.RuntimeIdentityError) as caught: + await headers.prepare() + assert "sensitive-key" not in str(caught.value) + assert caught.value.category == "authentication_failed" + with pytest.raises(runtime_identity.RuntimeIdentityError): + runtime_identity.MCPIdentityHeaders( + {"type": "api_key", "api_key": "raw-secret"} + ) + + +def test_identity_client_forwards_user_pool_to_api(): + from veadk.integrations.ve_identity.identity_client import IdentityClient + + requests = [] + + class API: + def get_resource_api_key(self, request): + requests.append(request.to_dict()) + return SimpleNamespace(api_key="synthetic-key") + + client = SimpleNamespace(_api_client=API()) + key = IdentityClient.get_api_key.__wrapped__( + client, + provider_name="user-provider", + pool_name="team-pool", + agent_identity_token="token", + ) + assert key == "synthetic-key" + assert requests == [ + { + "provider_name": "user-provider", + "pool_name": "team-pool", + "identity_token": "token", + } + ] + + +@pytest.mark.asyncio +async def test_local_worker_requires_explicit_model_source_and_never_bypasses_identity( + monkeypatch, +): + from veadk.runtime.managed_agents.worker import resolve_session_model_key + + monkeypatch.setenv("MODEL_AGENT_API_KEY", "local-model-fixture") + monkeypatch.delenv("MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE", raising=False) + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="missing_provider_reference" + ): + await resolve_session_model_key({}, "default") + monkeypatch.setenv("MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE", "environment") + assert await resolve_session_model_key({}, "default") == "local-model-fixture" + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="missing_provider_reference" + ): + await resolve_session_model_key({}, "agentkit_runtime") + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="invalid_provider_reference" + ): + await resolve_session_model_key({"ark_api_key_provider": "invalid"}, "default") + monkeypatch.delenv("MODEL_AGENT_API_KEY") + with pytest.raises( + runtime_identity.RuntimeIdentityError, match="missing_model_key" + ): + await resolve_session_model_key({}, "default") + + +@pytest.mark.asyncio +async def test_invalid_model_credential_source_fails_closed(monkeypatch): + from veadk.runtime.managed_agents.worker import resolve_session_model_key + + monkeypatch.setenv("MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE", "unknown") + with pytest.raises(ValueError, match="credential source"): + await resolve_session_model_key({}, "default") diff --git a/tests/runtime/managed_agents/test_cross_worker_completion.py b/tests/runtime/managed_agents/test_cross_worker_completion.py new file mode 100644 index 000000000..b5a4cc523 --- /dev/null +++ b/tests/runtime/managed_agents/test_cross_worker_completion.py @@ -0,0 +1,221 @@ +"""Two independent Pod caches consuming the same canonical event ledger.""" + +import asyncio +from types import SimpleNamespace + +from .test_managed_agent_loop import ManagedAgentsLoop, _Runner, _SDK, _AsyncPage + + +class Ledger: + def __init__(self): + self.events = [] + + def append(self, kind, **fields): + seq = len(self.events) + 1 + item = {"seq": seq, "id": f"event-{seq}", "type": kind, **fields} + self.events.append(item) + return item + + def list(self, _session, **options): + cursor = int(options.get("page", 0)) + return _AsyncPage([e for e in self.events if e["seq"] > cursor]) + + +def test_two_pods_do_not_repeat_other_pods_completed_input(): + async def check(): + ledger = Ledger() + sdk = _SDK() + sdk.beta.sessions.events.list = ledger.list + + # Preserve actual _run_turn and persisted completion, not a fake guard. + async def send(_session, *, events): + data = [] + for event in events: + fields = dict(event) + kind = fields.pop("type") + data.append(SimpleNamespace(**ledger.append(kind, **fields))) + return SimpleNamespace(data=data) + + sdk.beta.sessions.events.send = send + pods = [ + ManagedAgentsLoop(runner=_Runner([]), session_id="session") + for _ in range(2) + ] + ledger.append("user.message", content=[{"type": "text", "text": "first"}]) + assert await pods[0].run_pending(sdk) == 1 + ledger.append("user.message", content=[{"type": "text", "text": "recovery"}]) + assert await pods[1].run_pending(sdk) == 1 + # The stale Pod returns on a delayed successor, after the other Pod's idle. + assert await pods[0].run_pending(sdk) == 0 + fresh = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + assert await fresh.run_pending(sdk) == 0 + ledger.append("user.message", content=[{"type": "text", "text": "next"}]) + assert await pods[0].run_pending(sdk) == 1 + assert await pods[1].run_pending(sdk) == 0 + assert [len(p.runner.calls) for p in pods] == [2, 1] + + asyncio.run(check()) + + +def test_interrupt_cancels_all_earlier_inputs_without_consuming_next_turn(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + loop._restore_completed_inputs( + [ + {"id": "a", "type": "user.message"}, + {"id": "b", "type": "user.message"}, + {"id": "interrupt", "type": "user.interrupt"}, + {"id": "next", "type": "user.message"}, + ] + ) + assert loop._completed_inputs == {"a", "b"} + + +def test_duplicate_terminal_replay_does_not_consume_next_queued_input(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + one = {"id": "one", "type": "user.message", "seq": 1} + two = {"id": "two", "type": "user.message", "seq": 2} + idle = { + "id": "idle", + "type": "session.status_idle", + "seq": 3, + "stop_reason": {"type": "end_turn"}, + } + loop._restore_completed_inputs([one, two, idle, idle]) + assert loop._completed_inputs == {"one"} + loop._restore_completed_inputs([one, two, idle]) + assert loop._completed_inputs == {"one"} + + +def test_fresh_workers_reconcile_actual_sdk_http_event_ledger(): + import json + import anthropic + import httpx2 + + ledger = [] + requests = [] + + def append(event): + seq = len(ledger) + 1 + persisted = { + "id": f"event-{seq}", + "seq": seq, + "processed_at": "2026-10-10T00:00:00Z", + **event, + } + ledger.append(persisted) + return persisted + + def transport(request): + requests.append(request) + assert request.url.path == "/v1/sessions/session/events" + if request.method == "POST": + body = json.loads(request.content) + return httpx2.Response( + 200, json={"data": [append(event) for event in body["events"]]} + ) + cursor = int(request.url.params.get("page", "0")) + return httpx2.Response( + 200, + json={ + "data": [event for event in ledger if event["seq"] > cursor], + "has_more": False, + }, + ) + + async def check(): + pods = [ + ManagedAgentsLoop(runner=_Runner([]), session_id="session") + for _ in range(2) + ] + async with anthropic.AsyncAnthropic( + auth_token="synthetic-token", + base_url="http://ledger.invalid", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(transport)), + ) as sdk: + for prompt, pod in (("first", 0), ("second", 1), ("third", 0)): + append( + { + "type": "user.message", + "content": [{"type": "text", "text": prompt}], + } + ) + assert await pods[pod].run_pending(sdk) == 1 + assert await pods[1 - pod].run_pending(sdk) == 0 + fresh = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + assert await fresh.run_pending(sdk) == 0 + assert [len(pod.runner.calls) for pod in pods] == [2, 1] + + asyncio.run(check()) + assert any(request.url.params.get("page") for request in requests) + + +def test_transient_event_publish_uses_actual_sdk_transport(caplog): + import json + import anthropic + import httpx2 + + captured = [] + + def transport(request): + captured.append(json.loads(request.content)) + return httpx2.Response(200, json={"data": []}) + + async def check(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + async with anthropic.AsyncAnthropic( + auth_token="synthetic-token", + base_url="http://ledger.invalid", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(transport)), + ) as sdk: + await loop._publish_transient( + sdk, {"type": "agent.message_stream_start", "message_id": "message"} + ) + await loop._publish_transient( + sdk, + { + "type": "agent.message_chunk", + "message_id": "message", + "delta": "hello", + }, + ) + + asyncio.run(check()) + assert [entry["events"][0]["type"] for entry in captured] == [ + "event_start", + "event_delta", + ] + assert "publish failed" not in caplog.text + + +def test_transient_transport_failure_propagates_without_secret_logs(caplog): + import anthropic + import httpx2 + import pytest + + def transport(request): + return httpx2.Response( + 500, + json={ + "type": "error", + "error": {"type": "api_error", "message": "secret-provider-marker"}, + }, + ) + + async def check(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session") + async with anthropic.AsyncAnthropic( + auth_token="synthetic-token", + base_url="http://ledger.invalid", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(transport)), + ) as sdk: + with pytest.raises(anthropic.APIStatusError): + await loop._publish_transient( + sdk, {"type": "agent.message_stream_start", "message_id": "message"} + ) + + asyncio.run(check()) + assert "publish failed" in caplog.text + assert "secret-provider-marker" not in caplog.text diff --git a/tests/runtime/managed_agents/test_dispatcher.py b/tests/runtime/managed_agents/test_dispatcher.py new file mode 100644 index 000000000..ddcb02c95 --- /dev/null +++ b/tests/runtime/managed_agents/test_dispatcher.py @@ -0,0 +1,524 @@ +"""Protocol tests use the published Anthropic SDK, not dispatcher doubles.""" + +import asyncio +import base64 +import json + +import anthropic +import httpx2 +import pytest + +from veadk.runtime.managed_agents.dispatcher import EnvironmentWorkDispatcher + + +def item(number=1, session="session-1", secret=True): + return { + "id": f"work-{number}", + "created_at": "2026-10-10T00:00:00Z", + "data": {"id": session, "type": "session"}, + "environment_id": "env-1", + "metadata": {}, + "secret": base64.urlsafe_b64encode( + json.dumps({"sessions_token": f"scoped-{number}"}).encode() + ).decode() + if secret + else None, + "state": "starting", + "type": "work", + } + + +def heartbeat(stamp="stamp-1", **kwargs): + return { + "last_heartbeat": stamp, + "lease_extended": True, + "state": "active", + "ttl_seconds": 90, + "type": "work_heartbeat", + **kwargs, + } + + +class Protocol: + def __init__(self, items=(), *, heartbeat_status=200, poll_status=None): + self.items = list(items) + self.requests = [] + self.heartbeat_status = heartbeat_status + self.poll_status = poll_status + self.poll_block = None + + async def __call__(self, request): + self.requests.append(request) + path = request.url.path + if path.endswith("/poll"): + if self.poll_block: + await self.poll_block.wait() + if self.poll_status: + return httpx2.Response( + self.poll_status, + json={ + "error": {"type": "authentication_error", "message": "rejected"} + }, + ) + return ( + httpx2.Response(200, json=self.items.pop(0)) + if self.items + else httpx2.Response(204) + ) + if path.endswith("/heartbeat"): + return httpx2.Response( + self.heartbeat_status, + json=heartbeat() + if self.heartbeat_status == 200 + else {"error": {"type": "lease_error", "message": "rejected"}}, + ) + return httpx2.Response(200, json=item()) + + def paths(self, suffix): + return [r for r in self.requests if r.url.path.endswith(suffix)] + + +def client(protocol): + return anthropic.AsyncAnthropic( + api_key="parent-key", + auth_token="parent-token", + base_url="https://ma.example", + max_retries=0, + default_headers={ + "authorization": "Bearer parent-custom", + "x-api-key": "parent-custom-key", + "X-Custom": "preserved", + "X-Top-Account-Id": "test-account", + "X-Runtime-Type": "test-runtime", + }, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(protocol)), + ) + + +def dispatcher(c, handler, **kwargs): + return EnvironmentWorkDispatcher( + c, + handler=handler, + environment_id="env-1", + environment_key="environment-key", + worker_id="worker-1", + **kwargs, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("account_work", [False, True]) +async def test_wire_scope_ack_fence_and_stop(account_work): + protocol = Protocol([item()]) + async with client(protocol) as c: + + async def handle(work, scoped): + await scoped.beta.sessions.retrieve("session-1") + + d = dispatcher(c, handle, account_work=account_work) + assert await d.run(max_items=1) == 1 + assert d.ready.is_set() + assert len(protocol.paths("/ack")) == len(protocol.paths("/stop")) == 1 + poll = protocol.paths("/poll")[0] + assert poll.headers["authorization"] == "Bearer environment-key" + assert poll.url.params["block_ms"] == "999" + assert poll.headers["Anthropic-Worker-ID"] == "worker-1" + if account_work: + assert poll.url.path == "/v1/model-work/poll" + for request in protocol.requests: + assert "x-api-key" not in request.headers + assert request.headers["X-Custom"] == "preserved" + assert request.headers["X-Top-Account-Id"] == "test-account" + assert request.headers["X-Runtime-Type"] == "test-runtime" + if not request.url.path.endswith("/poll"): + assert request.headers["authorization"] == "Bearer scoped-1" + assert ( + protocol.paths("/heartbeat")[0].url.params["expected_last_heartbeat"] + == "NO_HEARTBEAT" + ) + assert json.loads(protocol.paths("/stop")[0].content)["force"] is True + + +@pytest.mark.asyncio +async def test_same_session_fifo_and_concurrency_bound(): + protocol = Protocol([item(1), item(2), item(3, "session-2")]) + first = asyncio.Event() + started = [] + async with client(protocol) as c: + + async def handle(work, scoped): + started.append(work.id) + if work.id == "work-1": + await first.wait() + + d = dispatcher(c, handle) + task = asyncio.create_task(d.run(max_items=3, max_concurrency=2)) + for _ in range(50): + if len(protocol.paths("/ack")) >= 2: + break + await asyncio.sleep(0.01) + assert started == ["work-1"] + assert len(protocol.paths("/poll")) == 2 + first.set() + assert await asyncio.wait_for(task, 2) == 3 + assert started.index("work-1") < started.index("work-2") + + +@pytest.mark.asyncio +async def test_drain_interrupts_poll_wait_and_preserves_accepted_handler(): + protocol = Protocol([item()]) + finished = asyncio.Event() + started = asyncio.Event() + async with client(protocol) as c: + + async def handle(*args): + started.set() + await finished.wait() + + d = dispatcher(c, handle) + task = asyncio.create_task(d.run(max_concurrency=2)) + await asyncio.wait_for(started.wait(), 1) + protocol.poll_block = asyncio.Event() + await asyncio.sleep(0.3) + d.drain() + assert not task.done() + finished.set() + assert await asyncio.wait_for(task, 1) == 1 + assert len(protocol.paths("/stop")) == 1 + + +@pytest.mark.asyncio +async def test_lease_loss_cancels_handler_without_stopping_new_owner(): + protocol = Protocol([item()], heartbeat_status=412) + ran = [] + async with client(protocol) as c: + + async def handle(*args): + ran.append(True) + await asyncio.Event().wait() + + assert await dispatcher(c, handle).run(max_items=1) == 1 + assert not ran + assert not protocol.paths("/stop") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["poll", "heartbeat"]) +async def test_authentication_failure_propagates_for_credential_refresh(phase): + protocol = Protocol( + [item()], + heartbeat_status=401 if phase == "heartbeat" else 200, + poll_status=401 if phase == "poll" else None, + ) + async with client(protocol) as c: + + async def handle(*args): + await asyncio.Event().wait() + + d = dispatcher(c, handle) + with pytest.raises(anthropic.AuthenticationError): + await asyncio.wait_for(d.run(max_items=1), 1) + if phase == "poll": + assert not d.ready.is_set() + + +@pytest.mark.asyncio +async def test_external_cancellation_waits_for_handler_cleanup_and_stop(): + protocol = Protocol([item()]) + started = asyncio.Event() + cleaned = asyncio.Event() + async with client(protocol) as c: + + async def handle(*args): + started.set() + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0.01) + cleaned.set() + + d = dispatcher(c, handle) + task = asyncio.create_task(d.run(max_items=1)) + await asyncio.wait_for(started.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert cleaned.is_set() + assert len(protocol.paths("/stop")) == 1 + + +@pytest.mark.asyncio +async def test_handler_failure_isolated_and_secret_details_not_logged(caplog): + protocol = Protocol([item(1), item(2)]) + ran = [] + async with client(protocol) as c: + + async def handle(work, scoped): + ran.append(work.id) + if work.id == "work-1": + raise RuntimeError("private-bearer-secret") + + assert await dispatcher(c, handle).run(max_items=2) == 2 + assert ran == ["work-1", "work-2"] + assert len(protocol.paths("/stop")) == 2 + assert "private-bearer-secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_account_poll_rejects_malformed_secret_without_logging_body(caplog): + protocol = Protocol([item(secret=False)]) + async with client(protocol) as c: + d = dispatcher(c, lambda *args: None, account_work=True) + with pytest.raises(ValueError, match="malformed"): + await d.run(max_items=1) + assert not d.ready.is_set() + assert not protocol.paths("/ack") + + +@pytest.mark.asyncio +async def test_heartbeat_echoes_timestamp_and_cancels_on_control_plane_stop( + monkeypatch, +): + from veadk.runtime.managed_agents import dispatcher as module + + monkeypatch.setattr(module, "HEARTBEAT_INTERVAL", 0.01) + protocol = Protocol([item()]) + beats = 0 + original = protocol.__call__ + + async def transport(request): + nonlocal beats + if request.url.path.endswith("/heartbeat"): + protocol.requests.append(request) + beats += 1 + return httpx2.Response( + 200, + json=heartbeat( + f"stamp-{beats}", state="stopping" if beats == 2 else "active" + ), + ) + return await original(request) + + started, cleaned = asyncio.Event(), asyncio.Event() + async with client(transport) as c: + + async def handle(*args): + started.set() + try: + await asyncio.Event().wait() + finally: + cleaned.set() + + await dispatcher(c, handle).run(max_items=1) + assert started.is_set() and cleaned.is_set() + assert ( + protocol.paths("/heartbeat")[1].url.params["expected_last_heartbeat"] + == "stamp-1" + ) + assert len(protocol.paths("/stop")) == 1 + + +@pytest.mark.asyncio +async def test_transient_heartbeat_failures_expire_lease_without_stop(monkeypatch): + from veadk.runtime.managed_agents import dispatcher as module + + monkeypatch.setattr(module, "HEARTBEAT_INTERVAL", 0.01) + monkeypatch.setattr(module, "LEASE_TTL", 0.03) + protocol = Protocol([item()], heartbeat_status=503) + async with client(protocol) as c: + + async def handle(*args): + pytest.fail("unfenced lease must not run a handler") + + await asyncio.wait_for(dispatcher(c, handle).run(max_items=1), 1) + assert len(protocol.paths("/heartbeat")) > 1 + assert not protocol.paths("/stop") + + +@pytest.mark.asyncio +async def test_empty_poll_marks_ready_but_claims_no_work(): + protocol = Protocol() + async with client(protocol) as c: + + async def handle(*args): + pytest.fail("204 is not Work") + + d = dispatcher(c, handle) + task = asyncio.create_task(d.run()) + await asyncio.wait_for(d.ready.wait(), 1) + d.drain() + assert await asyncio.wait_for(task, 1) == 0 + assert not protocol.paths("/ack") + + +@pytest.mark.asyncio +async def test_safe_poll_trace_reports_metadata_only(monkeypatch, caplog): + monkeypatch.setenv("MANAGED_AGENT_POLL_TRACE", "true") + caplog.set_level("INFO", logger="anthropic.managed_agent_poll") + protocol = Protocol([item()]) + async with client(protocol) as c: + + async def handle(*args): + pass + + await dispatcher(c, handle, account_work=True).run(max_items=1) + assert "poll_started" in caplog.text and "poll_completed" in caplog.text + assert "scoped-1" not in caplog.text and "environment-key" not in caplog.text + assert item()["secret"] not in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_items,max_concurrency", [(0, 1), (1, 0), (-1, 2)]) +async def test_invalid_limits_fail_before_poll(max_items, max_concurrency): + protocol = Protocol() + async with client(protocol) as c: + with pytest.raises(ValueError): + await dispatcher(c, lambda *args: None).run( + max_items=max_items, max_concurrency=max_concurrency + ) + assert not protocol.requests + + +@pytest.mark.asyncio +async def test_independent_sessions_execute_concurrently(): + protocol = Protocol([item(1, "session-1"), item(2, "session-2")]) + release = asyncio.Event() + both = asyncio.Event() + running = [] + async with client(protocol) as c: + + async def handle(work, scoped): + running.append(work.id) + if len(running) == 2: + both.set() + await release.wait() + + task = asyncio.create_task( + dispatcher(c, handle).run(max_items=2, max_concurrency=2) + ) + await asyncio.wait_for(both.wait(), 1) + release.set() + assert await task == 2 + + +@pytest.mark.asyncio +async def test_heartbeat_auth_failure_interrupts_blocked_poll(): + protocol = Protocol([item()]) + original = protocol.__call__ + + async def transport(request): + if request.url.path.endswith("/heartbeat"): + await asyncio.sleep(0.05) + return httpx2.Response(401, json={"error": {"message": "rejected"}}) + response = await original(request) + if request.url.path.endswith("/poll") and response.status_code == 200: + protocol.poll_block = asyncio.Event() + return response + + async with client(transport) as c: + + async def handle(*args): + await asyncio.Event().wait() + + with pytest.raises(anthropic.AuthenticationError): + await asyncio.wait_for(dispatcher(c, handle).run(max_concurrency=2), 1) + + +@pytest.mark.asyncio +async def test_nonterminal_unextended_lease_does_not_force_stop(): + protocol = Protocol([item()]) + original = protocol.__call__ + + async def transport(request): + if request.url.path.endswith("/heartbeat"): + return httpx2.Response(200, json=heartbeat(lease_extended=False)) + return await original(request) + + async with client(transport) as c: + + async def handle(*args): + pytest.fail("unextended lease cannot run") + + await dispatcher(c, handle).run(max_items=1) + assert not protocol.paths("/stop") + + +@pytest.mark.asyncio +async def test_drain_interrupts_retry_backoff(): + protocol = Protocol(poll_status=503) + async with client(protocol) as c: + d = dispatcher(c, lambda *args: None) + task = asyncio.create_task(d.run()) + for _ in range(20): + if protocol.requests: + break + await asyncio.sleep(0.01) + d.drain() + assert await asyncio.wait_for(task, 0.5) == 0 + assert not d.ready.is_set() + + +@pytest.mark.asyncio +async def test_stop_auth_failure_propagates_and_releases_session_barrier(): + protocol = Protocol([item()]) + original = protocol.__call__ + + async def transport(request): + if request.url.path.endswith("/stop"): + return httpx2.Response(403, json={"error": {"message": "rejected"}}) + return await original(request) + + async with client(transport) as c: + + async def handle(*args): + pass + + d = dispatcher(c, handle) + with pytest.raises(anthropic.PermissionDeniedError): + await d.run(max_items=1) + assert not d._session_tails + + +@pytest.mark.asyncio +async def test_environment_poll_uses_official_generated_raw_api(monkeypatch): + protocol = Protocol([item()]) + async with client(protocol) as sdk: + resource_type = type(sdk.beta.environments.work.with_raw_response) + original = resource_type.__init__ + calls = [] + + def wrap(resource, work): + original(resource, work) + poll = resource.poll + + async def record(*args, **kwargs): + calls.append((args, kwargs)) + return await poll(*args, **kwargs) + + resource.poll = record + + monkeypatch.setattr(resource_type, "__init__", wrap) + + async def handle(*args): + pass + + await dispatcher(sdk, handle).run(max_items=1) + assert len(calls) == 1 + assert calls[0][0] == ("env-1",) + assert calls[0][1]["anthropic_worker_id"] == "worker-1" + assert calls[0][1]["block_ms"] == 999 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status,body", [(200, {}), (200, None), (202, {})]) +async def test_environment_poll_rejects_invalid_response_before_readiness(status, body): + async def respond(request): + if body is None: + return httpx2.Response(status, text="not Work") + return httpx2.Response(status, json=body) + + async with client(respond) as sdk: + d = dispatcher(sdk, lambda *args: None) + with pytest.raises(ValueError, match="Work poll returned"): + await d.run(max_items=1) + assert not d.ready.is_set() diff --git a/tests/runtime/managed_agents/test_health.py b/tests/runtime/managed_agents/test_health.py new file mode 100644 index 000000000..1a58d378d --- /dev/null +++ b/tests/runtime/managed_agents/test_health.py @@ -0,0 +1,30 @@ +"""Verify the installed listener's public HTTP health and invocation routes.""" + +import httpx +import pytest + +from veadk.runtime.managed_agents import health + + +@pytest.mark.asyncio +async def test_health_routes_without_business_identity(monkeypatch): + for name in ( + "ANTHROPIC_ENVIRONMENT_KEY", + "ANTHROPIC_BASE_URL", + "MODEL_AGENT_API_KEY", + ): + monkeypatch.delenv(name, raising=False) + async with health.app.router.lifespan_context(health.app): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=health.app), base_url="http://runtime" + ) as client: + for path in ("/ping", "/health", "/readiness", "/liveness"): + assert (await client.get(path)).status_code == 200 + response = await client.post( + "/invoke", + json={"work_id": "ignored"}, + headers={"x-api-key": "private-marker"}, + ) + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + assert "private-marker" not in response.text diff --git a/tests/runtime/managed_agents/test_managed_agent_loop.py b/tests/runtime/managed_agents/test_managed_agent_loop.py new file mode 100644 index 000000000..ce0bbf22b --- /dev/null +++ b/tests/runtime/managed_agents/test_managed_agent_loop.py @@ -0,0 +1,945 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import importlib +from types import SimpleNamespace +from typing import get_origin + +import pytest + +from google.adk.agents.run_config import StreamingMode + + +loop_module = importlib.import_module("veadk.runtime.managed_agents.loop") +ManagedAgentsLoop = loop_module.ManagedAgentsLoop +ManagedAgentEventState = loop_module.ManagedAgentEventState + + +def test_response_context_survives_fresh_worker_without_session_database(): + from google.adk.sessions import InMemorySessionService + + async def check(): + runner = _Runner([]) + runner.agent = SimpleNamespace(name="test_agent") + runner.session_service = InMemorySessionService() + loop = ManagedAgentsLoop(runner=runner, session_id="session-ark") + loop._restore_completed_inputs( + [ + {"type": "user.message", "id": "input-1"}, + { + "type": "span.model_request_end", + "metadata": {"ark_response_id": "resp-1"}, + }, + {"type": "session.status_idle"}, + ] + ) + await loop._ensure_local_session() + session = await runner.session_service.get_session( + app_name=runner.app_name, + user_id=loop.user_id, + session_id=loop.session_id, + ) + assert len(session.events) == 1 + assert session.events[0].interaction_id == "resp-1" + assert session.events[0].author == "test_agent" + assert session.events[0].content is None + assert "input-1" in loop._completed_inputs + + asyncio.run(check()) + + +def test_completed_turn_persists_response_reference(): + event = _event(_part("answer")) + event.interaction_id = "resp-completed" + runner = _Runner([event]) + sdk = _SDK() + loop = ManagedAgentsLoop(runner=runner, session_id="session-ark") + asyncio.run(loop._run_turn(sdk, {"content": [{"type": "text", "text": "hello"}]})) + events = [event for _, batch in sdk.beta.sessions.events.batches for event in batch] + end = next(event for event in events if event["type"] == "span.model_request_end") + assert end["metadata"] == {"ark_response_id": "resp-completed"} + + +class _SessionService: + def __init__(self, *, exists=True): + self.session = object() if exists else None + self.created = [] + + async def get_session(self, **kwargs): + return self.session + + async def create_session(self, **kwargs): + self.created.append(kwargs) + self.session = object() + return self.session + + +class _Runner: + app_name = "managed-agent-test" + + def __init__(self, events, *, session_exists=True): + self.events = events + self.session_service = _SessionService(exists=session_exists) + self.calls = [] + + async def run_async(self, **kwargs): + self.calls.append(kwargs) + for event in self.events: + yield event + + +class _EventsAPI: + def __init__(self): + self.batches = [] + + async def send(self, session_id, *, events): + batch = list(events) + self.batches.append((session_id, batch)) + data = [] + for event in batch: + event_id = ( + "span-start-1" + if event["type"] == "span.model_request_start" + else "event" + ) + data.append(SimpleNamespace(type=event["type"], id=event_id)) + return SimpleNamespace(data=data) + + +class _SDK: + def __init__(self): + self.transient = [] + self.beta = SimpleNamespace( + sessions=SimpleNamespace(events=_EventsAPI()), + ) + + async def post(self, path, *, cast_to, body): + self.transient.append((path, cast_to, body)) + return {"accepted": True} + + +class _SessionClient: + session_id = "managed-session-1" + + @staticmethod + def _event_seq(event): + return int(event.get("_seq", 0)) + + +class _AsyncPage: + def __init__(self, events): + self.events = events + + def __aiter__(self): + async def iterate(): + for event in self.events: + yield event + + return iterate() + + +class _EventPages: + def __init__(self, responses): + self.responses = iter(responses) + self.calls = [] + + def __call__(self, session_id, **kwargs): + self.calls.append((session_id, kwargs)) + response = next(self.responses) + if isinstance(response, BaseException): + raise response + return _AsyncPage(response) + + +def _event(*parts, usage=None, partial=False, function_calls=None): + return SimpleNamespace( + partial=partial, + content=SimpleNamespace(parts=list(parts)), + usage_metadata=usage, + get_function_calls=lambda: list(function_calls or []), + ) + + +def _part(text, *, thought=False): + return SimpleNamespace(text=text, thought=thought) + + +def test_debug_covers_runner_transient_and_persisted_events(monkeypatch, capsys): + import json + + from google.adk.events import Event + from google.genai import types + + monkeypatch.setenv("MA_DEBUG_EVENTS", "1") + runner = _Runner( + [ + Event( + author="agent", + partial=partial, + content=types.Content(parts=[types.Part(text=text)]), + ) + for partial, text in [(True, "hello"), (False, "hello world")] + ] + ) + sdk = _SDK() + loop = ManagedAgentsLoop(runner=runner, session_id="session-debug") + asyncio.run(loop._run_turn(sdk, {"content": [{"type": "text", "text": "hi"}]})) + logged = [ + json.loads(line.removeprefix("MA_DEBUG_EVENT ")) + for line in capsys.readouterr().out.splitlines() + if line.startswith("MA_DEBUG_EVENT ") + ] + assert all(record["session_id"] == "session-debug" for record in logged) + assert len([record for record in logged if record["source"] == "runner"]) == 2 + assert { + record["event"]["type"] for record in logged if record["source"] == "transient" + } == {"event_start", "event_delta"} + assert [record["event"] for record in logged if record["source"] == "send"] == [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + + +def test_debug_logs_sse_events_before_filtering(monkeypatch, capsys): + import json + from contextlib import asynccontextmanager + + monkeypatch.setenv("MA_DEBUG_EVENTS", "1") + events = [ + SimpleNamespace(type="user.tool_result", content="tool output"), + SimpleNamespace(type="session.status_terminated"), + ] + + class Stream(_AsyncPage): + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + async def stream(session_id): + assert session_id == "session-sse" + return Stream(events) + + @asynccontextmanager + async def client_context(): + yield SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=SimpleNamespace(stream=stream)) + ) + ) + + client = SimpleNamespace( + session_id="session-sse", + list_events=lambda: [], + create_async_client=client_context, + ) + loop = ManagedAgentsLoop(runner=_Runner([]), session_client=client) + + async def run(): + return await asyncio.wait_for(loop.run(), timeout=5) + + assert asyncio.run(run()) == 0 + logged = [ + json.loads(line.removeprefix("MA_DEBUG_EVENT ")) + for line in capsys.readouterr().out.splitlines() + if line.startswith("MA_DEBUG_EVENT ") + ] + assert [record["event"] for record in logged] == [vars(event) for event in events] + assert all(record["source"] == "sse" for record in logged) + + +def test_turn_maps_veadk_events_to_managed_agent_events(): + usage = SimpleNamespace( + prompt_token_count=3, + candidates_token_count=2, + cached_content_token_count=1, + ) + runner = _Runner( + [ + _event(_part("considering", thought=True)), + _event(_part("ignored"), function_calls=[object()]), + _event(_part("final answer"), usage=usage), + ], + session_exists=False, + ) + sdk = _SDK() + loop = ManagedAgentsLoop(runner=runner, session_client=_SessionClient()) + user_event = SimpleNamespace( + content=[SimpleNamespace(type="text", text="hello")], + ) + + asyncio.run(loop._run_turn(sdk, user_event)) + + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert [event["type"] for event in emitted] == [ + "session.status_running", + "span.model_request_start", + "agent.thinking", + "agent.message", + "span.model_request_end", + "session.status_idle", + ] + assert emitted[2]["text"] == "considering" + assert emitted[3]["content"] == [{"type": "text", "text": "final answer"}] + assert emitted[4] == { + "type": "span.model_request_end", + "model_request_start_id": "span-start-1", + "model_usage": { + "input_tokens": 3, + "output_tokens": 2, + "cache_read_input_tokens": 1, + "cache_creation_input_tokens": 0, + }, + } + assert runner.session_service.created == [ + { + "app_name": "managed-agent-test", + "user_id": "managed_agents_user", + "session_id": "managed-session-1", + } + ] + assert runner.calls[0]["session_id"] == "managed-session-1" + assert runner.calls[0]["new_message"].parts[0].text == "hello" + assert runner.calls[0]["run_config"].streaming_mode == StreamingMode.SSE + + +def test_turn_broadcasts_partial_text_before_persisting_the_final_message(): + runner = _Runner( + [ + _event(_part("hel"), partial=True), + _event(_part("lo"), partial=True), + _event(_part("hello")), + ] + ) + sdk = _SDK() + loop = ManagedAgentsLoop(runner=runner, session_client=_SessionClient()) + + asyncio.run( + loop._run_turn(sdk, SimpleNamespace(content=[{"type": "text", "text": "hi"}])) + ) + + transient = [call[2]["events"][0] for call in sdk.transient] + assert all(get_origin(call[1]) is dict for call in sdk.transient) + assert all(event["id"].startswith("transient-") for event in transient) + assert all(event["processed_at"] for event in transient) + assert all(call[0].endswith("/events") for call in sdk.transient) + assert [event["type"] for event in transient] == [ + "event_start", + "event_delta", + "event_delta", + ] + assert [event["delta"]["content"]["text"] for event in transient[1:]] == [ + "hel", + "lo", + ] + assert all( + event["delta"]["type"] == "content_delta" and event["delta"]["index"] == 0 + for event in transient[1:] + ) + message_id = transient[0]["event"]["id"] + assert all(event["event_id"] == message_id for event in transient[1:]) + canonical = [ + event + for _, batch in sdk.beta.sessions.events.batches + for event in batch + if event["type"] == "agent.message" + ] + assert canonical == [ + { + "type": "agent.message", + "message_id": message_id, + "content": [{"type": "text", "text": "hello"}], + } + ] + + +def test_turn_closes_partial_stream_as_aborted_when_cancelled(): + class _CancelledRunner(_Runner): + async def run_async(self, **kwargs): + self.calls.append(kwargs) + yield _event(_part("partial"), partial=True) + raise asyncio.CancelledError + + sdk = _SDK() + loop = ManagedAgentsLoop( + runner=_CancelledRunner([]), session_client=_SessionClient() + ) + + async def run(): + try: + await loop._run_turn( + sdk, SimpleNamespace(content=[{"type": "text", "text": "hi"}]) + ) + except asyncio.CancelledError: + return + raise AssertionError("cancelled turn did not propagate cancellation") + + asyncio.run(run()) + + transient = [call[2]["events"][0] for call in sdk.transient] + assert transient[-1]["type"] == "event_delta" + assert all(call[0].endswith("/events") for call in sdk.transient) + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert emitted[-1]["type"] == "span.model_request_end" + assert emitted[-1]["is_error"] is True + assert [event["type"] for event in emitted].count("session.status_idle") == 0 + + +def test_turn_failure_emits_span_error_and_session_error_without_idle(): + class _FailingRunner(_Runner): + async def run_async(self, **kwargs): + if False: + yield None + raise RuntimeError("model failed") + + sdk = _SDK() + loop = ManagedAgentsLoop( + runner=_FailingRunner([]), + session_client=_SessionClient(), + ) + + with pytest.raises(RuntimeError, match="model failed"): + asyncio.run( + loop._run_turn( + sdk, + SimpleNamespace(content=[{"type": "text", "text": "hello"}]), + ) + ) + + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert [event["type"] for event in emitted] == [ + "session.status_running", + "span.model_request_start", + "span.model_request_end", + "session.error", + ] + assert emitted[-1]["error"] == "Managed Agents turn failed" + + +def test_invalid_user_message_emits_session_error_without_starting_runner(): + runner = _Runner([]) + sdk = _SDK() + loop = ManagedAgentsLoop(runner=runner, session_client=_SessionClient()) + + with pytest.raises(ValueError): + asyncio.run( + loop._run_turn( + sdk, + SimpleNamespace(content=[{"type": "image", "source": "ignored"}]), + ) + ) + + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert emitted == [ + { + "type": "session.error", + "error": "Managed Agents turn failed", + } + ] + assert runner.calls == [] + + +def test_recovery_only_marks_user_messages_followed_by_terminal_event_complete(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_client=_SessionClient()) + + loop._restore_completed_inputs( + [ + {"_seq": 1, "id": "old", "type": "user.message"}, + {"_seq": 2, "type": "session.status_running"}, + {"_seq": 3, "type": "session.status_idle"}, + {"_seq": 4, "id": "pending", "type": "user.message"}, + {"_seq": 5, "type": "session.status_running"}, + ] + ) + + assert loop._completed_inputs == {"old"} + + +def test_recovery_pairs_one_terminal_with_one_queued_user_message(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_client=_SessionClient()) + + loop._restore_completed_inputs( + [ + {"id": "first", "type": "user.message"}, + {"id": "second", "type": "user.message"}, + {"id": "running", "type": "session.status_running"}, + {"id": "idle", "type": "session.status_idle"}, + ] + ) + + assert loop._completed_inputs == {"first"} + + +def test_requires_action_does_not_mark_turn_complete_during_recovery(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_client=_SessionClient()) + + loop._restore_completed_inputs( + [ + {"id": "pending", "type": "user.message"}, + { + "id": "waiting", + "type": "session.status_idle", + "stop_reason": {"type": "requires_action"}, + }, + ] + ) + + assert loop._completed_inputs == set() + + +def test_claimed_session_keeps_processing_after_idle(): + async def run(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session-1") + sdk = _SDK() + statuses = iter(["idle", "idle", "running", "terminated"]) + heartbeats = [] + polls = [] + + async def retrieve(_session_id): + return {"status": next(statuses)} + + async def heartbeat(work_id, **kwargs): + heartbeats.append((work_id, kwargs)) + return SimpleNamespace(lease_extended=True, last_heartbeat="hb-2") + + async def pending(_sdk, *, max_turns): + polls.append(max_turns) + return 0 if len(polls) == 2 else 1 + + sdk.beta.sessions.retrieve = retrieve + sdk.beta.environments = SimpleNamespace( + work=SimpleNamespace(heartbeat=heartbeat) + ) + loop.run_pending = pending + turns = await loop.run_claimed( + sdk, + environment_id="env-1", + work_id="work-1", + last_heartbeat="hb-1", + poll_interval=0, + ) + assert turns == 2 + assert polls == [None, None, None] + assert heartbeats[0][1]["expected_last_heartbeat"] == "hb-1" + + asyncio.run(run()) + + +def test_claimed_session_cancels_tools_when_lease_is_lost(): + async def run(): + loop = ManagedAgentsLoop(runner=_Runner([]), session_id="session-1") + sdk = _SDK() + cancelled = asyncio.Event() + + async def retrieve(_session_id): + return {"status": "running"} + + async def heartbeat(*_args, **_kwargs): + return SimpleNamespace(lease_extended=False) + + async def pending(_sdk, *, max_turns): + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + sdk.beta.sessions.retrieve = retrieve + sdk.beta.environments = SimpleNamespace( + work=SimpleNamespace(heartbeat=heartbeat) + ) + loop.run_pending = pending + try: + await loop.run_claimed(sdk, environment_id="env-1", work_id="work-1") + except RuntimeError as error: + assert "lease" in str(error) + else: + raise AssertionError("lost lease did not stop the session") + assert cancelled.is_set() + + asyncio.run(run()) + + +def test_claimed_work_processes_only_pending_message_from_durable_log(): + runner = _Runner([]) + sdk = _SDK() + sdk.beta.sessions.events.list = lambda *args, **kwargs: _AsyncPage( + [ + { + "id": "old", + "type": "user.message", + "content": [{"type": "text", "text": "old"}], + }, + {"id": "idle", "type": "session.status_idle"}, + { + "id": "pending", + "type": "user.message", + "content": [{"type": "text", "text": "new"}], + }, + ] + ) + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + turns = asyncio.run(loop.run_pending(sdk)) + + assert turns == 1 + assert len(runner.calls) == 1 + assert runner.calls[0]["new_message"].parts[0].text == "new" + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert [event["type"] for event in emitted] == [ + "session.status_running", + "span.model_request_start", + "span.model_request_end", + "session.status_idle", + ] + + +def test_claimed_work_reconciles_full_history_once_then_uses_seq_cursor(): + runner = _Runner([]) + sdk = _SDK() + history = [ + {"seq": index, "id": f"system-{index}", "type": "system.message"} + for index in range(1, 1005) + ] + history.append( + { + "seq": 1005, + "id": "first", + "type": "user.message", + "content": [{"type": "text", "text": "first"}], + } + ) + pages = _EventPages( + [ + history, + [ + { + "seq": 1006, + "id": "second", + "type": "user.message", + "content": [{"type": "text", "text": "second"}], + } + ], + ] + ) + sdk.beta.sessions.events.list = pages + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + async def run(): + assert await loop.run_pending(sdk) == 1 + assert await loop.run_pending(sdk) == 1 + + asyncio.run(run()) + + assert pages.calls == [ + ( + "managed-session-1", + {"limit": 100, "order": "asc"}, + ), + ( + "managed-session-1", + {"limit": 100, "order": "asc", "page": "1005"}, + ), + ] + assert [call["new_message"].parts[0].text for call in runner.calls] == [ + "first", + "second", + ] + assert loop._event_cursor_seq == 1006 + + +def test_separate_work_handlers_share_the_session_event_cursor(): + first_runner = _Runner([]) + second_runner = _Runner([]) + sdk = _SDK() + pages = _EventPages( + [ + [ + { + "seq": 1, + "id": "first", + "type": "user.message", + "content": [{"type": "text", "text": "first"}], + } + ], + [ + { + "seq": 2, + "id": "second", + "type": "user.message", + "content": [{"type": "text", "text": "second"}], + } + ], + ] + ) + sdk.beta.sessions.events.list = pages + state = ManagedAgentEventState() + first = ManagedAgentsLoop( + runner=first_runner, session_id="managed-session-1", event_state=state + ) + second = ManagedAgentsLoop( + runner=second_runner, session_id="managed-session-1", event_state=state + ) + + async def run(): + assert await first.run_pending(sdk) == 1 + assert await second.run_pending(sdk) == 1 + + asyncio.run(run()) + + assert "page" not in pages.calls[0][1] + assert pages.calls[1][1]["page"] == "1" + assert first_runner.calls[0]["new_message"].parts[0].text == "first" + assert second_runner.calls[0]["new_message"].parts[0].text == "second" + assert state.cursor_seq == 2 + + +def test_claimed_work_does_not_advance_cursor_past_unprocessed_message(): + runner = _Runner([]) + sdk = _SDK() + pages = _EventPages( + [ + [ + { + "seq": 1, + "id": "first", + "type": "user.message", + "content": [{"type": "text", "text": "first"}], + }, + { + "seq": 2, + "id": "second", + "type": "user.message", + "content": [{"type": "text", "text": "second"}], + }, + ], + [ + { + "seq": 2, + "id": "second", + "type": "user.message", + "content": [{"type": "text", "text": "second"}], + } + ], + ] + ) + sdk.beta.sessions.events.list = pages + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + async def run(): + assert await loop.run_pending(sdk, max_turns=1) == 1 + assert loop._event_cursor_seq == 1 + assert await loop.run_pending(sdk, max_turns=1) == 1 + + asyncio.run(run()) + + assert pages.calls[1][1]["page"] == "1" + assert [call["new_message"].parts[0].text for call in runner.calls] == [ + "first", + "second", + ] + assert loop._event_cursor_seq == 2 + + +def test_claimed_work_falls_back_to_one_full_reconcile_when_cursor_fails(): + runner = _Runner([]) + sdk = _SDK() + pages = _EventPages( + [ + [ + {"seq": 1, "id": "old", "type": "user.message"}, + {"seq": 2, "id": "idle", "type": "session.status_idle"}, + ], + RuntimeError("invalid cursor"), + [ + {"seq": 1, "id": "old", "type": "user.message"}, + {"seq": 2, "id": "idle", "type": "session.status_idle"}, + { + "seq": 3, + "id": "new", + "type": "user.message", + "content": [{"type": "text", "text": "new"}], + }, + ], + ] + ) + sdk.beta.sessions.events.list = pages + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + async def run(): + assert await loop.run_pending(sdk) == 0 + assert await loop.run_pending(sdk) == 1 + + asyncio.run(run()) + + assert pages.calls[1][1]["page"] == "2" + assert "page" not in pages.calls[2][1] + assert [call["new_message"].parts[0].text for call in runner.calls] == ["new"] + assert loop._event_cursor_seq == 3 + + +def test_claimed_work_falls_back_when_server_ignores_cursor(): + runner = _Runner([]) + sdk = _SDK() + pages = _EventPages( + [ + [ + {"seq": 1, "id": "old", "type": "user.message"}, + {"seq": 2, "id": "idle", "type": "session.status_idle"}, + ], + [ + {"seq": 1, "id": "old", "type": "user.message"}, + {"seq": 2, "id": "idle", "type": "session.status_idle"}, + { + "seq": 3, + "id": "new", + "type": "user.message", + "content": [{"type": "text", "text": "new"}], + }, + ], + [ + {"seq": 1, "id": "old", "type": "user.message"}, + {"seq": 2, "id": "idle", "type": "session.status_idle"}, + { + "seq": 3, + "id": "new", + "type": "user.message", + "content": [{"type": "text", "text": "new"}], + }, + ], + ] + ) + sdk.beta.sessions.events.list = pages + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + async def run(): + assert await loop.run_pending(sdk) == 0 + assert await loop.run_pending(sdk) == 1 + + asyncio.run(run()) + + assert pages.calls[1][1]["page"] == "2" + assert "page" not in pages.calls[2][1] + assert [call["new_message"].parts[0].text for call in runner.calls] == ["new"] + assert loop._event_cursor_seq == 3 + + +def test_claimed_work_without_event_sequences_stays_on_full_reconciliation(): + runner = _Runner([]) + sdk = _SDK() + pages = _EventPages( + [ + [ + {"id": "old", "type": "user.message"}, + {"id": "idle", "type": "session.status_idle"}, + ], + [ + {"id": "old", "type": "user.message"}, + {"id": "idle", "type": "session.status_idle"}, + { + "id": "new", + "type": "user.message", + "content": [{"type": "text", "text": "new"}], + }, + ], + ] + ) + sdk.beta.sessions.events.list = pages + loop = ManagedAgentsLoop(runner=runner, session_id="managed-session-1") + + async def run(): + assert await loop.run_pending(sdk) == 0 + assert await loop.run_pending(sdk) == 1 + + asyncio.run(run()) + + assert all("page" not in kwargs for _, kwargs in pages.calls) + assert [call["new_message"].parts[0].text for call in runner.calls] == ["new"] + + +def test_thinking_previews_are_start_only(): + sdk = _SDK() + loop = ManagedAgentsLoop(runner=_Runner([]), session_client=_SessionClient()) + + async def run(): + for typ in ( + "agent.thinking_stream_start", + "agent.thinking_chunk", + "agent.thinking_stream_end", + ): + await loop._publish_transient( + sdk, {"type": typ, "thinking_id": "thinking", "delta": "private"} + ) + + asyncio.run(run()) + assert len(sdk.transient) == 1 + path, _, body = sdk.transient[0] + assert path.endswith("/events") + assert body["events"][0]["type"] == "event_start" + assert body["events"][0]["event"] == {"type": "agent.thinking", "id": "thinking"} + + +def test_failed_turn_does_not_complete_input_or_leak_provider_secret(caplog): + class FailingRunner(_Runner): + async def run_async(self, **kwargs): + if False: + yield None + raise RuntimeError("provider secret-must-not-escape") + + sdk = _SDK() + loop = ManagedAgentsLoop(runner=FailingRunner([]), session_client=_SessionClient()) + with pytest.raises(RuntimeError): + asyncio.run( + loop._run_turn( + sdk, {"id": "failed", "content": [{"type": "text", "text": "hello"}]} + ) + ) + events = [event for _, batch in sdk.beta.sessions.events.batches for event in batch] + assert "secret-must-not-escape" not in str(events) + assert "secret-must-not-escape" not in caplog.text + loop._restore_completed_inputs([{"id": "failed", "type": "user.message"}, *events]) + assert "failed" not in loop._completed_inputs + + +def test_reporting_failure_preserves_original_turn_exception(monkeypatch): + class FailingRunner(_Runner): + async def run_async(self, **kwargs): + if False: + yield None + raise RuntimeError("original model failure") + + sdk = _SDK() + original_send = loop_module.send_session_events + + async def send(client, session_id, *, events): + if any(event["type"] == "session.error" for event in events): + raise OSError("reporting failed") + return await original_send(client, session_id, events=events) + + monkeypatch.setattr(loop_module, "send_session_events", send) + loop = ManagedAgentsLoop(runner=FailingRunner([]), session_client=_SessionClient()) + with pytest.raises(RuntimeError, match="original model failure"): + asyncio.run( + loop._run_turn(sdk, {"content": [{"type": "text", "text": "hello"}]}) + ) diff --git a/tests/runtime/managed_agents/test_managed_event_debug.py b/tests/runtime/managed_agents/test_managed_event_debug.py new file mode 100644 index 000000000..6d77bb027 --- /dev/null +++ b/tests/runtime/managed_agents/test_managed_event_debug.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import asyncio +import importlib +import json +from types import SimpleNamespace + +import pytest +from pydantic import BaseModel + + +debug = importlib.import_module("veadk.runtime.managed_agents.events") + + +def records(capsys): + return [ + json.loads(line.removeprefix("MA_DEBUG_EVENT ")) + for line in capsys.readouterr().out.splitlines() + if line.startswith("MA_DEBUG_EVENT ") + ] + + +@pytest.mark.parametrize("value", [None, "0", "false", "off"]) +def test_disabled_trace_does_not_serialize_events(monkeypatch, capsys, value): + monkeypatch.delenv("MA_DEBUG_EVENTS", raising=False) + if value is not None: + monkeypatch.setenv("MA_DEBUG_EVENTS", value) + + class Event: + def model_dump(self, **kwargs): + pytest.fail("Disabled tracing must not serialize events") + + debug.trace_event("received", "runner", "session-1", Event()) + assert capsys.readouterr().out == "" + + +def test_complete_payload_is_logged_without_mutation_or_credentials( + monkeypatch, capsys +): + monkeypatch.setenv("MA_DEBUG_EVENTS", "1") + + class Event(BaseModel): + type: str + content: str + metadata: dict + + event = Event( + type="user.message", + content="完整内容\n" * 10000, + metadata={ + "API-Key": "private-key", + "nested": [{"Authorization": "Bearer private-token", "text": "visible"}], + }, + ) + debug.trace_event("received", "sse", "session-1", event) + (record,) = records(capsys) + assert record["direction"] == "received" + assert record["source"] == "sse" + assert record["session_id"] == "session-1" + assert record["event"]["content"] == event.content + assert record["event"]["metadata"] == { + "API-Key": "[REDACTED]", + "nested": [{"Authorization": "[REDACTED]", "text": "visible"}], + } + assert event.metadata["API-Key"] == "private-key" + + +def test_trace_failure_does_not_fail_a_turn_or_print_error_secrets(monkeypatch, capsys): + monkeypatch.setenv("MA_DEBUG_EVENTS", "true") + + class Event: + def model_dump(self, **kwargs): + raise RuntimeError("private-error") + + debug.trace_event("received", "runner", "session-1", Event()) + assert "private-error" not in capsys.readouterr().out + + +@pytest.mark.asyncio +async def test_send_and_paginated_receive_preserve_payloads(monkeypatch, capsys): + monkeypatch.setenv("MA_DEBUG_EVENTS", "1") + events = [{"type": "session.status_running"}, {"type": "agent.message"}] + response = SimpleNamespace(data=[SimpleNamespace(type="agent.message", id="saved")]) + + class Events: + async def send(self, session_id, *, events): + assert session_id == "session-1" + assert events is outgoing + return response + + async def list(self, session_id, **options): + assert session_id == "session-1" + assert options == {"page": "5", "order": "asc"} + for event in events: + await asyncio.sleep(0) + yield event + + outgoing = events + sdk = SimpleNamespace( + beta=SimpleNamespace(sessions=SimpleNamespace(events=Events())) + ) + assert await debug.send_session_events(sdk, "session-1", events=events) is response + received = [ + event + async for event in debug.iter_session_events( + sdk, "session-1", page="5", order="asc" + ) + ] + assert all(actual is expected for actual, expected in zip(received, events)) + logged = records(capsys) + assert [(item["direction"], item["source"]) for item in logged] == [ + ("sent", "send"), + ("sent", "send"), + ("received", "send_response"), + ("received", "list"), + ("received", "list"), + ] + + +def test_sync_client_traces_send_response_and_list(monkeypatch, capsys): + from veadk.integrations.mpa.session_client import SelfHostSandboxClient + + monkeypatch.setenv("MA_DEBUG_EVENTS", "1") + event = {"type": "agent.tool_result", "content": "complete tool output"} + response = SimpleNamespace(data=[event]) + api = SimpleNamespace( + send=lambda **kwargs: response, list=lambda **kwargs: response + ) + client = SelfHostSandboxClient.__new__(SelfHostSandboxClient) + client.session_id = "session-sync" + client.client = SimpleNamespace( + beta=SimpleNamespace(sessions=SimpleNamespace(events=api)) + ) + client.post_events([event]) + assert client.list_events() == [event] + logged = records(capsys) + assert [record["source"] for record in logged] == ["send", "send_response", "list"] + assert all(record["event"] == event for record in logged) diff --git a/tests/runtime/managed_agents/test_official_tool_config.py b/tests/runtime/managed_agents/test_official_tool_config.py new file mode 100644 index 000000000..36a14a1d1 --- /dev/null +++ b/tests/runtime/managed_agents/test_official_tool_config.py @@ -0,0 +1,143 @@ +"""Official SDK response parsing must preserve native tool discriminants.""" + +import anthropic +import httpx2 +import pytest + +from veadk.runtime.managed_agents import worker + + +def sdk_snapshot(configs, *, default_enabled=False): + payload = { + "type": "agent", + "id": "fixture-agent", + "name": "fixture-agent", + "model": {"id": "fixture-model"}, + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": { + "enabled": default_enabled, + "permission_policy": {"type": "always_allow"}, + }, + "configs": configs, + } + ], + "skills": [], + } + with anthropic.Anthropic( + api_key="offline-fixture-key", + base_url="http://sdk.invalid", + http_client=httpx2.Client( + transport=httpx2.MockTransport(lambda _: httpx2.Response(200, json=payload)) + ), + ) as sdk: + return sdk.beta.agents.retrieve("fixture-agent") + + +def test_official_sdk_type_only_configs_enable_all_six_native_tools(): + names = {"bash", "read", "write", "edit", "glob", "grep"} + snapshot = sdk_snapshot( + [ + { + "type": name, + "enabled": True, + "permission_policy": {"type": "always_allow"}, + } + for name in sorted(names) + ] + ) + assert isinstance(snapshot, anthropic.types.beta.BetaManagedAgentsAgent) + # The server round-trips the official input schema without legacy `name`. + assert all( + not getattr(config, "name", None) for config in snapshot.tools[0].configs + ) + config = worker.managed_agent_config(snapshot) + assert {tool.__name__ for tool in config["tools"]} == names + policies = worker._tool_permission_map(snapshot) + assert all(policies[name] == "always_allow" for name in names) + + +def test_typed_disable_override_and_permission_share_the_same_discriminant(): + snapshot = sdk_snapshot( + [ + { + "type": "bash", + "enabled": False, + "permission_policy": {"type": "always_allow"}, + }, + { + "type": "read", + "enabled": True, + "permission_policy": {"type": "always_ask"}, + }, + ], + default_enabled=True, + ) + config = worker.managed_agent_config(snapshot) + names = {tool.__name__ for tool in config["tools"]} + assert "bash" not in names + assert "read" in names + assert worker._tool_permission_map(snapshot)["read"] == "always_ask" + + +def test_official_type_takes_precedence_over_legacy_name(): + snapshot = { + "model": "fixture-model", + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": {"enabled": False}, + "configs": [ + { + "type": "read", + "name": "bash", + "enabled": True, + "permission_policy": {"type": "always_ask"}, + } + ], + } + ], + } + assert [ + tool.__name__ for tool in worker.managed_agent_config(snapshot)["tools"] + ] == ["read"] + assert worker._tool_permission_map(snapshot)["read"] == "always_ask" + assert "bash" not in worker._tool_permission_map(snapshot) + + +def test_unknown_official_type_cannot_be_hidden_by_a_legacy_name(): + snapshot = { + "model": "fixture-model", + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": {"enabled": False}, + "configs": [{"type": "unsupported", "name": "read", "enabled": True}], + } + ], + } + with pytest.raises(ValueError, match="unsupported"): + worker.managed_agent_config(snapshot) + + +def test_mcp_configs_continue_to_use_names_not_native_type_discriminants(): + snapshot = { + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "fixture-server", + "configs": [ + { + "name": "fixture-tool", + "type": "custom", + "permission_policy": {"type": "always_ask"}, + } + ], + } + ] + } + assert ( + worker._tool_permission_map(snapshot)["mcp__fixture-server__fixture-tool"] + == "always_ask" + ) diff --git a/tests/runtime/managed_agents/test_package.py b/tests/runtime/managed_agents/test_package.py new file mode 100644 index 000000000..bb110bf16 --- /dev/null +++ b/tests/runtime/managed_agents/test_package.py @@ -0,0 +1,26 @@ +import os +import subprocess +import sys +import textwrap + + +def test_runtime_import_needs_no_example_or_credentials(): + code = textwrap.dedent(""" + import anthropic + def unexpected_client(*args, **kwargs): + raise AssertionError("Library import constructed a remote client") + anthropic.Anthropic.__init__ = unexpected_client + anthropic.AsyncAnthropic.__init__ = unexpected_client + from veadk.runtime.managed_agents.loop import ManagedAgentsLoop + from veadk.runtime.managed_agents import worker + assert "main" not in __import__("sys").modules + assert worker.get_default_agent.cache_info().currsize == 0 + """) + env = { + key: value + for key, value in os.environ.items() + if not key.startswith( + ("ANTHROPIC_", "SANDBOX_", "MODEL_AGENT_", "AGENTKIT_", "MA_") + ) + } + subprocess.run([sys.executable, "-c", code], check=True, env=env) diff --git a/tests/runtime/managed_agents/test_sdk_integrity.py b/tests/runtime/managed_agents/test_sdk_integrity.py new file mode 100644 index 000000000..7153ccff4 --- /dev/null +++ b/tests/runtime/managed_agents/test_sdk_integrity.py @@ -0,0 +1,55 @@ +"""The image must reject a modified SDK even when imports still work.""" + +import base64 +import hashlib +import importlib.metadata +import runpy +from pathlib import Path + +import pytest + + +CHECKER = ( + Path(__file__).resolve().parents[3] + / "docker/managed-agents/verify_anthropic_sdk.py" +) + + +def checker(): + return runpy.run_path(str(CHECKER))["verify_distribution"] + + +def test_installed_official_sdk_is_intact(): + assert checker()(importlib.metadata.distribution("anthropic")) > 0 + + +@pytest.mark.parametrize("change", ["modified", "missing", "unrecorded"]) +def test_record_validation_rejects_changed_python_files(tmp_path, change): + package = tmp_path / "anthropic" + package.mkdir() + source = package / "__init__.py" + source.write_bytes(b"# official fixture\n") + digest = ( + base64.urlsafe_b64encode(hashlib.sha256(source.read_bytes()).digest()) + .decode() + .rstrip("=") + ) + record = f"anthropic/__init__.py,sha256={digest},{source.stat().st_size}\n" + + class Distribution: + def read_text(self, name): + assert name == "RECORD" + return record + + def locate_file(self, path): + return tmp_path / path + + assert checker()(Distribution()) == 1 + if change == "modified": + source.write_bytes(b"# replaced fixture\n") + elif change == "missing": + source.unlink() + else: + (package / "overlay.py").write_bytes(b"# overlay fixture\n") + with pytest.raises(RuntimeError, match="SDK integrity"): + checker()(Distribution()) diff --git a/tests/runtime/managed_agents/test_self_host_sandbox_agent.py b/tests/runtime/managed_agents/test_self_host_sandbox_agent.py new file mode 100644 index 000000000..9c6dd9864 --- /dev/null +++ b/tests/runtime/managed_agents/test_self_host_sandbox_agent.py @@ -0,0 +1,1083 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import importlib +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest + +try: + from builtins import BaseExceptionGroup +except ImportError: # Python 3.10 + from exceptiongroup import BaseExceptionGroup + + +agent_module = importlib.import_module("veadk.runtime.managed_agents.sandbox") +SandboxSessionManager = agent_module.SandboxSessionManager + +main_module = importlib.import_module("veadk.runtime.managed_agents.worker") + + +def test_managed_memory_ignores_legacy_postgresql_configuration(monkeypatch): + from google.adk.sessions import InMemorySessionService + + monkeypatch.setenv( + "VEADK_MANAGED_SESSION_DB_URL", "postgresql://invalid.invalid/db" + ) + memory = main_module.managed_short_term_memory() + assert isinstance(memory.session_service, InMemorySessionService) + + +class _FakeClient: + def __init__(self, remote_session_id: str): + self.session_id = None + self.remote_session_id = remote_session_id + self.created_titles = [] + self.idle_count = 0 + + def create_session(self, title: str): + self.created_titles.append(title) + self.session_id = self.remote_session_id + + def post_status_idle(self): + self.idle_count += 1 + + +def test_each_veadk_session_creates_a_distinct_remote_session(monkeypatch): + manager = SandboxSessionManager() + clients = iter((_FakeClient("remote-1"), _FakeClient("remote-2"))) + monkeypatch.setattr(manager, "_new_client", lambda: next(clients)) + + assert manager.create_remote_session("veadk-1") == "remote-1" + assert manager.create_remote_session("veadk-1") == "remote-1" + assert manager.create_remote_session("veadk-2") == "remote-2" + + assert manager.get("veadk-1") is not manager.get("veadk-2") + assert manager.get("veadk-1").created_titles == [ + "VeADK Self-Hosted Sandbox Session veadk-1" + ] + assert manager.get("veadk-2").created_titles == [ + "VeADK Self-Hosted Sandbox Session veadk-2" + ] + + +def test_remote_session_writes_no_synthetic_event_and_idles_once_per_turn( + monkeypatch, +): + manager = SandboxSessionManager() + client = _FakeClient("remote-1") + monkeypatch.setattr(manager, "_new_client", lambda: client) + manager.create_remote_session("veadk-1") + + manager.begin_turn("veadk-1") + manager.begin_turn("veadk-1") + + manager.end_turn("veadk-1") + assert client.idle_count == 0 + manager.end_turn("veadk-1") + assert client.idle_count == 1 + + manager.begin_turn("veadk-1") + manager.begin_turn("veadk-1") + manager.end_turn("veadk-1") + manager.end_turn("veadk-1") + assert client.idle_count == 2 + + +def test_runner_wrapper_ends_remote_turn_after_failure(monkeypatch): + lifecycle_calls = [] + + class _FailingRunner: + async def run_async(self, **kwargs): + yield "started" + raise RuntimeError("turn failed") + + monkeypatch.setattr( + agent_module.sandbox_sessions, + "begin_turn", + lambda session_id: lifecycle_calls.append(("begin", session_id)), + ) + monkeypatch.setattr( + agent_module.sandbox_sessions, + "end_turn", + lambda session_id: lifecycle_calls.append(("end", session_id)), + ) + runner = agent_module.enable_sandbox_turn_lifecycle(_FailingRunner()) + + async def consume(): + async for _ in runner.run_async(session_id="veadk-1"): + pass + + try: + asyncio.run(consume()) + except RuntimeError as error: + assert str(error) == "turn failed" + else: + raise AssertionError("the wrapped runner must preserve turn failures") + + assert lifecycle_calls == [("begin", "veadk-1"), ("end", "veadk-1")] + + +def test_dispatch_task_sends_only_the_model_tool_call(monkeypatch): + dispatched = [] + client = SimpleNamespace( + dispatch_tool=lambda name, arguments, *, dispatch_id: ( + dispatched.append((name, arguments, dispatch_id)) or {"stdout": "ok"} + ) + ) + monkeypatch.setattr(agent_module.sandbox_sessions, "get", lambda session_id: client) + tool_call = SimpleNamespace( + session_id="veadk-session", + id="tool-call-1", + name="bash", + arguments={"command": "printf ok"}, + ) + + result = asyncio.run(agent_module.dispatch_task(tool_call)) + + assert result == {"stdout": "ok"} + assert dispatched == [("bash", {"command": "printf ok"}, "tool-call-1")] + assert not hasattr(agent_module.get_default_agent(), "run_turn") + + +def test_managed_agent_config_uses_frozen_snapshot(): + config = main_module.managed_agent_config( + { + "name": "Distributed Agent", + "description": "test", + "model": {"id": "model-from-session"}, + "system": "Remember prior turns.", + "tools": [], + } + ) + + assert config["name"] == "Distributed_Agent" + assert config["model_name"] == "model-from-session" + assert config["instruction"] == "Remember prior turns." + assert config["tools"] == [] + assert config["before_tool_callback"] is None + + +def test_nested_http_status_detection_handles_task_group_authentication(): + authentication = RuntimeError("redacted") + authentication.status_code = 401 + nested = BaseExceptionGroup("dispatcher", [authentication]) + + assert main_module._contains_http_status(nested, {401, 403}) + assert not main_module._contains_http_status(nested, {429, 503}) + + +def test_dispatcher_refreshes_identity_after_startup_401(monkeypatch): + runs = 0 + invalidations = 0 + sleeps = [] + + class Dispatcher: + async def run(self, **kwargs): + nonlocal runs + runs += 1 + assert kwargs == {"max_items": 1, "max_concurrency": 2} + if runs == 1: + error = RuntimeError("redacted") + error.status_code = 401 + raise BaseExceptionGroup("poll", [error]) + return 7 + + class Provider: + def invalidate(self): + nonlocal invalidations + invalidations += 1 + + async def no_wait(delay): + sleeps.append(delay) + + monkeypatch.setattr(main_module.asyncio, "sleep", no_wait) + result = asyncio.run( + main_module._run_dispatcher_with_identity_retry( + Dispatcher, Provider(), max_items=1, max_concurrency=2 + ) + ) + + assert result == 7 + assert runs == 2 + assert invalidations == 1 + assert sleeps == [1.0] + + +def test_dispatcher_does_not_restart_after_shutdown_during_auth_backoff(monkeypatch): + runs = 0 + stopping = False + + class Dispatcher: + async def run(self, **kwargs): + nonlocal runs + runs += 1 + error = RuntimeError("redacted") + error.status_code = 401 + raise error + + class Provider: + def invalidate(self): + pass + + async def request_stop(_delay): + nonlocal stopping + stopping = True + + monkeypatch.setattr(main_module.asyncio, "sleep", request_stop) + result = asyncio.run( + main_module._run_dispatcher_with_identity_retry( + Dispatcher, + Provider(), + max_items=None, + max_concurrency=1, + should_stop=lambda: stopping, + ) + ) + + assert result == 0 + assert runs == 1 + + +def test_dispatcher_drains_new_instance_when_shutdown_races_with_factory(): + stopping = False + runs = 0 + drains = 0 + + class Dispatcher: + async def run(self, **kwargs): + nonlocal runs + runs += 1 + return 1 + + def drain(self): + nonlocal drains + drains += 1 + + def factory(): + nonlocal stopping + dispatcher = Dispatcher() + stopping = True + return dispatcher + + result = asyncio.run( + main_module._run_dispatcher_with_identity_retry( + factory, + None, + max_items=None, + max_concurrency=1, + should_stop=lambda: stopping, + ) + ) + + assert result == 0 + assert drains == 1 + assert runs == 0 + + +def test_identity_retry_can_wrap_readiness_poll(monkeypatch): + polls = 0 + invalidations = 0 + + async def poll(): + nonlocal polls + polls += 1 + if polls == 1: + error = RuntimeError("redacted") + error.status_code = 403 + raise error + + class Provider: + def invalidate(self): + nonlocal invalidations + invalidations += 1 + + async def no_wait(_delay): + pass + + monkeypatch.setattr(main_module.asyncio, "sleep", no_wait) + assert asyncio.run(main_module._run_with_identity_retry(poll, Provider())) is None + assert polls == 2 + assert invalidations == 1 + + +def test_managed_agent_config_rejects_unknown_enabled_tool(): + snapshot = { + "name": "demo", + "model": "model", + "tools": [ + { + "type": "agent_toolset_20260401", + "configs": [{"name": "unknown", "enabled": True}], + } + ], + } + + try: + main_module.managed_agent_config(snapshot) + except ValueError as error: + assert "unknown" in str(error) + else: + raise AssertionError("unknown enabled tools must be rejected") + + +def test_managed_agent_config_applies_toolset_defaults_and_overrides(): + snapshot = { + "name": "demo", + "model": "model", + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": {"enabled": False}, + "configs": [ + {"name": "bash", "enabled": True}, + {"name": "read", "enabled": True}, + ], + } + ], + } + + config = main_module.managed_agent_config(snapshot) + + assert [tool.__name__ for tool in config["tools"]] == ["bash", "read"] + + +def test_managed_agent_config_requires_runtime_context_for_custom_tools(): + custom_snapshot = { + "name": "demo", + "model": "model", + "tools": [{"type": "custom", "name": "external_tool"}], + } + try: + main_module.managed_agent_config(custom_snapshot) + except ValueError as error: + assert "custom" in str(error) + else: + raise AssertionError("unsupported top-level tool definitions must be rejected") + + +def test_managed_agent_config_ignores_skills_without_metadata(): + snapshot = { + "name": "demo", + "model": "model", + "tools": [], + "skills": [{"type": "custom", "skill_id": "skill_x", "version": "v1"}], + } + config = main_module.managed_agent_config(snapshot) + assert "managed_agent_skills" not in config["instruction"] + + +def test_managed_agent_config_loads_materialized_skill_instructions(tmp_path): + skill_dir = tmp_path / "demo-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text( + "# demo-skill\n\nDemo instructions.", encoding="utf-8" + ) + snapshot = { + "name": "demo", + "model": "model", + "skills": [{"type": "custom", "skill_id": "skill_x", "version": "v1"}], + "tools": [ + { + "type": "custom", + "name": "approve", + "description": "Approve", + "input_schema": {"type": "object"}, + } + ], + } + config = main_module.managed_agent_config( + snapshot, + sdk=SimpleNamespace(), + session_id="session-1", + skill_dirs=[skill_dir], + ) + assert "name='demo-skill'" in config["instruction"] + assert "# demo-skill\n\nDemo instructions." in config["instruction"] + assert str(skill_dir / "SKILL.md") in config["instruction"] + assert [tool.name for tool in config["tools"]] == ["approve"] + + +@pytest.mark.parametrize("result_type", ["user.tool_result", "agent.tool_result"]) +def test_managed_worker_remote_tool_waits_for_matching_result(monkeypatch, result_type): + monkeypatch.setenv("MANAGED_AGENT_TOOL_EXECUTION", "remote") + batches = [] + + async def send(session_id, *, events): + batches.extend(events) + + async def list_events(session_id, **kwargs): + yield {"type": "user.tool_result", "tool_use_id": "other", "content": "wrong"} + yield { + "type": result_type, + "tool_use_id": "tool-1", + "content": "remote-ok", + "is_error": False, + } + + async def forbidden_local(*args, **kwargs): + raise AssertionError("remote tools must not execute inside Agent Loop") + + monkeypatch.setattr(main_module, "_run_managed_tool", forbidden_local) + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace( + events=SimpleNamespace(send=send, list=list_events) + ) + ) + ) + runtime = main_module.managed_work_tool_runtime( + sdk, "session-1", workdir=Path("/tmp/managed-work") + ) + result = asyncio.run( + runtime.execute( + SimpleNamespace( + id="tool-1", + name="bash", + arguments={"command": "printf remote-ok"}, + tool=SimpleNamespace(), + ) + ) + ) + assert result == {"result": "remote-ok"} + expected = ["agent.tool_use"] + if result_type == "user.tool_result": + expected.append("agent.tool_result") + assert batches[-1]["tool_use_id"] == "tool-1" + assert batches[-1]["content"] == "remote-ok" + assert [event["type"] for event in batches] == expected + + +def test_managed_worker_executes_tool_locally_and_publishes_events(monkeypatch): + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=SimpleNamespace(batches=[])) + ) + ) + + async def send(session_id, *, events): + sdk.beta.sessions.events.batches.append((session_id, events)) + + async def run_tool(name, arguments, *, workdir, native_tools=None): + assert name == "bash" + assert arguments == {"command": "printf ok"} + assert workdir == Path("/tmp/managed-work") + return { + "status": "completed", + "exit_code": 0, + "stdout": "ok", + "stderr": "", + } + + sdk.beta.sessions.events.send = send + monkeypatch.setattr(main_module, "_run_managed_tool", run_tool) + runtime = main_module.managed_work_tool_runtime( + sdk, "session-1", workdir=Path("/tmp/managed-work") + ) + result = asyncio.run( + runtime.execute( + SimpleNamespace( + id="tool-1", + name="bash", + arguments={"command": "printf ok"}, + tool=SimpleNamespace(), + ) + ) + ) + + assert result["stdout"] == "ok" + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert emitted == [ + { + "type": "agent.tool_use", + "id": "tool-1", + "name": "bash", + "input": {"command": "printf ok"}, + "evaluated_permission": "allow", + }, + { + "type": "agent.tool_result", + "tool_use_id": "tool-1", + "content": "exit=0\nok", + "is_error": False, + }, + ] + + +def test_always_ask_tool_waits_for_confirmation(monkeypatch, tmp_path): + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=SimpleNamespace(batches=[])) + ) + ) + + async def send(session_id, *, events): + sdk.beta.sessions.events.batches.append((session_id, list(events))) + + class Page: + def __aiter__(self): + async def iterate(): + yield SimpleNamespace( + type="user.tool_confirmation", + tool_use_id="tool-ask", + result="allow", + ) + + return iterate() + + sdk.beta.sessions.events.send = send + sdk.beta.sessions.events.list = lambda *args, **kwargs: Page() + + async def run_tool(*args, **kwargs): + return {"status": "completed", "exit_code": 0, "stdout": "ok", "stderr": ""} + + monkeypatch.setattr(main_module, "_run_managed_tool", run_tool) + runtime = main_module.managed_work_tool_runtime( + sdk, + "session-1", + workdir=tmp_path, + snapshot={ + "tools": [ + { + "type": "agent_toolset_20260401", + "default_config": {"enabled": False}, + "configs": [ + { + "name": "bash", + "enabled": True, + "permission_policy": {"type": "always_ask"}, + } + ], + } + ] + }, + ) + result = asyncio.run( + runtime.execute( + SimpleNamespace( + id="tool-ask", + name="bash", + arguments={"command": "printf ok"}, + tool=SimpleNamespace(), + ) + ) + ) + assert result["stdout"] == "ok" + batches = [batch for _, batch in sdk.beta.sessions.events.batches] + assert batches[0][0]["evaluated_permission"] == "ask" + assert batches[0][1]["stop_reason"]["type"] == "requires_action" + assert batches[1] == [{"type": "session.status_running"}] + + +def test_mcp_tool_executes_locally_and_publishes_mcp_events(tmp_path): + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=SimpleNamespace(batches=[])) + ) + ) + + async def send(session_id, *, events): + sdk.beta.sessions.events.batches.append((session_id, list(events))) + + class FakeMcpTool: + name = "mcp__orders__lookup" + _mcp_tool = object() + managed_agents_mcp_server_name = "orders" + + async def run_async(self, *, args, tool_context): + return {"order": args["id"]} + + sdk.beta.sessions.events.send = send + runtime = main_module.managed_work_tool_runtime( + sdk, "session-1", workdir=tmp_path, snapshot={"tools": []} + ) + result = asyncio.run( + runtime.execute( + SimpleNamespace( + id="mcp-call-1", + name=FakeMcpTool.name, + arguments={"id": "42"}, + tool=FakeMcpTool(), + context=SimpleNamespace(), + ) + ) + ) + assert result == {"order": "42"} + emitted = [ + event for _, batch in sdk.beta.sessions.events.batches for event in batch + ] + assert emitted[0] == { + "type": "agent.mcp_tool_use", + "id": "mcp-call-1", + "mcp_server_name": "orders", + "name": "lookup", + "input": {"id": "42"}, + "evaluated_permission": "allow", + } + assert emitted[1]["type"] == "agent.mcp_tool_result" + assert emitted[1]["mcp_tool_use_id"] == "mcp-call-1" + assert '"order": "42"' in emitted[1]["content"][0]["text"] + + +def test_managed_tool_environment_removes_control_plane_secrets(monkeypatch): + monkeypatch.setenv("ANTHROPIC_ENVIRONMENT_KEY", "secret") + monkeypatch.setenv("DATABASE_POSTGRESQL_PASSWORD", "secret") + monkeypatch.setenv("VEADK_MANAGED_SESSION_DB_URL", "secret") + monkeypatch.setenv("MODEL_API_KEY", "secret") + monkeypatch.setenv("SAFE_VALUE", "visible") + + environment = main_module._managed_tool_environment() + + assert "SAFE_VALUE" not in environment + assert "ANTHROPIC_ENVIRONMENT_KEY" not in environment + assert "DATABASE_POSTGRESQL_PASSWORD" not in environment + assert "VEADK_MANAGED_SESSION_DB_URL" not in environment + assert "MODEL_API_KEY" not in environment + + +def test_managed_worker_executes_real_local_tool(tmp_path): + result = asyncio.run( + main_module._run_managed_tool( + "bash", {"command": "printf local-tool-ok"}, workdir=tmp_path + ) + ) + + assert result == { + "status": "completed", + "exit_code": 0, + "stdout": "local-tool-ok", + "stderr": "", + "content": "local-tool-ok", + "is_error": False, + } + + +def test_managed_worker_times_out_local_tool(tmp_path): + result = asyncio.run( + main_module._run_managed_tool( + "bash", + {"command": "sleep 30", "timeout_ms": 20}, + workdir=tmp_path, + ) + ) + + assert result["status"] == "failed" + assert result["is_error"] is True + assert "timed out" in result["stderr"].lower() + + +def test_managed_tool_timeout_defaults_to_five_minutes(monkeypatch): + monkeypatch.delenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, raising=False) + + assert main_module._managed_tool_timeout_seconds() == 300.0 + + +def test_managed_tool_timeout_honours_env_and_model_override(monkeypatch): + monkeypatch.setenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, "420") + + assert main_module._managed_tool_timeout_seconds() == 420.0 + + +def test_action_timeout_never_truncates_the_tool_budget(monkeypatch): + monkeypatch.delenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, raising=False) + monkeypatch.setenv("MANAGED_AGENT_ACTION_TIMEOUT_SECONDS", "180") + + assert main_module._managed_action_timeout_seconds() == 360.0 + + monkeypatch.setenv("MANAGED_AGENT_ACTION_TIMEOUT_SECONDS", "900") + assert main_module._managed_action_timeout_seconds() == 900.0 + + monkeypatch.setenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, "420") + assert main_module._managed_action_timeout_seconds() == 900.0 + monkeypatch.setenv("MANAGED_AGENT_ACTION_TIMEOUT_SECONDS", "300") + assert main_module._managed_action_timeout_seconds() == 480.0 + + +def test_published_tool_use_carries_the_deployment_budget(monkeypatch): + monkeypatch.delenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, raising=False) + + annotated = main_module._with_default_tool_timeout( + {"type": "agent.tool_use", "name": "bash", "input": {"command": "sleep 200"}} + ) + assert annotated["input"]["timeout_ms"] == 300000 + assert annotated["input"]["command"] == "sleep 200" + + explicit = main_module._with_default_tool_timeout( + { + "type": "agent.tool_use", + "name": "bash", + "input": {"command": "ls", "timeout_ms": 1500}, + } + ) + assert explicit["input"]["timeout_ms"] == 1500 + + read_only = main_module._with_default_tool_timeout( + {"type": "agent.tool_use", "name": "read", "input": {"file_path": "/tmp/x"}} + ) + assert "timeout_ms" not in read_only["input"] + + monkeypatch.setenv(main_module.MANAGED_TOOL_TIMEOUT_ENV, "420") + tuned = main_module._with_default_tool_timeout( + {"type": "agent.tool_use", "name": "bash", "input": {"command": "true"}} + ) + assert tuned["input"]["timeout_ms"] == 420000 + + +def test_native_file_tools_preserve_content_and_do_not_spawn_shell( + monkeypatch, tmp_path +): + async def forbidden(*args, **kwargs): + raise AssertionError("file tools must not invoke a subprocess") + + monkeypatch.setattr(asyncio, "create_subprocess_exec", forbidden) + + async def run(): + tools = main_module.NativeAgentTools(tmp_path) + content = "one\n' \" $HOME `uname` $(touch nope) & 中文\nthree\n" + try: + written = await tools.run( + "write", {"file_path": "nested/notes.txt", "content": content} + ) + assert not written["is_error"] + edited = await tools.run( + "edit", + { + "file_path": "nested/notes.txt", + "old_string": "three", + "new_string": "THREE", + }, + ) + assert not edited["is_error"] + read = await tools.run( + "read", {"file_path": "nested/notes.txt", "view_range": [2, 3]} + ) + assert "中文" in read["content"] and "THREE" in read["content"] + assert (tmp_path / "nested/notes.txt").read_text() == content.replace( + "three", "THREE" + ) + assert ( + "notes.txt" + in (await tools.run("glob", {"pattern": "**/*.txt"}))["content"] + ) + assert "THREE" in (await tools.run("grep", {"pattern": "THREE"}))["content"] + assert not (tmp_path / "nope").exists() + finally: + await tools.close() + + asyncio.run(run()) + + +def test_native_file_tools_reject_symlink_escape_and_ambiguous_edit(tmp_path): + outside = tmp_path.parent / "outside.txt" + outside.write_text("private") + (tmp_path / "escape").symlink_to(outside) + + async def run(): + tools = main_module.NativeAgentTools(tmp_path) + try: + for name, arguments in [ + ("read", {"file_path": "escape"}), + ("write", {"file_path": "escape", "content": "changed"}), + ("glob", {"pattern": "../*"}), + ("grep", {"pattern": "private", "path": ".."}), + ]: + assert (await tools.run(name, arguments))["is_error"] + await tools.run("write", {"file_path": "repeat.txt", "content": "x x"}) + assert ( + await tools.run( + "edit", + {"file_path": "repeat.txt", "old_string": "x", "new_string": "y"}, + ) + )["is_error"] + assert (tmp_path / "repeat.txt").read_text() == "x x" + assert outside.read_text() == "private" + finally: + await tools.close() + + asyncio.run(run()) + + +def test_native_bash_persists_state_and_restart_resets_it(tmp_path): + async def run(): + tools = main_module.NativeAgentTools(tmp_path) + try: + await tools.run("bash", {"command": "export NATIVE_TEST_VALUE=kept"}) + assert ( + "kept" + in ( + await tools.run( + "bash", {"command": "printf '%s' \"$NATIVE_TEST_VALUE\""} + ) + )["content"] + ) + assert not (await tools.run("bash", {"restart": True}))["is_error"] + assert ( + "cleared" + in ( + await tools.run( + "bash", + {"command": "printf '%s' \"${NATIVE_TEST_VALUE-cleared}\""}, + ) + )["content"] + ) + finally: + await tools.close() + + asyncio.run(run()) + + +def test_feishu_channel_stays_up_until_stopped_and_shuts_down(monkeypatch): + calls = [] + + class _FakeRunner: + def __init__(self, **kwargs): + calls.append(("runner", kwargs)) + + async def run_async(self, **kwargs): + if False: + yield None + + class _FakeChannel: + def __init__(self, *, runner, **kwargs): + calls.append(("channel", runner, kwargs)) + + def start(self, loop): + calls.append(("start", loop)) + + async def shutdown(self): + calls.append(("shutdown", None)) + + monkeypatch.setattr(main_module, "Runner", _FakeRunner) + monkeypatch.setattr(main_module, "FeishuChannelExtension", _FakeChannel) + stop_event = asyncio.Event() + stop_event.set() + + asyncio.run(main_module.serve_feishu_channel(stop_event)) + + assert calls[0] == ( + "runner", + { + "agent": main_module.get_default_agent(), + "app_name": "self_host_sandbox_demo", + }, + ) + assert calls[1][0] == "channel" + assert isinstance(calls[1][1], _FakeRunner) + assert calls[1][2] == { + "streaming": True, + "show_thinking": True, + "show_tool_calls": True, + "show_tool_results": True, + "separate_tool_call_cards": True, + "separate_thinking_card": True, + "create_topic": True, + } + assert calls[2][0] == "start" + assert calls[3] == ("shutdown", None) + + +def test_managed_agent_model_base_url_is_taken_from_snapshot(): + config = main_module.managed_agent_config( + { + "name": "custom", + "model": {"id": "model-x", "base_url": "https://model.example/api/v3"}, + "tools": [], + } + ) + assert config["model_name"] == "model-x" + assert config["model_api_base"] == "https://model.example/api/v3" + assert "model_api_key" not in config + + +def test_mcp_result_preserves_error_and_converts_image(): + content, error = main_module.mcp_result_content( + { + "isError": True, + "content": [ + {"type": "text", "text": "failed"}, + {"type": "image", "data": "fixture", "mimeType": "image/png"}, + ], + } + ) + assert error is True + assert content == [ + {"type": "text", "text": "failed"}, + { + "type": "image", + "source": {"type": "base64", "data": "fixture", "media_type": "image/png"}, + }, + ] + assert main_module.mcp_result_content({"structuredContent": {"value": 42}}) == ( + [{"type": "text", "text": '{"value": 42}'}], + False, + ) + + +@pytest.mark.parametrize("verdict", ["allow", "deny"]) +def test_mcp_permission_gate_and_error_result(monkeypatch, tmp_path, verdict): + monkeypatch.setenv("MANAGED_AGENT_TOOL_EXECUTION", "remote") + emitted, calls = [], [] + + async def send(session_id, *, events): + emitted.extend(events) + + async def confirmed(*args, **kwargs): + assert emitted[0]["evaluated_permission"] == "ask" + assert calls == [] + return {"result": verdict} + + class Tool: + name = "mcp__orders__lookup" + _mcp_tool = object() + managed_agents_mcp_server_name = "orders" + + async def run_async(self, *, args, tool_context): + calls.append(args) + return {"isError": True, "content": [{"type": "text", "text": "not found"}]} + + sdk = SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=SimpleNamespace(send=send)) + ) + ) + monkeypatch.setattr(main_module, "_wait_for_session_event", confirmed) + snapshot = { + "tools": [ + { + "type": "mcp_toolset", + "mcp_server_name": "orders", + "default_config": {"permission_policy": {"type": "always_ask"}}, + } + ] + } + + async def execute(): + runtime = main_module.managed_work_tool_runtime( + sdk, "session", workdir=tmp_path, snapshot=snapshot + ) + try: + return await runtime.execute( + SimpleNamespace( + id="call", + name=Tool.name, + arguments={"id": "42"}, + tool=Tool(), + context=SimpleNamespace(), + ) + ) + finally: + await runtime.close() + + asyncio.run(execute()) + assert len(calls) == (1 if verdict == "allow" else 0) + assert sum(e["type"] == "agent.mcp_tool_use" for e in emitted) == 1 + result = emitted[-1] + assert result["type"] == "agent.mcp_tool_result" + assert result["mcp_tool_use_id"] == "call" and result["is_error"] is True + assert result["content"][0]["text"] == ( + "not found" if verdict == "allow" else "MCP tool call denied by user" + ) + + +def test_runner_cleanup_closes_all_toolsets_before_reporting_failure(): + closed = [] + + class Toolset: + def __init__(self, name, fail=False): + self.name, self.fail = name, fail + + async def close(self): + closed.append(self.name) + if self.fail: + raise RuntimeError("close failed") + + runner = SimpleNamespace( + agent=SimpleNamespace(tools=[Toolset("first", True), Toolset("second")]) + ) + with pytest.raises(RuntimeError, match="close failed"): + asyncio.run(main_module.close_managed_runner(runner)) + assert closed == ["first", "second"] + + +def test_native_cleanup_closes_context_after_tool_failure(monkeypatch): + closed = [] + tools = object.__new__(main_module.NativeAgentTools) + tools.tools = {"first": "first", "second": "second"} + + async def close_tool(tool): + closed.append(tool) + if tool == "first": + raise RuntimeError("close failed") + + class Context: + async def close(self): + closed.append("context") + + tools.context = Context() + monkeypatch.setattr(main_module, "aclose_runnable_tool", close_tool) + with pytest.raises(RuntimeError, match="close failed"): + asyncio.run(tools.close()) + assert closed == ["first", "second", "context"] + + +@pytest.mark.asyncio +async def test_native_tool_cleanup_sanitizes_both_failures_and_exits_context(caplog): + from types import SimpleNamespace + from veadk.runtime.managed_agents.worker import aclose_runnable_tool + + phases = [] + + def close(): + phases.append("close") + raise RuntimeError("private-cleanup-marker") + + class Context: + async def __aexit__(self, *args): + phases.append("exit") + raise ValueError("private-exit-marker") + + await aclose_runnable_tool(SimpleNamespace(close=close, _context_manager=Context())) + assert phases == ["close", "exit"] + assert "private-cleanup-marker" not in caplog.text + assert "private-exit-marker" not in caplog.text + assert "RuntimeError" in caplog.text and "ValueError" in caplog.text + + +@pytest.mark.asyncio +async def test_native_tool_cleanup_awaits_close_and_preserves_cancellation(): + from types import SimpleNamespace + from veadk.runtime.managed_agents.worker import aclose_runnable_tool + + phases = [] + + async def close(): + phases.append("close") + + class Context: + def __exit__(self, *args): + phases.append("exit") + + await aclose_runnable_tool(SimpleNamespace(close=close, _context_manager=Context())) + assert phases == ["close", "exit"] + + async def cancel(): + raise asyncio.CancelledError() + + with pytest.raises(asyncio.CancelledError): + await aclose_runnable_tool(SimpleNamespace(aclose=cancel)) diff --git a/tests/runtime/managed_agents/test_worker_supervision.py b/tests/runtime/managed_agents/test_worker_supervision.py new file mode 100644 index 000000000..d06037d3d --- /dev/null +++ b/tests/runtime/managed_agents/test_worker_supervision.py @@ -0,0 +1,102 @@ +"""Run the real shell supervisor and health server with a controlled worker.""" + +import os +import socket +import subprocess +import sys +import time +from pathlib import Path + +import httpx +import pytest + +ROOT = Path(__file__).resolve().parents[3] + + +@pytest.mark.parametrize("stop_mode", ["terminate", "worker_failure"]) +@pytest.mark.parametrize("entrypoint_mode", ["worker", "claimed"]) +def test_supervisor_starts_listener_and_propagates_worker_exit( + tmp_path, stop_mode, entrypoint_mode +): + executable = tmp_path / "python-controlled" + executable.write_text( + f"#!{sys.executable}\n" + "import os,sys,time\nfrom pathlib import Path\n" + "root=Path(os.environ['TEST_RECORD_DIR'])\n" + "role='listener' if sys.argv[2].endswith('.health') else 'worker'\n" + "(root / (role+'.pid')).write_text(str(os.getpid()))\n" + "(root / (role+'.args')).write_text(' '.join(sys.argv[1:]))\n" + "if role=='listener': os.execv(sys.executable,[sys.executable]+sys.argv[1:])\n" + "while not (root/'fail').exists(): time.sleep(0.02)\n" + "sys.exit(23)\n" + ) + executable.chmod(0o755) + with socket.socket() as reserved: + reserved.bind(("127.0.0.1", 0)) + port = reserved.getsockname()[1] + env = dict( + os.environ, + TEST_RECORD_DIR=str(tmp_path), + MANAGED_AGENTS_PYTHON=str(executable), + MA_RUNTIME_PORT=str(port), + MA_RUNTIME_HOST="127.0.0.1", + MANAGED_AGENT_ENTRYPOINT_MODE=entrypoint_mode, + ) + env.pop("AGENTKIT_RUNTIME_PORT", None) + env.pop("AGENTKIT_RUNTIME_HOST", None) + process = subprocess.Popen( + [ + "bash", + str(ROOT / "docker/managed-agents/run-runtime.sh"), + "--worker-id", + "test-worker", + ], + env=env, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + with httpx.Client(trust_env=False, timeout=0.5) as client: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + assert process.poll() is None, "Supervisor exited during startup" + try: + if ( + client.get(f"http://127.0.0.1:{port}/ping") + ).status_code == 200 and (tmp_path / "worker.args").exists(): + break + except httpx.TransportError: + pass + time.sleep(0.05) + else: + raise AssertionError("Listener and worker did not start") + expected_mode = ( + "--managed-agent-worker" + if entrypoint_mode == "worker" + else "--managed-agent-work-item" + ) + assert ( + (tmp_path / "worker.args").read_text() + == f"-m veadk.runtime.managed_agents.worker {expected_mode} --worker-id test-worker" + ) + if stop_mode == "terminate": + process.terminate() + expected = 143 + else: + (tmp_path / "fail").touch() + expected = 23 + assert process.wait(timeout=10) == expected + for role in ("listener", "worker"): + with pytest.raises(ProcessLookupError): + os.kill(int((tmp_path / f"{role}.pid").read_text()), 0) + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=10) + for role in ("listener", "worker"): + path = tmp_path / f"{role}.pid" + if path.exists(): + try: + os.kill(int(path.read_text()), 9) + except ProcessLookupError: + pass diff --git a/tests/runtime/test_anthropic_gateway_e2e.py b/tests/runtime/test_anthropic_gateway_e2e.py new file mode 100644 index 000000000..4ba2bd1ce --- /dev/null +++ b/tests/runtime/test_anthropic_gateway_e2e.py @@ -0,0 +1,79 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +import importlib.util +import asyncio +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + + +PATH = ( + Path(__file__).parents[2] / "examples/16_self_host_sandbox/anthropic_gateway_e2e.py" +) +SPEC = importlib.util.spec_from_file_location("gateway_e2e_config_test", PATH) +e2e = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = e2e +SPEC.loader.exec_module(e2e) + + +@pytest.fixture +def config(monkeypatch): + for key in ("MANAGED_AGENTS_GATEWAY", "MANAGED_AGENTS_GATEWAY_TOKEN"): + monkeypatch.delenv(key, raising=False) + monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://current.example") + monkeypatch.setenv("ANTHROPIC_ENVIRONMENT_KEY", "test-current-key") + monkeypatch.setenv("ANTHROPIC_ENVIRONMENT_ID", "env-test") + monkeypatch.setenv("MANAGED_AGENTS_TEST_AGENT_ID", "agent-test") + return monkeypatch + + +def test_uses_current_anthropic_configuration(config): + args = e2e.parser().parse_args(["--mode", "conversation"]) + assert args.base_url == "https://current.example" + assert args.auth_token == "test-current-key" + + +def test_explicit_and_legacy_overrides(config): + config.setenv("MANAGED_AGENTS_GATEWAY", "https://override.example") + config.setenv("MANAGED_AGENTS_GATEWAY_TOKEN", "test-override-key") + args = e2e.parser().parse_args([]) + assert (args.base_url, args.auth_token) == ( + "https://override.example", + "test-override-key", + ) + args = e2e.parser().parse_args( + ["--base-url", "https://cli.example", "--auth-token", "test-cli"] + ) + assert (args.base_url, args.auth_token) == ("https://cli.example", "test-cli") + + +def test_missing_credentials_fail_before_network(config): + config.delenv("ANTHROPIC_ENVIRONMENT_KEY") + with pytest.raises(SystemExit): + e2e.parser().parse_args([]) + + +def test_failed_conversation_stops_and_archives_owned_session(config): + sessions = SimpleNamespace( + events=SimpleNamespace(send=AsyncMock()), archive=AsyncMock() + ) + client = SimpleNamespace(beta=SimpleNamespace(sessions=sessions)) + context = AsyncMock() + context.__aenter__.return_value = client + config.setattr(e2e.anthropic, "AsyncAnthropic", lambda **kwargs: context) + config.setattr(e2e, "create_session", AsyncMock(return_value="sesn-owned")) + config.setattr(e2e, "run_turn", AsyncMock(side_effect=TimeoutError("turn timeout"))) + args = e2e.parser().parse_args(["--mode", "conversation"]) + with pytest.raises(TimeoutError, match="turn timeout"): + asyncio.run(e2e.async_main(args)) + assert sessions.events.send.await_args_list[0].args == ("sesn-owned",) + assert sessions.events.send.await_args_list[0].kwargs == { + "events": [{"type": "user.interrupt"}] + } + assert sessions.events.send.await_args_list[1].kwargs == { + "events": [{"type": "session.status_terminated"}] + } + sessions.archive.assert_awaited_once_with("sesn-owned") diff --git a/tests/runtime/test_completion_gate.py b/tests/runtime/test_completion_gate.py new file mode 100644 index 000000000..47ad69962 --- /dev/null +++ b/tests/runtime/test_completion_gate.py @@ -0,0 +1,206 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""The first idle is not acceptance: wait for the Work successor transaction.""" + +import asyncio +import json +import runpy +import sys +from types import SimpleNamespace + +import pytest + +from test_fault_recovery_boundary import Events, client, event, fault, turn +import soak_load_test as soak + + +class Page: + def __init__(self, items): + self.data = items + self.next_page = None + + +class WorkAPI: + def __init__(self, rows): + self.rows = rows + self.calls = 0 + + async def list(self, *_args, **_kwargs): + self.calls += 1 + return Page(self.rows) + + +def work(work_id, state, *, kind="agent_loop"): + return { + "id": work_id, + "state": state, + "data": {"id": "session", "type": "session"}, + "runtime_type": "default" if kind == "agent_loop" else None, + "metadata": {"managed_agent_worker_id": work_id}, + "environment_id": "env", + } + + +def setup(*, later=(), rows=None): + events = Events(turn()) + history = turn() + list(later) + + async def list_events(*_args, **_kwargs): + return Page(history) + + events.list = list_events + sdk = client(events) + sdk.beta.environments = SimpleNamespace( + work=WorkAPI(rows or [work("pod-a", "stopped")]) + ) + return sdk + + +def run(sdk): + return asyncio.run( + fault.run_verified_tool_turn( + sdk, + "session", + marker="mark", + prompt="prompt", + timeout=0.15, + environment_id="env", + ) + ) + + +def test_late_different_work_execution_after_first_idle_fails(): + sdk = setup( + later=[event(15, "agent.tool_use", id="second-call")], + rows=[work("pod-a", "stopped"), work("pod-b", "active")], + ) + result = run(sdk) + with pytest.raises(AssertionError, match="after.*idle"): + fault.assert_tool_turn(result, "mark") + assert result["completion"]["events"][-1]["id"] == "second-call" + assert len(result["completion"]["work"]) == 2 + assert result["input_event_id"] == "event-10" + + +def test_idle_with_queued_successor_never_passes_or_archives_early(): + result = run(setup(rows=[work("pod-a", "stopped"), work("pod-b", "queued")])) + with pytest.raises(AssertionError, match="TimeoutError"): + fault.assert_tool_turn(result, "mark") + assert not result["completion"]["settled"] + + +def test_terminal_work_graph_and_no_new_execution_passes(): + sdk = setup(rows=[work("pod-a", "stopped"), work("sandbox", "active", kind="tool")]) + result = run(sdk) + fault.assert_tool_turn(result, "mark") + assert result["completion"]["settled"] + assert sdk.beta.environments.work.calls >= 2 + + +def test_missing_work_evidence_cannot_pass(): + sdk = setup(rows=[work("sandbox", "active", kind="tool")]) + result = run(sdk) + with pytest.raises(AssertionError): + fault.assert_tool_turn(result, "mark") + + +def test_delayed_execution_is_seen_while_successor_is_pending(): + sdk = setup(rows=[work("pod-a", "stopped"), work("pod-b", "queued")]) + + async def delayed_history(*_args, **_kwargs): + items = turn() + if sdk.beta.environments.work.calls >= 2: + items += [event(15, "agent.tool_use", id="late-other-pod")] + return Page(items) + + sdk.beta.sessions.events.list = delayed_history + result = run(sdk) + assert "after end_turn idle" in result["error"] + assert result["completion"]["events"][-1]["id"] == "late-other-pod" + + +def test_soak_uses_same_gate_and_preserves_failure_evidence(): + async def check(): + sdk = setup(later=[event(15, "agent.tool_use", id="late-other-pod")]) + now = soak.time.monotonic() + record = await soak.run_one_turn( + sdk, + "session", + marker="mark", + prompt="prompt", + timeout=0.2, + require_tool=True, + worker=1, + request_id="request", + run_start=now, + deadline=now + 1, + environment_id="env", + ) + assert not record.ok + assert "after end_turn idle" in record.error_reason + assert record.completion["input_event_id"] == "event-10" + assert record.completion["work"] + + asyncio.run(check()) + + +def test_multiple_tools_in_one_input_are_legal(): + sdk = setup() + items = [ + *turn()[:3], + event(13, "agent.tool_use", id="second-call"), + event(14, "agent.tool_result", tool_use_id="second-call", content="mark"), + event(15, "agent.message", content="mark"), + event(16, "session.status_idle", stop_reason={"type": "end_turn"}), + ] + sdk.beta.sessions.events.events = items + + async def list_events(*_args, **_kwargs): + return Page(items) + + sdk.beta.sessions.events.list = list_events + result = run(sdk) + fault.assert_tool_turn(result, "mark") + assert len(result["tool_uses"]) == 2 + assert result["completion"]["settled"] + + +def test_fault_cli_exits_nonzero_and_writes_failed_gate(monkeypatch, tmp_path): + second = run(setup(later=[event(15, "agent.tool_use", id="duplicate")])) + result = { + "recovery_ok": False, + "second_turn": second, + "cleanup": "archived", + "physical_cleanup_reclaimed": True, + } + output = tmp_path / "result.json" + + def completed(coro): + coro.close() + return result + + monkeypatch.setattr(asyncio, "run", completed) + monkeypatch.setattr( + sys, + "argv", + [ + fault.__file__, + "--base-url", + "http://fixture", + "--agent-id", + "agent", + "--environment-id", + "env", + "--tool-id", + "tool", + "--output", + str(output), + ], + ) + with pytest.raises(SystemExit) as raised: + runpy.run_path(fault.__file__, run_name="__main__") + assert raised.value.code == 1 + assert ( + json.loads(output.read_text())["second_turn"]["completion"]["events"][-1]["id"] + == "duplicate" + ) diff --git a/tests/runtime/test_fault_recovery_boundary.py b/tests/runtime/test_fault_recovery_boundary.py new file mode 100644 index 000000000..f39b9300f --- /dev/null +++ b/tests/runtime/test_fault_recovery_boundary.py @@ -0,0 +1,349 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +"""Fault gate regressions: canonical turn identity, not prompt/marker heuristics.""" + +import asyncio +import importlib +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +sys.path.insert(0, str(Path(__file__).parents[2] / "examples/16_self_host_sandbox")) +fault = importlib.import_module("soak_fault_preflight") + + +def event(seq, kind, **fields): + return {"seq": seq, "id": f"event-{seq}", "type": kind, **fields} + + +def turn(start=10, *, error=False): + return [ + event(start, "user.message"), + event(start + 1, "agent.tool_use", id=f"call-{start}"), + event( + start + 2, + "agent.tool_result", + tool_use_id=f"call-{start}", + is_error=error, + content="failed" if error else "mark", + ), + event(start + 3, "agent.message", content="mark"), + event(start + 4, "session.status_idle", stop_reason={"type": "end_turn"}), + ] + + +class Events: + def __init__(self, events, *, acknowledged=None, stall=False): + self.events = events + self.acknowledged = acknowledged or event(10, "user.message") + self.stall = stall + self.closed = False + self.options = None + + async def send(self, *_args, **_kwargs): + return {"data": [self.acknowledged]} + + async def stream(self, *_args, **kwargs): + self.options = kwargs + return self + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + self.closed = True + + async def __aiter__(self): + for item in self.events: + if isinstance(item, Exception): + raise item + yield item + if self.stall: + await asyncio.Event().wait() + + +def client(events): + return SimpleNamespace( + beta=SimpleNamespace(sessions=SimpleNamespace(events=events)) + ) + + +def run(events): + return asyncio.run( + fault.run_verified_tool_turn( + client(events), "session", marker="mark", prompt="new prompt", timeout=0.1 + ) + ) + + +def test_old_error_idle_replay_cannot_finish_recovery(): + events = Events(turn(1, error=True) + turn()) + result = run(events) + fault.assert_tool_turn(result, "mark") + assert result["tool_uses"] == ["call-10"] + assert events.options["extra_headers"] == {"Last-Event-ID": "9"} + assert result["input_event_id"] == "event-10" + assert events.closed + + +def test_new_tool_error_still_fails_and_keeps_evidence(): + result = run(Events(turn(error=True))) + with pytest.raises(AssertionError, match="tool execution failed"): + fault.assert_tool_turn(result, "mark") + assert result["tool_error_ids"] == ["call-10"] + assert result["events"][-1]["type"] == "session.status_idle" + + +@pytest.mark.parametrize( + "items", + [ + turn(1), # No acknowledged new input, even if old final contains marker. + [ + event(10, "user.message"), + event(11, "agent.message", content="mark"), + event(12, "session.status_idle"), + ], # No new tool. + turn()[:-1], # EOF before idle. + [*turn()[:3], event(13, "session.status_idle"), turn()[3]], # Early idle. + [*turn()[:3], event(13, "session.error")], + [*turn()[:3], event(13, "user.message")], # Concurrent unrelated turn. + ], +) +def test_incomplete_or_interleaved_turn_fails(items): + result = run(Events(items)) + with pytest.raises(AssertionError): + fault.assert_tool_turn(result, "mark") + + +def test_timeout_retains_partial_turn_and_closes_stream(): + events = Events(turn()[:2], stall=True) + result = run(events) + with pytest.raises(AssertionError, match="TimeoutError"): + fault.assert_tool_turn(result, "mark") + assert result["tool_uses"] == ["call-10"] + assert events.closed + + +def test_duplicate_event_replay_is_idempotent(): + items = turn() + result = run(Events([items[0], items[1], items[1], *items[2:]])) + fault.assert_tool_turn(result, "mark") + assert result["tool_uses"] == ["call-10"] + + +def test_conflicting_duplicate_sequence_fails(): + items = turn() + result = run(Events([*items[:3], {**items[2], "is_error": True}, *items[3:]])) + with pytest.raises(AssertionError, match="sequence"): + fault.assert_tool_turn(result, "mark") + + +def test_pre_tool_commentary_is_not_the_final(): + items = turn() + result = run( + Events( + [ + items[0], + event(11, "agent.message", content="I will run the tool"), + *[{**item, "seq": item["seq"] + 1} for item in items[1:]], + ] + ) + ) + fault.assert_tool_turn(result, "mark") + assert result["final"] == "mark" + + +def test_no_fault_injection_cannot_pass(): + with pytest.raises(AssertionError, match="physical sandbox"): + asyncio.run( + fault.run_loss_during_active_tool( + client( + Events( + [event(10, "user.message"), event(11, "session.status_idle")] + ) + ), + object(), + "session", + marker="mark", + timeout=0.1, + ) + ) + + +def test_real_sdk_history_pagination_uses_decimal_cursor(): + import anthropic + import httpx2 + from managed_agent_loop import ManagedAgentsLoop + + seen = [] + + def handle(request): + page = request.url.params.get("page") + seen.append(page) + if page not in {"100", "102"}: + return httpx2.Response( + 400, + json={ + "error": { + "type": "invalid_request_error", + "message": "page is invalid", + } + }, + ) + seqs = (101, 102) if page == "100" else (103,) + return httpx2.Response( + 200, + json={ + "data": [event(n, "system.message") for n in seqs], + "has_more": page == "100", + "next_page": "102" if page == "100" else None, + }, + ) + + async def check(): + async with anthropic.AsyncAnthropic( + api_key="test", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handle)), + ) as sdk: + loop = ManagedAgentsLoop(runner=object(), session_id="session") + result = await loop._list_events_after(sdk, 100) + assert [loop._event_seq(item) for item in result] == [101, 102, 103] + + asyncio.run(check()) + assert seen == ["100", "102"] + + +def test_real_sdk_send_ack_and_durable_sse_replay(): + import json + + import anthropic + import httpx2 + + def handle(request): + if request.method == "POST": + return httpx2.Response(200, json={"data": [event(10, "user.message")]}) + assert request.headers["Last-Event-ID"] == "9" + frames = "".join( + f"id: {item['seq']}\nevent: {item['type']}\ndata: {json.dumps(item)}\n\n" + for item in turn() + ) + return httpx2.Response( + 200, text=frames, headers={"Content-Type": "text/event-stream"} + ) + + async def check(): + async with anthropic.AsyncAnthropic( + api_key="test", + max_retries=0, + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handle)), + ) as sdk: + result = await fault.run_verified_tool_turn( + sdk, "session", marker="mark", prompt="new prompt", timeout=1 + ) + fault.assert_tool_turn(result, "mark") + + asyncio.run(check()) + + +@pytest.mark.parametrize( + "failure", + [ + "tool", + "timeout", + "snapshot", + "cleanup_snapshot", + "fault_timeout", + "not_injected", + ], +) +def test_main_keeps_failed_second_turn_before_cleanup(monkeypatch, failure): + items = turn(error=failure == "tool") + items[3]["content"] = "fault-b-12345678" + monkeypatch.setattr(fault.uuid, "uuid4", lambda: SimpleNamespace(hex="12345678")) + events = Events( + items[:2] if failure == "timeout" else items, stall=failure == "timeout" + ) + + class Sessions: + async def create(self, **_kwargs): + return SimpleNamespace(id="session") + + async def archive(self, _session): + return None + + sessions = Sessions() + sessions.events = events + + class Client: + beta = SimpleNamespace(sessions=sessions) + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + pass + + class Inspector: + async def wait_absent(self, *_args, **_kwargs): + pass + + async def wait_present(self, *_args, **_kwargs): + return {"session_id": "replacement"} + + async def loss(*_args, **_kwargs): + if failure == "not_injected": + raise AssertionError("fault never injected") + if failure == "fault_timeout": + return {"physical": {"session_id": "deleted"}, "error": "TimeoutError"} + return {"physical": {"session_id": "deleted"}, "error": None} + + calls = 0 + + async def snapshot(*_args, **_kwargs): + nonlocal calls + calls += 1 + if failure == "snapshot" and calls == 2: + raise RuntimeError("snapshot unavailable") + if failure == "cleanup_snapshot" and calls == 3: + raise RuntimeError("cleanup snapshot unavailable") + return {"rows": [{"work_id": "work", "state": "stopped"}]} + + monkeypatch.setattr(fault.anthropic, "AsyncAnthropic", lambda **_kwargs: Client()) + monkeypatch.setattr(fault, "AgentKitInspector", lambda *_args: Inspector()) + monkeypatch.setattr(fault, "run_loss_during_active_tool", loss) + monkeypatch.setattr(fault, "session_work_snapshot", snapshot) + args = SimpleNamespace( + inspector_script="unused", + tool_id="tool", + base_url="http://fixture", + auth_token="test", + turn_timeout=1, + agent_id="agent", + environment_id="env", + recovery_timeout=0.01, + ) + result = asyncio.run(fault.main(args)) + assert result["recovery_ok"] is False + if failure in {"fault_timeout", "not_injected"}: + assert events.options is None # Never start the dependent recovery turn. + assert result["error"] + assert result["cleanup"] == "archived" + return + assert result["second_turn"]["tool_uses"] == ["call-10"] + assert result["second_turn"]["events"] + assert result["cleanup"] == "archived" + if failure == "tool": + assert result["second_turn"]["tool_error_ids"] == ["call-10"] + elif failure == "timeout": + assert "TimeoutError" in result["second_turn"]["error"] + elif failure == "snapshot": + assert "snapshot unavailable" in result["error"] + else: + assert result["work_after_cleanup_query"]["complete"] is False + assert ( + "cleanup snapshot unavailable" + in result["work_after_cleanup_query"]["error"] + ) diff --git a/tests/runtime/test_managed_agents_delivery.py b/tests/runtime/test_managed_agents_delivery.py new file mode 100644 index 000000000..ecdc02e4c --- /dev/null +++ b/tests/runtime/test_managed_agents_delivery.py @@ -0,0 +1,108 @@ +"""Exercise the public build and installed-package delivery boundaries.""" + +import os +import subprocess +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] + + +def test_image_builder_defaults_to_local_build_without_adjacent_sdk(tmp_path): + commands = tmp_path / "commands" + executable = tmp_path / "docker" + executable.write_text('#!/bin/sh\nprintf "%s\\n" "$*" >> "$BUILD_TEST_COMMANDS"\n') + executable.chmod(0o755) + env = dict(os.environ, BUILD_TEST_COMMANDS=str(commands)) + env["PATH"] = str(tmp_path) + os.pathsep + env["PATH"] + result = subprocess.run( + ["bash", str(ROOT / "docker/managed-agents/build.sh"), "test-worker:local"], + env=env, + text=True, + capture_output=True, + check=True, + ) + command = commands.read_text() + assert "--load" in command + assert "--push" not in command + assert "Dockerfile.worker" in command + assert "test-worker:local" in command + assert "SDK" not in result.stdout + result.stderr + + +def test_image_builder_keeps_explicit_indexes_and_publication_choice(tmp_path): + commands = tmp_path / "commands" + executable = tmp_path / "docker" + executable.write_text('#!/bin/sh\nprintf "%s\\n" "$*" >> "$BUILD_TEST_COMMANDS"\n') + executable.chmod(0o755) + env = dict( + os.environ, + BUILD_TEST_COMMANDS=str(commands), + PIP_INDEX_URL="https://packages.example.com/simple", + UV_DEFAULT_INDEX="https://uv.example.com/simple", + ) + env["PATH"] = str(tmp_path) + os.pathsep + env["PATH"] + subprocess.run( + [ + "bash", + str(ROOT / "docker/managed-agents/build.sh"), + "test-worker:local", + "--push", + ], + env=env, + check=True, + capture_output=True, + ) + command = commands.read_text() + assert "--push" in command + assert "--load" not in command + assert "PIP_INDEX_URL=https://packages.example.com/simple" in command + assert "UV_DEFAULT_INDEX=https://uv.example.com/simple" in command + + +def test_worker_build_context_excludes_ui_credentials_and_local_evidence(): + context = ( + ROOT / "docker/managed-agents/Dockerfile.worker.dockerignore" + ).read_text() + assert "**/.env" in context + assert "frontend" in context + assert "veadk/webui" in context + assert "deploy_records" in context + source = (ROOT / "docker/managed-agents/Dockerfile.worker").read_text() + assert "--require-hashes" in source + assert ".docker-dist" not in source + assert "anthropic-*.whl" not in source + assert "veadk.runtime.managed_agents" in source + + +@pytest.mark.parametrize( + "index", + [ + "https://user:private-test-marker@packages.example.com/simple", + "https://packages.example.com/simple?token=private-test-marker", + "https://packages.example.com/simple#private-test-marker", + ], +) +def test_image_builder_rejects_credentials_before_invoking_docker(tmp_path, index): + commands = tmp_path / "commands" + executable = tmp_path / "docker" + executable.write_text('#!/bin/sh\nprintf called > "$BUILD_TEST_COMMANDS"\n') + executable.chmod(0o755) + env = dict( + os.environ, + BUILD_TEST_COMMANDS=str(commands), + PIP_INDEX_URL=index, + UV_DEFAULT_INDEX="https://pypi.org/simple", + ) + env["PATH"] = str(tmp_path) + os.pathsep + env["PATH"] + result = subprocess.run( + ["bash", str(ROOT / "docker/managed-agents/build.sh"), "test-worker:local"], + env=env, + text=True, + capture_output=True, + ) + assert result.returncode != 0 + assert not commands.exists() + assert "private-test-marker" not in result.stdout + result.stderr diff --git a/tests/runtime/test_preflight_probe.py b/tests/runtime/test_preflight_probe.py new file mode 100644 index 000000000..356e86dbd --- /dev/null +++ b/tests/runtime/test_preflight_probe.py @@ -0,0 +1,83 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + + +MODULE_PATH = ( + Path(__file__).parents[2] + / "examples" + / "16_self_host_sandbox" + / "preflight_probe.py" +) +SPEC = importlib.util.spec_from_file_location("preflight_probe_tested", MODULE_PATH) +probe = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(probe) + + +def test_event_trace_whitelists_tool_linkage_and_error(tmp_path): + output = tmp_path / "events.jsonl" + probe.append_event_trace( + output, + turn="tool-123", + event=SimpleNamespace( + type="agent.tool_result", + id="event-1", + event_id="event-0", + tool_use_id="toolu-1", + is_error=True, + sequence=9, + content="secret tool output", + input={"token": "secret"}, + ), + ) + probe.append_event_trace(output, turn="tool-123", error_type="TimeoutError") + + lines = [json.loads(line) for line in output.read_text().splitlines()] + assert lines == [ + { + "event_id": "event-0", + "id": "event-1", + "is_error": True, + "sequence": 9, + "tool_use_id": "toolu-1", + "turn": "tool-123", + "type": "agent.tool_result", + }, + {"error_type": "TimeoutError", "turn": "tool-123", "type": "probe.error"}, + ] + rendered = output.read_text() + assert "secret tool output" not in rendered + assert "token" not in rendered + + +def test_controlled_tool_events_preserve_probe_correlation(tmp_path): + output = tmp_path / "controlled.jsonl" + tool_use = SimpleNamespace( + type="agent.tool_use", + id="toolu-safe-1", + event_id=None, + tool_use_id=None, + is_error=None, + sequence=3, + input={"command": "printf controlled"}, + ) + tool_result = SimpleNamespace( + type="agent.tool_result", + id="event-result-1", + event_id=None, + tool_use_id="toolu-safe-1", + is_error=False, + sequence=4, + content="controlled output", + ) + probe.append_event_trace(output, turn="tool-controlled", event=tool_use) + probe.append_event_trace(output, turn="tool-controlled", event=tool_result) + rows = [json.loads(line) for line in output.read_text().splitlines()] + assert rows[0]["id"] == rows[1]["tool_use_id"] == "toolu-safe-1" + assert rows[1]["is_error"] is False + assert "printf controlled" not in output.read_text() + assert "controlled output" not in output.read_text() diff --git a/tests/runtime/test_self_host_sandbox_agent.py b/tests/runtime/test_self_host_sandbox_agent.py deleted file mode 100644 index 10d246f8a..000000000 --- a/tests/runtime/test_self_host_sandbox_agent.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import asyncio -import importlib -import sys -import os -from pathlib import Path -from types import SimpleNamespace - - -EXAMPLE_DIR = Path(__file__).parents[2] / "examples" / "16_self_host_sandbox" -if str(EXAMPLE_DIR) not in sys.path: - sys.path.insert(0, str(EXAMPLE_DIR)) - -os.environ.setdefault("ANTHROPIC_BASE_URL", "https://sandbox.example.com") -os.environ.setdefault("ANTHROPIC_ENVIRONMENT_KEY", "test-token") -os.environ.setdefault("MODEL_AGENT_API_KEY", "test-model-api-key") - -agent_module = importlib.import_module("agents.self_host_sandbox_agent.agent") -SandboxSessionManager = agent_module.SandboxSessionManager - - -class _FakeClient: - def __init__(self, remote_session_id: str): - self.session_id = None - self.remote_session_id = remote_session_id - self.created_titles = [] - self.idle_count = 0 - - def create_session(self, title: str): - self.created_titles.append(title) - self.session_id = self.remote_session_id - - def post_status_idle(self): - self.idle_count += 1 - - -def test_each_veadk_session_creates_a_distinct_remote_session(monkeypatch): - manager = SandboxSessionManager() - clients = iter((_FakeClient("remote-1"), _FakeClient("remote-2"))) - monkeypatch.setattr(manager, "_new_client", lambda: next(clients)) - - assert manager.create_remote_session("veadk-1") == "remote-1" - assert manager.create_remote_session("veadk-1") == "remote-1" - assert manager.create_remote_session("veadk-2") == "remote-2" - - assert manager.get("veadk-1") is not manager.get("veadk-2") - assert manager.get("veadk-1").created_titles == [ - "VeADK Self-Hosted Sandbox Session veadk-1" - ] - assert manager.get("veadk-2").created_titles == [ - "VeADK Self-Hosted Sandbox Session veadk-2" - ] - - -def test_web_session_service_creation_provisions_remote_session(monkeypatch): - created = [] - monkeypatch.setattr( - agent_module.sandbox_sessions, - "create_remote_session", - lambda session_id: created.append(session_id) or "remote-web", - ) - - asyncio.run( - agent_module.short_term_memory.session_service.create_session( - app_name="self_host_sandbox_agent", - user_id="user", - session_id="web-session", - ) - ) - - assert created == ["web-session"] - - -def test_dispatch_task_sends_only_the_model_tool_call(monkeypatch): - dispatched = [] - client = SimpleNamespace( - dispatch_tool=lambda name, arguments, *, dispatch_id: dispatched.append( - (name, arguments, dispatch_id) - ) - or {"stdout": "ok"} - ) - monkeypatch.setattr(agent_module.sandbox_sessions, "get", lambda session_id: client) - tool_call = SimpleNamespace( - session_id="veadk-session", - id="tool-call-1", - name="bash", - arguments={"command": "printf ok"}, - ) - - result = asyncio.run(agent_module.dispatch_task(tool_call)) - - assert result == {"stdout": "ok"} - assert dispatched == [("bash", {"command": "printf ok"}, "tool-call-1")] - assert not hasattr(agent_module.agent, "run_turn") diff --git a/tests/runtime/test_soak_load_test.py b/tests/runtime/test_soak_load_test.py new file mode 100644 index 000000000..056bad811 --- /dev/null +++ b/tests/runtime/test_soak_load_test.py @@ -0,0 +1,1074 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + +import asyncio +import importlib.util +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +SCRIPT = ( + Path(__file__).parents[2] + / "examples" + / "16_self_host_sandbox" + / "soak_load_test.py" +) +SPEC = importlib.util.spec_from_file_location("soak_load_test", SCRIPT) +assert SPEC and SPEC.loader +soak = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = soak +SPEC.loader.exec_module(soak) + + +class _Stream: + def __init__(self, events): + self._events = events + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + def __aiter__(self): + return self._iterate() + + async def _iterate(self): + for event in self._events: + if isinstance(event, BaseException): + raise event + yield event + + +class _Events: + def __init__(self, events): + self._events = events + + async def stream(self, *_args, **_kwargs): + return _Stream(self._events) + + async def send(self, *_args, **_kwargs): + return None + + +def _client(events): + return SimpleNamespace( + beta=SimpleNamespace( + sessions=SimpleNamespace(events=_Events(events)), + ) + ) + + +def _event(event_type, **values): + return SimpleNamespace(type=event_type, **values) + + +def _fault_module(): + fault_path = SCRIPT.with_name("soak_fault_preflight.py") + spec = importlib.util.spec_from_file_location("soak_fault_preflight", fault_path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_jsonl_writer_keeps_every_success(tmp_path): + output = tmp_path / "turns.jsonl" + with soak.DurableJsonlWriter(output) as writer: + for index in range(125): + writer.append( + soak.RequestRecord( + request_id=f"request-{index}", + worker=index % 50, + session_id=f"session-{index}", + started_at="start", + completed_at="end", + started_offset_seconds=float(index), + completed_offset_seconds=float(index + 1), + completion_phase="formal", + elapsed_seconds=1.0, + first_delta_seconds=0.5, + ok=True, + require_tool=True, + tool_uses=1, + tool_results=1, + tool_use_ids=[f"call-{index}"], + tool_result_ids=[f"call-{index}"], + ) + ) + + rows = [json.loads(line) for line in output.read_text().splitlines()] + assert len(rows) == 125 + assert [row["request_id"] for row in rows] == [ + f"request-{index}" for index in range(125) + ] + + +def test_tool_success_requires_exact_use_result_id_match(): + good = [ + _event("agent.tool_use", id="call-1"), + _event("agent.tool_result", tool_use_id="call-1", content=[]), + _event("agent.message", content=[{"type": "text", "text": "mark"}]), + _event("session.status_idle"), + ] + bad = [ + _event("agent.tool_use", id="call-1"), + _event("agent.tool_result", tool_use_id="call-other", content=[]), + _event("agent.message", content=[{"type": "text", "text": "mark"}]), + _event("session.status_idle"), + ] + + async def run(events): + now = soak.time.monotonic() + return await soak.run_one_turn( + _client(events), + "session", + marker="mark", + prompt="prompt", + timeout=1, + require_tool=True, + worker=0, + request_id="request", + run_start=now, + deadline=now + 10, + ) + + matched = asyncio.run(run(good)) + mismatched = asyncio.run(run(bad)) + assert matched.ok is True + assert matched.tool_use_ids == matched.tool_result_ids == ["call-1"] + assert mismatched.ok is False + + +def test_tool_error_result_is_a_failed_turn(): + events = [ + _event("agent.tool_use", id="call-1"), + _event( + "agent.tool_result", + tool_use_id="call-1", + is_error=True, + content=[{"type": "text", "text": "command failed"}], + ), + _event("agent.message", content=[{"type": "text", "text": "mark"}]), + _event("session.status_idle"), + ] + + async def run(): + now = soak.time.monotonic() + return await soak.run_one_turn( + _client(events), + "session", + marker="mark", + prompt="prompt", + timeout=1, + require_tool=True, + worker=0, + request_id="request", + run_start=now, + deadline=now + 10, + ) + + result = asyncio.run(run()) + assert result.ok is False + assert result.tool_use_ids == result.tool_result_ids == ["call-1"] + assert result.tool_error_ids == ["call-1"] + assert result.error_category == "tool_result_error" + + +def test_error_record_keeps_redacted_detail_and_separate_category(): + now = soak.time.monotonic() + record = soak._error_record( + "request", + 0, + "session", + "start", + now, + RuntimeError("rate limit Authorization: Bearer abc123 token=secret-value"), + require_tool=True, + run_start=now, + deadline=now + 10, + ) + + assert record.error_category == "rate_limit" + assert "rate limit" in record.error_reason + assert "abc123" not in record.error_reason + assert "secret-value" not in record.error_reason + assert "rate limit" in record.error_reason + + +@pytest.mark.parametrize( + "detail,secrets", + [ + ("secret_key=synthetic_secret request_id=req-1", ["synthetic_secret"]), + ( + '{"secret_key":"synthetic_secret","request_id":"req-1"}', + ["synthetic_secret"], + ), + ("{'api-key': 'synthetic-api', 'context': 'useful'}", ["synthetic-api"]), + ('access_key = "synthetic-ak" status=429', ["synthetic-ak"]), + ('secret_access_key: "synthetic-sk" path=/v1/work', ["synthetic-sk"]), + ( + "client_secret=synthetic-client session_token=synthetic-token", + ["synthetic-client", "synthetic-token"], + ), + ( + 'private_key="-----BEGIN_PRIVATE_KEY-----" request-id=req-2', + ["BEGIN_PRIVATE_KEY"], + ), + ( + "Authorization: Bearer synthetic-bearer request_id=req-3", + ["synthetic-bearer"], + ), + ], +) +def test_sanitize_error_detail_redacts_credential_key_forms(detail, secrets): + redacted = soak.sanitize_error_detail(detail) + for secret in secrets: + assert secret not in redacted + assert "" in redacted + if "request_id" in detail: + assert "request_id" in redacted + if "status=429" in detail: + assert "status=429" in redacted + + +def test_turn_exception_preserves_partial_tool_ids_category_and_redacted_detail(): + events = [ + _event("agent.tool_use", id="call-1"), + TimeoutError("timed out token=secret-value"), + ] + now = soak.time.monotonic() + record = asyncio.run( + soak.run_one_turn( + _client(events), + "session", + marker="mark", + prompt="prompt", + timeout=1, + require_tool=True, + worker=0, + request_id="request", + run_start=now, + deadline=now + 10, + ) + ) + assert record.ok is False + assert record.tool_use_ids == ["call-1"] + assert record.tool_result_ids == [] + assert record.error_category == "timeout" + assert "timed out" in record.error_reason + assert "secret-value" not in record.error_reason + + +def test_turn_exception_after_tool_error_preserves_error_evidence(): + events = [ + _event("agent.tool_use", id="call-1"), + _event( + "agent.tool_result", + tool_use_id="call-1", + is_error=True, + content="failed token=secret-value", + ), + RuntimeError("terminal failure"), + ] + now = soak.time.monotonic() + record = asyncio.run( + soak.run_one_turn( + _client(events), + "session", + marker="mark", + prompt="prompt", + timeout=1, + require_tool=True, + worker=0, + request_id="request", + run_start=now, + deadline=now + 10, + ) + ) + assert record.ok is False + assert record.tool_error_ids == ["call-1"] + assert "secret-value" not in record.tool_error_details[0] + + +def test_physical_lookup_is_filtered_to_run_sessions(): + calls = [] + + class _Inspector: + tool_id = "tool" + + class _Client: + @staticmethod + def call(action, body): + calls.append((action, body)) + return {"SessionInfos": []} + + client = _Client() + + result = asyncio.run( + soak.run_scoped_physical_sessions(_Inspector(), {"session-b", "session-a"}) + ) + + assert result == [] + assert calls == [ + ( + "ListSessions", + { + "ToolId": "tool", + "MaxResults": 100, + "Filters": [ + { + "Name": "UserSessionId", + "Values": ["session-a", "session-b"], + } + ], + }, + ) + ] + + +def test_physical_sample_records_same_instant_in_flight_mapping(tmp_path): + output = tmp_path / "samples.jsonl" + stop = asyncio.Event() + + class _Inspector: + async def sessions_for_managed_sessions(self, managed_session_ids): + assert managed_session_ids == { + "session-a", + "session-b", + "session-completed-awaiting-reclaim", + } + stop.set() + return [ + { + "session_id": "physical-a", + "user_session_id": "session-a", + "status": "Ready", + "created_at": "now", + }, + { + "session_id": "physical-completed", + "user_session_id": "session-completed-awaiting-reclaim", + "status": "Ready", + "created_at": "now", + }, + ] + + samples = [] + with soak.DurableJsonlWriter(output) as writer: + asyncio.run( + soak.sample_physical( + _Inspector(), + {"session-a", "session-b", "session-completed-awaiting-reclaim"}, + {"request-a": "session-a", "request-b": "session-b"}, + asyncio.Lock(), + stop, + samples, + writer, + 30, + ) + ) + + assert len(samples) == 1 + assert samples[0]["in_flight_requests"] == 2 + assert samples[0]["in_flight"] == { + "request-a": "session-a", + "request-b": "session-b", + } + assert [row["session_id"] for row in samples[0]["physical_sessions"]] == [ + "physical-a", + "physical-completed", + ] + assert samples[0]["snapshot_started_at"] <= samples[0]["snapshot_completed_at"] + assert json.loads(output.read_text())["in_flight_requests"] == 2 + + +def test_cleanup_confirms_physical_reclaim_and_records_query_window(tmp_path): + class _Sessions: + async def archive(self, session_id, *, timeout): + assert session_id == "session-a" + assert timeout == 7 + + class _Inspector: + async def sessions_for_managed_sessions(self, session_ids): + assert session_ids == {"session-a"} + return [] + + record = soak.RequestRecord( + request_id="request-a", + worker=0, + session_id="session-a", + started_at="start", + completed_at="end", + started_offset_seconds=0, + completed_offset_seconds=1, + completion_phase="formal", + elapsed_seconds=1, + first_delta_seconds=0.5, + ok=True, + require_tool=True, + tool_uses=1, + tool_results=1, + ) + active = {"session-a"} + cleanup_records = [] + with soak.DurableJsonlWriter(tmp_path / "cleanup.jsonl") as writer: + asyncio.run( + soak.cleanup_session( + SimpleNamespace(beta=SimpleNamespace(sessions=_Sessions())), + "session-a", + record, + lock=asyncio.Lock(), + active=active, + cleanup_records=cleanup_records, + cleanup_writer=writer, + archive_timeout=7, + ) + ) + asyncio.run( + soak.confirm_physical_reclaims( + _Inspector(), + active, + [record], + cleanup_records, + writer, + timeout=7, + interval=0, + ) + ) + + assert active == set() + assert record.cleanup_status == "archived_physical_absent" + assert record.physical_reclaimed is True + assert record.physical_query_started_at <= record.physical_query_completed_at + assert cleanup_records[0]["physical_reclaimed"] is True + assert record.physical_query_count == 1 + + +def test_physical_lookup_chunks_large_run_and_paginates(): + calls = [] + + class _Inspector: + tool_id = "tool" + + class _Client: + @staticmethod + def call(_action, body): + calls.append(body) + if ( + body.get("NextToken") is None + and len(body["Filters"][0]["Values"]) == 50 + ): + return {"SessionInfos": [], "NextToken": "next"} + return {"SessionInfos": []} + + client = _Client() + + session_ids = {f"session-{index:02d}" for index in range(51)} + assert ( + asyncio.run(soak.run_scoped_physical_sessions(_Inspector(), session_ids)) == [] + ) + assert [len(call["Filters"][0]["Values"]) for call in calls] == [50, 50, 1] + assert calls[1]["NextToken"] == "next" + + +def test_physical_lookup_fails_closed_on_wrong_owner_or_schema(): + class _Inspector: + tool_id = "tool" + + class _Client: + response = {"SessionInfos": {}} + + @classmethod + def call(cls, _action, _body): + return cls.response + + client = _Client() + + with pytest.raises(RuntimeError, match="SessionInfos list"): + asyncio.run(soak.run_scoped_physical_sessions(_Inspector(), {"session"})) + _Inspector._Client.response = { + "SessionInfos": [{"SessionId": "physical", "UserSessionId": "other"}] + } + with pytest.raises(RuntimeError, match="unexpected UserSessionId"): + asyncio.run(soak.run_scoped_physical_sessions(_Inspector(), {"session"})) + + +def test_physical_sample_error_log_redacts_json_credential(tmp_path): + stop = asyncio.Event() + + class _Inspector: + async def sessions_for_managed_sessions(self, _session_ids): + stop.set() + raise RuntimeError('{"secret_key":"synthetic_secret","request_id":"req-1"}') + + samples = [] + output = tmp_path / "samples.jsonl" + with soak.DurableJsonlWriter(output) as writer: + asyncio.run( + soak.sample_physical( + _Inspector(), + {"session"}, + {"request": "session"}, + asyncio.Lock(), + stop, + samples, + writer, + 30, + ) + ) + error = samples[0]["error"] + assert "synthetic_secret" not in error + assert "request_id" in error + assert "RuntimeError" in error + assert "synthetic_secret" not in output.read_text() + + +def test_batch_reclaim_marks_remaining_session_unconfirmed(tmp_path): + records = [] + cleanup_records = [] + for session_id in ("gone", "present"): + record = soak.RequestRecord( + request_id=session_id, + worker=0, + session_id=session_id, + started_at="start", + completed_at="end", + started_offset_seconds=0, + completed_offset_seconds=1, + completion_phase="formal", + elapsed_seconds=1, + first_delta_seconds=0.5, + ok=True, + require_tool=True, + tool_uses=1, + tool_results=1, + cleanup_status="archived", + ) + records.append(record) + cleanup_records.append( + {"session_id": session_id, "status": "archived", "error": None} + ) + + class _Inspector: + async def sessions_for_managed_sessions(self, _session_ids): + return [ + { + "session_id": "physical", + "user_session_id": "present", + "status": "Ready", + "created_at": "now", + } + ] + + active = {"gone", "present"} + with soak.DurableJsonlWriter(tmp_path / "cleanup.jsonl") as writer: + asyncio.run( + soak.confirm_physical_reclaims( + _Inspector(), + active, + records, + cleanup_records, + writer, + timeout=0.01, + interval=0.02, + ) + ) + assert active == {"present"} + assert records[0].physical_reclaimed is True + assert records[1].physical_reclaimed is False + assert records[1].error_category is None + assert "still present" in records[1].cleanup_error + + +def test_archive_timeout_is_recorded_before_physical_confirmation(tmp_path): + class _Sessions: + async def archive(self, _session_id, *, timeout): + assert timeout == 0.01 + await asyncio.sleep(10) + + class _Inspector: + async def sessions_for_managed_sessions(self, _session_ids): + return [ + { + "session_id": "physical-a", + "user_session_id": "session-a", + "status": "Ready", + "created_at": "now", + } + ] + + record = soak.RequestRecord( + request_id="request-a", + worker=0, + session_id="session-a", + started_at="start", + completed_at="end", + started_offset_seconds=0, + completed_offset_seconds=1, + completion_phase="formal", + elapsed_seconds=1, + first_delta_seconds=0.5, + ok=True, + require_tool=True, + tool_uses=1, + tool_results=1, + cleanup_status="pending", + ) + active = {"session-a"} + cleanup_records = [] + output = tmp_path / "cleanup.jsonl" + with soak.DurableJsonlWriter(output) as writer: + asyncio.run( + soak.cleanup_session( + SimpleNamespace(beta=SimpleNamespace(sessions=_Sessions())), + "session-a", + record, + lock=asyncio.Lock(), + active=active, + cleanup_records=cleanup_records, + cleanup_writer=writer, + archive_timeout=0.01, + ) + ) + asyncio.run( + soak.confirm_physical_reclaims( + _Inspector(), + active, + [record], + cleanup_records, + writer, + timeout=0.01, + interval=0.02, + ) + ) + + rows = [json.loads(line) for line in output.read_text().splitlines()] + assert rows[0]["status"] == "archive_failed" + assert rows[-1]["status"] == "archive_failed_physical_unconfirmed" + assert record.physical_query_count >= 1 + assert active == {"session-a"} + + +def test_summary_schema_keeps_distinct_phase_and_latency_metrics(): + fields = soak.Summary.__dataclass_fields__ + assert "success_latency" in fields + assert "latency_including_failures" in fields + assert "first_text_event" in fields + assert "formal_completions" in fields + assert "drain_completions" in fields + assert "peak_in_flight_requests" in fields + assert "observed_peak_physical_sessions" in fields + assert "physical_reclaim_unverified" in fields + + +def test_cli_result_fails_when_cleanup_or_physical_reclaim_is_incomplete( + monkeypatch, tmp_path +): + result = { + "summary": {}, + "cleanup_records": [ + {"status": "archived_physical_unconfirmed", "physical_reclaimed": False} + ], + "unreclaimed_logical_session_ids": ["session-a"], + } + monkeypatch.setattr(soak, "_with_progress", lambda _args: _async_value(result)) + monkeypatch.setattr( + sys, + "argv", + [ + "soak_load_test.py", + "--base-url", + "https://example", + "--agent-id", + "agent", + "--environment-id", + "env", + "--with-inspector", + "--tool-id", + "tool", + "--output", + str(tmp_path / "result.json"), + ], + ) + assert soak.main() == 1 + + +async def _async_value(value): + return value + + +def test_fault_preflight_requires_exact_tool_match(): + module = _fault_module() + + module.assert_tool_turn( + { + "tool_uses": ["call-1"], + "tool_results": ["call-1"], + "final": "marker", + "types": ["agent.tool_use", "agent.tool_result", "session.status_idle"], + }, + "marker", + ) + with pytest.raises(AssertionError, match="tool IDs do not match"): + module.assert_tool_turn( + { + "tool_uses": ["call-1"], + "tool_results": ["call-2"], + "final": "marker", + "types": ["session.status_idle"], + }, + "marker", + ) + with pytest.raises(AssertionError, match="tool execution failed"): + module.assert_tool_turn( + { + "tool_uses": ["call-1"], + "tool_results": ["call-1"], + "tool_error_ids": ["call-1"], + "tool_error_details": ["command failed"], + "final": "marker", + "types": ["session.status_idle"], + }, + "marker", + ) + + +def _fault_client(events): + boundary = _event("user.message", id="new-input", seq=1) + for index, event in enumerate(events, 2): + event.seq = index + if event.type == "session.status_idle": + event.stop_reason = {"type": "end_turn"} + sdk = _client([boundary, *events]) + + async def send(*_args, **_kwargs): + return {"data": [boundary]} + + sdk.beta.sessions.events.send = send + return sdk + + +def test_fault_verified_turn_rejects_tool_is_error(): + module = _fault_module() + events = [ + _event("agent.tool_use", id="call-1"), + _event( + "agent.tool_result", tool_use_id="call-1", is_error=True, content="failed" + ), + _event("agent.message", content=[{"type": "text", "text": "mark"}]), + _event("session.status_idle"), + ] + result = asyncio.run( + module.run_verified_tool_turn( + _fault_client(events), "session", marker="mark", prompt="prompt", timeout=1 + ) + ) + assert result["tool_error_ids"] == ["call-1"] + assert result["tool_error_details"] == ["failed"] + with pytest.raises(AssertionError, match="tool execution failed"): + module.assert_tool_turn(result, "mark") + + +def test_fault_error_detail_is_redacted(): + module = _fault_module() + events = [ + _event("agent.tool_use", id="call-1"), + _event( + "agent.tool_result", + tool_use_id="call-1", + is_error=True, + content="Authorization: Bearer abc token=secret-value", + ), + _event("session.status_idle"), + ] + result = asyncio.run( + module.run_verified_tool_turn( + _fault_client(events), "session", marker="mark", prompt="prompt", timeout=1 + ) + ) + detail = result["tool_error_details"][0] + assert "abc" not in detail + assert "secret-value" not in detail + + +def test_agentkit_inspector_uses_paginated_session_listing(): + module = _fault_module() + calls = [] + inspector = module.AgentKitInspector.__new__(module.AgentKitInspector) + inspector.tool_id = "tool" + inspector.client = object() + inspector.module = SimpleNamespace( + list_all_sessions=lambda client, tool_id, max_results: ( + calls.append((client, tool_id, max_results)) + or [ + { + "SessionId": "physical", + "UserSessionId": "managed", + "Status": "Ready", + "CreatedAt": "now", + } + ] + ) + ) + + sessions = asyncio.run(inspector.sessions()) + + assert calls == [(inspector.client, "tool", 100)] + assert sessions == [ + { + "session_id": "physical", + "user_session_id": "managed", + "status": "Ready", + "created_at": "now", + } + ] + + +def test_fault_preflight_preserves_physical_lookup_error(): + module = _fault_module() + + class _Inspector: + async def wait_present(self, _session_id, *, timeout): + assert timeout == 120 + raise RuntimeError( + "ListSessions failed Authorization: Bearer abc token=secret-value" + ) + + with pytest.raises(RuntimeError, match="physical lookup failed") as raised: + asyncio.run( + module.run_loss_during_active_tool( + _client([_event("agent.tool_use", id="call-1")]), + _Inspector(), + "session", + marker="mark", + timeout=1, + ) + ) + assert "abc" not in str(raised.value) + assert "secret-value" not in str(raised.value) + + +def _work(work_id, session_id, *, environment_id="environment", state="queued"): + return { + "id": work_id, + "work_id": work_id, + "type": "work", + "work_type": "session", + "environment_id": environment_id, + "data": {"id": session_id, "type": "session"}, + "state": state, + "metadata": {}, + "created_at": "now", + } + + +def test_work_snapshot_traverses_pages_and_filters_known_session(monkeypatch): + module = _fault_module() + responses = [ + {"data": [_work("other", "other-session")], "next_page": "page-2"}, + {"data": [_work("target", "session")], "next_page": None}, + ] + + class _Response: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self.payload + + class _HTTP: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, _url, *, params, **_kwargs): + expected_page = None if len(responses) == 2 else "page-2" + assert params.get("page") == expected_page + return _Response(responses.pop(0)) + + monkeypatch.setattr(module.httpx, "AsyncClient", lambda **_kwargs: _HTTP()) + snapshot = asyncio.run( + module.session_work_snapshot( + "https://gateway.example", "token", "environment", "session" + ) + ) + assert snapshot["complete"] is True + assert snapshot["page_count"] == 2 + assert snapshot["query_started_at"] <= snapshot["query_completed_at"] + assert [row["work_id"] for row in snapshot["rows"]] == ["target"] + + +@pytest.mark.parametrize( + "payload,match", + [ + ({"data": {}, "next_page": None}, "data list"), + ({"data": [{}], "next_page": None}, "work_id/id"), + ( + { + "data": [{**_work("target", "session"), "work_type": "other"}], + "next_page": None, + }, + "type/work_type", + ), + ( + {"data": [_work("target", "session")], "next_page": "same"}, + "repeated next_page", + ), + ], +) +def test_work_snapshot_fails_closed_on_schema_or_pagination( + monkeypatch, payload, match +): + module = _fault_module() + + class _Response: + def raise_for_status(self): + return None + + def json(self): + return payload + + class _HTTP: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, *_args, **_kwargs): + return _Response() + + monkeypatch.setattr(module.httpx, "AsyncClient", lambda **_kwargs: _HTTP()) + with pytest.raises(RuntimeError, match=match): + asyncio.run( + module.session_work_snapshot( + "https://gateway.example", "token", "environment", "session" + ) + ) + + +def test_work_snapshot_requires_previously_seen_work_ids(monkeypatch): + module = _fault_module() + + class _Response: + def raise_for_status(self): + return None + + def json(self): + return {"data": [], "next_page": None} + + class _HTTP: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + async def get(self, *_args, **_kwargs): + return _Response() + + monkeypatch.setattr(module.httpx, "AsyncClient", lambda **_kwargs: _HTTP()) + with pytest.raises(RuntimeError, match="lost known IDs"): + asyncio.run( + module.session_work_snapshot( + "https://gateway.example", + "token", + "environment", + "session", + expected_work_ids={"known-work"}, + ) + ) + + +def test_async_main_keeps_all_success_records(monkeypatch, tmp_path): + async def fake_worker( + _client, worker, *, records, lock, in_flight, writer, peak_in_flight, **kwargs + ): + for index in range(25): + request_id = f"{worker}-{index}" + record = soak.RequestRecord( + request_id=request_id, + worker=worker, + session_id=f"session-{request_id}", + started_at="start", + completed_at="end", + started_offset_seconds=0.0, + completed_offset_seconds=1.0, + completion_phase="formal", + elapsed_seconds=1.0, + first_delta_seconds=0.5, + ok=True, + require_tool=True, + tool_uses=1, + tool_results=1, + tool_use_ids=[f"call-{request_id}"], + tool_result_ids=[f"call-{request_id}"], + cleanup_status="archived", + ) + async with lock: + in_flight[request_id] = record.session_id + peak_in_flight[0] = max(peak_in_flight[0], len(in_flight)) + records.append(record) + writer.append(record) + in_flight.pop(request_id) + + class _Client: + kwargs = None + + def __init__(self, **kwargs): + type(self).kwargs = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return None + + monkeypatch.setattr(soak, "worker_loop", fake_worker) + monkeypatch.setattr(soak.anthropic, "AsyncAnthropic", _Client) + output = tmp_path / "result.json" + args = SimpleNamespace( + base_url="https://gateway.example", + auth_token="secret", + agent_id="agent", + environment_id="environment", + concurrency=5, + duration=1.0, + turn_timeout=1.0, + request_gap=0.0, + require_tool=True, + with_inspector=False, + tool_id=None, + sample_interval=30.0, + physical_reclaim_timeout=1.0, + output=output, + records_output=None, + samples_output=None, + cleanup_output=None, + progress_interval=60.0, + run_id="test-run", + ) + + result = asyncio.run(soak.async_main(args)) + rows = [ + json.loads(line) + for line in output.with_suffix(".records.jsonl").read_text().splitlines() + ] + assert result["summary"]["total_requests"] == 125 + assert result["summary"]["successes"] == 125 + assert len(result["records"]) == 125 + assert len(rows) == 125 + assert {row["request_id"] for row in rows} == { + f"{worker}-{index}" for worker in range(5) for index in range(25) + } + assert _Client.kwargs["max_retries"] == 0 diff --git a/tests/skills/test_ma_infra_skills_client.py b/tests/skills/test_ma_infra_skills_client.py new file mode 100644 index 000000000..75b04856b --- /dev/null +++ b/tests/skills/test_ma_infra_skills_client.py @@ -0,0 +1,233 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import zipfile +from io import BytesIO + +import httpx +import pytest + +from veadk.skills.ma_infra import ( + MaInfraSkillConfig, + MaInfraSkillError, + MaInfraSkillsClient, +) + + +def _zip(content: bytes = b"SKILL") -> bytes: + buffer = BytesIO() + with zipfile.ZipFile(buffer, "w") as target: + target.writestr("SKILL.md", content) + return buffer.getvalue() + + +def _config(**kwargs: str) -> MaInfraSkillConfig: + values = { + "base_url": "https://ma-infra.example", + "api_key": "key-1", + "account_id": "account-1", + **kwargs, + } + return MaInfraSkillConfig(**values) + + +def _client(handler) -> MaInfraSkillsClient: + transport = httpx.MockTransport(handler) + return MaInfraSkillsClient(_config(), transport=transport) + + +def _json_response(payload: object, status: int = 200) -> httpx.Response: + return httpx.Response( + status, + json=payload, + headers={"Content-Type": "application/json"}, + ) + + +def test_config_from_env_requires_base_url(monkeypatch) -> None: + monkeypatch.delenv("MA_SKILLS_BASE_URL", raising=False) + assert MaInfraSkillConfig.from_env() is None + + monkeypatch.setenv("MA_SKILLS_BASE_URL", " https://ma-infra.example/ ") + monkeypatch.setenv("MA_SKILLS_API_KEY", "key-1") + monkeypatch.setenv("MA_SKILLS_ACCOUNT_ID", "account-1") + + config = MaInfraSkillConfig.from_env() + assert config is not None + assert config.base_url == "https://ma-infra.example" + assert config.api_key == "key-1" + assert config.account_id == "account-1" + + +def test_list_skills_sends_account_headers_and_parses_data() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert request.headers["x-api-key"] == "key-1" + assert request.headers["X-Top-Account-Id"] == "account-1" + assert request.url.params["source"] == "custom" + return _json_response( + { + "data": [ + { + "id": "s-1", + "type": "skill", + "display_name": "document-summary", + "description": "d", + "latest_version_id": "v1", + "source": {"type": "custom"}, + } + ], + "next_page": None, + } + ) + + client = _client(handler) + skills = client.list_skills() + assert skills[0]["id"] == "s-1" + assert skills[0]["display_name"] == "document-summary" + assert skills[0]["latest_version_id"] == "v1" + + +def test_get_version_resolves_latest_to_concrete_id() -> None: + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.url.path) + return _json_response( + { + "id": "v2", + "type": "skill_version", + "name": "document-summary", + "description": "d", + } + ) + + client = _client(handler) + meta = client.get_version("s-1", "latest") + + assert meta["id"] == "v2" + assert meta["name"] == "document-summary" + assert seen == ["/v1/skills/s-1/versions/latest"] + + +def test_download_content_resolves_latest_before_download() -> None: + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path.endswith("/content"): + return httpx.Response(200, content=_zip()) + return _json_response( + { + "id": "v2", + "type": "skill_version", + "name": "document-summary", + "description": "d", + } + ) + + client = _client(handler) + content = client.download_content("s-1", "latest") + + assert content.startswith(b"PK") + assert paths == [ + "/v1/skills/s-1/versions/latest", + "/v1/skills/s-1/versions/v2/content", + ] + + +def test_download_content_uses_pinned_version_directly() -> None: + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + return httpx.Response(200, content=_zip()) + + client = _client(handler) + content = client.download_content("s-1", "v1") + + assert content.startswith(b"PK") + assert paths == ["/v1/skills/s-1/versions/v1/content"] + + +def test_create_skill_sends_multipart_archive() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert "multipart/form-data" in request.headers["Content-Type"] + assert ( + b"document-summary" in request.content or b"display_name" in request.content + ) + return _json_response( + { + "id": "s-9", + "type": "skill", + "display_name": "document-summary", + "description": "d", + "latest_version_id": "v1", + } + ) + + client = _client(handler) + created = client.create_skill("document-summary", _zip()) + + assert created["id"] == "s-9" + + +def test_delete_skill_round_trips() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert request.method == "DELETE" + return _json_response({"id": "s-1", "type": "skill_deleted"}) + + client = _client(handler) + assert client.delete_skill("s-1")["type"] == "skill_deleted" + + +def test_error_body_is_surfaced() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return _json_response( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": "bad id"}, + }, + status=400, + ) + + client = _client(handler) + with pytest.raises(MaInfraSkillError) as excinfo: + client.get_skill("s-1") + assert excinfo.value.status_code == 400 + assert excinfo.value.code == "invalid_request_error" + assert not excinfo.value.retryable + + +def test_transport_failure_is_retryable() -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("boom") + + client = _client(handler) + with pytest.raises(MaInfraSkillError) as excinfo: + client.list_skills() + assert excinfo.value.retryable is True + + +@pytest.mark.asyncio +async def test_async_helpers_delegate_to_sync() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return _json_response({"data": [], "next_page": None}) + + client = _client(handler) + skills = await client.alist_skills() + assert skills == [] + await asyncio.to_thread(client.close) diff --git a/tests/test_feishu_channel_extension.py b/tests/test_feishu_channel_extension.py index 24c321cb9..b0fc9b88a 100644 --- a/tests/test_feishu_channel_extension.py +++ b/tests/test_feishu_channel_extension.py @@ -22,6 +22,7 @@ from veadk.extensions.feishu_channel import ( FeishuChannelExtension, _call_in_fresh_event_loop, + _format_tool_payload, ) @@ -58,7 +59,8 @@ def __init__(self): async def stream(self, chat_id, spec, options=None): controller = FakeStreamController() await spec["markdown"](controller) - self.stream_calls.append((chat_id, controller.chunks, options)) + if controller.chunks: + self.stream_calls.append((chat_id, controller.chunks, options)) class FakeBlockingChannel(FakeChannel): @@ -192,6 +194,120 @@ async def run_async(self, user_id, session_id, new_message, run_config=None): ) +def build_stream_event( + *, + partial=False, + parts=None, + calls=None, + responses=None, +): + return SimpleNamespace( + partial=partial, + content=SimpleNamespace(parts=parts or []), + get_function_calls=lambda: calls or [], + get_function_responses=lambda: responses or [], + ) + + +class FakeDetailedStreamingRunner(FakeStreamingRunner): + async def run_async(self, user_id, session_id, new_message, run_config=None): + self.run_async_calls.append( + { + "user_id": user_id, + "session_id": session_id, + "new_message": new_message, + "run_config": run_config, + } + ) + call = SimpleNamespace( + id="call-1", + name="bash", + args={"command": "printf ok", "api_key": "secret-value"}, + ) + response = SimpleNamespace( + id="call-1", + name="bash", + response={"output": "ok", "access_token": "secret-token"}, + ) + thinking = "The\noutput\nshows one item: skills.\nI'll report it." + yield build_stream_event( + partial=True, + parts=[SimpleNamespace(text="The\noutput\n", thought=True)], + ) + yield build_stream_event( + partial=True, + parts=[ + SimpleNamespace( + text="shows one item: skills.\nI'll report it.", thought=True + ) + ], + ) + # Providers may repeat the accumulated thinking in a completed event. + yield build_stream_event( + parts=[SimpleNamespace(text=thinking, thought=True)], + ) + yield build_stream_event(calls=[call]) + yield build_stream_event(calls=[call]) + yield build_stream_event(responses=[response]) + yield build_stream_event( + partial=True, + parts=[SimpleNamespace(text="完成", thought=False)], + ) + # Some providers insert an empty reasoning event immediately before + # repeating the completed answer. + yield build_stream_event( + partial=True, + parts=[SimpleNamespace(text="\n", thought=True)], + ) + yield build_stream_event( + parts=[SimpleNamespace(text="完成", thought=False)], + ) + yield build_stream_event( + partial=True, + parts=[SimpleNamespace(text=".", thought=True)], + ) + + +class FakeSeparateToolCardRunner(FakeStreamingRunner): + async def run_async(self, user_id, session_id, new_message, run_config=None): + calls = [ + SimpleNamespace(id="call-1", name="bash", args={"command": "one"}), + SimpleNamespace(id="call-2", name="python", args={"code": "two"}), + ] + responses = [ + SimpleNamespace( + id="call-2", + name="python", + response={"status": "failed", "error": "bad"}, + ), + SimpleNamespace( + id="call-1", + name="bash", + response={"output": "ok"}, + ), + ] + yield build_stream_event(calls=calls) + yield build_stream_event(calls=calls) + yield build_stream_event(responses=responses) + yield build_stream_event(responses=responses) + yield build_stream_event( + partial=True, + parts=[SimpleNamespace(text="done", thought=False)], + ) + + +class FakeToolOnlyRunner(FakeStreamingRunner): + async def run_async(self, user_id, session_id, new_message, run_config=None): + call = SimpleNamespace(id="call-1", name="bash", args={"command": "one"}) + response = SimpleNamespace( + id="call-1", + name="bash", + response={"output": "ok"}, + ) + yield build_stream_event(calls=[call]) + yield build_stream_event(responses=[response]) + + def build_message(**overrides): message = SimpleNamespace( id="om_001", @@ -219,6 +335,133 @@ def build_message(**overrides): return message +def test_tool_payload_is_redacted_and_truncated(): + rendered = _format_tool_payload( + { + "password": "do-not-show", + "nested": {"authorization": "Bearer secret"}, + "output": "x" * 200, + }, + max_length=100, + ) + + assert len(rendered) == 100 + assert rendered.endswith("…") + assert "do-not-show" not in rendered + assert "Bearer secret" not in rendered + assert '"password": "***"' in rendered + + +def test_tool_payload_fallback_preserves_redaction_and_escapes_fences(): + rendered = _format_tool_payload( + {("unsupported", "json-key"): "```", "password": "do-not-show"}, + max_length=4000, + ) + + assert "do-not-show" not in rendered + assert "***" in rendered + assert "```" not in rendered + + +@pytest.mark.anyio +async def test_shutdown_closes_channel_then_drains_accepted_message(): + entered = asyncio.Event() + finish = asyncio.Event() + closed = asyncio.Event() + + class DrainingChannel(FakeChannel): + async def disconnect(self): + closed.set() + + async def handler(context): + entered.set() + await finish.wait() + return "finished" + + channel = DrainingChannel() + extension = FeishuChannelExtension( + runner=FakeRunner(), channel=channel, message_handler=handler + ) + extension._loop = asyncio.get_running_loop() + await extension._on_message(build_message()) + await entered.wait() + shutdown = asyncio.create_task(extension.shutdown(drain_timeout=1)) + await closed.wait() + + assert extension.is_draining + assert not shutdown.done() + assert len(extension._inflight) == 1 + + finish.set() + await shutdown + assert channel.sent_messages[0][1] == {"text": "finished"} + assert not extension._inflight + + +@pytest.mark.anyio +async def test_separate_tool_card_is_settled_when_runner_fails(): + class FailingRunner(FakeStreamingRunner): + async def run_async(self, **kwargs): + yield build_stream_event( + calls=[SimpleNamespace(id="pending", name="bash", args={})] + ) + raise RuntimeError("runner failed") + + channel = FakeStreamChannel() + extension = FeishuChannelExtension( + runner=FailingRunner(), + channel=channel, + streaming=True, + show_tool_calls=True, + separate_tool_call_cards=True, + ) + + with pytest.raises(RuntimeError, match="runner failed"): + await asyncio.wait_for(extension._on_message(build_message()), timeout=10) + + assert len(channel.stream_calls) == 1 + assert "pending" not in "".join(channel.stream_calls[0][1]) + + +@pytest.mark.anyio +async def test_separate_card_failure_during_cleanup_is_propagated(): + class MissingResultRunner(FakeStreamingRunner): + async def run_async(self, **kwargs): + yield build_stream_event( + calls=[SimpleNamespace(id="pending", name="bash", args={})] + ) + + class FailingCardChannel(FakeStreamChannel): + async def stream(self, chat_id, spec, options=None): + controller = FakeStreamController() + await spec["markdown"](controller) + if controller.chunks: + raise RuntimeError("card failed") + + extension = FeishuChannelExtension( + runner=MissingResultRunner(), + channel=FailingCardChannel(), + streaming=True, + show_tool_calls=True, + separate_tool_call_cards=True, + ) + + with pytest.raises(RuntimeError, match="card failed"): + await extension._on_message(build_message()) + + +def test_separate_cards_and_topic_can_be_enabled_from_environment(monkeypatch): + monkeypatch.setenv("TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS", "true") + monkeypatch.setenv("TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD", "true") + monkeypatch.setenv("TOOL_FEISHU_CHANNEL_CREATE_TOPIC", "true") + + extension = FeishuChannelExtension(runner=FakeRunner(), channel=FakeChannel()) + + assert extension.separate_tool_call_cards is True + assert extension.separate_thinking_card is True + assert extension.create_topic is True + + @pytest.mark.anyio async def test_extension_uses_union_id_and_thread_id(): runner = FakeRunner() @@ -244,7 +487,11 @@ async def test_extension_uses_union_id_and_thread_id(): } ] assert channel.sent_messages == [ - ("oc_chat", {"text": "echo:你好"}, {"reply_to": "om_001"}) + ( + "oc_chat", + {"text": "echo:你好"}, + {"reply_to": "om_001", "reply_in_thread": True}, + ) ] @@ -301,6 +548,137 @@ async def test_extension_streaming_uses_markdown_producer_controller(): assert channel.stream_calls == [("oc_chat", ["hel", "lo"], {"reply_to": "om_001"})] +@pytest.mark.anyio +async def test_extension_streams_thinking_tool_calls_results_and_answer(): + runner = FakeDetailedStreamingRunner() + channel = FakeStreamChannel() + extension = FeishuChannelExtension( + runner=runner, + channel=channel, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, + ) + + await extension._on_message(build_message()) + + assert len(channel.stream_calls) == 1 + rendered = "".join(channel.stream_calls[0][1]) + assert rendered.count("💭 **Thinking**") == 1 + assert rendered.count("🔧 **调用工具 `bash`**") == 1 + assert rendered.count("✅ **工具 `bash` 返回**") == 1 + assert rendered.count("💬 **回答**") == 1 + assert "The output shows one item: skills. I'll report it." in rendered + assert "The\noutput" not in rendered + assert rendered.count("完成") == 1 + assert rendered.endswith("完成") + assert '"api_key": "***"' in rendered + assert '"access_token": "***"' in rendered + assert "secret-value" not in rendered + assert "secret-token" not in rendered + assert rendered.index("Thinking") < rendered.index("调用工具") + assert rendered.index("调用工具") < rendered.index("工具 `bash` 返回") + assert rendered.index("工具 `bash` 返回") < rendered.index("回答") + + +@pytest.mark.anyio +async def test_extension_streams_each_tool_call_and_result_in_separate_card(): + runner = FakeSeparateToolCardRunner() + channel = FakeStreamChannel() + extension = FeishuChannelExtension( + runner=runner, + channel=channel, + streaming=True, + show_tool_calls=True, + show_tool_results=True, + separate_tool_call_cards=True, + ) + + await extension._on_message(build_message()) + + assert len(channel.stream_calls) == 3 + rendered_cards = ["".join(call[1]) for call in channel.stream_calls] + bash_card = next(card for card in rendered_cards if "`bash`" in card) + python_card = next(card for card in rendered_cards if "`python`" in card) + answer_card = next(card for card in rendered_cards if card == "done") + + assert bash_card.count("🔧 **调用工具 `bash`**") == 1 + assert bash_card.count("✅ **工具 `bash` 返回**") == 1 + assert '"command": "one"' in bash_card + assert '"output": "ok"' in bash_card + assert "python" not in bash_card + + assert python_card.count("🔧 **调用工具 `python`**") == 1 + assert python_card.count("❌ **工具 `python` 返回**") == 1 + assert '"code": "two"' in python_card + assert '"error": "bad"' in python_card + assert "bash" not in python_card + assert answer_card == "done" + + +@pytest.mark.anyio +async def test_extension_streams_thinking_tools_and_answer_as_topic_cards(): + runner = FakeDetailedStreamingRunner() + channel = FakeStreamChannel() + extension = FeishuChannelExtension( + runner=runner, + channel=channel, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, + separate_tool_call_cards=True, + separate_thinking_card=True, + create_topic=True, + ) + + await extension._on_message(build_message()) + + assert len(channel.stream_calls) == 3 + rendered_cards = ["".join(call[1]) for call in channel.stream_calls] + thinking_card = next(card for card in rendered_cards if "Thinking" in card) + tool_card = next(card for card in rendered_cards if "`bash`" in card) + answer_card = next(card for card in rendered_cards if card == "完成") + + assert "The output shows one item: skills. I'll report it." in thinking_card + assert "调用工具" not in thinking_card + assert "回答" not in thinking_card + assert tool_card.count("🔧 **调用工具 `bash`**") == 1 + assert tool_card.count("✅ **工具 `bash` 返回**") == 1 + assert "Thinking" not in tool_card + assert answer_card == "完成" + assert all( + call[2] == {"reply_to": "om_001", "reply_in_thread": True} + for call in channel.stream_calls + ) + + +@pytest.mark.anyio +async def test_extension_does_not_create_cards_for_missing_stages(): + channel = FakeStreamChannel() + extension = FeishuChannelExtension( + runner=FakeToolOnlyRunner(), + channel=channel, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, + separate_tool_call_cards=True, + separate_thinking_card=True, + create_topic=True, + ) + + await extension._on_message(build_message()) + + assert len(channel.stream_calls) == 1 + rendered = "".join(channel.stream_calls[0][1]) + assert "🔧 **调用工具 `bash`**" in rendered + assert "✅ **工具 `bash` 返回**" in rendered + assert "Thinking" not in rendered + assert "回答" not in rendered + + @pytest.mark.anyio async def test_extension_runs_sync_channel_connect_in_worker_thread(): runner = FakeRunner() diff --git a/tests/test_ve_identity_api_key.py b/tests/test_ve_identity_api_key.py new file mode 100644 index 000000000..029a2843b --- /dev/null +++ b/tests/test_ve_identity_api_key.py @@ -0,0 +1,42 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from veadk.integrations.ve_identity.identity_client import IdentityClient + + +@pytest.mark.parametrize("pool_name", [None, "customer-pool"]) +def test_api_key_request_preserves_optional_pool_scope(pool_name): + client = IdentityClient( + access_key="test-ak", secret_key="test-sk", enable_vefaas_iam_fallback=False + ) + api = Mock() + api.api_client.configuration = SimpleNamespace() + api.get_resource_api_key.return_value = SimpleNamespace(api_key="test-key") + client._api_client = api + arguments = { + "provider_name": "model-provider", + "agent_identity_token": "test-token", + } + if pool_name is not None: + arguments["pool_name"] = pool_name + assert client.get_api_key(**arguments) == "test-key" + request = api.get_resource_api_key.call_args.args[0] + assert request.pool_name == pool_name + assert request.provider_name == "model-provider" + assert request.identity_token == "test-token" diff --git a/uv.lock b/uv.lock index 757e83711..384e6cc2b 100644 --- a/uv.lock +++ b/uv.lock @@ -5896,7 +5896,7 @@ requires-dist = [ { name = "agentkit-sdk-python", specifier = ">=0.8.5" }, { name = "agentkit-sdk-python", marker = "extra == 'harness-sidecar'", specifier = ">=0.8.1,<0.9.0" }, { name = "aiomysql", specifier = "==0.3.2" }, - { name = "anthropic", marker = "extra == 'sandbox'", specifier = ">=0.40.0" }, + { name = "anthropic", marker = "extra == 'sandbox'", specifier = ">=1.3.0" }, { name = "async-timeout", marker = "python_full_version < '3.11'", specifier = ">=5,<6" }, { name = "asyncpg", specifier = ">=0.29.0" }, { name = "cookiecutter", specifier = "==2.6.0" }, diff --git a/veadk/consts.py b/veadk/consts.py index a20c7f0d0..108dc7ee9 100644 --- a/veadk/consts.py +++ b/veadk/consts.py @@ -13,7 +13,6 @@ # limitations under the License. import os -import time from veadk.version import VERSION @@ -38,7 +37,6 @@ # "thinking": { # "type": "disabled" # }, - "expire_at": int(time.time()) + 3600, # expire after 1 hour }, } diff --git a/veadk/extensions/feishu_channel.py b/veadk/extensions/feishu_channel.py index a108a12cc..1dbf0c179 100644 --- a/veadk/extensions/feishu_channel.py +++ b/veadk/extensions/feishu_channel.py @@ -54,6 +54,69 @@ def _concat(prev: str, chunk: str) -> str: SessionIdFactory = Callable[[Any], str] UserIdFactory = Callable[[Any], str] +_SENSITIVE_FIELD_MARKERS = ( + "api_key", + "apikey", + "access_key", + "authorization", + "cookie", + "credential", + "password", + "secret", + "token", +) + + +def _env_enabled(name: str) -> bool: + return str(os.getenv(name, "")).strip().lower() in {"1", "true", "yes", "on"} + + +def _redact_sensitive(value: Any) -> Any: + """Return a display-safe copy of tool input/output data.""" + if hasattr(value, "model_dump"): + value = value.model_dump(exclude_none=True) + if isinstance(value, dict): + redacted: dict[Any, Any] = {} + for key, item in value.items(): + normalized_key = str(key).lower().replace("-", "_") + if any(marker in normalized_key for marker in _SENSITIVE_FIELD_MARKERS): + redacted[key] = "***" + else: + redacted[key] = _redact_sensitive(item) + return redacted + if isinstance(value, (list, tuple)): + return [_redact_sensitive(item) for item in value] + return value + + +def _format_tool_payload(value: Any, max_length: int) -> str: + """Serialize, redact and bound a tool payload for a Feishu Markdown card.""" + redacted = _redact_sensitive(value) + try: + rendered = json.dumps( + redacted, + ensure_ascii=False, + indent=2, + default=str, + ) + except (TypeError, ValueError): + rendered = str(redacted) + # Do not let a tool-provided fence terminate our Markdown code block. + rendered = rendered.replace("```", "` ` `") + if len(rendered) > max_length: + rendered = f"{rendered[: max(0, max_length - 1)]}…" + return rendered + + +def _tool_response_failed(payload: Any) -> bool: + if not isinstance(payload, dict): + return False + status = str(payload.get("status", "")).lower() + if status in {"error", "failed", "failure"} or payload.get("error"): + return True + exit_code = payload.get("exit_code") + return exit_code not in (None, 0, "0") + def _coalesce(*values: Any) -> str: for value in values: @@ -318,6 +381,13 @@ def __init__( channel_kwargs: dict[str, Any] | None = None, streaming: bool = False, reactions: bool = False, + show_thinking: bool = False, + show_tool_calls: bool = False, + show_tool_results: bool = False, + separate_tool_call_cards: bool = False, + separate_thinking_card: bool = False, + create_topic: bool = False, + tool_detail_max_length: int = 4000, include_parent_message: bool = True, include_thread_history: bool = True, thread_history_limit: int = 20, @@ -337,6 +407,25 @@ def __init__( streaming or str(os.getenv("TOOL_FEISHU_CHANNEL_STREAMING", "")).lower() == "true" ) + self.show_thinking = show_thinking or _env_enabled( + "TOOL_FEISHU_CHANNEL_SHOW_THINKING" + ) + self.show_tool_calls = show_tool_calls or _env_enabled( + "TOOL_FEISHU_CHANNEL_SHOW_TOOL_CALLS" + ) + self.show_tool_results = show_tool_results or _env_enabled( + "TOOL_FEISHU_CHANNEL_SHOW_TOOL_RESULTS" + ) + self.separate_tool_call_cards = separate_tool_call_cards or _env_enabled( + "TOOL_FEISHU_CHANNEL_SEPARATE_TOOL_CALL_CARDS" + ) + self.separate_thinking_card = separate_thinking_card or _env_enabled( + "TOOL_FEISHU_CHANNEL_SEPARATE_THINKING_CARD" + ) + self.create_topic = create_topic or _env_enabled( + "TOOL_FEISHU_CHANNEL_CREATE_TOPIC" + ) + self.tool_detail_max_length = max(100, int(tool_detail_max_length)) self.include_parent_message = include_parent_message self.include_thread_history = include_thread_history self.thread_history_limit = max(1, int(thread_history_limit)) @@ -644,6 +733,8 @@ async def _process_message(self, message: Any) -> None: send_options = {} if self.reply_in_thread and context.message_id: send_options["reply_to"] = context.message_id + if self.create_topic or context.thread_id: + send_options["reply_in_thread"] = True if self.message_handler is not None: response_text = await self._maybe_await(self.message_handler(context)) @@ -680,22 +771,310 @@ async def _process_message(self, message: Any) -> None: ) async def stream_to_feishu(stream): - for converted_message in converted_messages: - async for event in self.runner.run_async( - user_id=context.user_id, - session_id=context.session_id, - new_message=converted_message, - run_config=run_config, - ): - if not getattr(event, "partial", False): - continue - if not (event.content and event.content.parts): - continue - for part in event.content.parts: - if getattr(part, "thought", False): - continue - if part.text: - await stream.append(part.text) + seen_tool_call_ids: set[str] = set() + seen_tool_response_ids: set[str] = set() + tool_card_streams: dict[ + str, + tuple[asyncio.Future[Any], asyncio.Task[Any]], + ] = {} + thinking_queue: asyncio.Queue[str | None] | None = None + thinking_task: asyncio.Task[Any] | None = None + active_section = "" + current_thinking_text = "" + last_emitted_thinking_text = "" + current_answer_text = "" + + async def flush_thinking() -> None: + nonlocal current_thinking_text, last_emitted_thinking_text + nonlocal thinking_queue, thinking_task + if not current_thinking_text: + return + # A Markdown stream append may be rendered as a separate + # block by Feishu. Buffer reasoning deltas and collapse + # provider whitespace so token-sized chunks do not become + # one-word-per-line output. + rendered = " ".join(current_thinking_text.split()) + if rendered: + if self.separate_thinking_card: + if thinking_queue is None: + thinking_queue = asyncio.Queue() + thinking_queue.put_nowait(rendered) + started = asyncio.Event() + + async def stream_thinking(thinking_stream) -> None: + first_chunk = await thinking_queue.get() + if first_chunk is None: + started.set() + return + try: + await thinking_stream.append( + f"💭 **Thinking**\n\n{first_chunk}" + ) + finally: + started.set() + while True: + chunk = await thinking_queue.get() + if chunk is None: + return + await thinking_stream.append(f"\n\n{chunk}") + + thinking_task = asyncio.create_task( + self._maybe_await( + self.channel.stream( + context.chat_id, + {"markdown": stream_thinking}, + send_options, + ) + ) + ) + thinking_task.add_done_callback(lambda _: started.set()) + await started.wait() + if thinking_task.done(): + await thinking_task + else: + thinking_queue.put_nowait(rendered) + else: + await begin_section("thinking", "💭 **Thinking**\n\n") + await stream.append(rendered) + last_emitted_thinking_text = rendered + current_thinking_text = "" + + async def begin_section(section: str, heading: str = "") -> None: + nonlocal active_section + if active_section == section: + return + separator = "\n\n" if active_section else "" + rendered_heading = f"{separator}{heading}" + if rendered_heading: + await stream.append(rendered_heading) + active_section = section + + async def append_tool_call(call: Any) -> None: + call_id = str(getattr(call, "id", "") or "") + if call_id and call_id in seen_tool_call_ids: + return + if call_id: + seen_tool_call_ids.add(call_id) + name = str(getattr(call, "name", "") or "unknown") + args = getattr(call, "args", None) + await flush_thinking() + if self.separate_tool_call_cards: + response_future = asyncio.get_running_loop().create_future() + started = asyncio.Event() + + async def stream_tool_call(tool_stream) -> None: + try: + await tool_stream.append( + f"🔧 **调用工具 `{name}`**\n\n" + "```json\n" + f"{_format_tool_payload(args, self.tool_detail_max_length)}\n" + "```" + ) + finally: + started.set() + + response = await response_future + if response is None: + return + response_name = str(getattr(response, "name", "") or name) + payload = getattr(response, "response", None) + icon = "❌" if _tool_response_failed(payload) else "✅" + await tool_stream.append( + f"\n\n{icon} **工具 `{response_name}` 返回**\n\n" + "```json\n" + f"{_format_tool_payload(payload, self.tool_detail_max_length)}\n" + "```" + ) + + task = asyncio.create_task( + self._maybe_await( + self.channel.stream( + context.chat_id, + {"markdown": stream_tool_call}, + send_options, + ) + ) + ) + task.add_done_callback(lambda _: started.set()) + tool_card_streams[call_id or str(id(call))] = ( + response_future, + task, + ) + await started.wait() + if task.done(): + await task + return + await begin_section( + f"tool-call:{call_id or id(call)}", + f"🔧 **调用工具 `{name}`**\n\n", + ) + await stream.append( + "```json\n" + f"{_format_tool_payload(args, self.tool_detail_max_length)}\n" + "```" + ) + + async def append_tool_response(response: Any) -> None: + response_id = str(getattr(response, "id", "") or "") + if response_id and response_id in seen_tool_response_ids: + return + if response_id: + seen_tool_response_ids.add(response_id) + name = str(getattr(response, "name", "") or "unknown") + payload = getattr(response, "response", None) + icon = "❌" if _tool_response_failed(payload) else "✅" + await flush_thinking() + if self.separate_tool_call_cards: + state = tool_card_streams.get(response_id) + if state is not None: + response_future, task = state + if not response_future.done(): + response_future.set_result(response) + await task + return + + async def stream_tool_response(tool_stream) -> None: + await tool_stream.append( + f"{icon} **工具 `{name}` 返回**\n\n" + "```json\n" + f"{_format_tool_payload(payload, self.tool_detail_max_length)}\n" + "```" + ) + + await self._maybe_await( + self.channel.stream( + context.chat_id, + {"markdown": stream_tool_response}, + send_options, + ) + ) + return + await begin_section( + f"tool-result:{response_id or id(response)}", + f"{icon} **工具 `{name}` 返回**\n\n", + ) + await stream.append( + "```json\n" + f"{_format_tool_payload(payload, self.tool_detail_max_length)}\n" + "```" + ) + + stream_failed = False + try: + for converted_message in converted_messages: + async for event in self.runner.run_async( + user_id=context.user_id, + session_id=context.session_id, + new_message=converted_message, + run_config=run_config, + ): + get_calls = getattr(event, "get_function_calls", None) + if self.show_tool_calls and callable(get_calls): + for call in get_calls() or []: + await append_tool_call(call) + + get_responses = getattr( + event, "get_function_responses", None + ) + if self.show_tool_results and callable(get_responses): + for response in get_responses() or []: + await append_tool_response(response) + + content = getattr(event, "content", None) + parts = getattr(content, "parts", None) or [] + for part in parts: + part_text = getattr(part, "text", None) + if not part_text: + continue + if getattr(part, "thought", False): + if not self.show_thinking: + continue + # Some providers emit a trailing reasoning + # fragment (often punctuation) after the final + # answer. Do not open a new Thinking section + # unless a tool event has moved the stream into + # another reasoning round. + if ( + active_section == "answer" + and current_answer_text + ): + continue + is_partial = getattr(event, "partial", False) + if is_partial: + current_thinking_text += part_text + continue + # SSE providers commonly emit the accumulated + # thought once more in a non-partial event. + normalized_part = " ".join(part_text.split()) + normalized_current = " ".join( + current_thinking_text.split() + ) + if not normalized_part or normalized_part in { + normalized_current, + last_emitted_thinking_text, + }: + continue + current_thinking_text = part_text + continue + + # Preserve the existing delta-first streaming + # behavior while supporting providers that only + # produce a completed answer event. + await flush_thinking() + if ( + not getattr(event, "partial", False) + and part_text == current_answer_text + ): + continue + answer_heading = ( + "💬 **回答**\n\n" + if ( + self.show_thinking + and not self.separate_thinking_card + or ( + not self.separate_tool_call_cards + and ( + self.show_tool_calls + or self.show_tool_results + ) + ) + ) + else "" + ) + if ( + getattr(event, "partial", False) + and active_section != "answer" + ): + current_answer_text = "" + await begin_section("answer", answer_heading) + await stream.append(part_text) + if getattr(event, "partial", False): + current_answer_text += part_text + else: + current_answer_text = part_text + + await flush_thinking() + except BaseException: + stream_failed = True + raise + finally: + pending_tasks: list[asyncio.Task[Any]] = [] + if thinking_queue is not None: + thinking_queue.put_nowait(None) + if thinking_task is not None: + pending_tasks.append(thinking_task) + for response_future, task in tool_card_streams.values(): + if not response_future.done(): + response_future.set_result(None) + pending_tasks.append(task) + if pending_tasks: + results = await asyncio.gather( + *pending_tasks, return_exceptions=True + ) + if not stream_failed: + for result in results: + if isinstance(result, BaseException): + raise result await self._maybe_await( self.channel.stream( diff --git a/veadk/integrations/mpa/session_client.py b/veadk/integrations/mpa/session_client.py new file mode 100644 index 000000000..8e32e7131 --- /dev/null +++ b/veadk/integrations/mpa/session_client.py @@ -0,0 +1,508 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Self-Hosted Sandbox client using official Anthropic Python SDK (client.beta.sessions & events). + +Endpoints managed via Anthropic SDK: +1. client.beta.sessions.create: + Creates a session and enqueues work into the WorkQueue for the worker. +2. client.beta.sessions.events.send: + Publishes ``agent.tool_use`` and ``session.status_idle`` events. +3. client.beta.sessions.events.list: + Retrieves session events to wait for ``user.tool_result`` from workers. +""" + +import json +import logging +import os +import re +import time +import uuid +from typing import Any +from urllib.parse import urlsplit + +import anthropic + +logger = logging.getLogger("veadk.sandbox") + + +def _runtime_id_from_environment() -> str: + """Resolve the AgentKit Runtime identity without guessing platform IDs.""" + for name in ("AGENTKIT_RUNTIME_ID", "MA_RUNTIME_ID"): + if runtime_id := os.getenv(name, "").strip(): + return runtime_id + + client_request_id = os.getenv("MODEL_AGENT_CLIENT_REQ_ID", "").strip() + match = re.fullmatch(r"agentkit/([^/\s]+)", client_request_id) + return match.group(1) if match else "" + + +class SelfHostSandboxClient: + """Client for Managed Agents Self-Hosted Sandbox backed by official Anthropic SDK.""" + + def __init__( + self, + base_url: str | None = None, + environment_id: str | None = None, + agent_id: str | None = None, + session_id: str | None = None, + bearer_token: str | None = None, + account_id: str | None = None, + remote_bash_tool_name: str | None = None, + timeout: int = 300, + trust_env: bool | None = None, + runtime_type: str | None = None, + ): + self.base_url = ( + base_url + or os.getenv("ANTHROPIC_BASE_URL") + or os.getenv("SANDBOX_BASE_URL", "") + ).rstrip("/") + self.environment_id = ( + environment_id + or os.getenv("ANTHROPIC_ENVIRONMENT_ID") + or os.getenv("SANDBOX_ENVIRONMENT_ID", "") + ) + self.agent_id = agent_id or os.getenv("SANDBOX_AGENT_ID", "") + self.session_id = ( + session_id + or os.getenv("ANTHROPIC_SESSION_ID") + or os.getenv("SANDBOX_SESSION_ID") + ) + + raw_token = ( + bearer_token + or os.getenv("ANTHROPIC_ENVIRONMENT_KEY") + or os.getenv("ANTHROPIC_API_KEY") + or os.getenv("SANDBOX_BEARER_TOKEN") + or os.getenv("SANDBOX_API_KEY", "") + ) + runtime_id = _runtime_id_from_environment() + self._identity_credentials = None + if not raw_token and runtime_id: + from veadk.runtime.managed_agents.identity import ( + get_runtime_credential_provider, + ) + + self._identity_credentials = get_runtime_credential_provider( + runtime_id=runtime_id, + region=os.getenv("VOLCENGINE_REGION") + or os.getenv("REGION") + or "cn-beijing", + ) + raw_token = self._identity_credentials.get_sync() + if raw_token.startswith("Bearer "): + raw_token = raw_token[7:].strip() + self.bearer_token = raw_token + self.account_id = account_id or os.getenv("X_TOP_ACCOUNT_ID", "") + # RuntimeType routes this worker's AgentLoop poll to a specific + # execution pool. Empty means legacy behaviour: the server treats a + # typeless poll as the reserved ``default`` pool. + self.runtime_type = ( + runtime_type or os.getenv("MA_AGENT_LOOP_RUNTIME_TYPE") or "" + ).strip() + self.remote_bash_tool_name = remote_bash_tool_name or os.getenv( + "SANDBOX_BASH_TOOL_NAME", "bash" + ) + self.timeout = int(os.getenv("SANDBOX_TIMEOUT_SECONDS", str(timeout))) + self.trust_env = trust_env + + if not self.base_url: + raise ValueError( + "ANTHROPIC_BASE_URL (or SANDBOX_BASE_URL) must be configured in .env or arguments." + ) + if not self.bearer_token: + raise ValueError( + "ANTHROPIC_ENVIRONMENT_KEY (or SANDBOX_BEARER_TOKEN) must be configured in .env or arguments." + ) + + # The Runtime gateway needs the Managed Agents beta flag and, when + # supplied, the Volcengine top-account routing header. + default_headers = { + "anthropic-beta": "managed-agents-2026-04-01", + } + if self.account_id: + default_headers["X-Top-Account-Id"] = str(self.account_id) + # Inject the AgentLoop RuntimeType routing label alongside the existing + # auth/beta headers so every SDK helper poll carries it. ma-infra pairs + # this with the optional ``runtime_type`` poll query and rejects a + # conflicting pair with 400. + if self.runtime_type: + default_headers["X-Runtime-Type"] = self.runtime_type + + self._default_headers = default_headers + http_client = ( + anthropic.DefaultHttpxClient( + trust_env=False if trust_env is None else trust_env, + event_hooks={ + "request": [self._identity_request], + "response": [self._identity_response], + }, + ) + if self._identity_credentials is not None + else anthropic.DefaultHttpxClient(trust_env=trust_env) + if trust_env is not None + else None + ) + self.client = anthropic.Anthropic( + base_url=self.base_url, + auth_token=self.bearer_token, + default_headers=default_headers, + timeout=float(self.timeout), + http_client=http_client, + ) + + def _set_gateway_header(self, request: Any, key: str) -> None: + # SDK Work sub-clients deliberately strip parent auth headers. Apply + # gateway auth at transport send time so their scoped Bearer survives. + expected = urlsplit(self.base_url) + actual = urlsplit(str(request.url)) + if (actual.scheme, actual.hostname, actual.port) != ( + expected.scheme, + expected.hostname, + expected.port, + ): + from veadk.runtime.managed_agents.identity import RuntimeIdentityError + + raise RuntimeIdentityError("authentication_failed", False) + request.headers["x-api-key"] = key + if request.headers.get("Authorization") == "Bearer " + self.bearer_token: + del request.headers["Authorization"] + + def _identity_request(self, request: Any) -> None: + assert self._identity_credentials is not None + self._set_gateway_header(request, self._identity_credentials.get_sync()) + + async def _identity_async_request(self, request: Any) -> None: + assert self._identity_credentials is not None + self._set_gateway_header(request, await self._identity_credentials.get()) + + def _identity_response(self, response: Any) -> None: + if response.status_code in {401, 403}: + assert self._identity_credentials is not None + self._identity_credentials.invalidate() + + async def _identity_async_response(self, response: Any) -> None: + self._identity_response(response) + + def create_async_client(self) -> anthropic.AsyncAnthropic: + """Create an async SDK client for the Agent Loop event stream.""" + http_client = ( + anthropic.DefaultAsyncHttpxClient( + trust_env=False if self.trust_env is None else self.trust_env, + event_hooks={ + "request": [self._identity_async_request], + "response": [self._identity_async_response], + }, + ) + if self._identity_credentials is not None + else anthropic.DefaultAsyncHttpxClient(trust_env=self.trust_env) + if self.trust_env is not None + else None + ) + return anthropic.AsyncAnthropic( + base_url=self.base_url, + auth_token=self.bearer_token, + default_headers=self._default_headers, + timeout=float(self.timeout), + http_client=http_client, + ) + + def create_session( + self, + title: str = "VeADK Self-Hosted Sandbox Session", + ) -> dict[str, Any]: + """POST /v1/sessions via Anthropic SDK (client.beta.sessions.create). + + This only creates the session. A real ``user.message`` event causes the + server control plane to enqueue the corresponding WorkItem. + """ + if not self.agent_id or not self.environment_id: + raise ValueError( + "Session creation requires configured agent_id and environment_id" + ) + sess = self.client.beta.sessions.create( + agent=self.agent_id, + environment_id=self.environment_id, + title=title, + # Fix the session's AgentLoop RuntimeType at creation when this + # client is configured with one; ma-infra snapshots it and never + # lets successor/recovery Work drift to another type. + extra_body=( + {"runtime_type": self.runtime_type} if self.runtime_type else None + ), + ) + self.session_id = sess.id + logger.info( + "Session %s created on remote Runtime via Anthropic SDK.", self.session_id + ) + if hasattr(sess, "model_dump"): + return sess.model_dump(warnings=False) + return { + "id": sess.id, + "environment_id": self.environment_id, + "agent": self.agent_id, + } + + def post_events(self, events: list[dict[str, Any]]) -> dict[str, Any]: + """POST /v1/sessions/{session_id}/events via Anthropic SDK (client.beta.sessions.events.send). + + Dispatches tool calls and tool results into the Session event stream. + """ + if not self.session_id: + self.create_session() + + from veadk.runtime.managed_agents.events import ( + trace_events, + trace_send_response, + ) + + trace_events("sent", "send", self.session_id, events) + response = self.client.beta.sessions.events.send( + session_id=self.session_id, + events=events, + ) + trace_send_response(self.session_id, response) + if hasattr(response, "model_dump"): + return response.model_dump(warnings=False) + return {"status": "ok"} + + def send_tool_result( + self, tool_use_id: str, content: str, is_error: bool = False + ) -> dict[str, Any]: + """Send tool execution result to POST /v1/sessions/{session_id}/events.""" + event = { + "type": "user.tool_result", + "tool_use_id": tool_use_id, + "content": [{"type": "text", "text": content}], + "is_error": is_error, + } + return self.post_events([event]) + + def list_events( + self, + *, + limit: int = 100, + order: str = "asc", + ) -> list[dict[str, Any]]: + """Read and normalize events from GET /v1/sessions/{id}/events via Anthropic SDK.""" + if not self.session_id: + raise RuntimeError("Cannot read events before a Session is created.") + + page = self.client.beta.sessions.events.list( + session_id=self.session_id, + limit=limit, + order=order, + ) + events: list[dict[str, Any]] = [] + from veadk.runtime.managed_agents.events import trace_event + + for item in page.data: + trace_event("received", "list", self.session_id, item) + if hasattr(item, "model_dump"): + d = item.model_dump(warnings=False) + elif isinstance(item, dict): + d = dict(item) + else: + d = { + "id": getattr(item, "id", None), + "type": getattr(item, "type", None), + "name": getattr(item, "name", None), + "input": getattr(item, "input", None), + "content": getattr(item, "content", None), + "is_error": getattr(item, "is_error", False), + "tool_use_id": getattr(item, "tool_use_id", None), + } + events.append(d) + return events + + def post_status_idle(self, stop_reason: str = "end_turn") -> None: + """Publish a ``session.status_idle`` event so workers know the turn ended and can release their lease.""" + if not self.session_id: + return + try: + self.post_events( + [ + { + "type": "session.status_idle", + "stop_reason": { + "type": stop_reason, + }, + } + ] + ) + logger.info( + "Session %s marked status_idle (%s).", self.session_id, stop_reason + ) + except Exception as e: # noqa: BLE001 -- idle notification is best effort + logger.warning( + "Failed to post session.status_idle for %s: %s", self.session_id, e + ) + + def dispatch_tool( + self, + tool_name: str, + arguments: dict[str, Any], + *, + dispatch_id: str = "", + ) -> dict[str, Any]: + """Dispatch the registered tool name and JSON input without shell conversion.""" + if not isinstance(tool_name, str) or not tool_name.strip(): + raise ValueError("tool_name must be nonempty") + if not isinstance(arguments, dict): + raise TypeError("tool input must be an object") + if not self.session_id: + self.create_session() + payload = dict(arguments) + timeout_seconds = float(self.timeout) + if tool_name == "bash": + # Only the bash tool has an execution timeout in the native schema. + if payload.get("timeout_ms") is None: + payload["timeout_ms"] = int(timeout_seconds * 1000) + timeout_seconds = self._positive_timeout(payload["timeout_ms"]) / 1000 + tool_use_id = dispatch_id or f"toolu_{uuid.uuid4().hex}" + self.post_events( + [ + { + "type": "agent.tool_use", + "id": tool_use_id, + "name": tool_name, + "input": payload, + } + ] + ) + return self._wait_for_tool_result( + timeout=timeout_seconds + 60, + tool_use_id=tool_use_id, + dispatch_id=dispatch_id or tool_use_id, + ) + + def execute_command( + self, + command: str, + *, + timeout: float | None = None, + dispatch_id: str = "", + ) -> dict[str, Any]: + """Convenience entry point for a native bash invocation.""" + return self.dispatch_tool( + self.remote_bash_tool_name, + { + "command": command, + "timeout_ms": int( + self._positive_timeout( + float(self.timeout) if timeout is None else timeout + ) + * 1000 + ), + }, + dispatch_id=dispatch_id, + ) + + def _wait_for_tool_result( + self, + *, + timeout: float, + tool_use_id: str, + dispatch_id: str, + ) -> dict[str, Any]: + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + events = self.list_events(limit=100, order="asc") + for event in events: + event_type = str(event.get("type") or "") + if event_type not in {"agent.tool_result", "user.tool_result"}: + continue + result_tool_use_id = str(event.get("tool_use_id") or "") + if result_tool_use_id == tool_use_id: + result = self._normalize_tool_result( + event, + tool_use_id=tool_use_id, + dispatch_id=dispatch_id, + ) + if event_type == "user.tool_result": + self.post_events( + [ + { + "type": "agent.tool_result", + "tool_use_id": tool_use_id, + "content": event.get("content") or "", + "is_error": bool(event.get("is_error")), + } + ] + ) + return result + time.sleep(0.5) + + raise TimeoutError( + f"Timed out after {timeout:g}s waiting for tool result {tool_use_id}." + ) + + @staticmethod + def _positive_timeout(value: Any) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0: + raise ValueError("timeout must be greater than zero") + return float(value) + + def _normalize_tool_result( + self, + event: dict[str, Any], + *, + tool_use_id: str, + dispatch_id: str, + ) -> dict[str, Any]: + content = self._content_text(event.get("content")) + match = re.match(r"^exit=(-?\d+)\n([\s\S]*)$", content) + if match: + exit_code = int(match.group(1)) + stdout = match.group(2) + else: + exit_code = 1 if event.get("is_error") else 0 + stdout = content + return { + "dispatch_id": dispatch_id, + "tool_use_id": tool_use_id, + "environment_id": self.environment_id, + "session_id": self.session_id, + "status": "failed" if exit_code else "completed", + "exit_code": exit_code, + "stdout": stdout, + "stderr": "", + "content": event.get("content"), + "is_error": bool(event.get("is_error")), + } + + @staticmethod + def _event_seq(event: dict[str, Any]) -> int: + try: + return int(event.get("_seq", 0) or 0) + except (TypeError, ValueError): + return 0 + + @staticmethod + def _content_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text") or "")) + elif isinstance(block, str): + parts.append(block) + return "\n".join(part for part in parts if part) + if content is None: + return "" + return json.dumps(content, ensure_ascii=False) diff --git a/veadk/integrations/mpa/session_resources.py b/veadk/integrations/mpa/session_resources.py new file mode 100644 index 000000000..f05137a68 --- /dev/null +++ b/veadk/integrations/mpa/session_resources.py @@ -0,0 +1,305 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Materialize resources referenced by a frozen Managed Agent Session.""" + +from __future__ import annotations + +import asyncio +import atexit +import os +import re +import shutil +import tempfile +from pathlib import Path +from typing import Any + +import anthropic +import httpx2 +from anthropic.lib.tools._skills import _download_and_extract +from google.adk.tools.mcp_tool.mcp_session_manager import ( + StreamableHTTPConnectionParams, +) +from google.adk.tools.mcp_tool.mcp_toolset import McpToolset + +_ACTIVE_TEMP_DIRS: set[Path] = set() + + +def _cleanup_all_active_temp_dirs() -> None: + for path in list(_ACTIVE_TEMP_DIRS): + try: + shutil.rmtree(path, ignore_errors=True) + except Exception: + pass + _ACTIVE_TEMP_DIRS.discard(path) + + +atexit.register(_cleanup_all_active_temp_dirs) + + +class SessionSkillMaterializationError(RuntimeError): + """A Session cannot run until its frozen Skill package is available.""" + + def __init__(self, message: str, *, retryable: bool) -> None: + super().__init__(message) + self.retryable = retryable + + +class ManagedMcpToolset(McpToolset): + """Keep the owning server name on resolved tools for event attribution.""" + + def __init__( + self, *, server_name: str, identity_headers: Any = None, **kwargs: Any + ) -> None: + self.managed_agents_server_name = server_name + self.identity_headers = identity_headers + super().__init__(**kwargs) + + async def get_tools(self, readonly_context: Any = None) -> list[Any]: + if self.identity_headers is not None: + await self.identity_headers.prepare() + tools = await super().get_tools(readonly_context) + for tool in tools: + tool.managed_agents_mcp_server_name = self.managed_agents_server_name + return tools + + +def field(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, dict): + return value.get(name, default) + return getattr(value, name, default) + + +def session_workdir(root: Path, session_id: str) -> Path: + safe_id = re.sub(r"[^A-Za-z0-9_.-]", "_", session_id).strip("._") + if not safe_id or safe_id != session_id: + raise ValueError( + "Managed Session id must contain only letters, numbers, underscores, dots or hyphens" + ) + result = (root.resolve() / "sessions" / safe_id).resolve() + result.relative_to(root.resolve()) + result.mkdir(parents=True, exist_ok=True, mode=0o700) + return result + + +async def materialize_session_skills( + session: Any, + workdir: Path | None = None, + *, + client: anthropic.AsyncAnthropic, +) -> list[Path]: + """Materialize every pinned Session Skill through the Anthropic SDK.""" + references = list(field(field(session, "agent"), "skills", []) or []) + if not references: + return [] + + session_id = str(field(session, "id") or "session") + safe_prefix = re.sub(r"[^A-Za-z0-9_]", "_", session_id)[:8] or "skills" + if workdir is not None: + workdir = workdir.resolve() + workdir.mkdir(parents=True, exist_ok=True, mode=0o700) + # Keep each invocation isolated, but inside the file tools' workspace. + work_root = Path( + tempfile.mkdtemp(prefix=f"ma_session_skills_{safe_prefix}_", dir=workdir) + ).resolve() + _ACTIVE_TEMP_DIRS.add(work_root) + + try: + result = await download_session_skills( + client, workdir=work_root, session=session + ) + if len(result) != len(references): + raise RuntimeError( + f"Anthropic SDK materialized {len(result)} of {len(references)} Session Skills" + ) + for directory in result: + if not (directory / "SKILL.md").is_file(): + raise RuntimeError( + f"Anthropic SDK materialized Skill without SKILL.md: {directory}" + ) + return result + except (Exception, asyncio.CancelledError) as error: + await asyncio.to_thread(shutil.rmtree, work_root, True) + _ACTIVE_TEMP_DIRS.discard(work_root) + if isinstance(error, asyncio.CancelledError): + raise + if isinstance(error, SessionSkillMaterializationError): + raise + raise SessionSkillMaterializationError( + "failed to materialize pinned Session Skills", + retryable=_skill_failure_retryable(error), + ) from error + + +def _skill_failure_retryable(error: Exception) -> bool: + if isinstance( + error, + ( + anthropic.APIConnectionError, + httpx2.TransportError, + ConnectionError, + TimeoutError, + ), + ): + return True + status = getattr(error, "status_code", None) + if isinstance(status, int): + return status in (408, 429) or status >= 500 + return bool(getattr(error, "retryable", False)) + + +async def download_session_skills( + client: anthropic.AsyncAnthropic, *, workdir: Path, session: Any +) -> list[Path]: + """Download all pinned skills, preserving SDK failures for retry decisions. + + The SDK's session helper logs and skips failures. Use its streaming download + and safe archive extractor directly because every frozen skill is required. + """ + skills_root = (workdir / "skills").resolve() + downloaded: list[Path] = [] + for skill in field(field(session, "agent"), "skills", []) or []: + skill_id = field(skill, "skill_id") + version = await client.beta.skills.versions.retrieve( + field(skill, "version"), skill_id=skill_id + ) + dirname = os.path.basename(version.name.strip()) or skill_id + if dirname in ("", ".", ".."): + dirname = skill_id + dest = (skills_root / dirname).resolve() + if dest == skills_root or not dest.is_relative_to(skills_root): + raise ValueError("Skill name escapes the skills directory") + if dest in downloaded: + raise ValueError("Session Skills have conflicting directory names") + await _download_and_extract(client, skill_id, version.id, dest) + downloaded.append(dest) + return downloaded + + +async def cleanup_session_skills(skill_dirs: list[Path]) -> None: + """Remove owned download roots without following agent-mutated symlinks.""" + roots = { + root + for root in _ACTIVE_TEMP_DIRS + if any(directory.absolute().is_relative_to(root) for directory in skill_dirs) + } + for root in roots: + if root.is_symlink(): + await asyncio.to_thread(root.unlink, missing_ok=True) + else: + await asyncio.to_thread(shutil.rmtree, root, True) + _ACTIVE_TEMP_DIRS.discard(root) + + +def skill_instructions(skill_dirs: list[Path]) -> str: + """Build model context from materialized pinned Skill packages.""" + blocks: list[str] = [] + for directory in skill_dirs: + skill_file = directory / "SKILL.md" + if not skill_file.is_file(): + raise RuntimeError(f"materialized Skill has no SKILL.md: {directory}") + blocks.append( + f"\n" + f"{skill_file.read_text(encoding='utf-8')}\n" + ) + if not blocks: + return "" + return ( + "\n\n\n" + "The following Session-pinned skills were materialized by the Worker. " + "Follow their instructions when relevant. Supporting files are below each listed path.\n" + + "\n\n".join(blocks) + + "\n" + ) + + +def _enabled(config: Any, default: bool) -> bool: + value = field(config, "enabled") + return default if value is None else bool(value) + + +def mcp_toolsets(snapshot: Any, *, credential_keys: Any = None) -> list[McpToolset]: + """Translate frozen URL MCP definitions and tool policies into ADK toolsets.""" + servers = { + str(field(server, "name", "")): server + for server in field(snapshot, "mcp_servers", []) or [] + } + bindings = {} + for reference in credential_keys or []: + name = field(reference, "MCPServerName") + if name is None: + continue + if name not in servers or name in bindings: + raise ValueError("Invalid Session MCP credential binding") + bindings[name] = reference + result: list[McpToolset] = [] + for toolset in field(snapshot, "tools", []) or []: + if field(toolset, "type") != "mcp_toolset": + continue + server_name = str(field(toolset, "mcp_server_name", "")) + server = servers.get(server_name) + if server is None: + raise ValueError(f"MCP toolset references unknown server {server_name!r}") + if field(server, "type") != "url": + raise ValueError(f"MCP server {server_name!r} must use type=url") + url = str(field(server, "url", "")) + if not url.startswith(("http://", "https://")): + raise ValueError(f"MCP server {server_name!r} must use an HTTP(S) URL") + + default = field(toolset, "default_config", {}) or {} + default_enabled = _enabled(default, True) + configs = { + str(field(item, "name", "")): item + for item in field(toolset, "configs", []) or [] + } + + def selected( + tool: Any, + _context: Any, + *, + configs=configs, + default_enabled=default_enabled, + ) -> bool: + config = configs.get(str(getattr(tool, "name", "")), {}) + return _enabled(config, default_enabled) + + identity_headers = None + if field(server, "authorization") is not None: + raise ValueError("MCP credentials must be associated with the Session") + if reference := bindings.get(server_name): + from veadk.runtime.managed_agents.identity import MCPIdentityHeaders + + identity_headers = MCPIdentityHeaders(reference) + result.append( + ManagedMcpToolset( + server_name=server_name, + identity_headers=identity_headers, + header_provider=identity_headers, + connection_params=StreamableHTTPConnectionParams(url=url), + tool_filter=selected, + tool_name_prefix=f"mcp__{server_name}_", + ) + ) + return result + + +__all__ = [ + "SessionSkillMaterializationError", + "cleanup_session_skills", + "field", + "materialize_session_skills", + "mcp_toolsets", + "session_workdir", + "skill_instructions", +] diff --git a/veadk/integrations/ve_identity/identity_client.py b/veadk/integrations/ve_identity/identity_client.py index 56e6dd58f..79b4f623d 100644 --- a/veadk/integrations/ve_identity/identity_client.py +++ b/veadk/integrations/ve_identity/identity_client.py @@ -621,12 +621,19 @@ def get_oauth2_token_or_auth_url( ) @refresh_credentials - def get_api_key(self, *, provider_name: str, agent_identity_token: str) -> str: + def get_api_key( + self, + *, + provider_name: str, + agent_identity_token: str, + pool_name: Optional[str] = None, + ) -> str: """Retrieve an API key from the identity service. Args: provider_name: Name of the API key credential provider. agent_identity_token: Agent's workload access token for authentication. + pool_name: User-owned credential pool; omitted keeps the service default. Returns: The API key string. @@ -635,6 +642,7 @@ def get_api_key(self, *, provider_name: str, agent_identity_token: str) -> str: request = volcenginesdkid.GetResourceApiKeyRequest( provider_name=provider_name, identity_token=agent_identity_token, + pool_name=pool_name, ) response: volcenginesdkid.GetResourceApiKeyResponse = ( diff --git a/veadk/models/ark_llm.py b/veadk/models/ark_llm.py index 4c8588d56..0700543cd 100644 --- a/veadk/models/ark_llm.py +++ b/veadk/models/ark_llm.py @@ -17,7 +17,9 @@ import base64 import copy import json -import time +import logging +import os +import re from typing import Any, Dict, Union, AsyncGenerator, Tuple, List, Optional, Literal from typing_extensions import override @@ -68,10 +70,107 @@ from veadk.utils.logger import get_logger logger = get_logger(__name__) +_MODEL_TRACE_LOGGER = logging.getLogger("anthropic.managed_agent_model") +_SAFE_TRACE_VALUE = re.compile(r"^[A-Za-z0-9_.:/-]{1,128}$") _ARK_TEXT_FIELD_TYPES = {"json_object", "json_schema"} + +def _safe_trace_value(value: Any) -> str: + text = str(value or "").strip() + return text if _SAFE_TRACE_VALUE.fullmatch(text) else "" + + +def _trace_field(value: Any, name: str) -> Any: + if isinstance(value, dict): + return value.get(name) + return getattr(value, name, None) + + +def _normalized_tool_choice(value: Any) -> Any: + if value is None: + return None + if isinstance(value, str): + return _safe_trace_value(value) + result = {} + for name in ("type", "name"): + safe = _safe_trace_value(_trace_field(value, name)) + if safe: + result[name] = safe + function = _trace_field(value, "function") + safe_name = _safe_trace_value(_trace_field(function, "name")) + if safe_name: + result["name"] = safe_name + return result + + +def _model_trace(event: str, **fields: Any) -> None: + if os.getenv("MANAGED_AGENT_MODEL_TRACE", "").strip().lower() != "true": + return + payload = {"event": event} + payload.update(fields) + _MODEL_TRACE_LOGGER.info(json.dumps(payload, sort_keys=True, separators=(",", ":"))) + + +def configure_managed_agent_model_trace() -> None: + """Configure the isolated, credential-safe model trace channel.""" + if os.getenv("MANAGED_AGENT_MODEL_TRACE", "").strip().lower() != "true": + return + logging.getLogger("anthropic._base_client").setLevel(logging.WARNING) + logging.getLogger("httpx").setLevel(logging.WARNING) + _MODEL_TRACE_LOGGER.setLevel(logging.INFO) + _MODEL_TRACE_LOGGER.propagate = False + installed = not any( + getattr(existing, "_managed_agent_model_trace", False) + for existing in _MODEL_TRACE_LOGGER.handlers + ) + if installed: + handler = logging.StreamHandler() + handler.setLevel(logging.INFO) + handler.setFormatter(logging.Formatter("%(message)s")) + setattr(handler, "_managed_agent_model_trace", True) + _MODEL_TRACE_LOGGER.addHandler(handler) + _model_trace( + "model_trace_ready", + module=_safe_trace_value(__file__), + ) + + +def _trace_model_request(responses_args: dict[str, Any]) -> None: + tools = [] + for tool in responses_args.get("tools") or []: + name = _safe_trace_value(_trace_field(tool, "name")) + tool_type = _safe_trace_value(_trace_field(tool, "type")) + if name or tool_type: + tools.append({"name": name, "type": tool_type}) + _model_trace( + "model_request", + model=_safe_trace_value(responses_args.get("model")), + tools=tools, + tool_choice=_normalized_tool_choice(responses_args.get("tool_choice")), + has_previous_response_id=bool(responses_args.get("previous_response_id")), + ) + + +def _trace_model_stream_event(event: Any) -> None: + event_type = _safe_trace_value(_trace_field(event, "type")) + if not event_type: + event_type = _safe_trace_value(type(event).__name__) + _model_trace("model_response_event", response_event_type=event_type) + + +def _trace_model_response(response: Any) -> None: + output_types = [] + for output in _trace_field(response, "output") or []: + output_type = _safe_trace_value(_trace_field(output, "type")) + if not output_type: + output_type = _safe_trace_value(type(output).__name__) + if output_type: + output_types.append(output_type) + _model_trace("model_response", output_types=output_types) + + _FINISH_REASON_MAPPING = { "incomplete": { "length": types.FinishReason.MAX_TOKENS, @@ -508,12 +607,11 @@ def request_reorganization_by_ark( key: value for key, value in request_data.items() if key in ark_supported_fields } - # expire time 1 hour when send aresponses req + # Respect explicit retention, otherwise use Ark's server default. extra_body = request_data.get("extra_body") if not isinstance(extra_body, dict): extra_body = {} request_data["extra_body"] = extra_body - extra_body["expire_at"] = int(time.time()) + 3600 # [Note: Ark Limitations] caching and text # After enabling caching, output_schema(text) cannot be used. Caching must be disabled. @@ -523,15 +621,8 @@ def request_reorganization_by_ark( ) _remove_caching(request_data) - # [Note: Ark Limitations] tools and previous_response_id - # Remove tools in subsequent rounds (when previous_response_id is present) - if ( - "tools" in request_data - and "previous_response_id" in request_data - and request_data["previous_response_id"] is not None - ): - # Remove tools in subsequent rounds regardless of caching status - del request_data["tools"] + # Tool declarations are request configuration, not conversation history. + # Keep them on continuation requests, including after a worker restart. # [Note: Ark Limitations] caching and store # Ensure store field is true or default when caching is enabled @@ -706,6 +797,7 @@ class ArkLlm(Gemini): _additional_args: Dict[str, Any] = None use_interactions_api: bool = True enable_responses_cache: bool = True + retry_expired_response: bool = True def __init__(self, **kwargs): # adk version check @@ -725,6 +817,7 @@ def __init__(self, **kwargs): self._additional_args.pop("stream", None) self._additional_args.pop("fallbacks", None) self._additional_args.pop("enable_responses_cache", None) + self._additional_args.pop("retry_expired_response", None) if drop_params is not None: self._additional_args["drop_params"] = drop_params @@ -805,6 +898,11 @@ async def _generate_content_with_fallbacks( failure = error if self._is_previous_response_not_found(error): + if not self.retry_expired_response: + raise RuntimeError( + "Ark conversation context expired; start a new Session. " + "The previous conversation cannot be silently discarded." + ) from error logger.warning( f"Interaction expired for model `{model}` (PreviousResponseNotFound). Retrying without previous_response_id. Error: {error}" ) @@ -851,13 +949,16 @@ def _is_previous_response_not_found(error: Exception) -> bool: async def generate_content_via_responses( self, responses_args: dict, stream: bool = False ): + configure_managed_agent_model_trace() model = responses_args["model"] responses_args = request_reorganization_by_ark( responses_args, enable_responses_cache=self.enable_responses_cache ) + _trace_model_request(responses_args) if stream: responses_args["stream"] = True async for part in await self.llm_client.aresponses(**responses_args): + _trace_model_stream_event(part) llm_response = event_to_generate_content_response( event=part, is_partial=True, model_version=model ) @@ -865,6 +966,7 @@ async def generate_content_via_responses( yield llm_response else: raw_response = await self.llm_client.aresponses(**responses_args) + _trace_model_response(raw_response) llm_response = ark_response_to_generate_content_response(raw_response) yield llm_response diff --git a/veadk/runtime/managed_agents/__init__.py b/veadk/runtime/managed_agents/__init__.py new file mode 100644 index 000000000..a672560fd --- /dev/null +++ b/veadk/runtime/managed_agents/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Managed Agents execution. Import concrete modules to select capabilities.""" diff --git a/veadk/runtime/managed_agents/dispatcher.py b/veadk/runtime/managed_agents/dispatcher.py new file mode 100644 index 000000000..c7986949e --- /dev/null +++ b/veadk/runtime/managed_agents/dispatcher.py @@ -0,0 +1,413 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Work dispatcher implemented with the published Anthropic HTTP/Work APIs. + +Poll credentials belong to the environment. Claimed Work credentials belong to +one lease. Neither credential nor secret-bearing API error text is logged. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import logging +import os +import random +import time +from collections.abc import Awaitable, Callable, Mapping +from typing import Any +from uuid import uuid4 + +import anthropic +import httpx2 +from anthropic.types.beta.environments import BetaSelfHostedWork + +logger = logging.getLogger(__name__) +poll_trace_logger = logging.getLogger("anthropic.managed_agent_poll") +HEARTBEAT_INTERVAL = 30.0 +LEASE_TTL = 90.0 +CLEANUP_TIMEOUT = 10.0 + + +def sessions_token(secret: str | None) -> str | None: + """Extract a per-Work token; opaque malformed secrets never enter errors.""" + if not secret: + return None + try: + value = json.loads(base64.urlsafe_b64decode(secret + "=" * (-len(secret) % 4))) + token = value.get("sessions_token") if isinstance(value, dict) else None + return token if isinstance(token, str) and token else None + except (ValueError, TypeError, UnicodeError): + return None + + +def _non_auth_headers(headers: Mapping[str, Any]) -> dict[str, Any]: + return { + k: v + for k, v in headers.items() + if k.lower() not in {"authorization", "x-api-key"} + } + + +def scoped_client(client: Any, token: str) -> Any: + """Preserve transport/settings while replacing all inherited authentication.""" + scoped = client.copy( + auth_token=token, + credentials=None, + set_default_headers=_non_auth_headers(client.default_headers), + max_retries=0, + ) + # copy(api_key=None) inherits the parent key in SDK 1.3.0. + scoped.api_key = None + return scoped + + +def _retryable(error: Exception) -> bool: + status = getattr(error, "status_code", None) + return ( + status in (408, 409, 429) + or (isinstance(status, int) and status >= 500) + or isinstance( + error, (anthropic.APIConnectionError, httpx2.TransportError, TimeoutError) + ) + ) + + +def _authentication_error(error: Exception) -> bool: + return getattr(error, "status_code", None) in (401, 403) + + +def _poll_trace(event: str, **fields: Any) -> None: + if os.getenv("MANAGED_AGENT_POLL_TRACE", "").strip().lower() == "true": + poll_trace_logger.info( + json.dumps( + {"event": event, **fields}, separators=(",", ":"), sort_keys=True + ) + ) + + +class EnvironmentWorkDispatcher: + """Bounded concurrent leases, with FIFO execution for each Session. + + ``drain`` interrupts polling and waits for accepted Work. Cancelling ``run`` + cancels handlers and waits for their cleanup. ``ready`` is set only after a + validated successful poll, including an empty 204 response. + """ + + def __init__( + self, + client: Any, + *, + handler: Callable[[Any, Any], Awaitable[None]], + environment_id: str, + environment_key: str, + account_work: bool = False, + worker_id: str | None = None, + extra_headers: Mapping[str, Any] | None = None, + ): + if not environment_key or (not environment_id and not account_work): + raise ValueError( + "Work dispatcher requires configured credentials and scope" + ) + self.client = client + self.handler = handler + self.environment_id = environment_id + self.environment_key = environment_key + self.account_work = account_work + self.worker_id = worker_id or f"worker-{uuid4().hex}" + self.headers = _non_auth_headers(extra_headers or {}) + self.draining = False + self.ready = asyncio.Event() + self._drain_event = asyncio.Event() + self.poll_task: asyncio.Task | None = None + self._session_tails: dict[str, asyncio.Future] = {} + + def drain(self) -> None: + self.draining = True + self._drain_event.set() + if self.poll_task is not None: + self.poll_task.cancel() + + async def _poll(self, client: Any) -> BetaSelfHostedWork | None: + headers = { + "anthropic-beta": "managed-agents-2026-04-01", + **self.headers, + "Anthropic-Worker-ID": self.worker_id, + } + trace_id = f"poll_{uuid4().hex}" + if os.getenv("MANAGED_AGENT_POLL_TRACE", "").strip().lower() == "true": + headers["X-MA-Poll-Trace-ID"] = trace_id + started = time.monotonic() + _poll_trace( + "poll_started", request_id=trace_id, account_scope=self.account_work + ) + try: + # Inspect 204 before parsing: SDK 1.3's typed poll result otherwise + # constructs an empty Work. Both paths use public SDK interfaces. + if self.account_work: + # Account model Work is a VeADK gateway extension, not a Claude + # endpoint. The official SDK supports custom requests via get. + response = await client.get( + "/v1/model-work/poll", + cast_to=httpx2.Response, + options={"headers": headers, "params": {"block_ms": 999}}, + ) + else: + raw = await client.beta.environments.work.with_raw_response.poll( + self.environment_id, + block_ms=999, + anthropic_worker_id=self.worker_id, + extra_headers=headers, + ) + response = raw.http_response + if response.status_code == 204: + item = None + else: + media = ( + response.headers.get("content-type", "") + .split(";", 1)[0] + .strip() + .lower() + ) + if response.status_code != 200 or media != "application/json": + raise ValueError("Work poll returned an unexpected response") + try: + item = BetaSelfHostedWork.model_validate(response.json()) + if not item.id or not item.environment_id or not item.data.id: + raise ValueError() + if self.account_work and not sessions_token(item.secret): + raise ValueError() + if ( + not self.account_work + and item.environment_id != self.environment_id + ): + raise ValueError() + except Exception: + raise ValueError("Work poll returned malformed Work") from None + except Exception as error: + _poll_trace( + "poll_error", + request_id=trace_id, + error_type=type(error).__name__, + status=getattr(error, "status_code", None), + retry=_retryable(error), + ) + raise + self.ready.set() + _poll_trace( + "poll_completed", + request_id=trace_id, + status=response.status_code, + result="work" if item is not None else "empty", + duration_ms=round((time.monotonic() - started) * 1000, 3), + ) + return item + + async def _pause(self, delay: float) -> None: + try: + await asyncio.wait_for(self._drain_event.wait(), timeout=delay) + except asyncio.TimeoutError: + pass + + async def run( + self, *, max_items: int | None = None, max_concurrency: int = 1 + ) -> int: + if max_concurrency < 1 or (max_items is not None and max_items < 1): + raise ValueError("Work limits must be positive") + client = scoped_client(self.client, self.environment_key) + active: set[asyncio.Task] = set() + count = attempt = 0 + try: + while not self.draining and (max_items is None or count < max_items): + completed = {task for task in active if task.done()} + for task in completed: + active.remove(task) + task.result() + if len(active) >= max_concurrency: + await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED) + continue + self.poll_task = asyncio.create_task(self._poll(client)) + try: + while not self.poll_task.done(): + done, _ = await asyncio.wait( + {self.poll_task, *active}, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in done - {self.poll_task}: + active.remove(task) + task.result() + item = await self.poll_task + except asyncio.CancelledError: + if self.draining and self.poll_task.cancelled(): + break + raise + except Exception as error: + if not _retryable(error): + raise + attempt += 1 + await self._pause(min(2 ** min(attempt, 5), 30) + random.random()) + continue + finally: + if self.poll_task is not None and not self.poll_task.done(): + self.poll_task.cancel() + await asyncio.gather(self.poll_task, return_exceptions=True) + self.poll_task = None + attempt = 0 + if item is None: + await self._pause(0.25) + continue + session = item.data.id + previous = self._session_tails.get(session) + finished = asyncio.get_running_loop().create_future() + self._session_tails[session] = finished + active.add( + asyncio.create_task(self._handle_item(item, previous, finished)) + ) + count += 1 + if active: + await asyncio.gather(*active) + return count + finally: + for task in active: + if not task.done(): + task.cancel() + if active: + await asyncio.gather(*active, return_exceptions=True) + + async def _handle_item( + self, + item: BetaSelfHostedWork, + previous: asyncio.Future | None = None, + finished: asyncio.Future | None = None, + ) -> None: + token = sessions_token(item.secret) + client = scoped_client(self.client, token or self.environment_key) + work = client.beta.environments.work + lost = False + fenced = asyncio.Event() + beat = handler = None + + async def heartbeat_loop() -> None: + nonlocal lost + last, interval, ttl = "NO_HEARTBEAT", HEARTBEAT_INTERVAL, LEASE_TTL + success = time.monotonic() + while True: + try: + response = await asyncio.wait_for( + work.heartbeat( + item.id, + environment_id=item.environment_id, + expected_last_heartbeat=last, + extra_headers=self.headers, + ), + timeout=min( + interval, max(0.001, ttl - (time.monotonic() - success)) + ), + ) + except Exception as error: + if getattr(error, "status_code", None) == 412: + lost = True + return + if not _retryable(error): + raise + if time.monotonic() - success >= ttl: + lost = True + return + else: + last = response.last_heartbeat + success = time.monotonic() + if response.ttl_seconds > 0: + ttl = float(response.ttl_seconds) + interval = min(ttl / 2, HEARTBEAT_INTERVAL) + if response.state in {"stopping", "stopped"}: + return + if not response.lease_extended: + lost = True + return + fenced.set() + await asyncio.sleep( + min(interval, max(0.001, ttl - (time.monotonic() - success))) + ) + + async def ordered_handler() -> None: + await fenced.wait() + if previous is not None: + await asyncio.shield(previous) + await self.handler(item, client) + + try: + await work.ack( + item.id, environment_id=item.environment_id, extra_headers=self.headers + ) + beat = asyncio.create_task(heartbeat_loop()) + handler = asyncio.create_task(ordered_handler()) + done, _ = await asyncio.wait( + {beat, handler}, return_when=asyncio.FIRST_COMPLETED + ) + for task in done: + task.result() + except asyncio.CancelledError: + raise + except Exception as error: + if getattr(error, "status_code", None) == 412: + lost = True + if _authentication_error(error): + raise + logger.error("Managed Work failed: %s", type(error).__name__) + finally: + # Keep a live heartbeat through asynchronous handler teardown. + if handler is not None: + if not handler.done(): + handler.cancel() + await asyncio.gather(handler, return_exceptions=True) + if beat is not None: + if not beat.done(): + beat.cancel() + await asyncio.gather(beat, return_exceptions=True) + try: + if not lost: + + async def stop() -> None: + try: + await asyncio.wait_for( + work.stop( + item.id, + environment_id=item.environment_id, + force=True, + extra_headers=self.headers, + ), + CLEANUP_TIMEOUT, + ) + except Exception as error: + if _authentication_error(error): + raise + if getattr(error, "status_code", None) != 409: + logger.error( + "Managed Work stop failed: %s", type(error).__name__ + ) + + cleanup = asyncio.create_task(stop()) + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + await cleanup + raise + finally: + if finished is not None: + if not finished.done(): + finished.set_result(None) + if self._session_tails.get(item.data.id) is finished: + self._session_tails.pop(item.data.id, None) diff --git a/veadk/runtime/managed_agents/events.py b/veadk/runtime/managed_agents/events.py new file mode 100644 index 000000000..8fcec1038 --- /dev/null +++ b/veadk/runtime/managed_agents/events.py @@ -0,0 +1,134 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Opt-in, payload-level tracing for Managed Agent events.""" + +from __future__ import annotations + +import json +import os +import threading +from collections.abc import AsyncIterator, Iterable +from datetime import datetime, timezone +from typing import Any + +_OUTPUT_LOCK = threading.Lock() + +_SECRET_FIELDS = { + "authorization", + "apikey", + "token", + "accesstoken", + "refreshtoken", + "accesskey", + "accesskeyid", + "secretkey", + "secret", + "secretaccesskey", + "sessiontoken", + "password", + "clientsecret", + "credentials", + "credential", + "workloadaccesstoken", + "agentidentitytoken", + "worktoken", + "bearertoken", + "environmentkey", + "xapikey", +} + + +def _json_value(value: Any) -> Any: + if hasattr(value, "model_dump"): + return value.model_dump(mode="json", warnings=False) + if isinstance(value, datetime): + return value.isoformat() + if hasattr(value, "__dict__"): + return vars(value) + raise TypeError("Unsupported event value") + + +def _redact(value: Any) -> Any: + if isinstance(value, dict): + return { + key: "[REDACTED]" + if key.lower().replace("_", "").replace("-", "") in _SECRET_FIELDS + else _redact(item) + for key, item in value.items() + } + if isinstance(value, list): + return [_redact(item) for item in value] + return value + + +def trace_event( + direction: str, source: str, session_id: str | None, event: Any +) -> None: + """Print one complete JSON record without changing the event or log levels.""" + if os.getenv("MA_DEBUG_EVENTS", "").strip().lower() not in { + "1", + "true", + "yes", + "on", + }: + return + try: + # Normalize SDK/Runner models first so nested credential fields are redacted too. + payload = _redact(json.loads(json.dumps(event, default=_json_value))) + record = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "direction": direction, + "source": source, + "session_id": session_id, + "event": payload, + } + line = "MA_DEBUG_EVENT " + json.dumps(record, ensure_ascii=False) + with _OUTPUT_LOCK: + print(line, flush=True) + except Exception: + # Debugging must not fail a turn or expose a secret-bearing exception. + pass + + +def trace_events( + direction: str, source: str, session_id: str | None, events: Iterable[Any] +) -> None: + for event in events: + trace_event(direction, source, session_id, event) + + +def trace_send_response(session_id: str | None, response: Any) -> None: + events = ( + response.get("data", []) + if isinstance(response, dict) + else getattr(response, "data", []) + ) + trace_events("received", "send_response", session_id, events or []) + + +async def send_session_events(sdk: Any, session_id: str, *, events: list[Any]) -> Any: + # Outbound records describe attempted sends; persisted events appear in send_response. + trace_events("sent", "send", session_id, events) + response = await sdk.beta.sessions.events.send(session_id, events=events) + trace_send_response(session_id, response) + return response + + +async def iter_session_events( + sdk: Any, session_id: str, **options: Any +) -> AsyncIterator[Any]: + async for event in sdk.beta.sessions.events.list(session_id, **options): + trace_event("received", "list", session_id, event) + yield event diff --git a/veadk/runtime/managed_agents/health.py b/veadk/runtime/managed_agents/health.py new file mode 100644 index 000000000..1eed3a574 --- /dev/null +++ b/veadk/runtime/managed_agents/health.py @@ -0,0 +1,53 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""AgentKit HTTP listener; business execution stays in the existing Agent Loop.""" + +from __future__ import annotations + +import os + +from agentkit.apps import AgentkitSimpleApp + + +def _first_env(*names: str, default: str = "") -> str: + for name in names: + value = os.getenv(name, "").strip() + if value: + return value + return default + + +app = AgentkitSimpleApp() + + +@app.ping +def ping() -> dict[str, str]: + return {"status": "ok"} + + +@app.entrypoint +async def invoke(payload: dict, headers: dict) -> dict[str, str]: + return {"status": "ok"} + + +if __name__ == "__main__": + app.run( + host=_first_env("AGENTKIT_RUNTIME_HOST", "MA_RUNTIME_HOST", default="0.0.0.0"), + port=int( + _first_env( + "AGENTKIT_RUNTIME_PORT", "MA_RUNTIME_PORT", "PORT", default="8080" + ) + ), + ) diff --git a/veadk/runtime/managed_agents/identity.py b/veadk/runtime/managed_agents/identity.py new file mode 100644 index 000000000..e3353df81 --- /dev/null +++ b/veadk/runtime/managed_agents/identity.py @@ -0,0 +1,351 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-local AgentKit Identity credential resolution. + +The API key is deliberately never written to an environment variable, file, +database, log message, exception, or response. +""" + +from __future__ import annotations + +import asyncio +import os +import re +import threading +import time +from collections import OrderedDict +from collections.abc import Callable +from typing import Any + +from veadk.integrations.ve_identity.identity_client import IdentityClient + +SESSION_MODEL_CACHE_CAPACITY = 2000 +SESSION_MODEL_CACHE_TTL_SECONDS = 60 * 60 + + +class RuntimeIdentityError(RuntimeError): + """A sanitized Identity failure safe to expose to the Runtime caller.""" + + def __init__(self, category: str, retryable: bool) -> None: + super().__init__(f"AgentKit Identity {category}") + self.category = category + self.retryable = retryable + + +def _api_key_value(value: Any) -> str: + if isinstance(value, str): + key = value + elif isinstance(value, dict): + key = value.get("api_key") or value.get("ApiKey") or "" + else: + key = getattr(value, "api_key", None) or getattr(value, "ApiKey", "") + if not isinstance(key, str) or not key: + raise RuntimeIdentityError("invalid_response", False) + return key + + +def _category(error: Exception) -> tuple[str, bool]: + status = getattr(error, "status", None) or getattr(error, "status_code", None) + if str(status) in {"401", "403"}: + return "authentication_failed", False + if str(status) == "404": + return "provider_not_found", False + if str(status) in {"408", "429", "500", "502", "503", "504"}: + return "network_error", True + text = f"{type(error).__name__} {error}".lower() + if any( + token in text + for token in ( + "unauthorized", + "forbidden", + "accessdenied", + "authentication", + "invalididentitytoken", + "401", + "403", + ) + ): + return "authentication_failed", False + if any(token in text for token in ("notfound", "not found", "404")): + return "provider_not_found", False + if any( + token in text + for token in ( + "timeout", + "connection", + "temporar", + "429", + "500", + "502", + "503", + "504", + ) + ): + return "network_error", True + return "identity_error", False + + +class _CachedKey: + def __init__(self, value: str, expires_at: float | None) -> None: + self.value = value + self.expires_at = expires_at + + +class RuntimeCredentialProvider: + """Cache gateway keys until invalidated; model keys have a bounded TTL.""" + + def __init__( + self, + *, + runtime_id: str, + provider_name: str | None = None, + region: str = "cn-beijing", + ttl_seconds: float | None = None, + attempts: int = 3, + client_factory: Callable[[], Any] | None = None, + credential_kind: str = "model", + pool_name: str | None = None, + ) -> None: + if not runtime_id: + raise ValueError("Runtime identity references are required") + self.runtime_id = runtime_id + pattern = ( + r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}" + if credential_kind == "mcp" + else r"ma-ark-[a-z2-7]{52}" + ) + if credential_kind not in {"model", "mcp"}: + raise RuntimeIdentityError("invalid_provider_reference", False) + if provider_name is not None and not re.fullmatch(pattern, provider_name): + raise RuntimeIdentityError("invalid_provider_reference", False) + self.provider_name = provider_name or runtime_id + "-ma" + self.region = region + self.pool_name = pool_name + max_ttl = SESSION_MODEL_CACHE_TTL_SECONDS if provider_name is not None else 60.0 + if ttl_seconds is None and provider_name is not None: + ttl_seconds = SESSION_MODEL_CACHE_TTL_SECONDS + self.ttl_seconds = ( + None if ttl_seconds is None else min(max_ttl, max(0.0, ttl_seconds)) + ) + self.attempts = min(3, max(1, attempts)) + self.client_factory = client_factory or (lambda: IdentityClient(region=region)) + self._cached: _CachedKey | None = None + self._lock = threading.Lock() + + def invalidate(self) -> None: + with self._lock: + self._cached = None + + async def get(self, *, force_refresh: bool = False) -> str: + return await asyncio.to_thread(self.get_sync, force_refresh=force_refresh) + + def get_sync(self, *, force_refresh: bool = False) -> str: + """Resolve credentials for the existing synchronous sandbox client.""" + with self._lock: + now = time.monotonic() + if ( + not force_refresh + and self._cached + and (self._cached.expires_at is None or self._cached.expires_at > now) + ): + return self._cached.value + self._cached = None + last: RuntimeIdentityError | None = None + for attempt in range(self.attempts): + phase = "workload" + try: + client = self.client_factory() + token = client.get_workload_access_token( + workload_name=self.runtime_id, + ) + token_value = getattr(token, "workload_access_token", None) + if not isinstance(token_value, str) or not token_value: + raise RuntimeIdentityError("invalid_token_response", False) + phase = "provider" + kwargs = { + "provider_name": self.provider_name, + "agent_identity_token": token_value, + } + if self.pool_name is not None: + kwargs["pool_name"] = self.pool_name + raw = client.get_api_key(**kwargs) + key = _api_key_value(raw) + expires_at = ( + None + if self.ttl_seconds is None + else time.monotonic() + self.ttl_seconds + ) + self._cached = _CachedKey(key, expires_at) + return key + except RuntimeIdentityError as error: + last = error + except Exception as error: # noqa: BLE001 - never expose secret-bearing SDK errors. + category, retryable = _category(error) + if phase == "workload" and category == "provider_not_found": + category, retryable = "authentication_failed", False + last = RuntimeIdentityError(category, retryable) + if not last.retryable or attempt + 1 >= self.attempts: + break + time.sleep(0.2 * (2**attempt)) + raise last or RuntimeIdentityError("identity_error", False) + + +_gateway_providers: dict[tuple[str, str], RuntimeCredentialProvider] = {} +_gateway_providers_lock = threading.Lock() + + +def get_runtime_credential_provider( + *, runtime_id: str, region: str = "cn-beijing" +) -> RuntimeCredentialProvider: + """Share a gateway credential and its refresh lock across clients in this process.""" + cache_key = (runtime_id, region) + with _gateway_providers_lock: + provider = _gateway_providers.get(cache_key) + if provider is None: + provider = RuntimeCredentialProvider(runtime_id=runtime_id, region=region) + _gateway_providers[cache_key] = provider + return provider + + +class SessionModelCredentials: + """Bounded LRU of Session model credentials, separate from endpoint keys.""" + + def __init__( + self, + capacity: int = SESSION_MODEL_CACHE_CAPACITY, + client_factory=None, + ttl_seconds: float = SESSION_MODEL_CACHE_TTL_SECONDS, + ): + if capacity < 1: + raise ValueError("credential cache capacity must be positive") + self.capacity = capacity + self.client_factory = client_factory + self.ttl_seconds = ttl_seconds + self._providers: OrderedDict[tuple[str, str], RuntimeCredentialProvider] = ( + OrderedDict() + ) + self._lock = threading.Lock() + + async def get(self, session: Any) -> str: + name = ( + session.get("ark_api_key_provider") + if isinstance(session, dict) + else getattr(session, "ark_api_key_provider", None) + ) + if name is None: + raise RuntimeIdentityError("missing_provider_reference", False) + if not isinstance(name, str) or not re.fullmatch(r"ma-ark-[a-z2-7]{52}", name): + raise RuntimeIdentityError("invalid_provider_reference", False) + from veadk.integrations.mpa.session_client import _runtime_id_from_environment + + workload = _runtime_id_from_environment() + if not workload: + raise RuntimeIdentityError("missing_workload_identity", False) + cache_key = (workload, name) + with self._lock: + provider = self._providers.get(cache_key) + if provider is None: + provider = RuntimeCredentialProvider( + runtime_id=workload, + provider_name=name, + region=os.getenv("VOLCENGINE_REGION") + or os.getenv("REGION") + or "cn-beijing", + client_factory=self.client_factory, + ttl_seconds=self.ttl_seconds, + ) + self._providers[cache_key] = provider + if len(self._providers) > self.capacity: + self._providers.popitem(last=False) + self._providers.move_to_end(cache_key) + return await provider.get() + + +session_model_credentials = SessionModelCredentials() + + +class MCPIdentityHeaders: + """Keep MCP credentials out of snapshots, model state and tool arguments.""" + + def __init__(self, reference: Any, *, client_factory=None): + def field(name, default=None): + return ( + reference.get(name, default) + if isinstance(reference, dict) + else getattr(reference, name, default) + ) + + self.name = field("Name") + self.pool_name = field("PoolName") + self.header = field("HeaderName", "Authorization") + self.prefix = field("Prefix", "Bearer ") + if any( + field(k) is not None for k in ("api_key", "ApiKey", "authorization_token") + ): + raise RuntimeIdentityError("invalid_mcp_authorization", False) + if not isinstance(self.name, str) or not re.fullmatch( + r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}", self.name + ): + raise RuntimeIdentityError("invalid_provider_reference", False) + if not isinstance(self.pool_name, str) or not re.fullmatch( + r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}", self.pool_name + ): + raise RuntimeIdentityError("invalid_provider_reference", False) + if not isinstance(self.header, str) or not re.fullmatch( + r"[A-Za-z0-9-]{1,64}", self.header + ): + raise RuntimeIdentityError("invalid_mcp_authorization", False) + if ( + not isinstance(self.prefix, str) + or len(self.prefix) > 32 + or any(c in self.prefix for c in "\r\n\x00") + ): + raise RuntimeIdentityError("invalid_mcp_authorization", False) + if self.header.lower() in {"host", "cookie"} or self.header.lower().startswith( + "proxy-" + ): + raise RuntimeIdentityError("invalid_mcp_authorization", False) + self.client_factory = client_factory + self._provider = None + + def provider(self): + if self._provider is None: + from veadk.integrations.mpa.session_client import ( + _runtime_id_from_environment, + ) + + workload = _runtime_id_from_environment() + if not workload: + raise RuntimeIdentityError("missing_workload_identity", False) + self._provider = RuntimeCredentialProvider( + runtime_id=workload, + provider_name=self.name, + credential_kind="mcp", + pool_name=self.pool_name, + region=os.getenv("VOLCENGINE_REGION") + or os.getenv("REGION") + or "cn-beijing", + client_factory=self.client_factory, + ) + return self._provider + + async def prepare(self): + await self.provider().get() + + def __call__(self, _context): + key = self.provider().get_sync() + if any(c in key for c in "\r\n\x00"): + raise RuntimeIdentityError("invalid_response", False) + return {self.header: self.prefix + key} diff --git a/veadk/runtime/managed_agents/loop.py b/veadk/runtime/managed_agents/loop.py new file mode 100644 index 000000000..220c03164 --- /dev/null +++ b/veadk/runtime/managed_agents/loop.py @@ -0,0 +1,739 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Event-driven Managed Agents loop backed entirely by the Anthropic SDK.""" + +from __future__ import annotations + +import asyncio +from contextlib import aclosing +import logging +import os +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict + +from veadk.runtime.managed_agents.events import ( + iter_session_events, + send_session_events, + trace_event, +) +from google.adk.agents import RunConfig +from google.adk.agents.run_config import StreamingMode +from google.adk.events import Event +from google.genai import types +from veadk.integrations.mpa.session_client import SelfHostSandboxClient + +from veadk.utils.adk_compat import get_event_function_calls + +logger = logging.getLogger("veadk.managed_agent_loop") + +_TERMINAL_EVENT_TYPES = { + "session.status_idle", + "session.status_terminated", + "session.error", +} + + +@dataclass(slots=True) +class ManagedAgentEventState: + """Incremental event state shared by handlers for one Managed Session.""" + + completed_inputs: set[str] = field(default_factory=set) + pending_inputs: list[str] = field(default_factory=list) + reconciled_events: set[str] = field(default_factory=set) + cursor_seq: int | None = None + history_reconciled: bool = False + ark_response_id: str | None = None + lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) + + +class ManagedAgentsLoop: + """Consume one pending Managed Session turn and drive a VeADK Runner. + + A self-hosted sandbox is normally started for one WorkItem. Consequently + ``run(max_turns=1)`` handles one uncompleted ``user.message``, emits the + canonical Agent/Session events, and exits after ``session.status_idle``. + A later user message creates another WorkItem and starts a fresh sandbox. + """ + + def __init__( + self, + *, + runner: Any, + session_client: SelfHostSandboxClient | None = None, + session_id: str | None = None, + user_id: str = "managed_agents_user", + event_state: ManagedAgentEventState | None = None, + ) -> None: + resolved_session_id = session_id or getattr(session_client, "session_id", None) + if not resolved_session_id: + raise ValueError( + "ANTHROPIC_SESSION_ID (or SANDBOX_SESSION_ID) is required for Agent Loop mode." + ) + self.runner = runner + self.session_client = session_client + self.session_id = str(resolved_session_id) + self.user_id = user_id + self._event_state = event_state or ManagedAgentEventState() + + @property + def _completed_inputs(self) -> set[str]: + return self._event_state.completed_inputs + + @property + def _event_cursor_seq(self) -> int | None: + return self._event_state.cursor_seq + + @_event_cursor_seq.setter + def _event_cursor_seq(self, value: int | None) -> None: + self._event_state.cursor_seq = value + + @property + def _history_reconciled(self) -> bool: + return self._event_state.history_reconciled + + @_history_reconciled.setter + def _history_reconciled(self, value: bool) -> None: + self._event_state.history_reconciled = value + + async def run( + self, + *, + max_turns: int | None = 1, + stop_event: asyncio.Event | None = None, + ) -> int: + """Listen through SDK SSE, process pending user messages, and return turns run.""" + if self.session_client is None: + raise ValueError( + "run() requires session_client; use run_pending() with a scoped SDK client" + ) + self._restore_completed_inputs( + await asyncio.to_thread(self.session_client.list_events) + ) + turns = 0 + backoff = 0.25 + stop_after_terminal = False + + async with self.session_client.create_async_client() as sdk: + while stop_event is None or not stop_event.is_set(): + try: + stream = await sdk.beta.sessions.events.stream(self.session_id) + terminal_seen = False + async with stream: + async for event in stream: + trace_event("received", "sse", self.session_id, event) + event_type = str(getattr(event, "type", "") or "") + if event_type == "session.status_terminated": + return turns + if ( + stop_after_terminal + and event_type in _TERMINAL_EVENT_TYPES + ): + terminal_seen = True + continue + if event_type != "user.message": + continue + + event_id = str(getattr(event, "id", "") or "") + if event_id and event_id in self._completed_inputs: + continue + + await self._run_turn(sdk, event) + if event_id: + self._completed_inputs.add(event_id) + turns += 1 + if max_turns is not None and turns >= max_turns: + stop_after_terminal = True + if terminal_seen: + return turns + backoff = 0.25 + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "Managed Session event stream failed; reconnecting in %.2fs", + backoff, + ) + if stop_event is not None: + try: + await asyncio.wait_for(stop_event.wait(), timeout=backoff) + return turns + except TimeoutError: + pass + else: + await asyncio.sleep(backoff) + backoff = min(backoff * 2, 5.0) + return turns + + async def run_pending(self, sdk: Any, *, max_turns: int | None = 1) -> int: + """Process persisted, incomplete user messages for one claimed work item. + + The first pass reconciles the complete durable ledger so a replacement + worker can recover work that arrived before it started. Later passes + use the last observed event sequence as a cursor and read only newer + events. Handlers for the same Session serialize through their shared + event state so concurrent WorkItems cannot advance the cursor past an + unprocessed message. + """ + async with self._event_state.lock: + return await self._run_pending_locked(sdk, max_turns=max_turns) + + async def _run_pending_locked(self, sdk: Any, *, max_turns: int | None) -> int: + previous_cursor = self._event_cursor_seq if self._history_reconciled else None + events, full_reconciliation = await self._list_pending_events(sdk) + # Another Pod may have completed inputs since our last cursor. A + # process-local completed set is only a cache of the canonical ledger. + self._restore_completed_inputs(events, incremental=not full_reconciliation) + turns = 0 + consumed = 0 + for index, event in enumerate(events): + consumed = index + 1 + if self._value(event, "type") != "user.message": + continue + event_id = str(self._value(event, "id") or "") + if event_id and event_id in self._completed_inputs: + continue + await self._run_turn(sdk, event) + if event_id: + self._completed_inputs.add(event_id) + turns += 1 + if max_turns is not None and turns >= max_turns: + break + self._history_reconciled = True + self._advance_event_cursor(events[:consumed], previous=previous_cursor) + return turns + + async def _list_pending_events(self, sdk: Any) -> tuple[list[Any], bool]: + """Read a recovery snapshot once, then advance through new events.""" + cursor_seq = self._event_cursor_seq if self._history_reconciled else None + incremental = cursor_seq is not None + try: + events = await self._list_events_after(sdk, cursor_seq) + if incremental and any( + (seq := self._event_seq(event)) is not None and seq <= cursor_seq + for event in events + ): + raise RuntimeError("event server did not advance the requested cursor") + full_reconciliation = cursor_seq is None + except asyncio.CancelledError: + raise + except Exception: + if not incremental: + raise + logger.warning( + "Managed Session event cursor failed; reconciling full history once", + exc_info=True, + ) + full_history = await self._list_events_after(sdk, None) + self._restore_completed_inputs(full_history) + events = full_history + full_reconciliation = False + + if incremental: + # A server that ignores ``page`` must not make an already-consumed + # user.message visible to the execution loop again. Missing seq is + # retained so correctness falls back to the completed-input guard. + events = [ + event + for event in events + if (seq := self._event_seq(event)) is None or seq > cursor_seq + ] + + return events, full_reconciliation + + async def _list_events_after(self, sdk: Any, cursor_seq: int | None) -> list[Any]: + options: dict[str, Any] = {"limit": 100, "order": "asc"} + if cursor_seq is not None: + # ma-infra ListEvents/next_page uses decimal sequence cursors. + options["page"] = str(cursor_seq) + return [ + event + async for event in iter_session_events(sdk, self.session_id, **options) + ] + + def _advance_event_cursor(self, events: list[Any], *, previous: int | None) -> None: + if not events: + if previous is None: + self._event_cursor_seq = 0 + return + sequences = [self._event_seq(event) for event in events] + if any(sequence is None for sequence in sequences): + # Custom/older servers may omit sequence metadata. Keep correctness + # by using full reconciliation on the next pass instead of risking + # a cursor that skips an event. + self._event_cursor_seq = None + logger.warning( + "Managed Session events omitted sequence metadata; " + "next poll will reconcile full history" + ) + return + self._event_cursor_seq = max(previous or 0, *(int(seq) for seq in sequences)) + + async def run_claimed( + self, + sdk: Any, + *, + environment_id: str, + work_id: str, + last_heartbeat: str = "NO_HEARTBEAT", + poll_interval: float = 1.0, + heartbeat_interval: float | None = None, + ) -> int: + turns = 0 + stopped = asyncio.Event() + if heartbeat_interval is None: + heartbeat_interval = float( + os.getenv("MANAGED_AGENT_HEARTBEAT_INTERVAL_SECONDS", "2") + ) + if heartbeat_interval <= 0: + raise ValueError("heartbeat_interval must be positive") + + async def heartbeat() -> None: + nonlocal last_heartbeat + while not stopped.is_set(): + response = await sdk.beta.environments.work.heartbeat( + work_id, + environment_id=environment_id, + expected_last_heartbeat=last_heartbeat, + desired_ttl_seconds=90, + ) + if not response.lease_extended: + raise RuntimeError("Managed work lease is no longer active") + last_heartbeat = response.last_heartbeat + try: + await asyncio.wait_for(stopped.wait(), timeout=heartbeat_interval) + except TimeoutError: + pass + + async def consume() -> None: + nonlocal turns + while True: + session = await sdk.beta.sessions.retrieve(self.session_id) + if self._value(session, "status") in {"terminated", "deleted"}: + return + turns += await self.run_pending(sdk, max_turns=None) + await asyncio.sleep(poll_interval) + + tasks = [asyncio.create_task(consume()), asyncio.create_task(heartbeat())] + try: + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() + finally: + stopped.set() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return turns + + def _restore_completed_inputs( + self, events: list[Any], *, incremental: bool = False + ) -> None: + """Reconcile completion on both full recovery and cross-Pod handoff.""" + if not incremental: + self._event_state.pending_inputs.clear() + self._event_state.reconciled_events.clear() + pending = self._event_state.pending_inputs + for event in events: + event_id = str(self._value(event, "id") or "") + sequence = self._event_seq(event) + key = f"seq:{sequence}" if sequence is not None else event_id + if key: + if key in self._event_state.reconciled_events: + continue + self._event_state.reconciled_events.add(key) + event_type = str(self._value(event, "type") or "") + if event_type == "span.model_request_end": + metadata = self._value(event, "metadata") or {} + response_id = self._value(metadata, "ark_response_id") + if response_id and not self._value(event, "is_error"): + self._event_state.ark_response_id = str(response_id) + if event_type == "user.message": + if event_id: + pending.append(event_id) + elif event_type in {"user.interrupt", "session.status_terminated"}: + self._completed_inputs.update(pending) + pending.clear() + elif event_type in _TERMINAL_EVENT_TYPES and event_type != "session.error": + if event_type == "session.status_idle": + reason = self._value(event, "stop_reason") + if self._value(reason, "type") not in {None, "end_turn"}: + continue + # A terminal event closes one turn. More than one user message + # may have accumulated while no worker was running, so marking + # the entire pending list complete here would silently drop all + # but the first queued turn. + if pending: + self._completed_inputs.add(pending.pop(0)) + + async def _run_turn(self, sdk: Any, user_event: Any) -> None: + span_id = "" + usage = { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } + + try: + prompt = self._message_text(self._value(user_event, "content")) + await self._ensure_local_session() + started = await send_session_events( + sdk, + self.session_id, + events=[ + {"type": "session.status_running"}, + {"type": "span.model_request_start"}, + ], + ) + for event in getattr(started, "data", None) or []: + if getattr(event, "type", None) == "span.model_request_start": + span_id = str(getattr(event, "id", "") or "") + break + if not span_id: + raise RuntimeError( + "events.send did not return the persisted span.model_request_start id" + ) + + message = types.Content(role="user", parts=[types.Part(text=prompt)]) + stream_ids: dict[str, str] = {} + try: + async with aclosing( + self.runner.run_async( + user_id=self.user_id, + session_id=self.session_id, + new_message=message, + run_config=RunConfig(streaming_mode=StreamingMode.SSE), + ) + ) as runner_events: + async for event in runner_events: + trace_event("received", "runner", self.session_id, event) + self._merge_usage(usage, getattr(event, "usage_metadata", None)) + if not getattr(event, "partial", False): + response_id = getattr(event, "interaction_id", None) + if response_id: + self._event_state.ark_response_id = str(response_id) + await self._publish_runner_event( + sdk, event, stream_ids=stream_ids + ) + except (Exception, asyncio.CancelledError): + try: + await asyncio.shield( + self._close_transient_streams(sdk, stream_ids, status="aborted") + ) + except Exception as cleanup_error: + logger.warning( + "Managed stream cleanup failed: %s", + type(cleanup_error).__name__, + ) + raise + await self._close_transient_streams(sdk, stream_ids, status="completed") + + await send_session_events( + sdk, + self.session_id, + events=[ + { + "type": "span.model_request_end", + "model_request_start_id": span_id, + "model_usage": usage, + **( + { + "metadata": { + "ark_response_id": self._event_state.ark_response_id + } + } + if self._event_state.ark_response_id + else {} + ), + }, + { + "type": "session.status_idle", + "stop_reason": {"type": "end_turn"}, + }, + ], + ) + except asyncio.CancelledError: + events: list[dict[str, Any]] = [] + if span_id: + events.append( + { + "type": "span.model_request_end", + "model_request_start_id": span_id, + "model_usage": usage, + "is_error": True, + } + ) + if events: + try: + await asyncio.shield( + send_session_events(sdk, self.session_id, events=events) + ) + except Exception as cleanup_error: + logger.warning( + "Managed cancellation event failed: %s", + type(cleanup_error).__name__, + ) + raise + except Exception as error: + logger.error("Managed Agents turn failed: %s", type(error).__name__) + events: list[dict[str, Any]] = [] + if span_id: + events.append( + { + "type": "span.model_request_end", + "model_request_start_id": span_id, + "model_usage": usage, + "is_error": True, + } + ) + events.append( + {"type": "session.error", "error": "Managed Agents turn failed"} + ) + + try: + await send_session_events(sdk, self.session_id, events=events) + except Exception as reporting_error: + logger.warning( + "Managed failure event failed: %s", type(reporting_error).__name__ + ) + raise + + async def _ensure_local_session(self) -> None: + service = self.runner.session_service + session = await service.get_session( + app_name=self.runner.app_name, + user_id=self.user_id, + session_id=self.session_id, + ) + if session is None: + session = await service.create_session( + app_name=self.runner.app_name, + user_id=self.user_id, + session_id=self.session_id, + ) + # A new Work may run in another process. Only the opaque response + # reference is restored; conversation inputs/outputs stay in Ark. + if self._event_state.ark_response_id: + await service.append_event( + session, + Event( + author=self.runner.agent.name, + interaction_id=self._event_state.ark_response_id, + ), + ) + + async def _publish_runner_event( + self, sdk: Any, event: Any, *, stream_ids: dict[str, str] + ) -> None: + partial = bool(getattr(event, "partial", False)) + content = getattr(event, "content", None) + if partial: + for part in getattr(content, "parts", None) or []: + text = str(getattr(part, "text", "") or "") + if not text: + continue + kind = ( + "thinking" if bool(getattr(part, "thought", False)) else "message" + ) + stream_id = stream_ids.get(kind) + if stream_id is None: + stream_id = f"sevt-{uuid.uuid4().hex}" + stream_ids[kind] = stream_id + await self._publish_transient( + sdk, + { + "type": f"agent.{kind}_stream_start", + f"{kind}_id": stream_id, + }, + ) + await self._publish_transient( + sdk, + { + "type": f"agent.{kind}_chunk", + f"{kind}_id": stream_id, + "delta": text, + }, + ) + return + + # Remote tools publish agent.tool_use before execution and the client + # canonicalizes user.tool_result into agent.tool_result. Do not duplicate + # those events when ADK later exposes its function-call bookkeeping. + if get_event_function_calls(event): + return + + output_events: list[dict[str, Any]] = [] + for part in getattr(content, "parts", None) or []: + text = str(getattr(part, "text", "") or "") + if not text: + continue + if bool(getattr(part, "thought", False)): + thinking_id = stream_ids.pop("thinking", None) + if thinking_id is not None: + await self._publish_transient( + sdk, + { + "type": "agent.thinking_stream_end", + "thinking_id": thinking_id, + "status": "completed", + }, + ) + else: + thinking_id = f"sevt-{uuid.uuid4().hex}" + output_events.append( + { + "type": "agent.thinking", + "thinking_id": thinking_id, + "text": text, + } + ) + else: + message_id = stream_ids.pop("message", None) + if message_id is not None: + await self._publish_transient( + sdk, + { + "type": "agent.message_stream_end", + "message_id": message_id, + "status": "completed", + }, + ) + else: + message_id = f"sevt-{uuid.uuid4().hex}" + output_events.append( + { + "type": "agent.message", + "message_id": message_id, + "content": [{"type": "text", "text": text}], + } + ) + if output_events: + await send_session_events(sdk, self.session_id, events=output_events) + + async def _publish_transient(self, sdk: Any, event: dict[str, Any]) -> None: + # Previews use the existing events.send route and SSE event delta wire + # format. The canonical event closes the preview; stream-end markers + # are local bookkeeping and need no extra event or endpoint. + typ = event["type"] + if typ.endswith("_stream_end") or typ == "agent.thinking_chunk": + return + thinking = typ.startswith("agent.thinking") + stream_id = event["thinking_id" if thinking else "message_id"] + if typ.endswith("_stream_start"): + event = { + "type": "event_start", + "event": { + "type": "agent.thinking" if thinking else "agent.message", + "id": stream_id, + }, + } + else: + event = { + "type": "event_delta", + "event_id": stream_id, + "delta": { + "type": "content_delta", + "index": 0, + "content": { + "type": "text", + "text": event["delta"], + }, + }, + } + event = { + "id": f"transient-{uuid.uuid4().hex}", + "processed_at": datetime.now(timezone.utc).isoformat(), + **event, + } + try: + trace_event("sent", "transient", self.session_id, event) + await sdk.post( + f"/v1/sessions/{self.session_id}/events", + cast_to=Dict[str, Any], + body={"events": [event]}, + ) + except Exception as error: + logger.warning( + "Managed Agents transient stream publish failed type=%s error_type=%s", + event.get("type"), + type(error).__name__, + ) + raise + + async def _close_transient_streams( + self, sdk: Any, stream_ids: dict[str, str], *, status: str + ) -> None: + for kind, stream_id in list(stream_ids.items()): + await self._publish_transient( + sdk, + { + "type": f"agent.{kind}_stream_end", + f"{kind}_id": stream_id, + "status": status, + }, + ) + stream_ids.pop(kind, None) + + @staticmethod + def _value(value: Any, name: str) -> Any: + return ( + value.get(name) if isinstance(value, dict) else getattr(value, name, None) + ) + + @classmethod + def _event_seq(cls, event: Any) -> int | None: + value = cls._value(event, "seq") + if value is None: + value = cls._value(event, "_seq") + try: + sequence = int(value) + except (TypeError, ValueError): + return None + return sequence if sequence >= 0 else None + + @staticmethod + def _message_text(content: Any) -> str: + blocks = content if isinstance(content, list) else list(content or []) + parts: list[str] = [] + for block in blocks: + block_type = ( + block.get("type") + if isinstance(block, dict) + else getattr(block, "type", None) + ) + if block_type != "text": + continue + text = ( + block.get("text") + if isinstance(block, dict) + else getattr(block, "text", None) + ) + if text: + parts.append(str(text)) + if not parts: + raise ValueError("user.message must contain at least one text block") + return "\n".join(parts) + + @staticmethod + def _merge_usage(target: dict[str, int], usage: Any) -> None: + if usage is None: + return + target["input_tokens"] += int(getattr(usage, "prompt_token_count", 0) or 0) + target["output_tokens"] += int(getattr(usage, "candidates_token_count", 0) or 0) + target["cache_read_input_tokens"] += int( + getattr(usage, "cached_content_token_count", 0) or 0 + ) diff --git a/veadk/runtime/managed_agents/sandbox.py b/veadk/runtime/managed_agents/sandbox.py new file mode 100644 index 000000000..8519bb11c --- /dev/null +++ b/veadk/runtime/managed_agents/sandbox.py @@ -0,0 +1,206 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VeADK agent whose tool calls run through a Self-Hosted Sandbox Runtime.""" + +import asyncio +import functools +import logging +import threading +from typing import Any + +from functools import lru_cache +from veadk.integrations.mpa.session_client import SelfHostSandboxClient +from veadk import Agent +from veadk.memory.short_term_memory import ShortTermMemory +from veadk.runtime import DispatchRuntimeProvider, ToolCall + +logger = logging.getLogger("veadk.sandbox_agent") + + +class SandboxSessionManager: + """Bind exactly one remote Managed Session to each VeADK session.""" + + def __init__(self) -> None: + self._clients: dict[str, SelfHostSandboxClient] = {} + self._active_turns: dict[str, int] = {} + self._lock = threading.Lock() + + def _new_client(self) -> SelfHostSandboxClient: + client = SelfHostSandboxClient() + client.session_id = None + return client + + def get(self, veadk_session_id: str) -> SelfHostSandboxClient: + key = veadk_session_id or "__default__" + with self._lock: + client = self._clients.get(key) + if client is None: + client = self._new_client() + self._clients[key] = client + return client + + def bind(self, veadk_session_id: str, client: SelfHostSandboxClient) -> None: + key = veadk_session_id or "__default__" + with self._lock: + self._clients[key] = client + + def create_remote_session(self, veadk_session_id: str) -> str: + client = self.get(veadk_session_id) + if not client.session_id: + client.create_session( + title=f"VeADK Self-Hosted Sandbox Session {veadk_session_id}" + ) + return client.session_id or "" + + def begin_turn(self, veadk_session_id: str) -> None: + """Track a local turn without writing a synthetic Session event.""" + key = veadk_session_id or "__default__" + client = self.get(key) + if not client.session_id: + self.create_remote_session(key) + + with self._lock: + self._active_turns[key] = self._active_turns.get(key, 0) + 1 + + def end_turn(self, veadk_session_id: str) -> None: + """Mark the remote Session idle after the last overlapping local turn.""" + key = veadk_session_id or "__default__" + with self._lock: + active_turns = self._active_turns.get(key, 0) + if active_turns <= 0: + return + if active_turns > 1: + self._active_turns[key] = active_turns - 1 + return + self._active_turns.pop(key, None) + client = self._clients.get(key) + if client: + client.post_status_idle() + + def release(self, veadk_session_id: str) -> None: + key = veadk_session_id or "__default__" + with self._lock: + client = self._clients.pop(key, None) + self._active_turns.pop(key, None) + if client: + client.post_status_idle() + + +sandbox_sessions = SandboxSessionManager() + + +def enable_sandbox_turn_lifecycle(runner: Any) -> Any: + """Wrap ``runner.run_async`` to mark the remote Session idle after a turn.""" + if getattr(runner, "_sandbox_turn_lifecycle_enabled", False): + return runner + + original_run_async = runner.run_async + + @functools.wraps(original_run_async) + async def run_async_with_sandbox_turn(*args: Any, **kwargs: Any): + session_id = str(kwargs.get("session_id") or "") + if not session_id: + async for event in original_run_async(*args, **kwargs): + yield event + return + + await asyncio.to_thread(sandbox_sessions.begin_turn, session_id) + try: + async for event in original_run_async(*args, **kwargs): + yield event + finally: + await asyncio.to_thread(sandbox_sessions.end_turn, session_id) + + runner.run_async = run_async_with_sandbox_turn + runner._sandbox_turn_lifecycle_enabled = True + return runner + + +def bash( + command: str | None = None, + restart: bool | None = None, + timeout_ms: int | None = None, +) -> dict: + """Execute a command or restart the persistent remote bash session.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +def read(file_path: str, view_range: list[int] | None = None) -> dict: + """Read file contents, optionally selecting an inclusive line range.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +def write(file_path: str, content: str) -> dict: + """Atomically write file contents in the remote sandbox.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +def edit( + file_path: str, old_string: str, new_string: str, replace_all: bool = False +) -> dict: + """Replace an exact string in a remote file.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +def glob(pattern: str, path: str | None = None) -> dict: + """Find files matching a glob pattern.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +def grep(pattern: str, path: str | None = None) -> dict: + """Search file contents using a regular expression.""" + raise RuntimeError("Remote tools must be intercepted by DispatchRuntimeProvider") + + +async def after_create_session(session: Any) -> None: + veadk_session_id = str(getattr(session, "id", "") or "") + remote_id = await asyncio.to_thread( + sandbox_sessions.create_remote_session, veadk_session_id + ) + logger.info( + "Remote Managed Session created: veadk=%s remote=%s", + veadk_session_id, + remote_id, + ) + + +async def dispatch_task(tool_call: ToolCall) -> dict: + client = sandbox_sessions.get(tool_call.session_id) + return await asyncio.to_thread( + client.dispatch_tool, + tool_call.name, + tool_call.arguments, + dispatch_id=tool_call.id, + ) + + +dispatch_runtime = DispatchRuntimeProvider(dispatch_task, dispatchable_tools=None) + + +@lru_cache(maxsize=1) +def get_default_agent() -> Agent: + """Build the optional demo Agent only when a caller explicitly requests it.""" + memory = ShortTermMemory(after_create_session_callback=after_create_session) + return Agent( + name="self_host_sandbox_agent", + description="An engineering agent dispatching tools to a self-hosted Worker.", + instruction=( + "You are an engineering assistant. Handle conversation and model " + "reasoning in VeADK and use the provided tools for sandbox operations." + ), + tools=[bash, read, write, edit, glob, grep], + short_term_memory=memory, + before_tool_callback=dispatch_runtime.before_tool_callback, + ) diff --git a/veadk/runtime/managed_agents/worker.py b/veadk/runtime/managed_agents/worker.py new file mode 100644 index 000000000..9fb4335e0 --- /dev/null +++ b/veadk/runtime/managed_agents/worker.py @@ -0,0 +1,1511 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the VeADK agent with remotely dispatched sandbox tools.""" + +import argparse +import asyncio +import inspect +import json +import logging +import os +import re +import signal +import socket +import tempfile +import uuid +from pathlib import Path +from typing import Any, NoReturn +from urllib.parse import urlsplit + +from anthropic.lib.tools import ToolError +from anthropic.lib.tools.agent_toolset import ( + AgentToolContext, + beta_agent_toolset_20260401, +) +from google.adk.tools.base_tool import BaseTool +from google.genai import types + +from veadk import Agent, Runner +from veadk.extensions import FeishuChannelExtension +from veadk.integrations.mpa.session_client import SelfHostSandboxClient +from veadk.integrations.mpa.session_resources import mcp_toolsets, skill_instructions +from veadk.memory.short_term_memory import ShortTermMemory +from veadk.runtime import ( + LocalRuntimeProvider, + RuntimeProvider, + ToolCall, +) +from veadk.runtime.managed_agents.events import iter_session_events, send_session_events +from veadk.runtime.managed_agents.loop import ManagedAgentEventState, ManagedAgentsLoop +from veadk.runtime.managed_agents.sandbox import ( + dispatch_runtime, + enable_sandbox_turn_lifecycle, + get_default_agent, + sandbox_sessions, +) + +logger = logging.getLogger(__name__) + +APP_NAME = "self_host_sandbox_demo" +_POLL_TRACE_LOGGER = "anthropic.managed_agent_poll" +_MODEL_TRACE_LOGGER = "anthropic.managed_agent_model" +_POLL_PATH = "/v1/model-work/poll" + + +def configure_managed_agent_poll_trace(*, worker_id: str, base_url: str) -> None: + """Enable the dedicated safe poll channel and announce its readiness.""" + if os.getenv("MANAGED_AGENT_POLL_TRACE", "").strip().lower() != "true": + return + logging.getLogger("anthropic._base_client").setLevel(logging.WARNING) + logging.getLogger("httpx").setLevel(logging.WARNING) + logger = logging.getLogger(_POLL_TRACE_LOGGER) + logger.setLevel(logging.INFO) + logger.propagate = False + if not any( + getattr(existing, "_managed_agent_poll_trace", False) + for existing in logger.handlers + ): + handler = logging.StreamHandler() + handler.setLevel(logging.INFO) + handler.setFormatter(logging.Formatter("%(message)s")) + handler._managed_agent_poll_trace = True + logger.addHandler(handler) + parsed = urlsplit(base_url) + host = ( + f"[{parsed.hostname}]" + if parsed.hostname and ":" in parsed.hostname + else parsed.hostname + ) + port = f":{parsed.port}" if parsed.port is not None else "" + origin = f"{parsed.scheme}://{host}{port}" if parsed.scheme and host else "" + safe_worker_id = ( + worker_id + if len(worker_id) <= 128 and re.fullmatch(r"[A-Za-z0-9_.:-]+", worker_id) + else "" + ) + logger.info( + json.dumps( + { + "account_scope": True, + "event": "poll_trace_ready", + "origin": origin, + "path": _POLL_PATH, + "worker_id": safe_worker_id, + }, + sort_keys=True, + separators=(",", ":"), + ) + ) + + +def configure_managed_agent_model_trace() -> None: + """Configure the logger at its owner so every Ark hook shares it.""" + from veadk.models.ark_llm import configure_managed_agent_model_trace as configure + + configure() + + +def _managed_tool_placeholder() -> NoReturn: + raise RuntimeError( + "Managed Agent tools must be intercepted by managed_work_tool_runtime" + ) + + +def managed_bash( + command: str | None = None, + restart: bool | None = None, + timeout_ms: int | None = None, +) -> dict[str, Any]: + """Run a command in the Managed Agent worker.""" + _managed_tool_placeholder() + + +def managed_read(file_path: str, view_range: list[int] | None = None) -> dict[str, Any]: + """Read a file from the Managed Agent worker.""" + _managed_tool_placeholder() + + +def managed_write(file_path: str, content: str) -> dict[str, Any]: + """Write a file in the Managed Agent worker.""" + _managed_tool_placeholder() + + +def managed_edit( + file_path: str, + old_string: str, + new_string: str, + replace_all: bool = False, +) -> dict[str, Any]: + """Replace text in a Managed Agent worker file.""" + _managed_tool_placeholder() + + +def managed_glob(pattern: str, path: str | None = None) -> dict[str, Any]: + """List paths matching a glob in the Managed Agent worker.""" + _managed_tool_placeholder() + + +def managed_grep(pattern: str, path: str | None = None) -> dict[str, Any]: + """Search files in the Managed Agent worker.""" + _managed_tool_placeholder() + + +for _tool, _name in ( + (managed_bash, "bash"), + (managed_read, "read"), + (managed_write, "write"), + (managed_edit, "edit"), + (managed_glob, "glob"), + (managed_grep, "grep"), +): + _tool.__name__ = _name + +_MANAGED_TOOLS = { + "bash": managed_bash, + "read": managed_read, + "write": managed_write, + "edit": managed_edit, + "glob": managed_glob, + "grep": managed_grep, +} +_LOCAL_BUILTIN_TOOLS = {"web_fetch", "web_search"} +_CANONICAL_MANAGED_TOOLS = ( + "bash", + "read", + "write", + "edit", + "glob", + "grep", + "web_fetch", + "web_search", +) +# Backward-compatible export used by the example's existing integrations/tests. + + +def managed_short_term_memory() -> ShortTermMemory: + """Keep per-Work ADK bookkeeping local to the Worker process.""" + return ShortTermMemory(backend="local") + + +def _field(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, dict): + return value.get(name, default) + return getattr(value, name, default) + + +def _enabled(config: Any, default: bool) -> bool: + value = _field(config, "enabled") + return default if value is None else bool(value) + + +def _contains_http_status(error: BaseException, statuses: set[int]) -> bool: + """Inspect nested SDK/TaskGroup failures without exposing their bodies.""" + status = getattr(error, "status_code", None) or getattr(error, "status", None) + try: + if status is not None and int(status) in statuses: + return True + except (TypeError, ValueError): + pass + return any( + _contains_http_status(child, statuses) + for child in getattr(error, "exceptions", ()) + if isinstance(child, BaseException) + ) + + +async def _run_with_identity_retry( + operation: Any, + credential_provider: Any | None, + *, + should_stop: Any = lambda: False, +) -> Any: + """Retry only startup auth propagation failures with a refreshed key.""" + auth_attempt = 0 + while True: + if should_stop(): + return 0 + try: + return await operation() + except BaseException as error: + if ( + credential_provider is None + or not _contains_http_status(error, {401, 403}) + or auth_attempt >= 5 + ): + raise + auth_attempt += 1 + credential_provider.invalidate() + delay = min(8.0, float(2 ** (auth_attempt - 1))) + print( + "MANAGED_AGENT_IDENTITY_RETRY " + f"attempt={auth_attempt} delay_seconds={delay:g}", + flush=True, + ) + await asyncio.sleep(delay) + if should_stop(): + return 0 + + +async def _run_dispatcher_with_identity_retry( + dispatcher_factory: Any, + credential_provider: Any | None, + *, + max_items: int | None, + max_concurrency: int, + should_stop: Any = lambda: False, +) -> int: + async def run_dispatcher() -> int: + if should_stop(): + return 0 + dispatcher = dispatcher_factory() + if should_stop(): + dispatcher.drain() + return 0 + return await dispatcher.run( + max_items=max_items, max_concurrency=max_concurrency + ) + + return await _run_with_identity_retry( + run_dispatcher, credential_provider, should_stop=should_stop + ) + + +class ManagedCustomTool(BaseTool): + """Expose an app-owned custom tool and wait for its Session result event.""" + + def __init__(self, definition: Any, sdk: Any, session_id: str) -> None: + self.input_schema = dict(_field(definition, "input_schema", {}) or {}) + self.sdk = sdk + self.session_id = session_id + super().__init__( + name=str(_field(definition, "name", "")), + description=str(_field(definition, "description", "") or ""), + ) + + def _get_declaration(self) -> types.FunctionDeclaration: + return types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema=self.input_schema, + ) + + async def run_async(self, *, args: dict[str, Any], tool_context: Any) -> Any: + tool_use_id = str( + getattr(tool_context, "function_call_id", "") or f"toolu_{uuid.uuid4().hex}" + ) + await send_session_events( + self.sdk, + self.session_id, + events=[ + { + "type": "agent.custom_tool_use", + "id": tool_use_id, + "name": self.name, + "input": args, + }, + { + "type": "session.status_idle", + "stop_reason": { + "type": "requires_action", + "action_type": "custom_tool_result", + "event_ids": [tool_use_id], + }, + }, + ], + ) + result = await _wait_for_session_event( + self.sdk, + self.session_id, + event_type="user.custom_tool_result", + link_field="custom_tool_use_id", + link_id=tool_use_id, + ) + await send_session_events( + self.sdk, self.session_id, events=[{"type": "session.status_running"}] + ) + content = _event_content_text(_field(result, "content")) + return { + "error" if bool(_field(result, "is_error", False)) else "result": content + } + + +async def _wait_for_session_event( + sdk: Any, + session_id: str, + *, + event_type: str | tuple[str, ...], + link_field: str, + link_id: str, +) -> Any: + timeout = _managed_action_timeout_seconds() + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + async for event in iter_session_events( + sdk, session_id, limit=1000, order="desc" + ): + if ( + _field(event, "type") + in ((event_type,) if isinstance(event_type, str) else event_type) + and _field(event, link_field) == link_id + ): + return event + await asyncio.sleep(0.5) + raise TimeoutError(f"timed out waiting for {event_type} for {link_id}") + + +def _event_content_text(content: Any) -> str: + if isinstance(content, str): + return content + values = [] + for block in content or []: + text = _field(block, "text") + if text: + values.append(str(text)) + return "\n".join(values) or "(no output)" + + +def _agent_tool_config_name(config: Any) -> str: + """Official SDK configs use the type discriminator; accept legacy names.""" + return str(_field(config, "type") or _field(config, "name") or "") + + +def _tool_permission_map(snapshot: Any) -> dict[str, str]: + policies: dict[str, str] = {} + for toolset in _field(snapshot, "tools", []) or []: + toolset_type = str(_field(toolset, "type", "")) + if toolset_type not in {"agent_toolset_20260401", "mcp_toolset"}: + continue + default = _field(toolset, "default_config", {}) or {} + default_policy = str( + _field( + _field(default, "permission_policy", {}) or {}, "type", "always_allow" + ) + ) + prefix = ( + f"mcp__{_field(toolset, 'mcp_server_name', '')}__" + if toolset_type == "mcp_toolset" + else "" + ) + for config in _field(toolset, "configs", []) or []: + name = ( + _agent_tool_config_name(config) + if toolset_type == "agent_toolset_20260401" + else str(_field(config, "name", "")) + ) + policy = str( + _field( + _field(config, "permission_policy", {}) or {}, + "type", + default_policy, + ) + ) + policies[prefix + name] = policy + policies[prefix + "*"] = default_policy + return policies + + +def managed_agent_config( + snapshot: Any, + *, + sdk: Any | None = None, + session_id: str | None = None, + skill_dirs: list[Path] | None = None, + credential_keys: Any = None, +) -> dict[str, Any]: + """Translate a frozen Managed Agent snapshot to a VeADK Agent config.""" + model = _field(snapshot, "model", "") + model_name = _field(model, "id", model) + if not isinstance(model_name, str) or not model_name: + model_name = os.getenv("MODEL_AGENT_NAME", "") + if not model_name: + raise ValueError("Managed Session agent snapshot has no model id") + + enabled_names: list[str] = [] + unsupported: list[str] = [] + for toolset in _field(snapshot, "tools", []) or []: + toolset_type = str(_field(toolset, "type", "") or "") + if toolset_type == "mcp_toolset": + continue + if toolset_type == "custom": + if sdk is None or not session_id: + unsupported.append( + f"custom tool {_field(toolset, 'name', '')} (missing Session client)" + ) + continue + if toolset_type != "agent_toolset_20260401": + unsupported.append(toolset_type or "unknown tool definition") + continue + default_config = _field(toolset, "default_config", {}) or {} + default_enabled = _enabled(default_config, True) + default_policy = _field( + _field(default_config, "permission_policy", {}) or {}, "type" + ) + configs = _field(toolset, "configs", []) or [] + configured: dict[str, Any] = { + _agent_tool_config_name(config): config for config in configs + } + names = set(configured) | ( + set(_CANONICAL_MANAGED_TOOLS) if default_enabled else set() + ) + for config in configs: + name = _agent_tool_config_name(config) + if not _enabled(config, default_enabled): + names.discard(name) + for name in sorted(names): + config = configured.get(name, {}) + enabled = _enabled(config, default_enabled) + if not enabled: + continue + policy = ( + _field(_field(config, "permission_policy", {}) or {}, "type") + or default_policy + ) + if policy not in {None, "always_allow", "always_ask"}: + unsupported.append(f"{name} (permission_policy={policy})") + enabled_names.append(name) + + unsupported.extend( + name + for name in enabled_names + if name not in _MANAGED_TOOLS and name not in _LOCAL_BUILTIN_TOOLS + ) + if unsupported: + raise ValueError( + f"Unsupported Managed Agent tools for VeADK: {', '.join(sorted(set(unsupported)))}" + ) + tools = list( + dict.fromkeys( + _MANAGED_TOOLS[name] for name in enabled_names if name in _MANAGED_TOOLS + ) + ) + if "web_fetch" in enabled_names: + from veadk.tools.builtin_tools.web_fetch import web_fetch + + tools.append(web_fetch) + if "web_search" in enabled_names: + from veadk.tools.builtin_tools.web_search import web_search + + tools.append(web_search) + resolved_instruction = str(_field(snapshot, "system", "") or "") + if skill_dirs: + resolved_instruction += skill_instructions(skill_dirs) + tools.extend(mcp_toolsets(snapshot, credential_keys=credential_keys)) + if sdk is not None and session_id: + tools.extend( + ManagedCustomTool(tool, sdk, session_id) + for tool in _field(snapshot, "tools", []) or [] + if _field(tool, "type") == "custom" + ) + raw_name = str(_field(snapshot, "name", "managed_agent")) + name = re.sub(r"[^A-Za-z0-9_]", "_", raw_name).strip("_") or "managed_agent" + if name[0].isdigit(): + name = f"agent_{name}" + model_base_url = _field(model, "base_url", None) + return { + **({"model_api_base": model_base_url} if model_base_url else {}), + "name": name, + "description": str(_field(snapshot, "description", "") or ""), + "instruction": resolved_instruction, + "model_name": model_name, + "enable_responses": True, + "enable_responses_cache": True, + "model_extra_config": { + "store": True, + "retry_expired_response": False, + # Conversation storage is independent of the paid cache service. + "extra_body": {"caching": {"type": "disabled"}}, + }, + "tools": tools, + "before_tool_callback": dispatch_runtime.before_tool_callback + if tools + else None, + } + + +async def resolve_session_model_key(session: Any, runtime_type: str) -> str: + """Require Session Identity unless a local worker explicitly opts in.""" + from veadk.runtime.managed_agents.identity import ( + RuntimeIdentityError, + session_model_credentials, + ) + + source = os.getenv("MANAGED_AGENT_MODEL_CREDENTIAL_SOURCE", "session") + if source not in {"session", "environment"}: + raise ValueError("Unknown Managed Agent model credential source") + if ( + source == "environment" + and runtime_type != "agentkit_runtime" + and _field(session, "ark_api_key_provider") is None + ): + key = os.getenv("MODEL_AGENT_API_KEY", "") + if not key: + raise RuntimeIdentityError("missing_model_key", False) + return key + return await session_model_credentials.get(session) + + +def managed_runner( + short_term_memory: ShortTermMemory, + snapshot: Any | None = None, + *, + before_tool_callback: Any | None = None, + sdk: Any | None = None, + session_id: str | None = None, + skill_dirs: list[Path] | None = None, + credential_keys: Any = None, + model_api_key: str | None = None, +) -> Runner: + if snapshot is None: + managed_agent = get_default_agent() + else: + config = managed_agent_config( + snapshot, + sdk=sdk, + session_id=session_id, + skill_dirs=skill_dirs, + credential_keys=credential_keys, + ) + if before_tool_callback is not None and config["tools"]: + config["before_tool_callback"] = before_tool_callback + if model_api_key is not None: + config["model_api_key"] = model_api_key + managed_agent = Agent(**config) + return Runner( + agent=managed_agent, + app_name=APP_NAME, + short_term_memory=short_term_memory, + ) + + +def _managed_tool_environment() -> dict[str, str]: + """Return a subprocess environment without control-plane credentials.""" + allowed = {"LANG", "LC_ALL", "PATH", "TERM", "TZ"} + return {key: value for key, value in os.environ.items() if key in allowed} + + +# Real execution budget for one tool invocation when the model does not ask for a +# specific timeout. 300s is the deployed contract; the outer wait +# (MANAGED_AGENT_ACTION_TIMEOUT_SECONDS) must exceed it so a long tool is never +# cut off before it can finish and be cleaned up. +MANAGED_TOOL_TIMEOUT_ENV = "MANAGED_AGENT_TOOL_TIMEOUT_SECONDS" +MANAGED_TOOL_TIMEOUT_DEFAULT_SECONDS = 300.0 + + +def _managed_tool_timeout_seconds() -> float: + """Tool execution budget: env override, otherwise the 300s default.""" + raw = os.getenv(MANAGED_TOOL_TIMEOUT_ENV, "").strip() + try: + value = float(raw) + except ValueError: + value = 0.0 + return value if value > 0 else MANAGED_TOOL_TIMEOUT_DEFAULT_SECONDS + + +# Cancellation and process-group cleanup need to fit between the tool budget and +# the outer wait that would give up on the tool result. +MANAGED_TOOL_CLEANUP_MARGIN_SECONDS = 60.0 + + +def _managed_action_timeout_seconds() -> float: + """Outer wait for a tool result, never below budget + cleanup margin.""" + raw = os.getenv("MANAGED_AGENT_ACTION_TIMEOUT_SECONDS", "3600").strip() + try: + configured = float(raw) + except ValueError: + configured = 3600.0 + return max( + configured, + _managed_tool_timeout_seconds() + MANAGED_TOOL_CLEANUP_MARGIN_SECONDS, + ) + + +def _with_default_tool_timeout(event: dict[str, Any]) -> dict[str, Any]: + """Annotate a published tool_use event with the deployment execution budget. + + The remote executor (sandbox ``ant`` worker) has no timeout option of its own; + it uses the ``timeout_ms`` carried by the event. Without this the remote path + silently falls back to the executor default and a long tool is cut short. + """ + if event.get("name") != "bash": + return event + payload = event.get("input") + if not isinstance(payload, dict): + return event + if payload.get("timeout") is not None or payload.get("timeout_ms") is not None: + return event + annotated = dict(event) + annotated["input"] = { + **payload, + "timeout_ms": int(_managed_tool_timeout_seconds() * 1000), + } + return annotated + + +async def aclose_runnable_tool(tool: object) -> None: + """Preserve SDK cleanup protocols without logging secret-bearing errors.""" + closer = getattr(tool, "aclose", None) or getattr(tool, "close", None) + if closer is not None: + try: + result = closer() + if inspect.isawaitable(result): + await result + except Exception as error: + logger.warning("Native tool cleanup failed: %s", type(error).__name__) + + # Official beta_tool/beta_async_tool retain their entered context here. + # Context-manager cleanup is additive to close/aclose, not an alternative. + context = getattr(tool, "_context_manager", None) + if context is not None: + try: + aexit = getattr(context, "__aexit__", None) + if aexit is not None: + await aexit(None, None, None) + else: + context.__exit__(None, None, None) + except Exception as error: + logger.warning("Native tool context exit failed: %s", type(error).__name__) + + +class NativeAgentTools: + """Own one SDK toolset and its persistent shell for a runner's lifetime.""" + + def __init__(self, workdir: Path) -> None: + self.context = AgentToolContext( + workdir=workdir, env=_managed_tool_environment() + ) + self.tools = { + tool.name: tool for tool in beta_agent_toolset_20260401(self.context) + } + + async def close(self) -> None: + failures: list[BaseException] = [] + for tool in self.tools.values(): + try: + await aclose_runnable_tool(tool) + except (Exception, asyncio.CancelledError) as error: + failures.append(error) + try: + await self.context.close() + except (Exception, asyncio.CancelledError) as error: + failures.append(error) + if failures: + raise failures[0] + + async def run(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: + tool = self.tools.get(name) + if tool is None: + raise ValueError(f"No native tool registered for {name!r}") + payload = _with_default_tool_timeout({"name": name, "input": arguments})[ + "input" + ] + try: + content = await tool.call(payload) + is_error = False + except ToolError as error: + content, is_error = error.content, True + text = _event_content_text(content) + return { + "status": "failed" if is_error else "completed", + "exit_code": 1 if is_error else 0, + "stdout": "" if is_error else text, + "stderr": text if is_error else "", + "content": content, + "is_error": is_error, + } + + +async def _run_managed_tool( + tool_name: str, + arguments: dict[str, Any], + *, + workdir: Path, + native_tools: NativeAgentTools | None = None, +) -> dict[str, Any]: + tools = native_tools or NativeAgentTools(workdir) + try: + return await tools.run(tool_name, arguments) + finally: + if native_tools is None: + await tools.close() + + +def mcp_result_content(result: Any) -> tuple[list[dict[str, Any]], bool]: + """Convert MCP CallToolResult to Managed Agents content blocks.""" + if hasattr(result, "model_dump"): + result = result.model_dump(mode="json", exclude_none=True) + is_error = isinstance(result, dict) and bool( + result.get("isError") or result.get("error") + ) + blocks = result.get("content") if isinstance(result, dict) else None + content: list[dict[str, Any]] = [] + if isinstance(blocks, list): + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text": + content.append({"type": "text", "text": block.get("text", "")}) + elif isinstance(block, dict) and block.get("type") == "image": + content.append( + { + "type": "image", + "source": { + "type": "base64", + "data": block["data"], + "media_type": block["mimeType"], + }, + } + ) + else: + # MCP resource/audio/link blocks have no identical Session + # variant; retain their payload as text instead of dropping it. + content.append( + {"type": "text", "text": json.dumps(block, ensure_ascii=False)} + ) + if not content: + payload = ( + result.get("structuredContent", result) + if isinstance(result, dict) + else result + ) + content = [ + { + "type": "text", + "text": payload + if isinstance(payload, str) + else json.dumps(payload, ensure_ascii=False, default=str), + } + ] + return content, is_error + + +def managed_work_tool_runtime( + sdk: Any, session_id: str, *, workdir: Path, snapshot: Any | None = None +) -> RuntimeProvider: + """Execute tools locally or await the configured remote sandbox worker.""" + + policies = _tool_permission_map(snapshot) if snapshot is not None else {} + local_runtime = LocalRuntimeProvider() + native_tools = NativeAgentTools(workdir) + + async def confirm(tool_call: ToolCall, event: dict[str, Any]) -> bool: + wildcard = ( + tool_call.name.rsplit("__", 1)[0] + "__*" if "__" in tool_call.name else "*" + ) + policy = policies.get(tool_call.name, policies.get(wildcard, "always_allow")) + if policy != "always_ask": + return True + if event["type"] in {"agent.tool_use", "agent.mcp_tool_use"}: + event["evaluated_permission"] = "ask" + await send_session_events( + sdk, + session_id, + events=[ + event, + { + "type": "session.status_idle", + "stop_reason": { + "type": "requires_action", + "action_type": "tool_confirmation", + "event_ids": [event["id"]], + }, + }, + ], + ) + verdict = await _wait_for_session_event( + sdk, + session_id, + event_type="user.tool_confirmation", + link_field="tool_use_id", + link_id=str(event["id"]), + ) + await send_session_events( + sdk, session_id, events=[{"type": "session.status_running"}] + ) + return _field(verdict, "result") == "allow" + + async def publish_result( + kind: str, tool_use_id: str, result: Any, is_error: bool + ) -> None: + content = ( + result + if isinstance(result, (str, list)) + else json.dumps(result, ensure_ascii=False, default=str) + ) + await send_session_events( + sdk, + session_id, + events=[ + { + "type": f"agent.{kind}_result", + f"{kind}_use_id": tool_use_id, + "content": content, + "is_error": is_error, + } + ], + ) + + async def execute(tool_call: ToolCall) -> Any: + tool_use_id = tool_call.id or f"toolu_{uuid.uuid4().hex}" + if isinstance(tool_call.tool, ManagedCustomTool): + return await local_runtime.execute(tool_call) + is_mcp = hasattr(tool_call.tool, "_mcp_tool") or hasattr( + tool_call.tool, "_mcp_session_manager" + ) + if is_mcp: + server_name = str( + getattr(tool_call.tool, "managed_agents_mcp_server_name", "mcp") + ) + prefix = f"mcp__{server_name}__" + original_name = ( + tool_call.name.removeprefix(prefix) + if tool_call.name.startswith(prefix) + else tool_call.name + ) + event = { + "type": "agent.mcp_tool_use", + "id": tool_use_id, + "mcp_server_name": server_name, + "name": original_name or tool_call.name, + "input": tool_call.arguments, + "evaluated_permission": "allow", + } + if not await confirm(tool_call, event): + await publish_result( + "mcp_tool", + tool_use_id, + [{"type": "text", "text": "MCP tool call denied by user"}], + True, + ) + return {"error": "tool call denied by user"} + if ( + policies.get( + tool_call.name, policies.get("mcp__" + server_name + "__*") + ) + != "always_ask" + ): + await send_session_events(sdk, session_id, events=[event]) + try: + result = await local_runtime.execute(tool_call) + except Exception: + # A tool failure is a model-visible result, not a failed turn. + result = {"error": "MCP tool execution failed"} + content, is_error = mcp_result_content(result) + await publish_result("mcp_tool", tool_use_id, content, is_error) + return result + + if ( + tool_call.name not in _MANAGED_TOOLS + and tool_call.name not in _LOCAL_BUILTIN_TOOLS + ): + return await local_runtime.execute(tool_call) + + event = { + "type": "agent.tool_use", + "id": tool_use_id, + "name": tool_call.name, + "input": tool_call.arguments, + "evaluated_permission": "allow", + } + if os.getenv("MANAGED_AGENT_TOOL_EXECUTION", "local") == "remote": + event = _with_default_tool_timeout(event) + if not await confirm(tool_call, event): + return {"error": "tool call denied by user"} + if event["evaluated_permission"] != "ask": + await send_session_events(sdk, session_id, events=[event]) + if os.getenv("MANAGED_AGENT_TOOL_EXECUTION", "local") == "remote": + result_event = await _wait_for_session_event( + sdk, + session_id, + event_type=("user.tool_result", "agent.tool_result"), + link_field="tool_use_id", + link_id=tool_use_id, + ) + content = _field(result_event, "content") + is_error = bool(_field(result_event, "is_error", False)) + if _field(result_event, "type") == "user.tool_result": + await publish_result("tool", tool_use_id, content, is_error) + return {"error" if is_error else "result": content} + try: + if tool_call.name in _MANAGED_TOOLS: + result = await _run_managed_tool( + tool_call.name, + tool_call.arguments, + workdir=workdir, + native_tools=native_tools, + ) + else: + result = await local_runtime.execute(tool_call) + except Exception as error: # noqa: BLE001 -- tool failures become result events + result = { + "status": "failed", + "exit_code": -1, + "stdout": "", + "stderr": str(error) + if isinstance(error, ToolError) + else "Tool execution failed", + } + if isinstance(result, dict) and "content" in result and "is_error" in result: + content, is_error = result["content"], result["is_error"] + elif ( + isinstance(result, dict) + and {"exit_code", "stdout", "stderr"} <= result.keys() + ): + content = "\n".join( + part + for part in ( + f"exit={result['exit_code']}", + str(result["stdout"]), + str(result["stderr"]), + ) + if part + ) + is_error = result.get("status") == "failed" + else: + content = ( + result + if isinstance(result, str) + else json.dumps(result, ensure_ascii=False, default=str) + ) + is_error = False + await send_session_events( + sdk, + session_id, + events=[ + { + "type": "agent.tool_result", + "tool_use_id": tool_use_id, + "content": content, + "is_error": is_error, + } + ], + ) + return result + + class ManagedWorkRuntime(RuntimeProvider): + def __init__(self) -> None: + super().__init__(name="managed_agents_work_runtime") + + async def execute(self, tool_call: ToolCall) -> Any: + return await execute(tool_call) + + async def close(self) -> None: + await native_tools.close() + + return ManagedWorkRuntime() + + +async def close_managed_runner(runner: Runner) -> None: + failures: list[BaseException] = [] + for tool in getattr(runner.agent, "tools", []) or []: + close = getattr(tool, "close", None) + if not callable(close): + continue + try: + result = close() + if inspect.isawaitable(result): + await result + except (Exception, asyncio.CancelledError) as error: + failures.append(error) + if failures: + raise failures[0] + + +async def serve_managed_agent_worker( + *, + worker_id: str | None = None, + max_work_items: int | None = None, + readiness_probe: bool = False, +) -> int: + """Poll Task Server work and execute claimed sessions with VeADK.""" + session_client = SelfHostSandboxClient() + short_term_memory = managed_short_term_memory() + resolved_worker_id = worker_id or f"{socket.gethostname()}-{uuid.uuid4().hex[:12]}" + if session_client.runtime_type == "agentkit_runtime": + session_client._default_headers["Anthropic-Worker-ID"] = resolved_worker_id + default_workdir = Path(tempfile.gettempdir()) / "veadk-managed-agents" / "workspace" + workdir = Path(os.getenv("MANAGED_AGENT_WORKDIR", default_workdir)).resolve() + workdir.mkdir(parents=True, exist_ok=True) + ready_file = Path( + os.getenv("MANAGED_AGENT_READY_FILE", "/tmp/managed-agent-worker-ready") + ) + ready_file.unlink(missing_ok=True) + event_states: dict[str, ManagedAgentEventState] = {} + max_event_states = int(os.getenv("MANAGED_AGENT_EVENT_STATE_CACHE_SIZE", "4096")) + if max_event_states < 1: + raise ValueError("MANAGED_AGENT_EVENT_STATE_CACHE_SIZE must be positive") + + def event_state_for(session_id: str) -> ManagedAgentEventState: + state = event_states.get(session_id) + if state is not None: + # Dict insertion order provides a small LRU without another cache + # dependency. Moving the key never changes the shared state object. + event_states.pop(session_id) + event_states[session_id] = state + return state + + if len(event_states) >= max_event_states: + for cached_session_id, cached_state in tuple(event_states.items()): + if not cached_state.lock.locked(): + event_states.pop(cached_session_id) + break + state = ManagedAgentEventState() + event_states[session_id] = state + return state + + try: + async with session_client.create_async_client() as sdk: + configure_managed_agent_poll_trace( + worker_id=resolved_worker_id, base_url=str(sdk.base_url) + ) + configure_managed_agent_model_trace() + if readiness_probe: + + async def poll_readiness() -> Any: + return await sdk.beta.environments.work.poll( + session_client.environment_id, + block_ms=1, + anthropic_worker_id=resolved_worker_id, + extra_headers=session_client._default_headers, + ) + + pending = await _run_with_identity_retry( + poll_readiness, session_client._identity_credentials + ) + if pending is not None: + raise RuntimeError( + "readiness probe requires an empty isolated work queue" + ) + + async def handle(work_item: Any, scoped_sdk: Any) -> None: + if getattr(getattr(work_item, "data", None), "type", None) != "session": + return + session_id = str(work_item.data.id) + print( + f"MANAGED_AGENT_WORK_CLAIM worker_id={resolved_worker_id} " + f"work_id={work_item.id} session_id={session_id}", + flush=True, + ) + phase = "work_metadata" + try: + await scoped_sdk.beta.environments.work.update( + work_item.id, + environment_id=work_item.environment_id, + metadata={"managed_agent_worker_id": resolved_worker_id}, + ) + phase = "session_retrieve" + session = await scoped_sdk.beta.sessions.retrieve(session_id) + from veadk.integrations.mpa.session_resources import ( + cleanup_session_skills, + materialize_session_skills, + session_workdir, + ) + + phase = "session_resources" + item_workdir = session_workdir(workdir, session_id) + downloaded_skills: list[Path] = [] + try: + downloaded_skills = await materialize_session_skills( + session, item_workdir, client=scoped_sdk + ) + phase = "model_identity" + model_api_key = await resolve_session_model_key( + session, session_client.runtime_type + ) + phase = "runner_create" + tool_runtime = managed_work_tool_runtime( + scoped_sdk, + session_id, + workdir=item_workdir, + snapshot=session.agent, + ) + runner = None + try: + runner = managed_runner( + short_term_memory, + session.agent, + before_tool_callback=tool_runtime.before_tool_callback, + sdk=scoped_sdk, + session_id=session_id, + skill_dirs=downloaded_skills, + model_api_key=model_api_key, + credential_keys=_field(session, "credential_keys", []), + ) + phase = "events_run" + await ManagedAgentsLoop( + runner=runner, + session_id=session_id, + event_state=event_state_for(session_id), + ).run_pending(scoped_sdk, max_turns=1) + finally: + try: + if runner is not None: + await close_managed_runner(runner) + finally: + await tool_runtime.close() + finally: + await cleanup_session_skills(downloaded_skills) + except asyncio.CancelledError: + raise + except Exception as error: + if session_client.runtime_type == "agentkit_runtime": + from veadk.runtime.managed_agents.identity import ( + RuntimeIdentityError, + ) + + safe_error = ( + error + if isinstance(error, RuntimeIdentityError) + else RuntimeIdentityError( + "work_execution_failed", + bool(getattr(error, "retryable", False)), + ) + ) + print( + f"MANAGED_AGENT_WORK_FAILURE phase={phase} " + f"error_type={type(error).__name__} " + f"http_status={getattr(error, 'status_code', None)}", + flush=True, + ) + await send_session_events( + scoped_sdk, + session_id, + events=[ + {"type": "session.error", "error": str(safe_error)} + ], + ) + raise safe_error from None + await send_session_events( + scoped_sdk, + session_id, + events=[ + { + "type": "session.error", + "error": "Managed Work execution failed", + } + ], + ) + raise + + from veadk.runtime.managed_agents.dispatcher import ( + EnvironmentWorkDispatcher, + ) + + concurrency = int(os.getenv("MANAGED_AGENT_WORK_CONCURRENCY", "1")) + if concurrency < 1: + raise ValueError("MANAGED_AGENT_WORK_CONCURRENCY must be positive") + loop = asyncio.get_running_loop() + + dispatcher: Any | None = None + readiness_tasks: list[asyncio.Task] = [] + + async def announce_ready(instance: Any) -> None: + await instance.ready.wait() + if shutdown_requested: + return + ready_file.parent.mkdir(parents=True, exist_ok=True) + ready_file.write_text(resolved_worker_id + "\n") + print( + f"MANAGED_AGENT_WORKER_READY worker_id={resolved_worker_id} " + f"environment_id={session_client.environment_id} " + f"runtime_type={session_client.runtime_type or 'default'}", + flush=True, + ) + + shutdown_requested = False + + def drain_worker() -> None: + nonlocal shutdown_requested + shutdown_requested = True + ready_file.unlink(missing_ok=True) + if dispatcher is not None: + dispatcher.drain() + print( + f"MANAGED_AGENT_WORKER_DRAIN worker_id={resolved_worker_id}", + flush=True, + ) + + for signum in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(signum, drain_worker) + try: + + def dispatcher_factory() -> Any: + nonlocal dispatcher + dispatcher = EnvironmentWorkDispatcher( + sdk, + handler=handle, + account_work=os.getenv( + "MANAGED_AGENT_WORK_SCOPE", "environment" + ) + == "account", + environment_id=session_client.environment_id, + environment_key=session_client.bearer_token, + worker_id=resolved_worker_id, + extra_headers=session_client._default_headers, + ) + readiness_tasks.append( + asyncio.create_task(announce_ready(dispatcher)) + ) + return dispatcher + + return await _run_dispatcher_with_identity_retry( + dispatcher_factory, + getattr(session_client, "_identity_credentials", None), + max_items=max_work_items, + max_concurrency=concurrency, + should_stop=lambda: shutdown_requested, + ) + finally: + for task in readiness_tasks: + task.cancel() + await asyncio.gather(*readiness_tasks, return_exceptions=True) + for signum in (signal.SIGINT, signal.SIGTERM): + loop.remove_signal_handler(signum) + finally: + ready_file.unlink(missing_ok=True) + + +async def serve_claimed_managed_agent_work() -> int: + """Serve a claimed session until the external dispatcher reclaims the sandbox.""" + session_client = SelfHostSandboxClient() + work_id = os.getenv("MA_WORK_ID") or os.getenv("ANTHROPIC_WORK_ID") + session_id = os.getenv("MA_SESSION_ID") or os.getenv("ANTHROPIC_SESSION_ID") + if not work_id or not session_id: + raise ValueError( + "MA_WORK_ID/ANTHROPIC_WORK_ID and MA_SESSION_ID/ANTHROPIC_SESSION_ID " + "are required for claimed Work mode" + ) + + root = Path(os.getenv("MANAGED_AGENT_WORKDIR", "/workspace")).resolve() + root.mkdir(parents=True, exist_ok=True) + short_term_memory = managed_short_term_memory() + async with session_client.create_async_client() as sdk: + session = await sdk.beta.sessions.retrieve(session_id) + from veadk.integrations.mpa.session_resources import ( + cleanup_session_skills, + materialize_session_skills, + session_workdir, + ) + + item_workdir = session_workdir(root, session_id) + downloaded_skills: list[Path] = [] + try: + downloaded_skills = await materialize_session_skills( + session, item_workdir, client=sdk + ) + model_api_key = await resolve_session_model_key( + session, session_client.runtime_type + ) + tool_runtime = managed_work_tool_runtime( + sdk, session_id, workdir=item_workdir, snapshot=session.agent + ) + runner = None + try: + runner = managed_runner( + short_term_memory, + session.agent, + before_tool_callback=tool_runtime.before_tool_callback, + sdk=sdk, + session_id=session_id, + skill_dirs=downloaded_skills, + model_api_key=model_api_key, + credential_keys=_field(session, "credential_keys", []), + ) + turns = await ManagedAgentsLoop( + runner=runner, session_id=session_id + ).run_claimed( + sdk, + environment_id=session_client.environment_id, + work_id=work_id, + last_heartbeat=os.getenv("MA_LATEST_HEARTBEAT_AT") + or "NO_HEARTBEAT", + ) + finally: + try: + if runner is not None: + await close_managed_runner(runner) + finally: + await tool_runtime.close() + finally: + await cleanup_session_skills(downloaded_skills) + + print( + f"MANAGED_AGENT_CLAIMED_WORK_COMPLETE work_id={work_id} " + f"session_id={session_id} turns={turns}", + flush=True, + ) + return turns + + +async def serve_feishu_channel(stop_event: asyncio.Event | None = None) -> None: + """Serve Feishu conversations until the process receives a stop signal.""" + runner = enable_sandbox_turn_lifecycle( + Runner(agent=get_default_agent(), app_name=APP_NAME) + ) + channel = FeishuChannelExtension( + runner=runner, + streaming=True, + show_thinking=True, + show_tool_calls=True, + show_tool_results=True, + separate_tool_call_cards=True, + separate_thinking_card=True, + create_topic=True, + ) + loop = asyncio.get_running_loop() + shutdown_event = stop_event or asyncio.Event() + + if stop_event is None: + for signal_number in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(signal_number, shutdown_event.set) + except NotImplementedError: # pragma: no cover - Windows fallback + signal.signal(signal_number, lambda *_: shutdown_event.set()) + + channel.start(loop) + print("Feishu Channel is running. Press Ctrl+C to stop.") + try: + await shutdown_event.wait() + finally: + await channel.shutdown() + + +async def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--prompt", + default=( + "Use the bash tool to run: printf 'veadk-self-host-ok'. " + "Then reply with the exact output." + ), + ) + parser.add_argument("--session-id", default=None) + parser.add_argument( + "--managed-agent-loop", + action="store_true", + help=( + "Consume one pending user.message from the existing Managed Session with " + "local bookkeeping, then exit." + ), + ) + parser.add_argument( + "--managed-agent-worker", + action="store_true", + help="Continuously poll and execute Managed Agent session work.", + ) + parser.add_argument( + "--managed-agent-work-item", + action="store_true", + help="Execute one WorkItem already claimed by an external dispatcher.", + ) + parser.add_argument("--worker-id", default=None) + parser.add_argument("--max-work-items", type=int, default=None) + parser.add_argument( + "--managed-agent-readiness-probe", + action="store_true", + help="Require one successful empty work poll before logging worker readiness.", + ) + parser.add_argument( + "--managed-agent-model-trace-probe", + action="store_true", + help="Run one controlled Ark Responses request through the real worker import path.", + ) + parser.add_argument("--model-trace-probe-url", default=None) + parser.add_argument( + "--feishu", + action="store_true", + help="Keep running and serve conversations through the Feishu bot channel.", + ) + args = parser.parse_args() + + if args.managed_agent_model_trace_probe: + if not args.model_trace_probe_url: + parser.error( + "--managed-agent-model-trace-probe requires --model-trace-probe-url" + ) + from veadk.models.ark_llm import ArkLlm + + configure_managed_agent_model_trace() + model = ArkLlm( + model="openai/managed-agent-trace-probe", + api_base=args.model_trace_probe_url, + api_key="managed-agent-trace-probe", # gitleaks:allow controlled test credential + enable_responses_cache=False, + ) + responses = [ + response + async for response in model.generate_content_via_responses( + { + "model": "openai/managed-agent-trace-probe", + "api_base": args.model_trace_probe_url, + "api_key": "managed-agent-trace-probe", # gitleaks:allow controlled test credential + "input": [], + "tools": [{"name": "bash", "type": "function"}], + "tool_choice": {"type": "function", "name": "bash"}, + } + ) + ] + calls = [ + call + for response in responses + for call in response.content.parts + if getattr(call, "function_call", None) is not None + ] + if len(calls) != 1 or calls[0].function_call.name != "bash": + raise RuntimeError("controlled model trace probe did not return bash") + print("MANAGED_AGENT_MODEL_TRACE_PROBE_OK tool=bash", flush=True) + return + + if args.managed_agent_loop: + session_client = SelfHostSandboxClient(session_id=args.session_id) + if not session_client.session_id: + parser.error( + "--managed-agent-loop requires --session-id, ANTHROPIC_SESSION_ID, or SANDBOX_SESSION_ID" + ) + sandbox_sessions.bind(session_client.session_id, session_client) + runner = managed_runner(managed_short_term_memory()) + turns = await ManagedAgentsLoop( + runner=runner, + session_client=session_client, + ).run(max_turns=1) + print(f"Managed Agents loop completed {turns} turn(s).") + return + + if args.managed_agent_worker: + count = await serve_managed_agent_worker( + worker_id=args.worker_id, + max_work_items=args.max_work_items, + readiness_probe=args.managed_agent_readiness_probe, + ) + print(f"Managed Agents worker completed {count} work item(s).") + return + + if args.managed_agent_work_item: + turns = await serve_claimed_managed_agent_work() + print(f"Managed Agents claimed Work completed {turns} turn(s).") + return + + if args.feishu: + await serve_feishu_channel() + return + + session_id = args.session_id or f"veadk-{uuid.uuid4()}" + runner = enable_sandbox_turn_lifecycle( + Runner(agent=get_default_agent(), app_name=APP_NAME) + ) + output = await runner.run(messages=args.prompt, session_id=session_id) + print(output) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/veadk/skills/ma_infra.py b/veadk/skills/ma_infra.py new file mode 100644 index 000000000..b3d272cee --- /dev/null +++ b/veadk/skills/ma_infra.py @@ -0,0 +1,312 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Client for the ma-infra ``/v1/skills`` interface. + +The ma-infra control plane owns the customer skill registry (each authenticated +account maps to one default skill space) and exposes the Anthropic-compatible +``/v1/skills`` REST API. This client lets runtime consumers read, +download and publish skills through that interface instead of talking to the +AgentKit Skill TOP service directly. + +Connection parameters come from the environment, which the ma-infra dispatcher +injects when it starts the worker process: + +- ``MA_SKILLS_BASE_URL`` — ma-infra base URL (e.g. ``https://ma-infra.example``). +- ``MA_SKILLS_API_KEY`` — API key sent as ``x-api-key`` (or ``Authorization``). +- ``MA_SKILLS_ACCOUNT_ID`` — account id sent as ``X-Top-Account-Id``. + +All methods are synchronous; the async helpers wrap them with +:func:`asyncio.to_thread` so blocking HTTP stays off the event loop. +""" + +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from typing import Any + +import httpx + +# Mirror the ma-infra content size cap (internal/skills/top.go maxContentBytes). +MAX_CONTENT_BYTES = 32 * 1024 * 1024 + +DEFAULT_SKILL_SPACE_ID = "ma_infra_customer_skills" +DEFAULT_SKILL_SPACE_DESCRIPTION = "ma-infra customer skill space" + + +class MaInfraSkillError(RuntimeError): + """A failed or rejected ma-infra skills request.""" + + def __init__( + self, + message: str, + *, + status_code: int = 502, + code: str = "api_error", + retryable: bool = False, + ) -> None: + super().__init__(message) + self.status_code = status_code + self.code = code + self.retryable = retryable + + +@dataclass(frozen=True) +class MaInfraSkillConfig: + """Connection configuration for the ma-infra skills interface.""" + + base_url: str + api_key: str = "" + account_id: str = "" + timeout: float = 30.0 + + @classmethod + def from_env(cls, env: dict[str, str] | None = None) -> "MaInfraSkillConfig | None": + """Read configuration from the environment, or ``None`` when unset.""" + source = os.environ if env is None else env + base_url = (source.get("MA_SKILLS_BASE_URL") or "").strip().rstrip("/") + if not base_url: + return None + return cls( + base_url=base_url, + api_key=(source.get("MA_SKILLS_API_KEY") or "").strip(), + account_id=(source.get("MA_SKILLS_ACCOUNT_ID") or "").strip(), + ) + + +class MaInfraSkillsClient: + """Thin HTTP client over the ma-infra ``/v1/skills`` REST API.""" + + def __init__( + self, config: MaInfraSkillConfig, *, transport: Any | None = None + ) -> None: + self._config = config + self._http = httpx.Client( + base_url=config.base_url, + timeout=config.timeout, + follow_redirects=False, + transport=transport, + ) + + def close(self) -> None: + self._http.close() + + # -- auth ------------------------------------------------------------- + + def _headers(self) -> dict[str, str]: + headers: dict[str, str] = {"Content-Type": "application/json"} + if self._config.api_key: + headers["x-api-key"] = self._config.api_key + if self._config.account_id: + headers["X-Top-Account-Id"] = self._config.account_id + return headers + + def _request( + self, + method: str, + path: str, + *, + params: dict[str, Any] | None = None, + json: Any | None = None, + files: Any | None = None, + data: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + ) -> httpx.Response: + merged = self._headers() + if files is not None: + # multipart requests carry their own content type. + merged.pop("Content-Type", None) + if headers: + merged.update(headers) + try: + response = self._http.request( + method, + path, + params=params, + json=json, + files=files, + data=data, + headers=merged, + ) + except httpx.HTTPError as error: + raise MaInfraSkillError( + "ma-infra skills request failed", + status_code=502, + retryable=True, + ) from error + if response.status_code // 100 != 2: + raise self._error(response) + return response + + @staticmethod + def _error(response: httpx.Response) -> MaInfraSkillError: + status = response.status_code + code, message = "api_error", f"ma-infra skills request failed (HTTP {status})" + try: + payload = response.json() + error = payload.get("error") or payload.get("Error") or {} + if isinstance(error, dict): + code = str( + error.get("type") or error.get("code") or error.get("Code") or code + ) + message = str(error.get("message") or error.get("Message") or message) + except Exception: # noqa: BLE001 - never fail on a malformed error body + pass + return MaInfraSkillError( + message, + status_code=status, + code=code, + retryable=status >= 500 or status == 429, + ) + + # -- skills (v1 interface) -------------------------------------------- + + def list_skills( + self, + *, + limit: int = 100, + page: int = 1, + source: str = "custom", + ) -> list[dict[str, Any]]: + """Return custom skills for the configured account.""" + return self.list_skills_page(limit=limit, page=page, source=source)["data"] + + def list_skills_page( + self, + *, + limit: int = 100, + page: int = 1, + source: str = "custom", + ) -> dict[str, Any]: + """Return one skills page together with server pagination metadata.""" + response = self._request( + "GET", + "/v1/skills", + params={"limit": limit, "page": page, "source": source}, + ) + payload = response.json() + return { + "data": list(payload.get("data") or []), + "next_page": payload.get("next_page"), + "total_count": payload.get("total_count", payload.get("totalCount")), + } + + def get_skill(self, skill_id: str) -> dict[str, Any]: + """Return one skill's metadata (display_name, description, latest_version_id).""" + response = self._request("GET", f"/v1/skills/{skill_id}") + return response.json() + + def get_version(self, skill_id: str, version: str) -> dict[str, Any]: + """Return one skill version (name, description, concrete ``id``). + + ``version`` may be ``"latest"``; the returned ``id`` is always a + concrete version id usable for content download. + """ + response = self._request("GET", f"/v1/skills/{skill_id}/versions/{version}") + return response.json() + + def download_content(self, skill_id: str, version: str) -> bytes: + """Download a skill version's zip archive. + + ``version`` may be ``"latest"``; a concrete version id is resolved first + because the content endpoint only serves pinned versions. + """ + version = (version or "").strip() + if not version or version == "latest": + meta = self.get_version(skill_id, "latest") + version = str(meta.get("id") or meta.get("version") or "") + if not version: + raise MaInfraSkillError( + f"ma-infra skill {skill_id} has no downloadable version", + status_code=404, + code="skill_version_not_found", + ) + response = self._request( + "GET", f"/v1/skills/{skill_id}/versions/{version}/content" + ) + content = response.content + if len(content) > MAX_CONTENT_BYTES: + raise MaInfraSkillError( + f"ma-infra skill {skill_id} archive exceeds 32 MiB", + status_code=413, + code="skill_archive_too_large", + ) + return content + + def create_skill(self, display_name: str, archive: bytes) -> dict[str, Any]: + """Publish a new custom skill from a zip archive.""" + if not display_name: + raise MaInfraSkillError("skill display_name is required", status_code=400) + response = self._request( + "POST", + "/v1/skills", + files={"archive": ("skill.zip", archive, "application/zip")}, + data={"display_name": display_name}, + ) + return response.json() + + def create_version( + self, skill_id: str, display_name: str, archive: bytes + ) -> dict[str, Any]: + """Publish a new version of an existing custom skill.""" + response = self._request( + "POST", + f"/v1/skills/{skill_id}/versions", + files={"archive": ("skill.zip", archive, "application/zip")}, + data={"display_name": display_name}, + ) + return response.json() + + def delete_skill(self, skill_id: str) -> dict[str, Any]: + """Delete a custom skill.""" + response = self._request("DELETE", f"/v1/skills/{skill_id}") + return response.json() + + # -- async helpers ---------------------------------------------------- + + async def alist_skills(self, **kwargs: Any) -> list[dict[str, Any]]: + return await asyncio.to_thread(self.list_skills, **kwargs) + + async def aget_skill(self, skill_id: str) -> dict[str, Any]: + return await asyncio.to_thread(self.get_skill, skill_id) + + async def aget_version(self, skill_id: str, version: str) -> dict[str, Any]: + return await asyncio.to_thread(self.get_version, skill_id, version) + + async def adownload_content(self, skill_id: str, version: str) -> bytes: + return await asyncio.to_thread(self.download_content, skill_id, version) + + async def acreate_skill(self, display_name: str, archive: bytes) -> dict[str, Any]: + return await asyncio.to_thread(self.create_skill, display_name, archive) + + async def acreate_version( + self, skill_id: str, display_name: str, archive: bytes + ) -> dict[str, Any]: + return await asyncio.to_thread( + self.create_version, skill_id, display_name, archive + ) + + async def adelete_skill(self, skill_id: str) -> dict[str, Any]: + return await asyncio.to_thread(self.delete_skill, skill_id) + + +__all__ = [ + "DEFAULT_SKILL_SPACE_ID", + "DEFAULT_SKILL_SPACE_DESCRIPTION", + "MaInfraSkillConfig", + "MaInfraSkillError", + "MaInfraSkillsClient", +]