mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-11-11 04:31:34 -10:00
feat(backend/api): External API v2 (#12206)
Adds v2 of the platform's public API at `/api/v2`: REST endpoints for a user's graphs, library agents, folders, presets and webhook triggers, runs and their human-in-the-loop reviews, schedules, workspace files, credits, credentials and marketplace submissions, plus an MCP server exposing AutoPilot's tools, all callable with an API key or an OAuth token and gated by per-resource scopes. - Resolves #11678 - Partly based on #12284 ### Why - v1 can create and execute a graph, execute a block, read the marketplace and manage integration credentials, and nothing else. - An app signing users in through OAuth (#11617) cannot list their agents, run a saved preset, set up a webhook trigger, follow, stop or share a run, answer a review, schedule a run, or read their files and spend through v1, so it cannot use the platform as its backend. ### What New surface, in `backend.api.external.v2` (too much to review whole; read `mcp_server.py`, `mcp_calls.py`, `mcp_tenancy.py` and `library/presets.py`, and scan `models.py`): - Resources: `/graphs`, `/library/agents`, `/library/folders`, `/library/presets`, `/runs`, `/runs/reviews`, `/schedules`, `/files`, `/credits`, `/integrations/credentials`, `/marketplace`, `/blocks`, `/search`, `/me` - MCP server at `/api/v2/mcp`, with its RFC 9728 metadata at `/.well-known/oauth-protected-resource/api/v2/mcp`: AutoPilot tools that opt in through `allow_external_use`; `tools/list` shows only the tools the caller's scopes allow; a refused or failed call returns an `isError` result - 13 new `APIKeyPermission` scopes, one per resource and verb (`schema.prisma`, migration `20261008140000_add_v2_api_key_permissions`; `READ_STORE` was already on `dev`), each described on the OAuth consent screen and in the API key dialog - v1 stays at `/external-api/v1`, deprecated: it stops working on 2027-02-01, after its last day of service on 2027-01-31 - v2 guides under `docs/platform/integrating/`, and the v2 OpenAPI spec published to GitBook from `master` (`docs-gitbook-openapi.yml`) Who can call what: - Every route resolves its credential to one organization and team; a row in another organization is a 404, over REST and MCP alike (`mcp_tenancy.py` checks every id argument of every exposed tool before it runs). - A banned account's API key or OAuth token is refused on every path: `X-API-Key`, an API key sent as a Bearer token (how MCP clients send one), an OAuth token, and the MCP server's own token check. - A graph someone else owns, a colleague's or a marketplace agent's, is read as its shared copy, as on the internal routes: no owner ids, credential references or secret inputs. - A colleague's run or graph in the caller's organization is readable but not writable: `DELETE /runs/{id}`, `POST /runs/{id}/stop`, `PUT /graphs/{id}` and `PUT /graphs/{id}/versions/active` answer 403 before acting. - An account without an active plan gets a 402 (`payment_required`) from every route that starts work, stores files or calls out, and from every MCP tool call. Reads, billing, deletes, stopping a run, sharing and organising stay open, as on the internal API since #15676. - MCP tool calls run outside the action approval gate: an MCP session is an automation's, so nobody is asked to approve a card, and the credential's scopes are the authorization, as in REST. A run an MCP call starts is tagged an API run, as v1's are, so its reviews come back through `/runs/reviews` rather than a chat that was never saved. Runs, presets and triggers: - `POST /library/agents/{id}/runs` and `POST /library/presets/{id}/execute` take an `Idempotency-Key`: a repeat returns the run the first request started, a repeat while that one is in flight is a 409, and the key reused on another agent, preset or body is a 422. - Both answer 422 for an input the graph does not have, or a required one omitted or `null`, with one `details.errors` entry per input (for a preset, unless it is a trigger, whose inputs are its block's settings); `POST /schedules` does the same, and schedules the active version when the request names none. - Presets: list, get, create, update, delete and execute. A preset must name a graph in the key's organization. One created for a graph that a single hired expert installed belongs to that expert, as on the internal routes, and so do triggers and schedules. Executing one runs for the expert it belongs to (409 while that expert is paused) and lands in the preset's organization. - Webhook triggers: `POST /library/presets/setup-trigger` registers the webhook and returns `webhook_url`, where a hand-configured service sends its events. New trigger `inputs` or `credentials_inputs` re-register the webhook, and whichever a PATCH leaves out keeps its saved value. Deleting a trigger removes its webhook once nothing else uses it. - A trigger's secret settings (a generic webhook's `secret_token`, All Quiet's `signing_secret`) are returned only when the trigger is created. Later reads leave them out of `inputs` and list them in `hidden_inputs`, and an update that leaves them out keeps them. - An active trigger starts runs, so arming one needs `RUN_AGENT` as well as `WRITE_LIBRARY`: `setup-trigger`; a PATCH that sets `is_active` to true, or gives a trigger that stays active new inputs or credentials (which re-register its webhook); and creating an active preset of an agent with a webhook trigger. Over MCP, `update_preset` needs it for `is_active: true`, and for new `inputs` unless they come with `is_active: false`, since the tool can't see the graph before it runs. - Over MCP: `list_presets`, `update_preset`, `delete_preset`, `list_agent_triggers` and `setup_agent_webhook_trigger`, scoped to the key's organization as above. Everything else: - `POST /schedules` creates a schedule for the `graph_id` in its body. - `POST /files/upload` writes to the workspace `GET /files` lists, and returns the file's id and a `workspace://` URI. `GET /files` lists the user's files, not Otto's working files or skill packages, as the Files page does. - Every response carries `X-RateLimit-Limit`, `-Remaining` and `-Reset`, and a 429 adds `Retry-After` (none when Redis is unreachable). An API key is counted per client and key head before it is hashed, and a request is verified once; a limiter that fails open logs only its own name, never the bucket. - A v2 request body is capped at 10 MB, MCP included; an upload gets its file cap plus 1 MB of framing, and a marketplace media upload past 10 MB gets a 413 as soon as it passes the cap. - `POST /runs/{id}/reviews` takes `auto_approve_future`, and answers 404 for an unknown review and 409 for a run that is not awaiting review. - One spelling per concept: `inputs`, `credentials_inputs`, amounts in `_cents`, durations in `_seconds`, times as ISO datetimes; one `ErrorResponse` body on every error. MCP results and 500s carry no internal exception text. - `POST /graphs` and `PUT /graphs/{id}` validate a version's credentials only when it is saved active. Activating a version moves the graph's webhook presets onto it only if the credential also has `RUN_AGENT`. - The served `openapi.json` lists only the declared `servers` (`root_path_in_servers=False`), not the mount path FastAPI would put first. Changes to existing code: > [!NOTE] > For an overview of all changes to the internal API, check out [`autogpt_platform/frontend/src/app/api/openapi.json`](https://github.com/Significant-Gravitas/AutoGPT/pull/12206/changes#diff-a1bb1783d6023599f9aaba5414bfd0b49637b61dad8da09c7bed1b2ee01ba852) - The external middleware resolves every credential in one `resolve_auth_info`, verified at most once per request, and refuses a banned owner's on each form; `require_api_key` and `require_access_token` are gone. - The internal upload route and the internal review route call the services v2 uses: `workspace/service.py` and `graph_executions/review/service.py`. The review service carries #15564's step-card handling and #15062's chat-queue wake. - `workspace/service.py`'s size-capped reader is public (`read_within_size_limit`), and refuses before reading when the multipart parser already counted the size. - The copilot's `run_agent` checks inputs with `unknown_inputs` and `missing_inputs` (`data/graph.py`), which v2 shares; a required input sent as `null` now counts as missing there too. - `db.update_preset` stores a preset's inputs and credentials in the same rows, and now rewrites them whenever both are given, either of them possibly `{}`. Before, it refused new inputs beside empty credentials, which the v2 PATCH and the copilot's `update_preset` tool both send for a preset that needs none, and it silently kept the old inputs when every input was cleared. `update_triggered_preset` prunes a webhook it registered for an update that then fails, unless it is the one the trigger already had. - The copilot's preset and trigger tools scope their listings, and the library entry a trigger setup creates, to an External API caller's organization; in a chat they behave as before. - `UserTransaction` is renamed `UserCreditTransactionAdminView` (admin credit routes and the admin spending page). - The Redis fixed-window rate limiter is extracted from `search/rate_limit.py` into `api/utils/rate_limit.py`, which v2 shares; #15703's public store-search limiter sits beside it. Its Redis round trip is bounded at 0.25 s and fails open past it, because every v2 request waits on several of these checks. - The store cache's helpers are public (`get_cached_store_agents` and the rest), for v1 and v2. - New setting `TRUSTED_PROXY_COUNT` (default 3), commented in `backend/.env.default`. - The API key dialog lists permissions in three groups, "API v2", "API v1 and v2" and "API v1 only (stops working 2027-02-01)", each with the consent screen's description, and a key's info dialog marks its v1-only permissions. The descriptions and each permission's API version live in `lib/api-key-scopes.ts`, which the consent page reads too; both are typed `Record<APIKeyPermission, ...>`, so a new permission fails tsc until it is described and classified. - `oauth-tool` offers every permission, v2's included, read from the enum. - Library data layer: `list_library_agents` takes `published` and `favorite` filters, and matches `published` by the graph's live, approved listing; `get_library_agent` fetches the marketplace listing with its active version and creator profile in one query; `add_store_agent_to_library` drops unneeded queries. - HITL data layer: `get_reviews` with optional filters; `get_pending_reviews_for_user` and `get_pending_reviews_for_execution` wrap it; `graph_runs_only` leaves out AutoPilot chat reviews. - `get_graph_all_versions` takes `include_subgraphs` and `offset`; `get_active_graph_version` returns a graph's active version number, and `get_graph(version=None)` keeps returning the newest. - `delete_store_submission_or_raise` raises why a submission couldn't be deleted; `delete_store_submission` wraps it and keeps answering true or false. - `GraphSettings.from_graph()` and `update_library_agent_version_and_settings` are removed (replaced by `GraphSettings()` and `update_agent_version_in_library`). - `PATCH /api/store/profile` accepts partial updates; `Profile` is replaced by `ProfileUpdateRequest`. `PATCH /library/agents/{id}` with `graph_version` points an entry only at the caller's own version in the key's organization, or at another user's version approved in the marketplace. - The consent page's Deny button does nothing until the app's info has loaded, which happens only once the backend has accepted the `redirect_uri`; the rest of that fix reached `dev` with #15057. Two tests pin the app-info endpoint's `redirect_uri` check. - `add_auth_responses_to_openapi` works with any auth scheme; exception handler setup is extracted from `rest_api` into `api/utils/exceptions.py:add_exception_handlers`; `export-api-schema` exports the external APIs' OpenAPI schemas too. - `# pragma: allowlist secret` on existing lines the secret scanner flagged; fast `pre-commit` checks run before slow ones. ### How - v2 is its own FastAPI app, mounted on the REST app at `/api/v2` beside the internal routes; v1 stays under `/external-api` (`backend.api.external.fastapi_app`). Each has its own `/docs`, and neither serves metrics, which the REST process serves on its own port. `rest_api_test.py` keeps v2 and the internal API apart: no internal route may start with `/api/v<n>`, and every v2 path sent to the host app must reach the v2 mount rather than an internal route with a path parameter. - Every v2 route depends on `require_permission(<scope>)`, which resolves the API key or OAuth token to a `TenantContext` (user, organization, team). A route that needs a plan depends on `require_paid_permission(<scope>)` instead, which adds the internal API's `enforce_payment_paywall`. v2's own `paywall_routes_test.py` fails on a route that neither checks nor sits on its open list. - An MCP call is held to the credential before the tool runs (`mcp_calls.py`): unknown arguments are refused, an argument that does more than the tool's own permission needs the matching one (`external_permissions`), ids are checked against the organization, and a tool that bills the platform runs the paywall and usage-cap checks. - The MCP server is served through `MCPMount`, which the platform app's lifespan runs (a mounted app's lifespan never runs); it answers 503 before startup. - The preset routes call the internal API's trigger helpers (`library/triggers.py`), so webhook registration and cleanup are one implementation. The hidden trigger settings are read from the trigger blocks' own inputs marked `secret`, so a new trigger block is covered without a list to maintain. - The anonymous rate-limit bucket keys on the `X-Forwarded-For` entry our own proxies appended, `TRUSTED_PROXY_COUNT` (default 3: Cloudflare, then Google's load balancer) places from the right; `0` uses the socket peer. - Idempotency keys live in Redis, held per user, organization and team, where the run lands. A key that can't be claimed is a retryable 503; an in-flight key holds for 24 hours, so a request that dies after its run started never lets a retry start a second one. - `contract_shape_test.py` checks the published OpenAPI document, so a later model cannot reintroduce a retired spelling or an amount without its unit. ### Known gaps - `pause_schedule` and `resume_schedule` are not exposed over MCP pending a decision on exposing them under `WRITE_SCHEDULE`. - MCP OAuth authorization-server discovery (RFC 8414) is not published: without dynamic client registration, and with a token endpoint that takes JSON rather than a form, it would advertise a flow a conformant client cannot complete. - `/search` still counts the user's library agents in other organizations in `total_count`; it needs the organization in the shared hybrid-search query. - One valid key presented from many addresses still costs a hash per request, up to 300 a minute per address. - The internal API and v1 have no request body cap; this PR caps v2 only. - v2's requests are counted under the REST API's metrics (`autogpt_rest_api_*`, handlers `/api/v2/...`), not the external API's, so the dashboard's External API panels show v1 only until a separate infra change adds v2's handlers to them. That change also routes `/api/v2` on `dev-api.agpt.co`, which routes only `/external-api` today. ### Verified On `7b1b7cf192` I ran `backend/api/external/`, the library, review, search and store suites, all of `backend/copilot/tools/`, `architecture_test.py`, both paywall route tests and `test_block.py`: 8071 passed, 84 skipped. What did not pass is this machine's, not this branch's: the store skill suites refuse a database that already holds listings (62 errors at setup), the browser tests need Chromium (5), `test_ensure_content_embedding_blocks` fails on clean `dev` here, and `skill_file_routes_test.py` passes alone (9 of 9). pyright is clean over the backend, and the frontend type-checks and its auth page tests pass (45). The commits after it (two test-only fixes, the rate-limit warning, the file listing, the shared graph copy and the MCP run's tagging) were checked with `backend/api/external/`, `copilot/gate/`, the graph routes, `run_agent`'s tests and the executor's. Each new guard turned its test red when removed: the ban check on each of the four credential paths, the hidden trigger settings on read and on update, the `RUN_AGENT` requirement over REST and MCP, the rate-limit warning naming its bucket, Otto's files in the listing, a non-owner's graph read whole, and an MCP session made interactive, which parks an approval. None of this ran against a live stack. On `20e9247c22`, which moves v2 to `/api/v2`, groups the dialog's permissions and moves v1's end date: `backend/api/external/`, both paywall route tests, `rest_api_test.py`, `public_metrics_test.py` and `architecture_test.py` (472 passed), `test_block.py` (3714 passed, 84 skipped), and the API key, consent and profile frontend suites (135 passed), with tsc clean. The two route guards each went red on a mutation (an internal `/api/v3` route; an internal `/api/{x}/library/agents`, which answered `GET /api/v2/library/agents`), as did the dialog tests and the tsc check on the version map. A REST server started from this tree served `/api/v2/openapi.json` and `/api/v2/docs`, answered `/api/v2/me` and `/api/v2/mcp/` with 401s pointing at the new discovery document, served that document, still served v1 at `/external-api/v1`, and answered 404 at `/external-api/v2`. On `f4db212b6b`, which fixes the preset update: `backend/api/features/library`, `backend/api/external/`, the copilot's preset tool tests, the webhook ingress tests and `architecture_test.py` (680 passed), and `test_block.py` (3714 passed, 84 skipped). The two PATCH tests drive the real `db.update_preset` and went red with the old guard; the prune went red when removed and when it pruned the trigger's own webhook; and the tests added for a listing's and an update's withheld trigger secrets and a trigger's run each went red when that code was reverted. On `61dd619efb`, which merges `dev` at `ea588f7375` (#15057's redirect checks and client-id encoding, #16175's invitation-only membership): the consent page conflicted, and keeps the shared descriptions module, the Deny guard, `dev`'s `encodeURIComponent` on the app-info request and `dev`'s test file. The internal spec equals a fresh export. `backend/api/external/`, the library, OAuth, orgs, paywall, `rest_api` and public-metrics tests, the copilot's preset tool tests and `architecture_test.py` passed (860), as did `test_block.py` (3714 passed, 84 skipped) and the consent, setup-wizard, proxy and API key frontend suites (324), with tsc clean. On `76ea8cf9e9` through `0a63dac8cf`, which close the remaining review findings, each fix has a test that goes red without it, through the real route, or the real `db.update_preset` where the bug sat between layers: credentials-only and clear-all PATCHes; a limiter whose Redis never answers; an idempotency key reused under another team; expert attribution on presets, triggers and schedules; a key holding only `WRITE_LIBRARY` refused when it arms an active trigger by PATCH or creates an active trigger preset, and the MCP equivalent; a failed attach pruning its webhook; an MCP tool failure returning no internal text; and a team-pinned key stopping once its membership ends. On `0a63dac8cf`, `backend/api/external/`, the library suites, the copilot's preset and trigger tool tests, both paywall route tests and `architecture_test.py` passed (703), as did `test_block.py` (3714 passed, 84 skipped). ### Agents and large language models used - Claude Code with Claude Opus 5.5 (implementation, merges, tests, docs) - T3 Code delegated reviewers for #15376's security review rounds: Claude Mythos 5.1 and Claude Fable 5.1 (Claude), GPT Daybreak Blue (Codex) ### Checklist 📋 #### For code changes: - [x] I have clearly listed my changes in the PR description - [x] I have made a test plan - [x] I have tested my changes according to the test plan: - [x] The `backend/api/external/` package and the library, graphs, review, workspace and copilot tool suites, on this head - [x] Every v2 route's scope and tenant check (`permissions_test.py`, `tenancy_test.py`), and every MCP tool's id arguments (`mcp_tenancy_test.py`) - [x] The published OpenAPI contract (`contract_test.py`, `contract_shape_test.py`) #### For configuration changes: - [x] `.env.default` is updated or already compatible with my changes - [x] `docker-compose.yml` is updated or already compatible with my changes - [x] I have included a list of my configuration changes in the PR description (under **Changes**) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <[email protected]> Co-authored-by: Nicholas Tindle <[email protected]> Co-authored-by: Nicholas Tindle <[email protected]>
This commit is contained in:
247 files changed
+25390
-2038
No files matched your search
@@ -0,0 +1,97 @@
|
||||
name: Docs - Publish the API spec to GitBook
|
||||
|
||||
# The API reference on agpt.co/docs is generated by GitBook from an OpenAPI
|
||||
# spec stored in the GitBook organization as `autogpt-api-v2`, the slug that
|
||||
# docs/platform/SUMMARY.md names. This republishes the spec from master, the
|
||||
# commit docs-gitbook-sync.yml copies the guides from, so the reference and the
|
||||
# guides describe the same API.
|
||||
#
|
||||
# Needs the GITBOOK_TOKEN secret (a GitBook API token of an organization admin)
|
||||
# and the GITBOOK_ORGANIZATION_ID variable (the segment after /o/ in GitBook's
|
||||
# app URLs). Without them the run is skipped, not failed.
|
||||
on:
|
||||
push:
|
||||
branches: [master]
|
||||
paths:
|
||||
- "autogpt_platform/backend/backend/api/external/**"
|
||||
- ".github/workflows/docs-gitbook-openapi.yml"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
check:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 2
|
||||
outputs:
|
||||
ready: ${{ steps.check.outputs.ready }}
|
||||
steps:
|
||||
- id: check
|
||||
env:
|
||||
GITBOOK_TOKEN: ${{ secrets.GITBOOK_TOKEN }}
|
||||
GITBOOK_ORGANIZATION_ID: ${{ vars.GITBOOK_ORGANIZATION_ID }}
|
||||
run: |
|
||||
if [ -n "$GITBOOK_TOKEN" ] && [ -n "$GITBOOK_ORGANIZATION_ID" ]; then
|
||||
echo "ready=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "::notice::GITBOOK_TOKEN or GITBOOK_ORGANIZATION_ID is not set; not publishing the API spec."
|
||||
echo "ready=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
publish:
|
||||
needs: check
|
||||
# Only master's spec is what the published docs describe.
|
||||
if: >-
|
||||
needs.check.outputs.ready == 'true'
|
||||
&& github.repository == 'Significant-Gravitas/AutoGPT'
|
||||
&& github.ref == 'refs/heads/master'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
concurrency:
|
||||
group: docs-gitbook-openapi
|
||||
cancel-in-progress: true
|
||||
defaults:
|
||||
run:
|
||||
working-directory: autogpt_platform/backend
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Install Poetry
|
||||
run: |
|
||||
POETRY_VERSION=$(python ../../.github/workflows/scripts/get_package_version_from_lockfile.py poetry)
|
||||
curl -sSL https://install.python-poetry.org | POETRY_VERSION=$POETRY_VERSION python3 -
|
||||
|
||||
- uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/pypoetry
|
||||
key: poetry-${{ runner.os }}-${{ hashFiles('autogpt_platform/backend/poetry.lock') }}
|
||||
|
||||
- name: Install the backend
|
||||
run: poetry install --only main
|
||||
|
||||
- name: Generate the Prisma client
|
||||
run: poetry run prisma generate && poetry run gen-prisma-stub
|
||||
|
||||
- name: Export the v2 OpenAPI spec
|
||||
run: poetry run export-api-schema --api v2 --pretty --output "$RUNNER_TEMP/autogpt-api-v2.json"
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: "22"
|
||||
|
||||
- name: Publish to GitBook
|
||||
env:
|
||||
GITBOOK_TOKEN: ${{ secrets.GITBOOK_TOKEN }}
|
||||
GITBOOK_ORGANIZATION_ID: ${{ vars.GITBOOK_ORGANIZATION_ID }}
|
||||
run: |
|
||||
npx --yes @gitbook/[email protected] openapi publish \
|
||||
--spec autogpt-api-v2 \
|
||||
--organization "$GITBOOK_ORGANIZATION_ID" \
|
||||
"$RUNNER_TEMP/autogpt-api-v2.json"
|
||||
+74
-74
@@ -54,6 +54,79 @@ repos:
|
||||
files: ^autogpt_platform/(backend|frontend)/\.env\.default$
|
||||
pass_filenames: false
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.7.2
|
||||
hooks:
|
||||
- id: ruff
|
||||
name: Lint (Ruff) - AutoGPT Platform - Backend
|
||||
alias: ruff-lint-platform-backend
|
||||
files: ^autogpt_platform/backend/
|
||||
args: [--fix]
|
||||
|
||||
- id: ruff
|
||||
name: Lint (Ruff) - AutoGPT Platform - Libs
|
||||
alias: ruff-lint-platform-libs
|
||||
files: ^autogpt_platform/autogpt_libs/
|
||||
args: [--fix]
|
||||
|
||||
- id: ruff-format
|
||||
name: Format (Ruff) - AutoGPT Platform - Libs
|
||||
alias: ruff-lint-platform-libs
|
||||
files: ^autogpt_platform/autogpt_libs/
|
||||
|
||||
- repo: local
|
||||
# isort needs the context of which packages are installed to function, so we
|
||||
# can't use a vendored isort pre-commit hook (which runs in its own isolated venv).
|
||||
hooks:
|
||||
- id: isort
|
||||
name: Lint (isort) - AutoGPT Platform - Backend
|
||||
alias: isort-platform-backend
|
||||
entry: poetry -P autogpt_platform/backend run isort -p backend
|
||||
files: ^autogpt_platform/backend/
|
||||
# Vendored upstream packages are test data compared byte for byte
|
||||
# against what the code stored; sorting their imports rewrites the
|
||||
# fixture. isort's own skip settings do not reach a file named on the
|
||||
# command line, which is how this hook invokes it.
|
||||
exclude: ^autogpt_platform/backend/test/fixtures/
|
||||
types: [file, python]
|
||||
language: system
|
||||
|
||||
- id: isort
|
||||
name: Lint (isort) - Classic
|
||||
alias: isort-classic
|
||||
entry: bash -c 'cd classic && poetry run isort $(echo "$@" | sed "s|classic/||g")' --
|
||||
files: ^classic/(original_autogpt|forge|direct_benchmark)/
|
||||
types: [file, python]
|
||||
language: system
|
||||
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 24.10.0
|
||||
# Black has sensible defaults, doesn't need package context, and ignores
|
||||
# everything in .gitignore, so it works fine without any config or arguments.
|
||||
hooks:
|
||||
- id: black
|
||||
name: Format (Black)
|
||||
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 7.0.0
|
||||
# Use consolidated flake8 config at classic/.flake8
|
||||
hooks:
|
||||
- id: flake8
|
||||
name: Lint (Flake8) - Classic
|
||||
alias: flake8-classic
|
||||
files: ^classic/(original_autogpt|forge|direct_benchmark)/
|
||||
args: [--config=classic/.flake8]
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: prettier
|
||||
name: Format (Prettier) - AutoGPT Platform - Frontend
|
||||
alias: format-platform-frontend
|
||||
entry: bash -c 'cd autogpt_platform/frontend && npx prettier --write $(echo "$@" | sed "s|autogpt_platform/frontend/||g")' --
|
||||
files: ^autogpt_platform/frontend/
|
||||
types: [file]
|
||||
language: system
|
||||
|
||||
- repo: local
|
||||
# For proper type checking, all dependencies need to be up-to-date.
|
||||
# It's also a good idea to check that poetry.lock is consistent with pyproject.toml.
|
||||
@@ -156,7 +229,7 @@ repos:
|
||||
entry: >
|
||||
bash -c '
|
||||
cd autogpt_platform/backend
|
||||
&& poetry run export-api-schema --output ../frontend/src/app/api/openapi.json
|
||||
&& poetry run export-api-schema --api internal --output ../frontend/src/app/api/openapi.json
|
||||
&& cd ../frontend
|
||||
&& pnpm prettier --write ./src/app/api/openapi.json
|
||||
'
|
||||
@@ -182,79 +255,6 @@ repos:
|
||||
pass_filenames: false
|
||||
stages: [pre-commit, post-checkout]
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.7.2
|
||||
hooks:
|
||||
- id: ruff
|
||||
name: Lint (Ruff) - AutoGPT Platform - Backend
|
||||
alias: ruff-lint-platform-backend
|
||||
files: ^autogpt_platform/backend/
|
||||
args: [--fix]
|
||||
|
||||
- id: ruff
|
||||
name: Lint (Ruff) - AutoGPT Platform - Libs
|
||||
alias: ruff-lint-platform-libs
|
||||
files: ^autogpt_platform/autogpt_libs/
|
||||
args: [--fix]
|
||||
|
||||
- id: ruff-format
|
||||
name: Format (Ruff) - AutoGPT Platform - Libs
|
||||
alias: ruff-lint-platform-libs
|
||||
files: ^autogpt_platform/autogpt_libs/
|
||||
|
||||
- repo: local
|
||||
# isort needs the context of which packages are installed to function, so we
|
||||
# can't use a vendored isort pre-commit hook (which runs in its own isolated venv).
|
||||
hooks:
|
||||
- id: isort
|
||||
name: Lint (isort) - AutoGPT Platform - Backend
|
||||
alias: isort-platform-backend
|
||||
entry: poetry -P autogpt_platform/backend run isort -p backend
|
||||
files: ^autogpt_platform/backend/
|
||||
# Vendored upstream packages are test data compared byte for byte
|
||||
# against what the code stored; sorting their imports rewrites the
|
||||
# fixture. isort's own skip settings do not reach a file named on the
|
||||
# command line, which is how this hook invokes it.
|
||||
exclude: ^autogpt_platform/backend/test/fixtures/
|
||||
types: [file, python]
|
||||
language: system
|
||||
|
||||
- id: isort
|
||||
name: Lint (isort) - Classic
|
||||
alias: isort-classic
|
||||
entry: bash -c 'cd classic && poetry run isort $(echo "$@" | sed "s|classic/||g")' --
|
||||
files: ^classic/(original_autogpt|forge|direct_benchmark)/
|
||||
types: [file, python]
|
||||
language: system
|
||||
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 24.10.0
|
||||
# Black has sensible defaults, doesn't need package context, and ignores
|
||||
# everything in .gitignore, so it works fine without any config or arguments.
|
||||
hooks:
|
||||
- id: black
|
||||
name: Format (Black)
|
||||
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 7.0.0
|
||||
# Use consolidated flake8 config at classic/.flake8
|
||||
hooks:
|
||||
- id: flake8
|
||||
name: Lint (Flake8) - Classic
|
||||
alias: flake8-classic
|
||||
files: ^classic/(original_autogpt|forge|direct_benchmark)/
|
||||
args: [--config=classic/.flake8]
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: prettier
|
||||
name: Format (Prettier) - AutoGPT Platform - Frontend
|
||||
alias: format-platform-frontend
|
||||
entry: bash -c 'cd autogpt_platform/frontend && npx prettier --write $(echo "$@" | sed "s|autogpt_platform/frontend/||g")' --
|
||||
files: ^autogpt_platform/frontend/
|
||||
types: [file]
|
||||
language: system
|
||||
|
||||
- repo: local
|
||||
# To have watertight type checking, we check *all* the files in an affected
|
||||
# project. To trigger on poetry.lock we also reset the file `types` filter.
|
||||
|
||||
+1
-11
@@ -276,7 +276,7 @@
|
||||
"filename": "autogpt_platform/backend/backend/blocks/replicate/replicate_block.py",
|
||||
"hashed_secret": "8bbdd6f26368f58ea4011d13d7f763cb662e66f0",
|
||||
"is_verified": false,
|
||||
"line_number": 100
|
||||
"line_number": 109
|
||||
}
|
||||
],
|
||||
"autogpt_platform/backend/backend/blocks/slant3d/webhook.py": [
|
||||
@@ -436,16 +436,6 @@
|
||||
"is_secret": false
|
||||
}
|
||||
],
|
||||
"autogpt_platform/docker-compose.platform.yml": [
|
||||
{
|
||||
"type": "Basic Auth Credentials",
|
||||
"filename": "autogpt_platform/docker-compose.platform.yml",
|
||||
"hashed_secret": "de6b16a1aa2c97b8963a4f966774ae09fbe47300",
|
||||
"is_verified": false,
|
||||
"line_number": 35,
|
||||
"is_secret": false
|
||||
}
|
||||
],
|
||||
"autogpt_platform/docker-compose.yml": [
|
||||
{
|
||||
"type": "Secret Keyword",
|
||||
|
||||
@@ -19,10 +19,13 @@ class APIKeySmith:
|
||||
PREFIX: str = "agpt_"
|
||||
HEAD_LENGTH: int = 8
|
||||
TAIL_LENGTH: int = 8
|
||||
# The platform's OAuth access and refresh tokens start with these. No key
|
||||
# does, so a credential that starts with one is never taken for a key.
|
||||
RESERVED_PREFIXES: tuple[str, ...] = ("agpt_xt_", "agpt_rt_")
|
||||
|
||||
def generate_key(self) -> APIKeyContainer:
|
||||
"""Generate a new API key with secure hashing."""
|
||||
raw_key = f"{self.PREFIX}{secrets.token_urlsafe(32)}"
|
||||
raw_key = self._new_raw_key()
|
||||
hash, salt = self.hash_key(raw_key)
|
||||
|
||||
return APIKeyContainer(
|
||||
@@ -33,6 +36,13 @@ class APIKeySmith:
|
||||
salt=salt,
|
||||
)
|
||||
|
||||
def _new_raw_key(self) -> str:
|
||||
"""A fresh random key that doesn't start like an OAuth token."""
|
||||
while True:
|
||||
raw_key = f"{self.PREFIX}{secrets.token_urlsafe(32)}"
|
||||
if not raw_key.startswith(self.RESERVED_PREFIXES):
|
||||
return raw_key
|
||||
|
||||
def verify_key(
|
||||
self, provided_key: str, known_hash: str, known_salt: str | None = None
|
||||
) -> bool:
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
import hashlib
|
||||
|
||||
import pytest
|
||||
|
||||
from autogpt_libs.api_key import keysmith as keysmith_module
|
||||
from autogpt_libs.api_key.keysmith import APIKeySmith
|
||||
|
||||
|
||||
@@ -77,3 +80,18 @@ def test_invalid_salt_format():
|
||||
|
||||
# Invalid salt format should fail gracefully
|
||||
assert keysmith.verify_key(key.key, key.hash, "invalid_hex") is False
|
||||
|
||||
|
||||
def test_a_generated_key_never_starts_like_an_oauth_token(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""A key headed like a token would be looked up as one, or cost a hash as both."""
|
||||
bodies = iter(["xt_" + "a" * 40, "rt_" + "b" * 40, "c" * 43])
|
||||
monkeypatch.setattr(
|
||||
keysmith_module.secrets, "token_urlsafe", lambda _nbytes: next(bodies)
|
||||
)
|
||||
|
||||
key = APIKeySmith().generate_key()
|
||||
|
||||
assert key.key == "agpt_" + "c" * 43
|
||||
assert not key.head.startswith(APIKeySmith.RESERVED_PREFIXES)
|
||||
@@ -8,7 +8,7 @@ from .dependencies import (
|
||||
requires_team_permission,
|
||||
requires_user,
|
||||
)
|
||||
from .helpers import add_auth_responses_to_openapi
|
||||
from .jwt_utils import add_auth_responses_to_openapi
|
||||
from .models import RequestContext, User
|
||||
from .permissions import (
|
||||
OrgAction,
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from fastapi import FastAPI
|
||||
|
||||
from .jwt_utils import bearer_jwt_auth
|
||||
|
||||
|
||||
def add_auth_responses_to_openapi(app: FastAPI) -> None:
|
||||
def add_auth_responses_to_openapi(
|
||||
app: FastAPI, supported_auth_schemes: list[str]
|
||||
) -> None:
|
||||
"""
|
||||
Patch a FastAPI instance's `openapi()` method to add 401 responses
|
||||
to all authenticated endpoints.
|
||||
@@ -29,7 +29,7 @@ def add_auth_responses_to_openapi(app: FastAPI) -> None:
|
||||
for auth_option in details.get("security", [])
|
||||
for schema in auth_option.keys()
|
||||
]
|
||||
if bearer_jwt_auth.scheme_name not in security_schemas:
|
||||
if not any(s in security_schemas for s in supported_auth_schemes):
|
||||
continue
|
||||
|
||||
if "responses" not in details:
|
||||
|
||||
@@ -5,8 +5,7 @@ Tests OpenAPI schema generation and authentication response handling.
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from autogpt_libs.auth.helpers import add_auth_responses_to_openapi
|
||||
from autogpt_libs.auth.jwt_utils import bearer_jwt_auth
|
||||
from autogpt_libs.auth.jwt_utils import add_auth_responses_to_openapi, bearer_jwt_auth
|
||||
|
||||
|
||||
def test_add_auth_responses_to_openapi_basic():
|
||||
|
||||
@@ -4,7 +4,7 @@ import threading
|
||||
from typing import Any
|
||||
|
||||
import jwt
|
||||
from fastapi import HTTPException, Security
|
||||
from fastapi import FastAPI, HTTPException, Security
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
|
||||
from .config import get_settings
|
||||
@@ -162,3 +162,12 @@ def verify_user(jwt_payload: dict | None, admin_only: bool) -> User:
|
||||
raise HTTPException(status_code=403, detail="Admin access required")
|
||||
|
||||
return User.from_payload(jwt_payload)
|
||||
|
||||
|
||||
def add_auth_responses_to_openapi(app: FastAPI) -> None:
|
||||
"""
|
||||
Add 401 responses to all endpoints that use the bearer JWT authentication scheme.
|
||||
"""
|
||||
from .helpers import add_auth_responses_to_openapi
|
||||
|
||||
add_auth_responses_to_openapi(app, [bearer_jwt_auth.scheme_name])
|
||||
@@ -85,6 +85,11 @@ FRONTEND_BASE_URL=http://localhost:3000
|
||||
# links may target — JSON list of full origins or "regex:"-prefixed patterns,
|
||||
# same format as BACKEND_CORS_ALLOW_ORIGINS. Self-hosting needs nothing here.
|
||||
# TRUSTED_FRONTEND_ORIGINS=["regex:https://autogpt-pr-\\d+\\.vercel\\.app"]
|
||||
# X-Forwarded-For entries this app's own proxies append; the anonymous
|
||||
# rate-limit bucket keys on that entry from the right. The default, 3, is the
|
||||
# hosted platform's Cloudflare + Google load balancer. Set 0 when nothing
|
||||
# fronts the app, or the header is a caller-chosen bucket key.
|
||||
# TRUSTED_PROXY_COUNT=3
|
||||
|
||||
# Optional split GCS buckets. Public media contains only explicitly published
|
||||
# marketplace assets and OAuth app logos. Private storage contains user media,
|
||||
|
||||
@@ -22,6 +22,7 @@ _PAYWALL_CHECK_SITES = (
|
||||
"backend.copilot.rate_limit.is_user_paywalled",
|
||||
"backend.executor.utils.is_user_paywalled",
|
||||
"backend.api.features.integrations.router.is_user_paywalled",
|
||||
"backend.api.external.v2.mcp_calls.is_user_paywalled",
|
||||
"backend.copilot.gate.workflow.is_user_paywalled",
|
||||
)
|
||||
_BAN_CHECK_SITES = (
|
||||
|
||||
Whitespace-only changes.
+28
-29
@@ -1,40 +1,39 @@
|
||||
from fastapi import FastAPI
|
||||
from fastapi.requests import Request
|
||||
from fastapi.responses import JSONResponse
|
||||
"""
|
||||
External API Application
|
||||
|
||||
This module defines the FastAPI application mounted at /external-api, which
|
||||
serves v1. The host app mounts v2 at /api/v2 itself.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
from backend.api.middleware.security import SecurityHeadersMiddleware
|
||||
from backend.copilot.rate_limit import UserPaywalledError
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import GraphActivationError
|
||||
from backend.monitoring.instrumentation import instrument_fastapi
|
||||
|
||||
from .v1.routes import v1_router
|
||||
from .v1.app import v1_app
|
||||
|
||||
external_api = FastAPI(
|
||||
title="AutoGPT External API",
|
||||
description="External API for AutoGPT integrations",
|
||||
docs_url="/docs",
|
||||
version="1.0",
|
||||
)
|
||||
|
||||
external_api.add_middleware(SecurityHeadersMiddleware)
|
||||
external_api.include_router(v1_router, prefix="/v1")
|
||||
# Each version publishes its own spec and docs; this app only routes to them.
|
||||
external_api = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
|
||||
|
||||
|
||||
# UserPaywalledError raised at deep enqueue paths (e.g. add_graph_execution)
|
||||
# maps to HTTP 402 here too — the external API is a separate FastAPI app
|
||||
# so the handler in rest_api.py doesn't reach it.
|
||||
@external_api.exception_handler(UserPaywalledError)
|
||||
async def _user_paywalled_handler(_request: Request, exc: UserPaywalledError):
|
||||
return JSONResponse(status_code=402, content={"detail": str(exc)})
|
||||
@external_api.get("/", include_in_schema=False)
|
||||
async def root_redirect() -> RedirectResponse:
|
||||
"""Redirect root to the current version's documentation."""
|
||||
# Relative, so it keeps a proxy prefix such as the single container's /_agpt.
|
||||
return RedirectResponse(url="../api/v2/docs")
|
||||
|
||||
|
||||
# Mirror of rest_api.py's mapping — bad/revoked node credentials during graph
|
||||
# activation surface as a clear 400 instead of an opaque 500.
|
||||
@external_api.exception_handler(GraphActivationError)
|
||||
async def _graph_activation_error_handler(_request: Request, exc: GraphActivationError):
|
||||
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
||||
@external_api.get("/openapi.json", include_in_schema=False)
|
||||
async def v1_spec_redirect(request: Request) -> RedirectResponse:
|
||||
"""v1's spec, at the address it had before the API was versioned, where
|
||||
published docs and generated clients still fetch it."""
|
||||
return RedirectResponse(
|
||||
url=f"{request.scope.get('root_path', '')}/v1/openapi.json", status_code=308
|
||||
)
|
||||
|
||||
|
||||
# Add Prometheus instrumentation. The REST API process this app is mounted in
|
||||
# serves the metrics on their own port, never on this public app.
|
||||
external_api.mount("/v1", v1_app)
|
||||
|
||||
# The REST API process this app is mounted in serves the metrics on their own
|
||||
# port, never on this public app.
|
||||
instrument_fastapi(external_api, service_name="external-api", expose_endpoint=False)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""The external API's own addresses, as the host app mounts it."""
|
||||
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pytest
|
||||
|
||||
from backend.api.external.fastapi_app import external_api
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> fastapi.testclient.TestClient:
|
||||
host = fastapi.FastAPI()
|
||||
host.mount("/external-api", external_api)
|
||||
return fastapi.testclient.TestClient(host, follow_redirects=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ["", "/_agpt"])
|
||||
def test_the_root_redirects_to_v2s_docs_behind_any_proxy_prefix(
|
||||
client: fastapi.testclient.TestClient, prefix: str
|
||||
) -> None:
|
||||
response = client.get("/external-api/")
|
||||
|
||||
page = f"https://example.com{prefix}/external-api/"
|
||||
assert urljoin(page, response.headers["location"]) == (
|
||||
f"https://example.com{prefix}/api/v2/docs"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/docs", "/redoc", "/versions.json"])
|
||||
def test_only_the_versions_publish_a_spec_and_docs(
|
||||
client: fastapi.testclient.TestClient, path: str
|
||||
) -> None:
|
||||
"""A spec of the wrapper app documents nothing a client can call."""
|
||||
assert client.get(f"/external-api{path}").status_code == 404
|
||||
|
||||
|
||||
def test_v1s_spec_is_still_found_at_its_old_address(
|
||||
client: fastapi.testclient.TestClient,
|
||||
) -> None:
|
||||
"""Published docs and generated clients fetch it from there."""
|
||||
response = client.get("/external-api/openapi.json")
|
||||
|
||||
assert response.status_code == 308
|
||||
assert response.headers["location"] == "/external-api/v1/openapi.json"
|
||||
+111
-52
@@ -1,14 +1,18 @@
|
||||
from typing import Optional
|
||||
|
||||
from autogpt_libs.auth.dependencies import ACCOUNT_BANNED_DETAIL
|
||||
from fastapi import HTTPException, Security, status
|
||||
from fastapi import FastAPI, HTTPException, Security, status
|
||||
from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from starlette.types import Scope
|
||||
|
||||
from backend.data.auth.api_key import APIKeyInfo, validate_api_key
|
||||
from backend.data.auth.api_key import validate_api_key
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import (
|
||||
InvalidClientError,
|
||||
InvalidTokenError,
|
||||
OAuthAccessTokenInfo,
|
||||
is_access_token,
|
||||
validate_access_token,
|
||||
)
|
||||
from backend.data.user import is_user_banned
|
||||
@@ -16,59 +20,29 @@ from backend.data.user import is_user_banned
|
||||
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||
bearer_auth = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
async def require_api_key(api_key: str | None = Security(api_key_header)) -> APIKeyInfo:
|
||||
"""Middleware for API key authentication only"""
|
||||
if api_key is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing API key"
|
||||
)
|
||||
|
||||
api_key_obj = await validate_api_key(api_key)
|
||||
|
||||
if not api_key_obj:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key"
|
||||
)
|
||||
|
||||
await _refuse_banned_owner(api_key_obj)
|
||||
return api_key_obj
|
||||
_REQUEST_AUTH = "external_api_auth"
|
||||
|
||||
|
||||
async def require_access_token(
|
||||
bearer: HTTPAuthorizationCredentials | None = Security(bearer_auth),
|
||||
) -> OAuthAccessTokenInfo:
|
||||
"""Middleware for OAuth access token authentication only"""
|
||||
if bearer is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing Authorization header",
|
||||
)
|
||||
class VerifiedCredential(BaseModel):
|
||||
"""What this request's credential verified as: who, or rejected."""
|
||||
|
||||
try:
|
||||
token_info, _ = await validate_access_token(bearer.credentials)
|
||||
except (InvalidClientError, InvalidTokenError) as e:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=str(e))
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
await _refuse_banned_owner(token_info)
|
||||
return token_info
|
||||
credential: Optional[str]
|
||||
auth: Optional[APIAuthorizationInfo] = None
|
||||
rejection: Optional[HTTPException] = None
|
||||
|
||||
|
||||
async def require_auth(
|
||||
async def resolve_auth_info(
|
||||
api_key: str | None = Security(api_key_header),
|
||||
bearer: HTTPAuthorizationCredentials | None = Security(bearer_auth),
|
||||
) -> APIAuthorizationInfo:
|
||||
) -> APIAuthorizationInfo | None:
|
||||
"""
|
||||
Unified authentication middleware supporting both API keys and OAuth tokens.
|
||||
Resolve authentication from API key or Bearer token headers.
|
||||
|
||||
Supports two authentication methods, which are checked in order:
|
||||
1. X-API-Key header (existing API key authentication)
|
||||
2. Authorization: Bearer <token> header (OAuth access token)
|
||||
|
||||
Returns:
|
||||
APIAuthorizationInfo: base class of both APIKeyInfo and OAuthAccessTokenInfo.
|
||||
Returns the auth info if valid credentials are provided, or None if no
|
||||
credentials are present. Raises HTTPException on *invalid* credentials.
|
||||
"""
|
||||
# Try API key first
|
||||
if api_key is not None:
|
||||
api_key_info = await validate_api_key(api_key)
|
||||
if api_key_info:
|
||||
@@ -78,8 +52,16 @@ async def require_auth(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key"
|
||||
)
|
||||
|
||||
# Try OAuth bearer token
|
||||
if bearer is not None:
|
||||
# MCP clients can only send credentials as a Bearer token, so an API key
|
||||
# arrives here too; without this the request resolves to anonymous and
|
||||
# gets the per-IP rate limit. Order matches v2's MCP TokenVerifier.
|
||||
if not is_access_token(bearer.credentials) and (
|
||||
api_key_info := await validate_api_key(bearer.credentials)
|
||||
):
|
||||
await _refuse_banned_owner(api_key_info)
|
||||
return api_key_info
|
||||
|
||||
try:
|
||||
token_info, _ = await validate_access_token(bearer.credentials)
|
||||
except (InvalidClientError, InvalidTokenError) as e:
|
||||
@@ -87,20 +69,83 @@ async def require_auth(
|
||||
await _refuse_banned_owner(token_info)
|
||||
return token_info
|
||||
|
||||
# No credentials provided
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing authentication. Provide API key or access token.",
|
||||
return None
|
||||
|
||||
|
||||
async def resolve_request_auth(
|
||||
scope: Scope,
|
||||
api_key: str | None,
|
||||
bearer: HTTPAuthorizationCredentials | None,
|
||||
) -> APIAuthorizationInfo | None:
|
||||
"""`resolve_auth_info`, verified at most once per request.
|
||||
|
||||
v2's rate limiter and route dependency both need the caller, and an API key
|
||||
costs a Scrypt hash per check. A rejection is remembered too, so an invalid
|
||||
key costs one hash, not two; an error that isn't a rejection (the database
|
||||
unreachable) is not, so the route's own check tries again. The result is
|
||||
kept with the credential it was for, and only reused for that credential.
|
||||
"""
|
||||
credential = (
|
||||
api_key if api_key is not None else bearer.credentials if bearer else None
|
||||
)
|
||||
state = scope.setdefault("state", {})
|
||||
verified = state.get(_REQUEST_AUTH)
|
||||
if not isinstance(verified, VerifiedCredential) or (
|
||||
verified.credential != credential
|
||||
):
|
||||
try:
|
||||
verified = VerifiedCredential(
|
||||
credential=credential,
|
||||
auth=await resolve_auth_info(api_key=api_key, bearer=bearer),
|
||||
)
|
||||
except HTTPException as rejection:
|
||||
verified = VerifiedCredential(credential=credential, rejection=rejection)
|
||||
state[_REQUEST_AUTH] = verified
|
||||
if verified.rejection is not None:
|
||||
raise HTTPException(
|
||||
status_code=verified.rejection.status_code,
|
||||
detail=verified.rejection.detail,
|
||||
headers=verified.rejection.headers,
|
||||
)
|
||||
return verified.auth
|
||||
|
||||
|
||||
def verified_credential(scope: Scope) -> Optional[VerifiedCredential]:
|
||||
"""What this request's credential already verified as, if it was checked."""
|
||||
verified = scope.get("state", {}).get(_REQUEST_AUTH)
|
||||
return verified if isinstance(verified, VerifiedCredential) else None
|
||||
|
||||
|
||||
async def require_auth(
|
||||
auth: APIAuthorizationInfo | None = Security(resolve_auth_info),
|
||||
) -> APIAuthorizationInfo:
|
||||
"""
|
||||
Unified authentication dependency that requires valid credentials.
|
||||
|
||||
Depends on `resolve_auth_info` (which accepts API key or Bearer token)
|
||||
and rejects requests with no credentials.
|
||||
"""
|
||||
if auth is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing authentication. Provide API key or access token.",
|
||||
)
|
||||
return auth
|
||||
|
||||
|
||||
async def _refuse_banned_owner(auth: APIAuthorizationInfo) -> None:
|
||||
if await is_user_banned(auth.user_id):
|
||||
if await owner_is_banned(auth):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN, detail=ACCOUNT_BANNED_DETAIL
|
||||
)
|
||||
|
||||
|
||||
async def owner_is_banned(auth: APIAuthorizationInfo) -> bool:
|
||||
"""Whether the account a credential belongs to is banned; for verifiers that
|
||||
resolve credentials outside `resolve_auth_info`, such as v2's MCP server."""
|
||||
return await is_user_banned(auth.user_id)
|
||||
|
||||
|
||||
def require_permission(*permissions: APIKeyPermission):
|
||||
"""
|
||||
Dependency function for checking required permissions.
|
||||
@@ -109,7 +154,9 @@ def require_permission(*permissions: APIKeyPermission):
|
||||
"""
|
||||
|
||||
async def check_permissions(
|
||||
auth: APIAuthorizationInfo = Security(require_auth),
|
||||
auth: APIAuthorizationInfo = Security(
|
||||
require_auth, scopes=[p.value for p in permissions]
|
||||
),
|
||||
) -> APIAuthorizationInfo:
|
||||
missing = [p for p in permissions if p not in auth.scopes]
|
||||
if missing:
|
||||
@@ -121,3 +168,15 @@ def require_permission(*permissions: APIKeyPermission):
|
||||
return auth
|
||||
|
||||
return check_permissions
|
||||
|
||||
|
||||
def add_auth_responses_to_openapi(app: FastAPI) -> None:
|
||||
"""
|
||||
Add 401 responses to all endpoints secured with `require_auth`,
|
||||
`require_api_key`, or `require_access_token` middleware.
|
||||
"""
|
||||
from autogpt_libs.auth.helpers import add_auth_responses_to_openapi
|
||||
|
||||
add_auth_responses_to_openapi(
|
||||
app, [api_key_header.scheme_name, bearer_auth.scheme_name]
|
||||
)
|
||||
+102
-19
@@ -1,17 +1,36 @@
|
||||
"""An API key or OAuth access token is refused when its owner is banned."""
|
||||
"""Which lookup a credential gets, what checking it costs, and that a banned
|
||||
owner's credential is refused.
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
An API key is matched on its head and then hashed with Scrypt; an OAuth access
|
||||
token is looked up by its digest. A bearer in the token format must never be
|
||||
tried as a key: one old key headed like a token would otherwise make every
|
||||
forged token cost a hash, ahead of every rate limit.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from autogpt_libs.api_key.keysmith import APIKeySmith
|
||||
from fastapi import HTTPException
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.enums import APIKeyPermission, APIKeyStatus
|
||||
|
||||
from backend.api.external import middleware
|
||||
from backend.api.external.middleware import resolve_auth_info
|
||||
from backend.data.auth.api_key import APIKeyInfo
|
||||
from backend.data.auth.oauth import OAuthAccessTokenInfo
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import (
|
||||
ACCESS_TOKEN_PREFIX,
|
||||
REFRESH_TOKEN_PREFIX,
|
||||
InvalidTokenError,
|
||||
OAuthAccessTokenInfo,
|
||||
)
|
||||
|
||||
_NOW = datetime.now(timezone.utc)
|
||||
_NOW = datetime.now(UTC)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -19,16 +38,17 @@ _NOW = datetime.now(timezone.utc)
|
||||
[("banned", 403), ("until_tomorrow", 403), ("lapsed", 200), ("unbanned", 200)],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"path, headers",
|
||||
"headers",
|
||||
[
|
||||
("/auth", {"X-API-Key": "key"}),
|
||||
("/auth", {"Authorization": "Bearer token"}),
|
||||
("/api-key", {"X-API-Key": "key"}),
|
||||
("/access-token", {"Authorization": "Bearer token"}),
|
||||
{"X-API-Key": "key"},
|
||||
# MCP clients can only send an API key this way.
|
||||
{"Authorization": "Bearer key"},
|
||||
{"Authorization": f"Bearer {ACCESS_TOKEN_PREFIX}token"},
|
||||
],
|
||||
ids=["api-key", "api-key-as-bearer", "access-token"],
|
||||
)
|
||||
def test_a_credential_answers_by_its_owners_ban_state(
|
||||
mocker, auth_row, state, status, path, headers
|
||||
mocker, auth_row, state, status, headers
|
||||
):
|
||||
user_id = auth_row(state)
|
||||
mocker.patch.object(middleware, "validate_api_key", return_value=_key(user_id))
|
||||
@@ -36,7 +56,7 @@ def test_a_credential_answers_by_its_owners_ban_state(
|
||||
middleware, "validate_access_token", return_value=(_token(user_id), None)
|
||||
)
|
||||
|
||||
response = _client().get(path, headers=headers)
|
||||
response = _client().get("/auth", headers=headers)
|
||||
|
||||
assert response.status_code == status
|
||||
if status == 403:
|
||||
@@ -45,6 +65,77 @@ def test_a_credential_answers_by_its_owners_ban_state(
|
||||
assert response.json() == {"user_id": user_id}
|
||||
|
||||
|
||||
def _bearer(token: str) -> HTTPAuthorizationCredentials:
|
||||
return HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)
|
||||
|
||||
|
||||
def _principal() -> APIAuthorizationInfo:
|
||||
return APIAuthorizationInfo(
|
||||
user_id="user-1", scopes=[], type="api_key", created_at=datetime.now(UTC)
|
||||
)
|
||||
|
||||
|
||||
async def test_a_token_shaped_bearer_is_never_tried_as_an_api_key(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
validate_key = mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key", new_callable=AsyncMock
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.middleware.validate_access_token",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=InvalidTokenError("not found"),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as rejected:
|
||||
await resolve_auth_info(api_key=None, bearer=_bearer("agpt_xt_forged"))
|
||||
|
||||
assert rejected.value.status_code == 401
|
||||
validate_key.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_an_api_key_sent_as_a_bearer_still_authenticates(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""MCP clients can only send a key as a bearer."""
|
||||
principal = _principal()
|
||||
mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=principal,
|
||||
)
|
||||
validate_token = mocker.patch(
|
||||
"backend.api.external.middleware.validate_access_token",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
assert await resolve_auth_info(api_key=None, bearer=_bearer("agpt_key")) is (
|
||||
principal
|
||||
)
|
||||
validate_token.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_token_shaped_value_sent_as_an_api_key_is_tried_as_one(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The header says what it is; the v2 limiter counts it as a key too."""
|
||||
validate_key = mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await resolve_auth_info(api_key="agpt_xt_value", bearer=None)
|
||||
|
||||
validate_key.assert_awaited_once_with("agpt_xt_value")
|
||||
|
||||
|
||||
def test_no_api_key_is_issued_with_an_oauth_token_prefix() -> None:
|
||||
assert ACCESS_TOKEN_PREFIX in APIKeySmith.RESERVED_PREFIXES
|
||||
assert REFRESH_TOKEN_PREFIX in APIKeySmith.RESERVED_PREFIXES
|
||||
|
||||
|
||||
def _client() -> TestClient:
|
||||
app = fastapi.FastAPI()
|
||||
|
||||
@@ -52,14 +143,6 @@ def _client() -> TestClient:
|
||||
async def by_auth(auth=fastapi.Security(middleware.require_auth)):
|
||||
return {"user_id": auth.user_id}
|
||||
|
||||
@app.get("/api-key")
|
||||
async def by_api_key(auth=fastapi.Security(middleware.require_api_key)):
|
||||
return {"user_id": auth.user_id}
|
||||
|
||||
@app.get("/access-token")
|
||||
async def by_access_token(auth=fastapi.Security(middleware.require_access_token)):
|
||||
return {"user_id": auth.user_id}
|
||||
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
V1 External API Application
|
||||
|
||||
This module defines the FastAPI application for the v1 external API.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from backend.api.external.middleware import add_auth_responses_to_openapi
|
||||
from backend.api.middleware.security import SecurityHeadersMiddleware
|
||||
from backend.api.utils.exceptions import add_exception_handlers
|
||||
from backend.api.utils.openapi import sort_openapi
|
||||
|
||||
from .routes import v1_router
|
||||
|
||||
DESCRIPTION = """
|
||||
The v1 API provides access to core AutoGPT functionality for external integrations.
|
||||
|
||||
**Deprecated**: v1 stops working on 2027-02-01; its last day of service is
|
||||
2027-01-31. New integrations should use [v2](../v2/docs).
|
||||
|
||||
For authentication details and usage examples, see the
|
||||
[API Integration Guide](https://agpt.co/docs/platform/api-and-integrations/api-guide).
|
||||
"""
|
||||
|
||||
v1_app = FastAPI(
|
||||
title="AutoGPT Platform API",
|
||||
summary="External API for AutoGPT Platform integrations (v1)",
|
||||
description=DESCRIPTION,
|
||||
version="1.0.0",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc",
|
||||
openapi_url="/openapi.json",
|
||||
openapi_tags=[
|
||||
{"name": "user", "description": "User information"},
|
||||
{"name": "blocks", "description": "Block operations"},
|
||||
{"name": "graphs", "description": "Graph execution"},
|
||||
{"name": "store", "description": "Marketplace agents and creators"},
|
||||
{"name": "integrations", "description": "OAuth credential management"},
|
||||
{"name": "tools", "description": "AI assistant tools"},
|
||||
],
|
||||
)
|
||||
|
||||
v1_app.add_middleware(SecurityHeadersMiddleware)
|
||||
v1_app.include_router(v1_router)
|
||||
|
||||
# Mounted sub-apps do NOT inherit exception handlers from the parent app.
|
||||
# Third parties call v1, so a 500 says nothing about what failed inside.
|
||||
add_exception_handlers(v1_app, server_error_detail=False)
|
||||
|
||||
# Add 401 responses to authenticated endpoints in OpenAPI spec
|
||||
add_auth_responses_to_openapi(v1_app)
|
||||
# Sort OpenAPI schema to eliminate diff on refactors
|
||||
sort_openapi(v1_app)
|
||||
@@ -354,7 +354,7 @@ async def get_store_agents(
|
||||
if page_size < 1:
|
||||
raise HTTPException(status_code=422, detail="Page size must be greater than 0")
|
||||
|
||||
agents = await store_cache._get_cached_store_agents(
|
||||
agents = await store_cache.get_cached_store_agents(
|
||||
featured=featured,
|
||||
creator=creator,
|
||||
sorted_by=sorted_by,
|
||||
@@ -388,7 +388,7 @@ async def get_store_agent(
|
||||
"""
|
||||
username = urllib.parse.unquote(username).lower()
|
||||
agent_name = urllib.parse.unquote(agent_name).lower()
|
||||
agent = await store_cache._get_cached_agent_details(
|
||||
agent = await store_cache.get_cached_agent_details(
|
||||
username=username, agent_name=agent_name
|
||||
)
|
||||
return agent
|
||||
@@ -426,7 +426,7 @@ async def get_store_creators(
|
||||
if page_size < 1:
|
||||
raise HTTPException(status_code=422, detail="Page size must be greater than 0")
|
||||
|
||||
creators = await store_cache._get_cached_store_creators(
|
||||
creators = await store_cache.get_cached_store_creators(
|
||||
featured=featured,
|
||||
search_query=search_query,
|
||||
sorted_by=sorted_by,
|
||||
@@ -455,5 +455,5 @@ async def get_store_creator(
|
||||
CreatorDetails: Detailed information about the creator
|
||||
"""
|
||||
username = urllib.parse.unquote(username).lower()
|
||||
creator = await store_cache._get_cached_creator_details(username=username)
|
||||
creator = await store_cache.get_cached_creator_details(username=username)
|
||||
return creator
|
||||
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
V2 External API
|
||||
|
||||
This module provides the v2 external API for programmatic access to the AutoGPT Platform.
|
||||
"""
|
||||
|
||||
from .routes import v2_router
|
||||
|
||||
__all__ = ["v2_router"]
|
||||
@@ -0,0 +1,191 @@
|
||||
"""
|
||||
V2 External API Application
|
||||
|
||||
This module defines the FastAPI application for the v2 external API.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from backend.api.middleware.security import SecurityHeadersMiddleware
|
||||
from backend.api.utils.openapi import sort_openapi
|
||||
|
||||
from .body_limit import BodySizeLimitMiddleware
|
||||
from .errors import add_v2_exception_handlers
|
||||
from .global_rate_limit import GlobalRateLimitMiddleware
|
||||
from .mcp_server import mcp_mount
|
||||
from .routes import v2_router
|
||||
|
||||
DESCRIPTION = """
|
||||
The v2 API provides comprehensive access to the AutoGPT Platform for building
|
||||
integrations, automations, and custom applications.
|
||||
|
||||
### Key Improvements over v1
|
||||
|
||||
- **Consistent naming**: Uses `graph_id`/`graph_version` consistently
|
||||
- **One list shape**: every list endpoint takes `limit`/`cursor` and returns
|
||||
`{"items": [...], "next_cursor": ..., "total_count": ...}`
|
||||
- **One error shape**: every non-2xx response is
|
||||
`{"error": {"code", "message", "details"}}`
|
||||
- **Comprehensive coverage**: Access to library, runs, schedules, credits, and more
|
||||
- **Human-in-the-loop**: Review and approve agent decisions via the API
|
||||
|
||||
For authentication details and usage examples, see the
|
||||
[API Integration Guide](https://agpt.co/docs/platform/api-and-integrations/api-guide).
|
||||
|
||||
### Rate Limits
|
||||
|
||||
All endpoints are subject to a global rate limit per authenticated user:
|
||||
**200 requests per minute**. Unauthenticated requests are limited to
|
||||
**5 requests per minute** per client IP.
|
||||
|
||||
Some endpoints have additional per-endpoint rate limits (documented on each
|
||||
endpoint). When a rate limit is exceeded, the API returns `429 Too Many Requests`.
|
||||
|
||||
### Pagination
|
||||
|
||||
List endpoints take `limit` (1-100, default 20) and `cursor`, and return
|
||||
`{"items": [...], "next_cursor": ..., "total_count": ...}`. Pass `next_cursor`
|
||||
back as `cursor` for the next page; it is `null` on the last page. Cursors are
|
||||
opaque. `total_count` counts the matches across all pages, and is `null` only
|
||||
where the source cannot report one.
|
||||
|
||||
### Status codes
|
||||
|
||||
`201` when a resource is created, `202` when work is queued, `204` on delete,
|
||||
`200` on everything else. An endpoint that starts work, stores files or calls
|
||||
out answers `402` (`payment_required`) for an account without an active plan.
|
||||
|
||||
### Errors
|
||||
|
||||
Every non-2xx response is `{"error": {"code", "message", "details"}}`. Branch
|
||||
on `code`, a stable snake_case identifier; `message` is for humans.
|
||||
""".strip()
|
||||
|
||||
v2_app = FastAPI(
|
||||
title="AutoGPT Platform External API",
|
||||
summary="External API for AutoGPT Platform integrations (v2)",
|
||||
description=DESCRIPTION,
|
||||
version="2.0.0",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc",
|
||||
openapi_url="/openapi.json",
|
||||
# Otherwise the mount path is served as a relative first server, which is
|
||||
# wrong behind a proxy prefix such as the single container's `/_agpt`.
|
||||
root_path_in_servers=False,
|
||||
servers=[
|
||||
{
|
||||
"url": "https://backend.agpt.co/api/v2",
|
||||
"description": "AutoGPT Cloud",
|
||||
},
|
||||
{
|
||||
"url": "{origin}/api/v2",
|
||||
"description": "Self-hosted with Docker Compose",
|
||||
"variables": {"origin": {"default": "http://localhost:8006"}},
|
||||
},
|
||||
{
|
||||
"url": "{origin}/_agpt/api/v2",
|
||||
"description": "Self-hosted single container",
|
||||
"variables": {"origin": {"default": "http://localhost:3000"}},
|
||||
},
|
||||
],
|
||||
openapi_tags=[
|
||||
{
|
||||
"name": "graphs",
|
||||
"x-page-title": "Graphs",
|
||||
"description": "Create, update, and manage agent graphs",
|
||||
},
|
||||
{
|
||||
"name": "schedules",
|
||||
"x-page-title": "Schedules",
|
||||
"description": "Manage scheduled graph executions",
|
||||
},
|
||||
{
|
||||
"name": "blocks",
|
||||
"x-page-title": "Blocks",
|
||||
"description": "Discover available building blocks",
|
||||
},
|
||||
{
|
||||
"name": "search",
|
||||
"x-page-title": "Search",
|
||||
"description": "Cross-domain hybrid search across agents, blocks, and docs",
|
||||
},
|
||||
{
|
||||
"name": "marketplace",
|
||||
"x-page-title": "Marketplace",
|
||||
"description": "Browse agents and creators, manage submissions",
|
||||
},
|
||||
{
|
||||
"name": "library",
|
||||
"x-page-title": "Library",
|
||||
"description": "Manage your agent library, run agents, organize with folders",
|
||||
},
|
||||
{
|
||||
"name": "presets",
|
||||
"x-page-title": "Presets",
|
||||
"description": "Saved run inputs for agents, and webhook triggers that run them",
|
||||
},
|
||||
{
|
||||
"name": "runs",
|
||||
"x-page-title": "Runs",
|
||||
"description": "Monitor, stop, delete, and share agent runs; manage human-in-the-loop reviews",
|
||||
},
|
||||
{
|
||||
"name": "credits",
|
||||
"x-page-title": "Credits",
|
||||
"description": "Check balance and view transaction history",
|
||||
},
|
||||
{
|
||||
"name": "integrations",
|
||||
"x-page-title": "Integrations",
|
||||
"description": "List, create, and delete integration credentials",
|
||||
},
|
||||
{
|
||||
"name": "files",
|
||||
"x-page-title": "Files",
|
||||
"description": "Upload, list, download, and delete workspace files",
|
||||
},
|
||||
{
|
||||
"name": "identity",
|
||||
"x-page-title": "Identity",
|
||||
"description": "Who the credentials act as, and in which organization and team",
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
v2_app.add_middleware(GlobalRateLimitMiddleware)
|
||||
# Inside the security headers, so its 413 carries them; outside the rate
|
||||
# limiter, so a body declared too long is refused before any credential check.
|
||||
v2_app.add_middleware(BodySizeLimitMiddleware)
|
||||
v2_app.add_middleware(SecurityHeadersMiddleware)
|
||||
v2_app.include_router(v2_router)
|
||||
|
||||
# Mounted sub-apps do NOT inherit exception handlers from the parent app,
|
||||
# so we must register them here for the v2 API specifically.
|
||||
add_v2_exception_handlers(v2_app)
|
||||
|
||||
# MCP server (Copilot tools via Streamable HTTP). The platform app's lifespan
|
||||
# runs it; see `MCPMount`.
|
||||
v2_app.mount("/mcp", mcp_mount)
|
||||
|
||||
# Sort OpenAPI schema to eliminate diff on refactors
|
||||
sort_openapi(v2_app)
|
||||
|
||||
|
||||
def _route_try_it_through_the_docs_proxy(app: FastAPI) -> None:
|
||||
"""The docs site calls the API from the reader's browser to "Test it",
|
||||
and the API does not allow that origin; GitBook's proxy makes the call
|
||||
server-side instead, and only to the hosts listed in `servers`."""
|
||||
wrapped = app.openapi
|
||||
|
||||
def openapi():
|
||||
if app.openapi_schema:
|
||||
return app.openapi_schema
|
||||
schema = wrapped()
|
||||
schema["x-enable-proxy"] = True
|
||||
app.openapi_schema = schema
|
||||
return schema
|
||||
|
||||
app.openapi = openapi
|
||||
|
||||
|
||||
_route_try_it_through_the_docs_proxy(v2_app)
|
||||
@@ -0,0 +1,346 @@
|
||||
"""
|
||||
Tests for the v2 API error envelope.
|
||||
|
||||
The v2 app registers its own exception handlers (since mounted sub-apps don't
|
||||
inherit handlers from the parent app). These tests verify that exceptions from
|
||||
the DB/service layer are correctly mapped to HTTP status codes.
|
||||
|
||||
We construct a lightweight test app rather than importing the full v2_app,
|
||||
because the latter eagerly loads the MCP server, block registry, and other
|
||||
heavy dependencies that are irrelevant for error handling tests.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import aiohttp
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pytest_snapshot.plugin import Snapshot
|
||||
|
||||
from backend.data.workspace import WorkspaceFile
|
||||
from backend.util.exceptions import DatabaseError, NotFoundError
|
||||
|
||||
from .errors import add_v2_exception_handlers
|
||||
from .files import file_workspace_router
|
||||
from .library.agents import agents_router
|
||||
from .marketplace import marketplace_router
|
||||
from .tenancy import TenantContext, require_auth
|
||||
|
||||
TEST_USER_ID = "test-user-id"
|
||||
TEST_ORG_ID = "test-org-id"
|
||||
|
||||
_mock_auth = TenantContext(
|
||||
user_id=TEST_USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=TEST_ORG_ID,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Build a lightweight test app with the shared exception handlers
|
||||
# but only the routers we need for testing.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(agents_router, prefix="/library")
|
||||
app.include_router(marketplace_router, prefix="/marketplace")
|
||||
app.include_router(file_workspace_router, prefix="/files")
|
||||
add_v2_exception_handlers(app)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _override_auth():
|
||||
"""Bypass API key / OAuth auth for all tests in this module."""
|
||||
|
||||
async def fake_auth() -> TenantContext:
|
||||
return _mock_auth
|
||||
|
||||
app.dependency_overrides[require_auth] = fake_auth
|
||||
yield
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
client = fastapi.testclient.TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# NotFoundError → 404
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_not_found_error_returns_404(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
snapshot: Snapshot,
|
||||
) -> None:
|
||||
"""NotFoundError raised by the DB layer should become a 404 response."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=NotFoundError("Agent #nonexistent not found"),
|
||||
)
|
||||
|
||||
response = client.get("/library/agents/nonexistent")
|
||||
|
||||
assert response.status_code == 404
|
||||
body = response.json()
|
||||
assert body["error"]["code"] == "not_found"
|
||||
assert body["error"]["message"] == "Agent #nonexistent not found"
|
||||
|
||||
snapshot.snapshot_dir = "snapshots"
|
||||
snapshot.assert_match(
|
||||
json.dumps(body, indent=2, sort_keys=True),
|
||||
"v2_not_found_error_404",
|
||||
)
|
||||
|
||||
|
||||
def test_not_found_error_on_delete_returns_404(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""NotFoundError on DELETE should return 404, not 204 or 500."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=TEST_ORG_ID),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.delete_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=NotFoundError("Agent #gone not found"),
|
||||
)
|
||||
|
||||
response = client.delete("/library/agents/gone")
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.json()["error"]["message"] == "Agent #gone not found"
|
||||
|
||||
|
||||
def test_not_found_error_on_marketplace_returns_404(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""NotFoundError from store DB layer should become a 404."""
|
||||
mocker.patch(
|
||||
"backend.api.features.store.db.get_store_agent_by_version_id",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=NotFoundError("Store listing not found"),
|
||||
)
|
||||
|
||||
response = client.get("/marketplace/agents/by-version/nonexistent")
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.json()["error"]["message"] == "Store listing not found"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ValueError → 400
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_value_error_returns_400(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
snapshot: Snapshot,
|
||||
) -> None:
|
||||
"""ValueError raised by the service layer should become a 400 response."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=TEST_ORG_ID),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.agents.writable_organization_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=frozenset({TEST_ORG_ID}),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.update_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=ValueError("Invalid graph version: -1"),
|
||||
)
|
||||
|
||||
response = client.patch(
|
||||
"/library/agents/some-id",
|
||||
json={"graph_version": -1},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
body = response.json()
|
||||
assert body["error"]["code"] == "bad_request"
|
||||
assert body["error"]["message"] == "Invalid graph version: -1"
|
||||
|
||||
snapshot.snapshot_dir = "snapshots"
|
||||
snapshot.assert_match(
|
||||
json.dumps(body, indent=2, sort_keys=True),
|
||||
"v2_value_error_400",
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# NotFoundError is a ValueError subclass — verify specificity wins
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_not_found_error_takes_precedence_over_value_error(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""
|
||||
NotFoundError(ValueError) should match the NotFoundError handler (404),
|
||||
not the ValueError handler (400).
|
||||
"""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=NotFoundError("Specific not found"),
|
||||
)
|
||||
|
||||
response = client.get("/library/agents/test-id")
|
||||
|
||||
# Must be 404, not 400
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Unhandled Exception → 500
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_unhandled_exception_returns_500(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
snapshot: Snapshot,
|
||||
) -> None:
|
||||
"""
|
||||
Unexpected exceptions should return a generic 500 without leaking
|
||||
internal details.
|
||||
"""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=DatabaseError("connection refused"),
|
||||
)
|
||||
|
||||
response = client.get("/library/agents/some-id")
|
||||
|
||||
assert response.status_code == 500
|
||||
body = response.json()
|
||||
assert body["error"]["code"] == "internal_error"
|
||||
assert "connection refused" not in json.dumps(body)
|
||||
|
||||
snapshot.snapshot_dir = "snapshots"
|
||||
snapshot.assert_match(
|
||||
json.dumps(body, indent=2, sort_keys=True),
|
||||
"v2_unhandled_exception_500",
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_error_returns_500(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""RuntimeError (not ValueError) should hit the catch-all 500 handler."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=TEST_ORG_ID),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.delete_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("something broke"),
|
||||
)
|
||||
|
||||
response = client.delete("/library/agents/some-id")
|
||||
|
||||
assert response.status_code == 500
|
||||
assert response.json()["error"]["code"] == "internal_error"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Response format consistency
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_all_error_responses_have_consistent_format(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Every error response is `{"error": {"code", "message", "details"}}`."""
|
||||
cases = [
|
||||
(NotFoundError("not found"), 404),
|
||||
(ValueError("bad value"), 400),
|
||||
(RuntimeError("boom"), 500),
|
||||
]
|
||||
|
||||
for exc, expected_status in cases:
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=exc,
|
||||
)
|
||||
|
||||
response = client.get("/library/agents/test-id")
|
||||
|
||||
assert response.status_code == expected_status, (
|
||||
f"Expected {expected_status} for {type(exc).__name__}, "
|
||||
f"got {response.status_code}"
|
||||
)
|
||||
body = response.json()
|
||||
assert set(body) == {
|
||||
"error"
|
||||
}, f"Extra top-level keys for {type(exc).__name__}: {body}"
|
||||
assert set(body["error"]) == {
|
||||
"code",
|
||||
"message",
|
||||
"details",
|
||||
}, f"Wrong error keys for {type(exc).__name__}: {body}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# A file download whose content cannot be read
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"storage_error, expected_status",
|
||||
[
|
||||
(FileNotFoundError("blob missing"), 404),
|
||||
(aiohttp.ClientError("connection reset"), 502),
|
||||
],
|
||||
ids=["missing-blob", "storage-failure"],
|
||||
)
|
||||
def test_a_download_that_cannot_read_its_content(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
storage_error: Exception,
|
||||
expected_status: int,
|
||||
) -> None:
|
||||
now = datetime.now(timezone.utc)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(id="ws-1"),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace_file",
|
||||
new_callable=AsyncMock,
|
||||
return_value=WorkspaceFile(
|
||||
id="file-1",
|
||||
workspace_id="ws-1",
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
name="report.pdf",
|
||||
path="/report.pdf",
|
||||
storage_path="local://ws-1/report.pdf",
|
||||
mime_type="application/pdf",
|
||||
size_bytes=3,
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.workspace.routes.get_workspace_storage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(retrieve=AsyncMock(side_effect=storage_error)),
|
||||
)
|
||||
|
||||
response = client.get("/files/file-1/download")
|
||||
|
||||
assert response.status_code == expected_status
|
||||
assert set(response.json()) == {"error"}
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
V2 External API - Blocks Endpoints
|
||||
|
||||
Provides read-only access to available building blocks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, Security
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.blocks import get_blocks
|
||||
from backend.util.cache import cached
|
||||
|
||||
from .models import BlockInfo
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .tenancy import TenantContext, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
blocks_router = APIRouter(tags=["blocks"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Internal Functions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _compute_blocks_sync() -> list[BlockInfo]:
|
||||
"""
|
||||
Synchronous function to compute blocks data.
|
||||
This does the heavy lifting: instantiate 226+ blocks, compute costs, serialize.
|
||||
"""
|
||||
return [
|
||||
BlockInfo.from_internal(block)
|
||||
for block_class in get_blocks().values()
|
||||
if not (block := block_class()).disabled
|
||||
]
|
||||
|
||||
|
||||
@cached(ttl_seconds=3600)
|
||||
async def _get_cached_blocks() -> list[BlockInfo]:
|
||||
"""
|
||||
Async cached function with thundering herd protection.
|
||||
On cache miss: runs heavy work in thread pool
|
||||
On cache hit: returns cached list immediately
|
||||
"""
|
||||
return await run_in_threadpool(_compute_blocks_sync)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@blocks_router.get(
|
||||
path="",
|
||||
summary="List available blocks",
|
||||
operation_id="listAvailableBlocks",
|
||||
)
|
||||
async def list_available_blocks(
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_BLOCK)),
|
||||
) -> Page[BlockInfo]:
|
||||
"""List all available blocks with their input/output schemas and cost information."""
|
||||
return page.slice(await _get_cached_blocks())
|
||||
@@ -0,0 +1,124 @@
|
||||
"""
|
||||
V2 External API - Request body size limit
|
||||
|
||||
ASGI middleware that bounds how much of a request body the v2 app reads.
|
||||
FastAPI reads a JSON body whole before the route runs, and FastMCP reads a
|
||||
JSON-RPC body whole before it dispatches the tool, so without a bound one
|
||||
caller holding a few large requests open could exhaust a server's memory.
|
||||
|
||||
A body declared longer than its endpoint's cap is refused before anything
|
||||
reads it. One sent in chunks is counted as the app reads it, and refused as
|
||||
soon as it crosses the cap. This middleware sends that 413 itself and drops
|
||||
whatever the app sends after it: FastAPI turns a failed body read into a 400,
|
||||
and FastMCP into a 500.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from starlette import status
|
||||
from starlette.datastructures import Headers
|
||||
from starlette.responses import Response
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from backend.util.settings import Settings
|
||||
|
||||
from .errors import error_response
|
||||
from .global_rate_limit import app_path
|
||||
from .marketplace import MEDIA_MAX_SIZE_MB
|
||||
|
||||
settings = Settings()
|
||||
|
||||
_MiB = 1024 * 1024
|
||||
# Every request that is not an upload, an MCP tool call included.
|
||||
MAX_BODY_BYTES = 10 * _MiB
|
||||
# The multipart framing around an uploaded file.
|
||||
_MULTIPART_OVERHEAD = _MiB
|
||||
|
||||
|
||||
def max_body_bytes(path: str) -> int:
|
||||
"""The most body a request to *path*, inside the v2 app, may send."""
|
||||
if path == "/files/upload":
|
||||
return settings.config.max_file_size_mb * _MiB + _MULTIPART_OVERHEAD
|
||||
if path == "/marketplace/submissions/media":
|
||||
return MEDIA_MAX_SIZE_MB * _MiB + _MULTIPART_OVERHEAD
|
||||
return MAX_BODY_BYTES
|
||||
|
||||
|
||||
class BodySizeLimitMiddleware:
|
||||
def __init__(self, app: ASGIApp):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
limit = max_body_bytes(app_path(scope))
|
||||
declared = _declared_length(scope)
|
||||
if declared is not None and declared > limit:
|
||||
await _too_large(limit)(scope, receive, send)
|
||||
return
|
||||
|
||||
guard = _BodyGuard(scope, receive, send, limit)
|
||||
try:
|
||||
await self.app(scope, guard.receive, guard.send)
|
||||
except RequestBodyTooLarge:
|
||||
# Nothing caught the failed read on the way out. The 413 has gone
|
||||
# out already, unless the app had started its own response.
|
||||
pass
|
||||
|
||||
|
||||
class RequestBodyTooLarge(Exception):
|
||||
"""Raised into the app when the body it is reading crosses the cap."""
|
||||
|
||||
|
||||
class _BodyGuard:
|
||||
"""Counts the body as the app reads it, and answers 413 when it crosses
|
||||
the cap."""
|
||||
|
||||
def __init__(self, scope: Scope, receive: Receive, send: Send, limit: int):
|
||||
self._scope = scope
|
||||
self._receive = receive
|
||||
self._send = send
|
||||
self._limit = limit
|
||||
self._received = 0
|
||||
self._response_started = False
|
||||
self._answered = False
|
||||
|
||||
async def receive(self) -> Message:
|
||||
message = await self._receive()
|
||||
if message["type"] != "http.request":
|
||||
return message
|
||||
self._received += len(message.get("body", b""))
|
||||
if self._received <= self._limit:
|
||||
return message
|
||||
if not (self._response_started or self._answered):
|
||||
self._answered = True
|
||||
await _too_large(self._limit)(self._scope, self._receive, self._send)
|
||||
raise RequestBodyTooLarge()
|
||||
|
||||
async def send(self, message: Message) -> None:
|
||||
if self._answered:
|
||||
# The app's own answer to the failed read; the 413 went out first.
|
||||
return
|
||||
if message["type"] == "http.response.start":
|
||||
self._response_started = True
|
||||
await self._send(message)
|
||||
|
||||
|
||||
def _declared_length(scope: Scope) -> Optional[int]:
|
||||
value = Headers(scope=scope).get("content-length")
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _too_large(limit: int) -> Response:
|
||||
return error_response(
|
||||
status.HTTP_413_CONTENT_TOO_LARGE,
|
||||
f"The request body is larger than this endpoint accepts ({limit:,} bytes)",
|
||||
details={"limit_bytes": limit},
|
||||
)
|
||||
@@ -0,0 +1,231 @@
|
||||
"""A request body is bounded before the app holds it, whoever reads it.
|
||||
|
||||
FastAPI and FastMCP both read a body whole before acting on it, so the bound
|
||||
has to sit in front of them, and its 413 has to survive whatever each does
|
||||
with a read that fails.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi.testclient import TestClient
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
from backend.api.external.v2 import body_limit, global_rate_limit
|
||||
from backend.api.external.v2.app import v2_app
|
||||
from backend.api.external.v2.body_limit import (
|
||||
MAX_BODY_BYTES,
|
||||
BodySizeLimitMiddleware,
|
||||
max_body_bytes,
|
||||
)
|
||||
|
||||
MiB = 1024 * 1024
|
||||
|
||||
|
||||
async def test_a_body_declared_over_the_cap_is_refused_before_the_app_runs() -> None:
|
||||
called = []
|
||||
|
||||
async def app(scope, receive, send):
|
||||
called.append(scope)
|
||||
|
||||
sent = await _call(
|
||||
BodySizeLimitMiddleware(app),
|
||||
_scope("/mcp", headers=[(b"content-length", str(MAX_BODY_BYTES + 1).encode())]),
|
||||
_chunks(b""),
|
||||
)
|
||||
|
||||
assert not called
|
||||
assert _statuses(sent) == [413]
|
||||
error = _body(sent)["error"]
|
||||
assert error["code"] == "payload_too_large"
|
||||
assert error["details"] == {"limit_bytes": MAX_BODY_BYTES}
|
||||
|
||||
|
||||
async def test_a_body_sent_in_chunks_is_refused_when_it_crosses_the_cap() -> None:
|
||||
"""No length to check up front: FastAPI would answer the failed read 400."""
|
||||
reached = []
|
||||
api = fastapi.FastAPI()
|
||||
|
||||
@api.post("/graphs")
|
||||
async def create(body: dict) -> dict:
|
||||
reached.append(body)
|
||||
return {}
|
||||
|
||||
piece = b" " * MiB
|
||||
sent = await _call(
|
||||
BodySizeLimitMiddleware(api),
|
||||
_scope("/graphs"),
|
||||
_chunks(b"{", *[piece] * 11, b"}"),
|
||||
)
|
||||
|
||||
assert not reached
|
||||
assert _statuses(sent) == [413]
|
||||
assert _body(sent)["error"]["code"] == "payload_too_large"
|
||||
|
||||
|
||||
async def test_an_app_that_answers_a_failed_read_itself_cannot_replace_the_413() -> (
|
||||
None
|
||||
):
|
||||
"""FastMCP catches the read error and answers 500; the client sees 413."""
|
||||
reads = []
|
||||
|
||||
async def swallowing_app(scope, receive, send):
|
||||
try:
|
||||
while (message := await receive())["type"] == "http.request":
|
||||
reads.append(len(message.get("body", b"")))
|
||||
if not message.get("more_body"):
|
||||
break
|
||||
except Exception as error:
|
||||
await JSONResponse({"error": str(error)}, status_code=500)(
|
||||
scope, receive, send
|
||||
)
|
||||
# Reading on after the failure must not send a second 413.
|
||||
with contextlib.suppress(Exception):
|
||||
await receive()
|
||||
return
|
||||
await JSONResponse({}, status_code=200)(scope, receive, send)
|
||||
|
||||
sent = await _call(
|
||||
BodySizeLimitMiddleware(swallowing_app),
|
||||
_scope("/mcp"),
|
||||
_chunks(*[b"x" * MiB] * 11),
|
||||
)
|
||||
|
||||
assert _statuses(sent) == [413]
|
||||
# What the app saw stops short of the cap.
|
||||
assert sum(reads) <= MAX_BODY_BYTES
|
||||
|
||||
|
||||
async def test_a_body_within_the_cap_reaches_the_app_whole() -> None:
|
||||
received = []
|
||||
|
||||
async def app(scope, receive, send):
|
||||
while True:
|
||||
message = await receive()
|
||||
received.append(message.get("body", b""))
|
||||
if not message.get("more_body"):
|
||||
break
|
||||
await JSONResponse({}, status_code=200)(scope, receive, send)
|
||||
|
||||
parts = [b"a" * MiB, b"b" * MiB, b"c"]
|
||||
sent = await _call(BodySizeLimitMiddleware(app), _scope("/graphs"), _chunks(*parts))
|
||||
|
||||
assert _statuses(sent) == [200]
|
||||
assert received == parts
|
||||
|
||||
|
||||
def test_uploads_get_room_for_their_file(mocker: pytest_mock.MockFixture) -> None:
|
||||
mocker.patch.object(body_limit.settings.config, "max_file_size_mb", 100)
|
||||
|
||||
assert max_body_bytes("/files/upload") == 101 * MiB
|
||||
assert max_body_bytes("/marketplace/submissions/media") == 11 * MiB
|
||||
assert max_body_bytes("/graphs") == MAX_BODY_BYTES
|
||||
assert max_body_bytes("/mcp") == MAX_BODY_BYTES
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/files/upload", "/marketplace/submissions/media"])
|
||||
def test_the_upload_caps_name_real_upload_routes(path: str) -> None:
|
||||
"""A moved route would fall back to the JSON cap and refuse its uploads."""
|
||||
routes = {
|
||||
(route.path, method)
|
||||
for route in v2_app.routes
|
||||
if isinstance(route, fastapi.routing.APIRoute)
|
||||
for method in route.methods
|
||||
}
|
||||
|
||||
assert (path, "POST") in routes
|
||||
|
||||
|
||||
async def test_the_cap_follows_the_path_inside_the_mounted_app(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
mocker.patch.object(body_limit.settings.config, "max_file_size_mb", 100)
|
||||
called = []
|
||||
|
||||
async def app(scope, receive, send):
|
||||
called.append(scope)
|
||||
await JSONResponse({}, status_code=201)(scope, receive, send)
|
||||
|
||||
upload = _scope(
|
||||
"/api/v2/files/upload",
|
||||
headers=[(b"content-length", str(50 * MiB).encode())],
|
||||
)
|
||||
upload["root_path"] = "/api/v2"
|
||||
sent = await _call(BodySizeLimitMiddleware(app), upload, _chunks(b""))
|
||||
|
||||
assert called
|
||||
assert _statuses(sent) == [201]
|
||||
|
||||
|
||||
def test_the_v2_app_refuses_an_oversized_body_with_its_security_headers(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Registered inside the security headers, and ahead of the rate limiter."""
|
||||
mocker.patch.object(
|
||||
global_rate_limit._anonymous_limiter,
|
||||
"check",
|
||||
side_effect=AssertionError("rate-limited before the size was checked"),
|
||||
)
|
||||
|
||||
response = TestClient(v2_app).post(
|
||||
"/graphs",
|
||||
content=b"x" * (MAX_BODY_BYTES + 1),
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
assert response.status_code == 413
|
||||
assert response.json()["error"]["code"] == "payload_too_large"
|
||||
assert response.headers["x-content-type-options"] == "nosniff"
|
||||
|
||||
|
||||
def _scope(path: str, headers: list[tuple[bytes, bytes]] | None = None) -> dict:
|
||||
return {
|
||||
"type": "http",
|
||||
"asgi": {"version": "3.0"},
|
||||
"http_version": "1.1",
|
||||
"method": "POST",
|
||||
"scheme": "http",
|
||||
"path": path,
|
||||
"raw_path": path.encode(),
|
||||
"root_path": "",
|
||||
"query_string": b"",
|
||||
"headers": [(b"content-type", b"application/json"), *(headers or [])],
|
||||
"client": ("10.0.0.9", 4242),
|
||||
"server": ("testserver", 80),
|
||||
}
|
||||
|
||||
|
||||
def _chunks(*parts: bytes):
|
||||
"""A receive that delivers *parts* as one body, then reports a disconnect."""
|
||||
pending = list(parts)
|
||||
|
||||
async def receive() -> dict:
|
||||
if not pending:
|
||||
return {"type": "http.disconnect"}
|
||||
body = pending.pop(0)
|
||||
return {"type": "http.request", "body": body, "more_body": bool(pending)}
|
||||
|
||||
return receive
|
||||
|
||||
|
||||
async def _call(app, scope: dict, receive) -> list[dict]:
|
||||
sent: list[dict] = []
|
||||
|
||||
async def send(message: dict) -> None:
|
||||
sent.append(message)
|
||||
|
||||
await app(scope, receive, send)
|
||||
return sent
|
||||
|
||||
|
||||
def _statuses(sent: list[dict]) -> list[int]:
|
||||
return [m["status"] for m in sent if m["type"] == "http.response.start"]
|
||||
|
||||
|
||||
def _body(sent: list[dict]) -> dict:
|
||||
return json.loads(
|
||||
b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")
|
||||
)
|
||||
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
Common utilities for V2 External API
|
||||
"""
|
||||
|
||||
# Constants for pagination
|
||||
MAX_PAGE_SIZE = 100
|
||||
DEFAULT_PAGE_SIZE = 20
|
||||
@@ -0,0 +1,323 @@
|
||||
"""
|
||||
Rules the v2 contract holds across every model, checked against the schema it
|
||||
publishes rather than the source it is written in.
|
||||
|
||||
The v1 surface accumulated four spellings of "credentials for a run" and money
|
||||
in three units because each model was reviewed alone. These tests read the whole
|
||||
OpenAPI document, so a new model cannot reintroduce a spelling the API retired.
|
||||
"""
|
||||
|
||||
from typing import Any, Iterator, get_args
|
||||
|
||||
import prisma.enums
|
||||
import pytest
|
||||
|
||||
import backend.blocks._base as block_types
|
||||
from backend.data.model import CREDENTIALS_ADAPTER
|
||||
|
||||
# `node_exec_id` is deliberately absent: v2 renames graph executions to runs,
|
||||
# and a node execution is not one.
|
||||
RETIRED_FIELD_NAMES = {
|
||||
"input_data": "inputs",
|
||||
"output_data": "outputs",
|
||||
"credentials": "credentials_inputs",
|
||||
"credential_inputs": "credentials_inputs",
|
||||
"agent_credentials": "credentials_inputs",
|
||||
"execution_id": "run_id",
|
||||
"cost_amount": "cost_cents",
|
||||
"duration": "duration_seconds",
|
||||
}
|
||||
|
||||
# A property whose name contains one of these holds an amount of money.
|
||||
MONEY_WORDS = ("cost", "balance", "amount", "price")
|
||||
|
||||
# ...unless it is one of these, which are not amounts.
|
||||
NOT_AMOUNTS = {"cost_type", "cost_filter", "costs"}
|
||||
|
||||
# A property whose name ends with one of these is a point in time.
|
||||
TIME_SUFFIXES = ("_at", "_time", "_end", "_start")
|
||||
|
||||
|
||||
def test_no_model_uses_a_retired_field_name() -> None:
|
||||
offenders = [
|
||||
f"{model}.{field} (use {RETIRED_FIELD_NAMES[field]})"
|
||||
for model, field, _ in _properties()
|
||||
if field in RETIRED_FIELD_NAMES
|
||||
]
|
||||
assert not offenders, f"retired field names in the v2 contract: {offenders}"
|
||||
|
||||
|
||||
def test_every_money_amount_is_named_in_cents() -> None:
|
||||
"""Three units and four spellings is how v1's billing fields ended up."""
|
||||
offenders = [
|
||||
f"{model}.{field}"
|
||||
for model, field, _ in _properties()
|
||||
if field not in NOT_AMOUNTS
|
||||
and any(word in field for word in MONEY_WORDS)
|
||||
and not field.endswith("_cents")
|
||||
and not field.endswith("_count")
|
||||
]
|
||||
assert not offenders, f"money fields without their unit: {offenders}"
|
||||
|
||||
|
||||
def test_every_duration_is_named_in_seconds() -> None:
|
||||
offenders = [
|
||||
f"{model}.{field}"
|
||||
for model, field, _ in _properties()
|
||||
if "duration" in field and not field.endswith(("_seconds", "_count"))
|
||||
]
|
||||
assert not offenders, f"durations without their unit: {offenders}"
|
||||
|
||||
|
||||
def test_every_point_in_time_is_a_timestamp_not_a_number() -> None:
|
||||
"""`current_period_end` was unix seconds beside ISO datetimes in one model."""
|
||||
offenders = [
|
||||
f"{model}.{field}: {schema}"
|
||||
for model, field, schema in _properties()
|
||||
if field.endswith(TIME_SUFFIXES) and _is_numeric(schema)
|
||||
]
|
||||
assert not offenders, f"times published as numbers: {offenders}"
|
||||
|
||||
|
||||
def test_the_schema_has_models_to_check() -> None:
|
||||
"""Guards every test above against passing on an empty document."""
|
||||
assert len(list(_properties())) > 300
|
||||
|
||||
|
||||
def _properties() -> Iterator[tuple[str, str, dict[str, Any]]]:
|
||||
"""Every (model, field, schema) the v2 OpenAPI document publishes."""
|
||||
from .app import v2_app
|
||||
|
||||
schemas = v2_app.openapi()["components"]["schemas"]
|
||||
for model, schema in schemas.items():
|
||||
for field, field_schema in (schema.get("properties") or {}).items():
|
||||
yield model, field, field_schema
|
||||
|
||||
|
||||
def _is_numeric(schema: dict[str, Any]) -> bool:
|
||||
variants = schema.get("anyOf") or schema.get("oneOf") or [schema]
|
||||
return any(variant.get("type") in ("integer", "number") for variant in variants)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def _quiet_mcp_registration() -> None:
|
||||
"""Importing the app registers the MCP tools; nothing here needs the log."""
|
||||
import logging
|
||||
|
||||
logging.getLogger("backend.api.external.v2.mcp_server").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Enums: v2 publishes its own, and they cannot drift from what they mirror
|
||||
# ============================================================================
|
||||
|
||||
# v2 value set -> the internal value set it must equal, and why anything differs.
|
||||
MIRRORED_ENUMS: dict[str, tuple[set[str], set[str], str]] = {}
|
||||
|
||||
|
||||
def _register() -> None:
|
||||
from . import models
|
||||
|
||||
MIRRORED_ENUMS.update(
|
||||
{
|
||||
"RunStatus": (
|
||||
set(get_args(models.RunStatus)),
|
||||
{s.value for s in prisma.enums.AgentExecutionStatus},
|
||||
"",
|
||||
),
|
||||
"AgentRunReviewStatus": (
|
||||
set(get_args(models.AgentRunReviewStatus)),
|
||||
{s.value for s in prisma.enums.ReviewStatus},
|
||||
"",
|
||||
),
|
||||
"TransactionType": (
|
||||
set(get_args(models.TransactionType)),
|
||||
{t.value for t in prisma.enums.CreditTransactionType},
|
||||
"",
|
||||
),
|
||||
"SubscriptionTierValue": (
|
||||
set(get_args(models.SubscriptionTierValue)),
|
||||
{t.value for t in prisma.enums.SubscriptionTier},
|
||||
"",
|
||||
),
|
||||
"SubmissionStatus": (
|
||||
set(get_args(models.SubmissionStatus)),
|
||||
{s.value for s in prisma.enums.SubmissionStatus},
|
||||
"",
|
||||
),
|
||||
"CredentialType": (
|
||||
set(get_args(models.CredentialType)),
|
||||
_credential_union_types(),
|
||||
"",
|
||||
),
|
||||
"SearchContentType": (
|
||||
{t.value for t in models.SearchContentType},
|
||||
{t.value for t in prisma.enums.ContentType} - {"CHAT_SESSION"},
|
||||
"chat sessions have no v2 surface and no scope that grants them",
|
||||
),
|
||||
"BlockCostType": (
|
||||
{t.value for t in models.BlockCostType},
|
||||
{t.value for t in block_types.BlockCostType},
|
||||
"",
|
||||
),
|
||||
"BlockType": (
|
||||
{t.name for t in models.BlockType},
|
||||
{t.name for t in block_types.BlockType},
|
||||
"compared by member name: the internal values are UI labels",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_no_v2_enum_has_drifted_from_the_enum_it_mirrors() -> None:
|
||||
"""A hand-copied enum is silent when its source grows: the new value simply
|
||||
fails response validation in production. These pin every copy."""
|
||||
_register()
|
||||
drifted = {
|
||||
name: {
|
||||
"missing from v2": sorted(source - published),
|
||||
"not in the source": sorted(published - source),
|
||||
"deliberately": reason,
|
||||
}
|
||||
for name, (published, source, reason) in MIRRORED_ENUMS.items()
|
||||
if published != source
|
||||
}
|
||||
assert not drifted, f"v2 enums out of step with their sources: {drifted}"
|
||||
|
||||
|
||||
def test_every_mirrored_enum_is_registered() -> None:
|
||||
"""Guards the test above against an empty table."""
|
||||
_register()
|
||||
assert len(MIRRORED_ENUMS) == 9
|
||||
|
||||
|
||||
def _credential_union_types() -> set[str]:
|
||||
"""The `type` discriminator values a `Credentials` object can actually carry.
|
||||
|
||||
Not `data.model.CredentialsType`, which also lists `device_code` — no member
|
||||
of the union declares it, and a device-code grant stores an
|
||||
`OAuth2Credentials`.
|
||||
"""
|
||||
return set(CREDENTIALS_ADAPTER.core_schema["choices"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Credential requirements: read the schema graphs actually publish
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def test_credential_requirements_are_read_from_a_real_graph_schema() -> None:
|
||||
"""The parser looked for `provider` or an `anyOf` of them; graphs publish
|
||||
`credentials_provider` as a list, so both requirement endpoints answered 200
|
||||
with an empty list for every agent that needs a credential."""
|
||||
from unittest import mock
|
||||
|
||||
from .integrations import helpers
|
||||
|
||||
schema = _one_block_graph_credentials_schema()
|
||||
assert schema["properties"], "fixture graph publishes no credentials field"
|
||||
|
||||
with mock.patch.object(
|
||||
helpers.creds_manager.store,
|
||||
"get_all_creds",
|
||||
new=mock.AsyncMock(return_value=[]),
|
||||
):
|
||||
requirements = await helpers.get_credential_requirements(schema, "user-1")
|
||||
|
||||
assert requirements, f"no requirements parsed from {schema}"
|
||||
assert requirements[0].provider
|
||||
assert requirements[0].field_name in schema["properties"]
|
||||
assert requirements[0].supported_types
|
||||
|
||||
|
||||
def _one_block_graph_credentials_schema() -> dict[str, Any]:
|
||||
"""`credentials_input_schema` of a graph holding one credential-taking block."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from backend.blocks import get_blocks
|
||||
from backend.data.graph import GraphModel, NodeModel
|
||||
|
||||
block_id = next(
|
||||
block_id
|
||||
for block_id, block in get_blocks().items()
|
||||
if any(
|
||||
"credentials" in field
|
||||
for field in block().input_schema.jsonschema().get("properties", {})
|
||||
)
|
||||
)
|
||||
graph = GraphModel(
|
||||
id="g",
|
||||
version=1,
|
||||
name="fixture",
|
||||
description="",
|
||||
user_id="u",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
nodes=[
|
||||
NodeModel(
|
||||
id="n1",
|
||||
block_id=block_id,
|
||||
input_default={},
|
||||
graph_id="g",
|
||||
graph_version=1,
|
||||
)
|
||||
],
|
||||
links=[],
|
||||
)
|
||||
return graph.credentials_input_schema
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# A write says what it replaces
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_a_submission_edit_cannot_silently_blank_what_it_omits() -> None:
|
||||
"""`PUT` writes every field, so a defaulted one is data loss: a caller that
|
||||
sent `{"name": "..."}` erased the listing's description, images and
|
||||
categories. Required fields make the omission a 422 instead."""
|
||||
from .app import v2_app
|
||||
|
||||
edit = v2_app.openapi()["components"]["schemas"][
|
||||
"MarketplaceAgentSubmissionEditRequest"
|
||||
]
|
||||
written_by_the_route = {
|
||||
"name",
|
||||
"sub_heading",
|
||||
"description",
|
||||
"image_urls",
|
||||
"categories",
|
||||
}
|
||||
assert written_by_the_route <= set(edit.get("required", []))
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# An absent value is null, not an empty string
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_a_listed_agent_without_an_image_reports_null() -> None:
|
||||
"""The store fills a missing image with `""`; the field is `Optional[str]`,
|
||||
and the details endpoint already answers `None` for the same agent."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from backend.api.features.store.model import StoreAgent
|
||||
|
||||
from .models import MarketplaceAgent
|
||||
|
||||
def listed(images: list[str]) -> MarketplaceAgent:
|
||||
row = SimpleNamespace(
|
||||
slug="agent",
|
||||
agent_name="Agent",
|
||||
agent_image=images,
|
||||
creator_username="creator",
|
||||
creator_avatar=None,
|
||||
sub_heading="",
|
||||
description="",
|
||||
runs=0,
|
||||
rating=0.0,
|
||||
graph_id="graph-1",
|
||||
)
|
||||
return MarketplaceAgent.from_internal(StoreAgent.from_db(row)) # type: ignore[arg-type]
|
||||
|
||||
assert listed([]).image_url is None
|
||||
assert listed(["https://img/1.png"]).image_url == "https://img/1.png"
|
||||
@@ -0,0 +1,975 @@
|
||||
"""
|
||||
Regression tests for the v2 external API contract.
|
||||
|
||||
Each test here pins a defect that shipped once and was invisible to the rest of
|
||||
the suite: a schema that only fails when FastAPI lazily builds it, a route that
|
||||
only fails once a sibling route is registered, and an auth path only the MCP
|
||||
transport exercises.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Optional, get_origin, get_type_hints
|
||||
from unittest import mock
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pydantic
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException
|
||||
from fastapi.dependencies.utils import get_flat_params
|
||||
from fastapi.routing import APIRoute
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
from prisma.enums import APIKeyPermission, ReviewStatus
|
||||
from starlette.requests import Request
|
||||
from starlette.routing import Match, Route
|
||||
|
||||
from backend.api.external.middleware import resolve_auth_info
|
||||
from backend.api.external.v2.errors import add_v2_exception_handlers
|
||||
from backend.api.external.v2.pagination import (
|
||||
MAX_PAGE,
|
||||
Page,
|
||||
PageRequest,
|
||||
encode_page_cursor,
|
||||
encode_token_cursor,
|
||||
single_page_request,
|
||||
)
|
||||
from backend.api.external.v2.tenancy import TenantContext, require_auth
|
||||
from backend.api.features.graph_executions.review.model import PendingHumanReviewModel
|
||||
from backend.api.features.library.exceptions import (
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
from backend.api.utils.rate_limit import RateLimiter
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.util.exceptions import NotAuthorizedError, NotFoundError
|
||||
|
||||
TEST_USER_ID = "test-user-id"
|
||||
|
||||
_api_key_auth = APIAuthorizationInfo(
|
||||
user_id=TEST_USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
created_at=datetime.now(tz=timezone.utc),
|
||||
)
|
||||
|
||||
_tenant = TenantContext(
|
||||
user_id=TEST_USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="test-org-id",
|
||||
)
|
||||
|
||||
|
||||
def test_openapi_schema_builds():
|
||||
"""Every response model must be resolvable at runtime.
|
||||
|
||||
`/docs`, `/redoc` and `/openapi.json` all 500 when one is not, and FastAPI
|
||||
builds response models lazily, so importing the app is not enough.
|
||||
"""
|
||||
from backend.api.external.v2.app import v2_app
|
||||
|
||||
schema = v2_app.openapi()
|
||||
|
||||
assert schema["paths"], "v2 OpenAPI schema has no paths"
|
||||
assert "/search" in schema["paths"]
|
||||
|
||||
|
||||
async def test_the_served_spec_lists_only_the_declared_servers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""FastAPI otherwise serves the mount path as the first server, a relative URL
|
||||
that is wrong behind a proxy prefix such as the single container's `/_agpt`."""
|
||||
from backend.api.external.v2.app import v2_app
|
||||
|
||||
declared = [server["url"] for server in v2_app.servers]
|
||||
monkeypatch.setattr(v2_app, "openapi_schema", None)
|
||||
monkeypatch.setattr(v2_app, "servers", list(v2_app.servers))
|
||||
(served,) = [
|
||||
route
|
||||
for route in v2_app.routes
|
||||
if isinstance(route, Route) and route.path == v2_app.openapi_url
|
||||
]
|
||||
|
||||
response = await served.endpoint(
|
||||
Request({"type": "http", "root_path": "/_agpt/api/v2"})
|
||||
)
|
||||
|
||||
assert [s["url"] for s in json.loads(response.body)["servers"]] == declared
|
||||
|
||||
|
||||
def test_static_routes_are_not_shadowed_by_path_params():
|
||||
"""Starlette matches in registration order.
|
||||
|
||||
A static path registered after a sibling `/{id}` route is unreachable: the
|
||||
parameterised route matches the literal segment first.
|
||||
"""
|
||||
from backend.api.external.v2.routes import v2_router
|
||||
|
||||
for method, path, expected_endpoint in [
|
||||
("GET", "/runs/reviews", "list_reviews"),
|
||||
("GET", "/runs", "list_runs"),
|
||||
]:
|
||||
endpoint = _first_matching_endpoint(v2_router, method, path)
|
||||
assert (
|
||||
endpoint == expected_endpoint
|
||||
), f"{method} {path} resolves to {endpoint}, not {expected_endpoint}"
|
||||
|
||||
|
||||
async def test_reviews_are_reachable_over_http():
|
||||
"""`GET /runs/reviews` end-to-end, not just through route matching.
|
||||
|
||||
Before the ordering fix this answered `404 Run #reviews not found` from
|
||||
`get_run`.
|
||||
"""
|
||||
from backend.util.models import Pagination
|
||||
|
||||
from .runs import runs_router
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(runs_router, prefix="/runs")
|
||||
app.dependency_overrides[require_auth] = lambda: _tenant
|
||||
add_v2_exception_handlers(app)
|
||||
|
||||
async def no_reviews(**kwargs) -> tuple[list, Pagination]:
|
||||
return [], Pagination(
|
||||
total_items=0, total_pages=0, current_page=1, page_size=20
|
||||
)
|
||||
|
||||
async def no_such_run(**kwargs) -> None:
|
||||
return None
|
||||
|
||||
# `get_run` is stubbed too: when the ordering regresses the request lands
|
||||
# there, and this keeps the failure a legible 404 instead of a DB error.
|
||||
with (
|
||||
mock.patch("backend.data.human_review.get_reviews", new=no_reviews),
|
||||
mock.patch("backend.data.execution.get_graph_execution", new=no_such_run),
|
||||
):
|
||||
response = fastapi.testclient.TestClient(app).get("/runs/reviews")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["items"] == []
|
||||
|
||||
|
||||
async def test_submitting_reviews_resumes_the_run(mocker: pytest_mock.MockFixture):
|
||||
"""One approval and one rejection, after which nothing is left pending.
|
||||
|
||||
The resume is the point of the endpoint: the graph sits parked until the
|
||||
last review lands, so a failure here strands the run rather than erroring.
|
||||
"""
|
||||
approved = _review("node-a", ReviewStatus.APPROVED)
|
||||
rejected = _review("node-b", ReviewStatus.REJECTED)
|
||||
close_step_cards = _mock_review_db(
|
||||
mocker,
|
||||
pending=[approved, rejected],
|
||||
processed={"node-a": approved, "node-b": rejected},
|
||||
still_pending=False,
|
||||
)
|
||||
add_graph_execution = _mock_resume_path(mocker)
|
||||
|
||||
result = await _submit(
|
||||
{"node_exec_id": "node-a", "approved": True},
|
||||
{"node_exec_id": "node-b", "approved": False, "message": "no"},
|
||||
)
|
||||
|
||||
assert (result.run_id, result.approved_count, result.rejected_count) == (
|
||||
"run-1",
|
||||
1,
|
||||
1,
|
||||
)
|
||||
add_graph_execution.assert_awaited_once()
|
||||
assert add_graph_execution.await_args.kwargs["graph_exec_id"] == "run-1"
|
||||
# Answered here, a step's approval card in a chat is closed, as on the run page.
|
||||
assert set(close_step_cards.await_args.args[1]) == {"node-a", "node-b"}
|
||||
|
||||
|
||||
async def test_reviews_still_pending_leave_the_run_parked(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
approved = _review("node-a", ReviewStatus.APPROVED)
|
||||
_mock_review_db(
|
||||
mocker,
|
||||
pending=[approved],
|
||||
processed={"node-a": approved},
|
||||
still_pending=True,
|
||||
)
|
||||
add_graph_execution = _mock_resume_path(mocker)
|
||||
|
||||
result = await _submit({"node_exec_id": "node-a", "approved": True})
|
||||
|
||||
assert result.approved_count == 1
|
||||
add_graph_execution.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_review_id_from_another_run_is_rejected(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
_mock_review_db(
|
||||
mocker,
|
||||
pending=[_review("node-a", ReviewStatus.WAITING)],
|
||||
processed={},
|
||||
still_pending=True,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _submit({"node_exec_id": "node-from-another-run", "approved": True})
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "node-from-another-run" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
async def test_deleting_an_unknown_schedule_is_a_404_by_type_not_by_message(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The route matched "not found" in the message; a reworded error broke it."""
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
from . import schedules
|
||||
|
||||
mocker.patch.object(
|
||||
schedules,
|
||||
"_assert_schedule_in_tenant",
|
||||
new=mock.AsyncMock(),
|
||||
)
|
||||
mocker.patch.object(
|
||||
schedules,
|
||||
"get_scheduler_client",
|
||||
return_value=mock.Mock(
|
||||
delete_schedule=mock.AsyncMock(side_effect=NotFoundError("Job #s-1 gone."))
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await schedules.delete_schedule(schedule_id="s-1", auth=_tenant)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
async def test_listing_runs_passes_the_status_and_time_filters_through(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
from backend.data.execution import ExecutionStatus
|
||||
|
||||
from . import runs
|
||||
|
||||
paginated = mocker.patch.object(
|
||||
runs.execution_db,
|
||||
"get_graph_executions_paginated",
|
||||
new=mock.AsyncMock(
|
||||
return_value=mock.Mock(executions=[], pagination=mock.Mock(total_items=0))
|
||||
),
|
||||
)
|
||||
since = datetime(2026, 9, 1, tzinfo=timezone.utc)
|
||||
|
||||
await runs.list_runs(
|
||||
graph_id=None,
|
||||
statuses=["RUNNING", "FAILED"],
|
||||
started_after=since,
|
||||
started_before=None,
|
||||
page=PageRequest(limit=25),
|
||||
auth=_tenant,
|
||||
)
|
||||
|
||||
kwargs = paginated.await_args.kwargs
|
||||
assert kwargs["statuses"] == [ExecutionStatus.RUNNING, ExecutionStatus.FAILED]
|
||||
assert kwargs["created_time_gte"] == since
|
||||
|
||||
|
||||
async def test_the_review_status_filter_keeps_its_wire_name(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The parameter is `status` on the wire; naming it that in Python shadowed
|
||||
the `starlette.status` the rest of the module raises through."""
|
||||
from . import runs
|
||||
|
||||
parameter = next(
|
||||
p
|
||||
for p in inspect.signature(runs.list_reviews).parameters.values()
|
||||
if p.name == "review_status"
|
||||
)
|
||||
assert parameter.default.alias == "status"
|
||||
|
||||
|
||||
async def test_the_marketplace_listing_lookup_goes_through_the_store_db(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The route held a raw Prisma query, so the store's own view was bypassed."""
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
from . import graphs
|
||||
|
||||
mocker.patch.object(
|
||||
graphs.store_db,
|
||||
"get_store_agent_by_graph_id",
|
||||
new=mock.AsyncMock(side_effect=NotFoundError("nope")),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await graphs.get_marketplace_listing_for_graph(graph_id="graph-1", auth=_tenant)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
assert "#graph-1" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
async def test_a_review_belonging_to_another_run_is_not_reachable_through_this_one(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The URL's run is tenancy-checked; the reviews must be that run's own.
|
||||
|
||||
`get_reviews_by_node_exec_ids` scopes to the user, not to the org, so a
|
||||
review of a run in another tenant would otherwise be decidable here.
|
||||
"""
|
||||
elsewhere = _review("node-a", ReviewStatus.WAITING)
|
||||
elsewhere.graph_exec_id = "run-2"
|
||||
_mock_review_db(mocker, pending=[elsewhere], processed={}, still_pending=True)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _submit({"node_exec_id": "node-a", "approved": True})
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
async def test_reviews_are_refused_on_a_run_that_is_not_awaiting_them(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""v2 had no status check where the internal route answers 409."""
|
||||
from backend.api.features.graph_executions.review import service as review_service
|
||||
from backend.data.execution import ExecutionStatus
|
||||
|
||||
_mock_review_db(
|
||||
mocker,
|
||||
pending=[_review("node-a", ReviewStatus.WAITING)],
|
||||
processed={},
|
||||
still_pending=True,
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_graph_execution_meta",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(status=ExecutionStatus.COMPLETED)),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _submit({"node_exec_id": "node-a", "approved": True})
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
|
||||
|
||||
async def test_an_auto_approved_review_records_the_node_and_ignores_edits(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The internal route has always supported this; v2 dropped it silently."""
|
||||
from backend.api.features.graph_executions.review import service as review_service
|
||||
|
||||
approved = _review("node-a", ReviewStatus.APPROVED)
|
||||
_mock_review_db(
|
||||
mocker,
|
||||
pending=[approved],
|
||||
processed={"node-a": approved},
|
||||
still_pending=True,
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_node_executions",
|
||||
new=mock.AsyncMock(
|
||||
return_value=[mock.Mock(node_exec_id="node-a", node_id="node-def-1")]
|
||||
),
|
||||
)
|
||||
record = mocker.patch.object(
|
||||
review_service, "create_auto_approval_record", new=mock.AsyncMock()
|
||||
)
|
||||
|
||||
await _submit(
|
||||
{
|
||||
"node_exec_id": "node-a",
|
||||
"approved": True,
|
||||
"edited_payload": {"ignored": True},
|
||||
"auto_approve_future": True,
|
||||
}
|
||||
)
|
||||
|
||||
record.assert_awaited_once()
|
||||
assert record.await_args.kwargs["node_id"] == "node-def-1"
|
||||
decisions = review_service.process_all_reviews_for_execution.await_args.kwargs[
|
||||
"review_decisions"
|
||||
]
|
||||
assert decisions["node-a"][1] is None
|
||||
|
||||
|
||||
async def test_a_failed_resume_does_not_fail_the_submission(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The decisions are already persisted; a resume failure must not undo them."""
|
||||
approved = _review("node-a", ReviewStatus.APPROVED)
|
||||
_mock_review_db(
|
||||
mocker,
|
||||
pending=[approved],
|
||||
processed={"node-a": approved},
|
||||
still_pending=False,
|
||||
)
|
||||
_mock_resume_path(mocker).side_effect = RuntimeError("executor unreachable")
|
||||
|
||||
result = await _submit({"node_exec_id": "node-a", "approved": True})
|
||||
|
||||
assert result.approved_count == 1
|
||||
|
||||
|
||||
async def test_a_string_without_the_key_prefix_is_not_an_api_key():
|
||||
"""The prefix check short-circuits before any database lookup."""
|
||||
from backend.data.auth.api_key import validate_api_key
|
||||
|
||||
assert await validate_api_key("not-an-agpt-key") is None
|
||||
|
||||
|
||||
async def test_api_key_sent_as_bearer_resolves_to_its_user(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""MCP clients can only send credentials in the `Authorization` header.
|
||||
|
||||
Resolving those to `None` silently drops the request into the anonymous
|
||||
per-IP rate-limit bucket (5 req/min), which an MCP handshake exhausts.
|
||||
"""
|
||||
mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
return_value=_api_key_auth,
|
||||
)
|
||||
validate_access_token = mocker.patch(
|
||||
"backend.api.external.middleware.validate_access_token"
|
||||
)
|
||||
|
||||
auth = await resolve_auth_info(
|
||||
api_key=None,
|
||||
bearer=HTTPAuthorizationCredentials(scheme="Bearer", credentials="agpt_test"),
|
||||
)
|
||||
|
||||
assert auth is not None and auth.user_id == TEST_USER_ID
|
||||
validate_access_token.assert_not_called()
|
||||
|
||||
|
||||
async def test_invalid_bearer_still_raises_401(mocker: pytest_mock.MockFixture):
|
||||
from backend.data.auth.oauth import InvalidTokenError
|
||||
|
||||
mocker.patch("backend.api.external.middleware.validate_api_key", return_value=None)
|
||||
mocker.patch(
|
||||
"backend.api.external.middleware.validate_access_token",
|
||||
side_effect=InvalidTokenError("nope"),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await resolve_auth_info(
|
||||
api_key=None,
|
||||
bearer=HTTPAuthorizationCredentials(scheme="Bearer", credentials="junk"),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
async def test_bearer_api_key_gets_the_authenticated_rate_limit(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""The consequence of the above, at the layer a client actually feels."""
|
||||
from backend.api.external.v2 import global_rate_limit
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
return_value=_api_key_auth,
|
||||
)
|
||||
authenticated = mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
anonymous = mocker.patch.object(
|
||||
global_rate_limit._anonymous_limiter, "check", return_value=None
|
||||
)
|
||||
# The pre-verification counters, which this test isn't about.
|
||||
mocker.patch.object(RateLimiter, "exhausted", return_value=False)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._presented_credential_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
await _call_rate_limit_middleware(headers=[(b"authorization", b"Bearer agpt_test")])
|
||||
|
||||
authenticated.assert_awaited_once_with(TEST_USER_ID)
|
||||
anonymous.assert_not_awaited()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Contract shape: one envelope, one error body, deterministic status codes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Collection endpoints whose name does not start with `list_`.
|
||||
_EXTRA_COLLECTION_ENDPOINTS = {"search", "get_folder_tree"}
|
||||
|
||||
# Non-2xx responses every operation must document with the error envelope.
|
||||
_DOCUMENTED_ERROR_STATUSES = {"400", "401", "403", "404", "413", "422", "429", "500"}
|
||||
|
||||
|
||||
def test_every_collection_endpoint_returns_the_page_envelope():
|
||||
"""`{"items": [...], "next_cursor": ...}` is the only list shape in v2."""
|
||||
offenders = []
|
||||
for route in _v2_routes():
|
||||
returns = get_type_hints(route.endpoint).get("return")
|
||||
is_collection = (
|
||||
route.name.startswith("list_") or route.name in _EXTRA_COLLECTION_ENDPOINTS
|
||||
)
|
||||
if is_collection and not _is_page(returns):
|
||||
offenders.append(f"{route.name} returns {returns}, not Page[...]")
|
||||
if get_origin(returns) is list:
|
||||
offenders.append(f"{route.name} returns a bare {returns}")
|
||||
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
||||
|
||||
def test_every_page_response_carries_all_three_fields():
|
||||
"""`total_count` is present on every list, even where its value is null."""
|
||||
from backend.api.external.v2.app import v2_app
|
||||
|
||||
schemas = v2_app.openapi()["components"]["schemas"]
|
||||
page_schemas = {k: v for k, v in schemas.items() if k.startswith("Page_")}
|
||||
|
||||
assert page_schemas
|
||||
for name, schema in page_schemas.items():
|
||||
assert set(schema["required"]) == {
|
||||
"items",
|
||||
"next_cursor",
|
||||
"total_count",
|
||||
}, f"{name} requires {schema['required']}"
|
||||
|
||||
|
||||
def test_every_collection_endpoint_takes_limit_and_cursor():
|
||||
for route in _v2_routes():
|
||||
if not _is_page(get_type_hints(route.endpoint).get("return")):
|
||||
continue
|
||||
params = {p.alias for p in get_flat_params(route.dependant)}
|
||||
assert {"limit", "cursor"} <= params, f"{route.name} query params: {params}"
|
||||
|
||||
|
||||
def test_status_codes_follow_one_rule_per_method():
|
||||
"""200 read/update, 204 delete. POST is disambiguated by the table below."""
|
||||
allowed = {
|
||||
"GET": {200},
|
||||
"POST": {200, 201, 202},
|
||||
"PUT": {200, 204},
|
||||
"PATCH": {200},
|
||||
"DELETE": {204},
|
||||
}
|
||||
offenders = [
|
||||
f"{sorted(r.methods)[0]} {r.path} -> {r.status_code or 200} ({r.name})"
|
||||
for r in _v2_routes()
|
||||
if (r.status_code or 200) not in allowed[sorted(r.methods)[0]]
|
||||
]
|
||||
assert not offenders, "\n".join(offenders)
|
||||
|
||||
|
||||
def test_every_post_declares_whether_it_creates_or_enqueues():
|
||||
"""201 when a resource comes into being, 202 when work is queued, else 200.
|
||||
|
||||
Listed one by one so a new POST has to make the choice, not inherit a default.
|
||||
"""
|
||||
expected = {
|
||||
"create_credential": 201,
|
||||
"create_folder": 201,
|
||||
"create_graph": 201,
|
||||
"create_preset": 201,
|
||||
"create_schedule": 201,
|
||||
"create_submission": 201,
|
||||
"setup_trigger": 201,
|
||||
"fork_library_agent": 201,
|
||||
"add_agent_to_library": 201,
|
||||
"enable_sharing": 201,
|
||||
"upload_file": 201,
|
||||
"upload_submission_media": 201,
|
||||
"execute_agent": 202,
|
||||
"execute_preset": 202,
|
||||
"stop_run": 202,
|
||||
"submit_reviews": 202,
|
||||
"move_folder": 200,
|
||||
}
|
||||
actual = {r.name: r.status_code or 200 for r in _v2_routes() if "POST" in r.methods}
|
||||
assert actual == expected
|
||||
|
||||
|
||||
def test_every_operation_documents_the_error_envelope():
|
||||
from backend.api.external.v2.app import v2_app
|
||||
from backend.api.external.v2.errors import ErrorResponse
|
||||
|
||||
expected_ref = f"#/components/schemas/{ErrorResponse.__name__}"
|
||||
offenders = []
|
||||
for path, operations in v2_app.openapi()["paths"].items():
|
||||
for method, operation in operations.items():
|
||||
for code, response in operation["responses"].items():
|
||||
if code not in _DOCUMENTED_ERROR_STATUSES:
|
||||
continue
|
||||
schema = response.get("content", {}).get("application/json", {})
|
||||
if schema.get("schema", {}).get("$ref") != expected_ref:
|
||||
offenders.append(f"{method.upper()} {path} {code}: {schema}")
|
||||
|
||||
assert not offenders, "\n".join(offenders)
|
||||
assert _DOCUMENTED_ERROR_STATUSES <= set(
|
||||
next(iter(v2_app.openapi()["paths"]["/runs"].values()))["responses"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("code", ["402", "409", "503"])
|
||||
def test_run_start_documents_its_payment_and_idempotency_refusals(code: str):
|
||||
"""402: no balance or no plan. 409: the `Idempotency-Key`'s run is not recorded
|
||||
yet. 503: the key cannot be checked, so the run is not started."""
|
||||
from backend.api.external.v2.app import v2_app
|
||||
from backend.api.external.v2.errors import ErrorResponse
|
||||
|
||||
run_start = v2_app.openapi()["paths"]["/library/agents/{agent_id}/runs"]["post"]
|
||||
|
||||
schema = run_start["responses"][code]["content"]["application/json"]["schema"]
|
||||
assert schema["$ref"] == f"#/components/schemas/{ErrorResponse.__name__}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exception,expected_status,expected_code",
|
||||
[
|
||||
(NotFoundError("gone"), 404, "not_found"),
|
||||
(NotAuthorizedError("nope"), 403, "forbidden"),
|
||||
(ValueError("bad"), 400, "bad_request"),
|
||||
(RuntimeError("boom"), 500, "internal_error"),
|
||||
(HTTPException(status_code=409, detail="taken"), 409, "conflict"),
|
||||
(LibraryAgentInAnotherOrganizationError("elsewhere"), 409, "conflict"),
|
||||
],
|
||||
)
|
||||
def test_error_bodies_all_have_the_envelope(
|
||||
exception: Exception, expected_status: int, expected_code: str
|
||||
):
|
||||
response = _error_client(exception).get("/boom")
|
||||
|
||||
assert response.status_code == expected_status
|
||||
body = response.json()
|
||||
assert set(body) == {"error"}
|
||||
assert set(body["error"]) == {"code", "message", "details"}
|
||||
assert body["error"]["code"] == expected_code
|
||||
|
||||
|
||||
def test_validation_errors_have_the_envelope():
|
||||
response = _error_client(None).get("/typed?number=not-a-number")
|
||||
|
||||
assert response.status_code == 422
|
||||
body = response.json()
|
||||
assert body["error"]["code"] == "validation_error"
|
||||
assert body["error"]["details"]["errors"]
|
||||
|
||||
|
||||
def test_field_validator_errors_have_the_envelope():
|
||||
"""A `@field_validator` leaves a live exception in the error's `ctx`.
|
||||
|
||||
Serializing that without a fallback raises a `ValueError` subclass, which
|
||||
the 400 handler then catches — turning a 422 into a 400 whose message is a
|
||||
serializer complaint.
|
||||
"""
|
||||
response = _error_client(None).post("/validated", json={"host": ""})
|
||||
|
||||
assert response.status_code == 422
|
||||
body = response.json()
|
||||
assert body["error"]["code"] == "validation_error"
|
||||
assert "Invalid hostname" in json.dumps(body["error"]["details"])
|
||||
|
||||
|
||||
def test_a_paywalled_request_is_not_logged_as_a_server_warning(caplog):
|
||||
"""402 means the caller lacks credits, which is not a server event."""
|
||||
from backend.copilot.rate_limit import UserPaywalledError
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="backend.api.external.v2.errors"):
|
||||
response = _error_client(UserPaywalledError("no credits")).get("/boom")
|
||||
|
||||
assert response.status_code == 402
|
||||
assert not caplog.records
|
||||
|
||||
|
||||
def test_http_exception_headers_are_forwarded():
|
||||
response = _error_client(
|
||||
HTTPException(
|
||||
status_code=429, detail="slow down", headers={"Retry-After": "30"}
|
||||
)
|
||||
).get("/boom")
|
||||
|
||||
assert response.status_code == 429
|
||||
assert response.headers["Retry-After"] == "30"
|
||||
|
||||
|
||||
async def test_rate_limit_body_has_the_envelope(mocker: pytest_mock.MockFixture):
|
||||
"""The limiter is ASGI middleware, outside the handlers — easy to miss."""
|
||||
from backend.api.external.v2 import global_rate_limit
|
||||
|
||||
mocker.patch("backend.api.external.middleware.validate_api_key", return_value=None)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._anonymous_limiter,
|
||||
"check",
|
||||
side_effect=HTTPException(status_code=429, detail="Rate limit exceeded"),
|
||||
)
|
||||
|
||||
sent: list[dict] = []
|
||||
await _call_rate_limit_middleware(headers=[], capture=sent)
|
||||
|
||||
start = next(m for m in sent if m["type"] == "http.response.start")
|
||||
body = json.loads(
|
||||
b"".join(m.get("body", b"") for m in sent if m["type"] == "http.response.body")
|
||||
)
|
||||
assert start["status"] == 429
|
||||
assert body["error"]["code"] == "rate_limit_exceeded"
|
||||
assert body["error"]["message"] == "Rate limit exceeded"
|
||||
|
||||
|
||||
def test_cursors_round_trip_and_reject_garbage():
|
||||
assert PageRequest(limit=10, cursor=None).page == 1
|
||||
assert PageRequest(limit=10, cursor=encode_page_cursor(4)).page == 4
|
||||
assert PageRequest(limit=10, cursor=encode_token_cursor("t")).token == "t"
|
||||
|
||||
for bad in ["not-a-cursor", _forge({"v": 99, "k": "p", "p": 2}), _forge({"p": 2})]:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
PageRequest(limit=10, cursor=bad)
|
||||
assert exc_info.value.status_code == 400, bad
|
||||
|
||||
|
||||
def test_a_cursor_from_another_endpoint_is_rejected():
|
||||
"""A wrong-kind cursor must not silently reset the caller to page 1."""
|
||||
page_cursor = PageRequest(limit=10, cursor=encode_page_cursor(3))
|
||||
token_cursor = PageRequest(limit=10, cursor=encode_token_cursor("2026-01-01"))
|
||||
|
||||
for request, read in [(page_cursor, "token"), (token_cursor, "page")]:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
getattr(request, read)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
page_cursor.uncounted(["a"])
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
page_cursor.keyset(["a"], next_token=None)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
def test_a_single_page_endpoint_rejects_a_cursor_before_it_does_any_work():
|
||||
"""The guard runs in dependency resolution, not the route body.
|
||||
|
||||
`/credits/invoices` calls Stripe; a cursor must not buy an outbound request
|
||||
the endpoint then throws away.
|
||||
"""
|
||||
request = PageRequest(limit=10, cursor=encode_page_cursor(2))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
single_page_request(request)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert single_page_request(PageRequest(limit=10, cursor=None)).limit == 10
|
||||
|
||||
|
||||
def test_the_documented_cursor_is_one_the_encoder_emits():
|
||||
"""A reader copying the API guide's example must not get a 400."""
|
||||
documented = "eyJ2IjoxLCJrIjoicCIsInAiOjJ9"
|
||||
|
||||
assert encode_page_cursor(2) == documented
|
||||
assert PageRequest(limit=10, cursor=documented).page == 2
|
||||
|
||||
|
||||
def test_an_absurd_page_is_rejected_rather_than_passed_to_the_database():
|
||||
"""An offset that deep is a forged cursor; unbounded it becomes a 500."""
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
PageRequest(limit=100, cursor=encode_page_cursor(MAX_PAGE + 1)).page
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
assert PageRequest(limit=100, cursor=encode_page_cursor(MAX_PAGE)).page == MAX_PAGE
|
||||
|
||||
|
||||
def test_next_cursor_is_null_on_the_last_page():
|
||||
request = PageRequest(limit=10, cursor=None)
|
||||
|
||||
assert request.paged(["a"] * 10, total_count=25).next_cursor is not None
|
||||
assert request.paged(["a"] * 10, total_count=10).next_cursor is None
|
||||
assert request.paged([], total_count=0).next_cursor is None
|
||||
|
||||
|
||||
def test_total_count_reports_the_whole_result_set():
|
||||
request = PageRequest(limit=10, cursor=None)
|
||||
|
||||
assert request.paged(["a"] * 10, total_count=25).total_count == 25
|
||||
assert request.slice(["a"] * 25).total_count == 25
|
||||
assert request.keyset(["a"], next_token=None).total_count is None
|
||||
assert request.uncounted(["a"]).total_count is None
|
||||
|
||||
|
||||
def _forge(payload: dict) -> str:
|
||||
return base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
def _is_page(annotation: Any) -> bool:
|
||||
return isinstance(annotation, type) and issubclass(annotation, Page)
|
||||
|
||||
|
||||
def _v2_routes() -> list[APIRoute]:
|
||||
from backend.api.external.v2.routes import v2_router
|
||||
|
||||
return [r for r in v2_router.routes if isinstance(r, APIRoute)]
|
||||
|
||||
|
||||
def _error_client(exception: Optional[Exception]) -> fastapi.testclient.TestClient:
|
||||
from backend.api.external.v2.errors import add_v2_exception_handlers
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
|
||||
@app.get("/boom")
|
||||
async def boom() -> None:
|
||||
assert exception is not None
|
||||
raise exception
|
||||
|
||||
@app.get("/typed")
|
||||
async def typed(number: int) -> int:
|
||||
return number
|
||||
|
||||
@app.post("/validated")
|
||||
async def validated(body: _ValidatedBody) -> str:
|
||||
return body.host
|
||||
|
||||
add_v2_exception_handlers(app)
|
||||
return fastapi.testclient.TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
async def _submit(*decisions: dict):
|
||||
"""Call the endpoint directly: a TestClient portal would put the resume
|
||||
path's async clients on a different event loop than the test."""
|
||||
from .models import AgentRunReviewsSubmitRequest
|
||||
from .runs import submit_reviews
|
||||
|
||||
return await submit_reviews(
|
||||
request=AgentRunReviewsSubmitRequest.model_validate({"reviews": decisions}),
|
||||
run_id="run-1",
|
||||
auth=_tenant,
|
||||
)
|
||||
|
||||
|
||||
def _review(node_exec_id: str, status: ReviewStatus) -> PendingHumanReviewModel:
|
||||
return PendingHumanReviewModel(
|
||||
node_exec_id=node_exec_id,
|
||||
user_id=TEST_USER_ID,
|
||||
graph_exec_id="run-1",
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
payload={},
|
||||
editable=True,
|
||||
status=status,
|
||||
created_at=datetime.now(tz=timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
def _mock_review_db(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
pending: list[PendingHumanReviewModel],
|
||||
processed: dict[str, PendingHumanReviewModel],
|
||||
still_pending: bool,
|
||||
) -> mock.AsyncMock:
|
||||
"""Stub the review rows; returns the step-card close to assert on."""
|
||||
from backend.api.features.graph_executions.review import service as review_service
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data.execution import ExecutionStatus
|
||||
|
||||
# The caller's own run, carrying the tenancy the reviews are checked against.
|
||||
mocker.patch.object(
|
||||
execution_db,
|
||||
"get_graph_execution",
|
||||
new=mock.AsyncMock(
|
||||
return_value=mock.Mock(
|
||||
organization_id=_tenant.organization_id, user_id=_tenant.user_id
|
||||
)
|
||||
),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_reviews_by_node_exec_ids",
|
||||
new=mock.AsyncMock(return_value={r.node_exec_id: r for r in pending}),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_graph_execution_meta",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(status=ExecutionStatus.REVIEW)),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"process_all_reviews_for_execution",
|
||||
new=mock.AsyncMock(return_value=processed),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"has_pending_reviews_for_graph_exec",
|
||||
new=mock.AsyncMock(return_value=still_pending),
|
||||
)
|
||||
close_step_cards = mock.AsyncMock()
|
||||
mocker.patch.object(review_service, "close_step_cards", new=close_step_cards)
|
||||
return close_step_cards
|
||||
|
||||
|
||||
def _mock_resume_path(mocker: pytest_mock.MockFixture) -> mock.AsyncMock:
|
||||
"""Stub everything the resume needs; returns the enqueue mock to assert on."""
|
||||
from backend.api.features.graph_executions.review import service as review_service
|
||||
from backend.data.graph import GraphSettings
|
||||
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_user_by_id",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(timezone="Europe/Amsterdam")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_graph_settings",
|
||||
new=mock.AsyncMock(return_value=GraphSettings()),
|
||||
)
|
||||
mocker.patch.object(
|
||||
review_service,
|
||||
"get_or_create_workspace",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(id="workspace-1")),
|
||||
)
|
||||
enqueue = mock.AsyncMock()
|
||||
mocker.patch.object(review_service, "add_graph_execution", new=enqueue)
|
||||
return enqueue
|
||||
|
||||
|
||||
def _first_matching_endpoint(router: Any, method: str, path: str) -> str | None:
|
||||
"""Name of the endpoint Starlette would dispatch `method path` to."""
|
||||
scope = {"type": "http", "method": method, "path": path, "headers": []}
|
||||
for route in router.routes:
|
||||
match, _ = route.matches(scope)
|
||||
if match == Match.FULL:
|
||||
return route.endpoint.__name__
|
||||
return None
|
||||
|
||||
|
||||
async def _call_rate_limit_middleware(
|
||||
headers: list[tuple[bytes, bytes]], capture: Optional[list[dict]] = None
|
||||
) -> None:
|
||||
from backend.api.external.v2.global_rate_limit import GlobalRateLimitMiddleware
|
||||
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b""})
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message):
|
||||
if capture is not None:
|
||||
capture.append(message)
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/runs",
|
||||
"headers": headers,
|
||||
"client": ("1.2.3.4", 0),
|
||||
}
|
||||
await GlobalRateLimitMiddleware(app)(scope, receive, send)
|
||||
|
||||
|
||||
class _ValidatedBody(pydantic.BaseModel):
|
||||
host: str
|
||||
|
||||
@pydantic.field_validator("host")
|
||||
@classmethod
|
||||
def _reject_empty(cls, value: str) -> str:
|
||||
if not value:
|
||||
raise ValueError("Invalid hostname: expected a domain like api.example.com")
|
||||
return value
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
V2 External API - Credits Endpoints
|
||||
|
||||
Provides read-only access to credit balance, transaction history,
|
||||
subscription status, invoices, and execution cost summary.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
import stripe
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security
|
||||
from prisma.enums import APIKeyPermission, SubscriptionTier
|
||||
from starlette.concurrency import run_in_threadpool
|
||||
|
||||
from backend.copilot.rate_limit import get_tier_multipliers
|
||||
from backend.data.credit import (
|
||||
PendingChangeUnknown,
|
||||
get_active_subscription_period_end,
|
||||
get_credit_model,
|
||||
get_pending_subscription_change,
|
||||
get_proration_credit_cents,
|
||||
get_subscription_price_id,
|
||||
get_user_billing_cycle,
|
||||
)
|
||||
from backend.data.execution_cost_summary import get_user_cost_summary
|
||||
from backend.data.user import get_user_by_id
|
||||
from backend.util.cache import cached
|
||||
|
||||
from .models import (
|
||||
AutomationCostSummary,
|
||||
CreditBalance,
|
||||
CreditTransaction,
|
||||
InvoiceItem,
|
||||
SubscriptionStatus,
|
||||
TransactionType,
|
||||
)
|
||||
from .pagination import Page, PageRequest, page_request, single_page_request
|
||||
from .rate_limit import enforce, subscription_limiter
|
||||
from .tenancy import TenantContext, require_billing_permission, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
credits_router = APIRouter(tags=["credits"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@credits_router.get(
|
||||
path="",
|
||||
summary="Get credit balance",
|
||||
operation_id="getCreditBalance",
|
||||
)
|
||||
async def get_balance(
|
||||
auth: TenantContext = Security(
|
||||
require_billing_permission(APIKeyPermission.READ_CREDITS)
|
||||
),
|
||||
) -> CreditBalance:
|
||||
"""Get the current credit balance for the authenticated user."""
|
||||
user_credit_model = await get_credit_model(auth.user_id, auth.organization_id)
|
||||
balance = await user_credit_model.get_credits(auth.user_id)
|
||||
|
||||
return CreditBalance(balance_cents=balance)
|
||||
|
||||
|
||||
@credits_router.get(
|
||||
path="/transactions",
|
||||
summary="Get credit transaction history",
|
||||
operation_id="listCreditTransactions",
|
||||
)
|
||||
async def get_transactions(
|
||||
transaction_type: Optional[TransactionType] = Query(
|
||||
default=None, description="Filter by transaction type"
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_billing_permission(APIKeyPermission.READ_CREDITS)
|
||||
),
|
||||
) -> Page[CreditTransaction]:
|
||||
"""Get credit transaction history for the authenticated user.
|
||||
|
||||
`total_count` is `null`: the history groups the raw rows (every node charge
|
||||
of one run collapses into one item), so a row count would not match what
|
||||
paging through this endpoint yields.
|
||||
"""
|
||||
user_credit_model = await get_credit_model(auth.user_id, auth.organization_id)
|
||||
|
||||
# The history's own cursor, as the internal route pages: a time ceiling
|
||||
# alone skipped the groups that tied on time at a page boundary.
|
||||
history = await user_credit_model.get_transaction_history(
|
||||
user_id=auth.user_id,
|
||||
transaction_count_limit=page.limit,
|
||||
transaction_type=transaction_type,
|
||||
cursor=page.token,
|
||||
viewer_organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.keyset(
|
||||
[CreditTransaction.from_internal(t) for t in history.transactions],
|
||||
next_token=history.next_cursor,
|
||||
)
|
||||
|
||||
|
||||
@credits_router.get(
|
||||
path="/subscription",
|
||||
summary="Get subscription status",
|
||||
operation_id="getSubscriptionStatus",
|
||||
)
|
||||
async def get_subscription_status(
|
||||
response: Response,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_CREDITS)),
|
||||
) -> SubscriptionStatus:
|
||||
"""Get the current subscription tier, pricing, and pending changes.
|
||||
|
||||
**Rate limit:** 60 requests per minute per user.
|
||||
"""
|
||||
await enforce(subscription_limiter, auth.user_id, response)
|
||||
|
||||
user = await get_user_by_id(auth.user_id)
|
||||
tier = user.subscription_tier or SubscriptionTier.NO_TIER
|
||||
|
||||
priceable_tiers = [
|
||||
SubscriptionTier.BASIC,
|
||||
SubscriptionTier.PRO,
|
||||
SubscriptionTier.MAX,
|
||||
SubscriptionTier.BUSINESS,
|
||||
]
|
||||
monthly_price_ids, yearly_price_ids = await asyncio.gather(
|
||||
asyncio.gather(
|
||||
*[get_subscription_price_id(t, "monthly") for t in priceable_tiers]
|
||||
),
|
||||
asyncio.gather(
|
||||
*[get_subscription_price_id(t, "yearly") for t in priceable_tiers]
|
||||
),
|
||||
)
|
||||
|
||||
monthly_costs, yearly_costs = await asyncio.gather(
|
||||
asyncio.gather(*[_get_stripe_price_amount(pid) for pid in monthly_price_ids]),
|
||||
asyncio.gather(*[_get_stripe_price_amount(pid) for pid in yearly_price_ids]),
|
||||
)
|
||||
|
||||
tier_costs: dict[str, int] = {}
|
||||
tier_costs_yearly: dict[str, int] = {}
|
||||
for t, m_pid, y_pid, m_cost, y_cost in zip(
|
||||
priceable_tiers,
|
||||
monthly_price_ids,
|
||||
yearly_price_ids,
|
||||
monthly_costs,
|
||||
yearly_costs,
|
||||
):
|
||||
if m_pid or y_pid:
|
||||
tier_costs[t.value] = (m_cost or 0) if m_pid else 0
|
||||
if y_pid:
|
||||
tier_costs_yearly[t.value] = y_cost or 0
|
||||
|
||||
multipliers = await get_tier_multipliers()
|
||||
tier_multipliers: dict[str, float] = {
|
||||
t.value: multipliers.get(t.value, 1.0)
|
||||
for t in priceable_tiers
|
||||
if t.value in tier_costs
|
||||
}
|
||||
|
||||
user_cycle = await get_user_billing_cycle(auth.user_id) or "monthly"
|
||||
if user_cycle == "yearly":
|
||||
current_monthly_cost = tier_costs_yearly.get(tier.value, 0)
|
||||
else:
|
||||
current_monthly_cost = tier_costs.get(tier.value, 0)
|
||||
|
||||
proration_credit, current_period_end = await asyncio.gather(
|
||||
get_proration_credit_cents(auth.user_id, current_monthly_cost),
|
||||
get_active_subscription_period_end(auth.user_id),
|
||||
)
|
||||
|
||||
try:
|
||||
pending = await get_pending_subscription_change(auth.user_id)
|
||||
except (stripe.StripeError, PendingChangeUnknown):
|
||||
logger.exception(
|
||||
"get_subscription_status: failed to resolve pending change for user %s",
|
||||
auth.user_id,
|
||||
)
|
||||
pending = None
|
||||
|
||||
result = SubscriptionStatus(
|
||||
tier=tier.value,
|
||||
monthly_cost_cents=current_monthly_cost,
|
||||
tier_costs_cents=tier_costs,
|
||||
tier_costs_yearly_cents=tier_costs_yearly,
|
||||
billing_cycle=user_cycle,
|
||||
tier_multipliers=tier_multipliers,
|
||||
proration_credit_cents=proration_credit,
|
||||
has_active_stripe_subscription=current_period_end is not None,
|
||||
# Stripe reports the period end as a unix timestamp; every other v2
|
||||
# timestamp is an ISO datetime, so it is converted at the boundary.
|
||||
current_period_end=(
|
||||
datetime.fromtimestamp(current_period_end, tz=timezone.utc)
|
||||
if current_period_end is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
if pending is not None:
|
||||
pending_tier_enum, pending_effective_at, pending_cycle = pending
|
||||
if pending_tier_enum in (
|
||||
SubscriptionTier.NO_TIER,
|
||||
SubscriptionTier.BASIC,
|
||||
SubscriptionTier.PRO,
|
||||
SubscriptionTier.MAX,
|
||||
SubscriptionTier.BUSINESS,
|
||||
):
|
||||
result.pending_tier = pending_tier_enum.value
|
||||
result.pending_tier_effective_at = pending_effective_at
|
||||
result.pending_billing_cycle = pending_cycle
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@credits_router.get(
|
||||
path="/invoices",
|
||||
summary="List Stripe invoices",
|
||||
operation_id="listCreditInvoices",
|
||||
)
|
||||
async def list_invoices(
|
||||
page: PageRequest = Depends(single_page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_billing_permission(APIKeyPermission.READ_CREDITS)
|
||||
),
|
||||
) -> Page[InvoiceItem]:
|
||||
"""Recent Stripe invoices for the current user.
|
||||
|
||||
Stripe returns at most `limit` invoices and offers neither a cursor nor a
|
||||
total here, so `next_cursor` and `total_count` are always `null`; raise
|
||||
`limit` to see further back.
|
||||
"""
|
||||
user_credit_model = await get_credit_model(auth.user_id, auth.organization_id)
|
||||
invoices = await user_credit_model.list_invoices(auth.user_id, limit=page.limit)
|
||||
|
||||
return page.uncounted([InvoiceItem.from_internal(inv) for inv in invoices])
|
||||
|
||||
|
||||
@credits_router.get(
|
||||
path="/cost-summary",
|
||||
summary="Get execution cost summary",
|
||||
operation_id="getAutomationCostSummary",
|
||||
)
|
||||
async def get_cost_summary(
|
||||
since: Optional[datetime] = Query(
|
||||
default=None,
|
||||
description="Window start (UTC). Defaults to start of current calendar month.",
|
||||
),
|
||||
until: Optional[datetime] = Query(
|
||||
default=None,
|
||||
description="Window end (UTC). Defaults to now.",
|
||||
),
|
||||
top_runs_limit: int = Query(
|
||||
10,
|
||||
ge=1,
|
||||
le=50,
|
||||
description="Maximum number of top-cost runs to return.",
|
||||
),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_CREDITS)),
|
||||
) -> AutomationCostSummary:
|
||||
"""Aggregated cost breakdown for the user's graph executions."""
|
||||
# A time without an offset is UTC, as documented; comparing it with one
|
||||
# that has an offset would otherwise fail.
|
||||
since = since.replace(tzinfo=timezone.utc) if since and not since.tzinfo else since
|
||||
until = until.replace(tzinfo=timezone.utc) if until and not until.tzinfo else until
|
||||
if since is not None and until is not None and since > until:
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail="`since` must be earlier than or equal to `until`.",
|
||||
)
|
||||
summary = await get_user_cost_summary(
|
||||
user_id=auth.user_id,
|
||||
since=since,
|
||||
until=until,
|
||||
top_runs_limit=top_runs_limit,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
return AutomationCostSummary.from_internal(summary)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@cached(ttl_seconds=300, maxsize=32, cache_none=False)
|
||||
async def _get_stripe_price_amount(price_id: str | None) -> int | None:
|
||||
"""Return the unit_amount (cents) for a Stripe Price ID, cached 5 minutes.
|
||||
|
||||
``None`` on a Stripe error, so ``cache_none=False`` keeps it out of the cache.
|
||||
"""
|
||||
if not price_id:
|
||||
return 0
|
||||
try:
|
||||
price = await run_in_threadpool(stripe.Price.retrieve, price_id)
|
||||
return price.unit_amount or 0
|
||||
except stripe.StripeError:
|
||||
logger.warning(
|
||||
"Failed to retrieve Stripe price %s — returning None (not cached)",
|
||||
price_id,
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,150 @@
|
||||
import inspect
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
import stripe
|
||||
from fastapi import HTTPException
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.data.model import TransactionHistory
|
||||
|
||||
from .credits import (
|
||||
_get_stripe_price_amount,
|
||||
get_balance,
|
||||
get_cost_summary,
|
||||
get_transactions,
|
||||
list_invoices,
|
||||
)
|
||||
from .pagination import PageRequest
|
||||
from .tenancy import TenantContext, require_billing_permission
|
||||
|
||||
|
||||
async def test_a_stripe_error_is_not_cached_as_a_zero_price(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
_get_stripe_price_amount.cache_clear()
|
||||
mocker.patch.object(
|
||||
stripe.Price,
|
||||
"retrieve",
|
||||
side_effect=[stripe.StripeError("network"), SimpleNamespace(unit_amount=2000)],
|
||||
)
|
||||
|
||||
await _get_stripe_price_amount("price_pro")
|
||||
assert await _get_stripe_price_amount("price_pro") == 2000
|
||||
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def credit_model(mocker: pytest_mock.MockFixture) -> Mock:
|
||||
model = Mock(
|
||||
get_transaction_history=AsyncMock(
|
||||
return_value=TransactionHistory(
|
||||
transactions=[],
|
||||
next_transaction_time=datetime(2026, 10, 1, tzinfo=timezone.utc),
|
||||
next_cursor="history-cursor-2",
|
||||
)
|
||||
)
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.credits.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=model,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
async def test_transactions_page_by_the_historys_own_cursor(credit_model: Mock):
|
||||
"""A time ceiling alone skipped groups that tied on time at a page boundary."""
|
||||
first = await get_transactions(
|
||||
transaction_type=None, page=PageRequest(limit=10), auth=_AUTH
|
||||
)
|
||||
assert first.next_cursor is not None
|
||||
|
||||
await get_transactions(
|
||||
transaction_type=None,
|
||||
page=PageRequest(limit=10, cursor=first.next_cursor),
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
kwargs = credit_model.get_transaction_history.await_args.kwargs
|
||||
assert kwargs["cursor"] == "history-cursor-2"
|
||||
assert kwargs["viewer_organization_id"] == "org-1"
|
||||
assert "transaction_time_ceiling" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"member, allowed",
|
||||
[
|
||||
(Mock(isOwner=True, isBillingManager=False), True),
|
||||
(Mock(isOwner=False, isBillingManager=True), True),
|
||||
(Mock(isOwner=False, isBillingManager=False), False),
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
async def test_only_the_owner_or_a_billing_manager_reads_the_orgs_credits(
|
||||
mocker: pytest_mock.MockFixture, member: Mock | None, allowed: bool
|
||||
):
|
||||
"""The web app shows a pooled wallet only to those roles; keys follow suit."""
|
||||
prisma = mocker.patch("backend.api.external.v2.tenancy.prisma")
|
||||
prisma.orgmember.find_unique = AsyncMock(return_value=member)
|
||||
check = require_billing_permission(APIKeyPermission.READ_CREDITS)
|
||||
|
||||
if allowed:
|
||||
assert await check(tenant=_AUTH) is _AUTH
|
||||
else:
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await check(tenant=_AUTH)
|
||||
assert raised.value.status_code == 403
|
||||
|
||||
|
||||
def test_the_org_wallet_routes_check_the_billing_role():
|
||||
for route in (get_balance, get_transactions, list_invoices):
|
||||
source = inspect.getsource(route)
|
||||
assert "require_billing_permission" in source, route.__name__
|
||||
|
||||
|
||||
async def test_a_cost_summary_takes_a_time_without_an_offset_as_utc(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
summary = mocker.patch(
|
||||
"backend.api.external.v2.credits.get_user_cost_summary",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
mocker.patch("backend.api.external.v2.credits.AutomationCostSummary.from_internal")
|
||||
|
||||
await get_cost_summary(
|
||||
since=datetime(2026, 10, 1),
|
||||
until=datetime(2026, 10, 2, tzinfo=timezone.utc),
|
||||
top_runs_limit=10,
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
assert summary.await_args.kwargs["since"] == datetime(
|
||||
2026, 10, 1, tzinfo=timezone.utc
|
||||
)
|
||||
|
||||
|
||||
async def test_a_cost_summary_counts_only_this_organizations_runs(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""`GET /runs` hides another organization's runs; their costs, graph ids
|
||||
and names must not come back here either."""
|
||||
summary = mocker.patch(
|
||||
"backend.api.external.v2.credits.get_user_cost_summary",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
mocker.patch("backend.api.external.v2.credits.AutomationCostSummary.from_internal")
|
||||
|
||||
await get_cost_summary(since=None, until=None, top_runs_limit=10, auth=_AUTH)
|
||||
|
||||
assert summary.await_args.kwargs["organization_id"] == "org-1"
|
||||
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
V2 External API - Error responses
|
||||
|
||||
One body shape for every non-2xx: `{"error": {"code", "message", "details"}}`.
|
||||
|
||||
The shared handlers in `api/utils/exceptions` emit `{"message", "detail",
|
||||
"hint"}` and FastAPI's own `HTTPException` path emits `{"detail": ...}`, so v2
|
||||
registers its own handlers over both rather than inheriting either.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Mapping, Optional
|
||||
|
||||
import fastapi
|
||||
import fastapi.responses
|
||||
import pydantic
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from prisma.errors import PrismaError
|
||||
from prisma.errors import RecordNotFoundError as PrismaRecordNotFoundError
|
||||
from pydantic import BaseModel, Field, JsonValue
|
||||
from starlette import status
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
|
||||
from backend.api.features.library.exceptions import (
|
||||
FolderAlreadyExistsError,
|
||||
FolderValidationError,
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
from backend.api.features.store.exceptions import UnauthorizedError as StoreForbidden
|
||||
from backend.copilot.rate_limit import UserPaywalledError
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import GraphActivationError
|
||||
from backend.util.exceptions import (
|
||||
ConflictError,
|
||||
MissingConfigError,
|
||||
NotAuthorizedError,
|
||||
NotFoundError,
|
||||
PreconditionFailed,
|
||||
WebhookRegistrationError,
|
||||
WebhookSetupUnavailableError,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ErrorDetail(BaseModel):
|
||||
code: str = Field(description="Stable machine-readable identifier for the failure")
|
||||
message: str = Field(description="Human-readable explanation")
|
||||
details: Optional[dict[str, JsonValue]] = Field(
|
||||
default=None, description="Structured context, when the failure has any"
|
||||
)
|
||||
|
||||
|
||||
class ErrorResponse(BaseModel):
|
||||
"""The body of every v2 response that is not 2xx."""
|
||||
|
||||
error: ErrorDetail
|
||||
|
||||
|
||||
def add_v2_exception_handlers(app: fastapi.FastAPI) -> None:
|
||||
"""Register the v2 error envelope over every exception that can reach a route."""
|
||||
for exception, status_code in {
|
||||
NotFoundError: status.HTTP_404_NOT_FOUND,
|
||||
PrismaRecordNotFoundError: status.HTTP_404_NOT_FOUND,
|
||||
# 404 not 403: whether a listing exists in another tenant is not this
|
||||
# credential's business, which is how `in_tenant` treats every other row.
|
||||
StoreForbidden: status.HTTP_404_NOT_FOUND,
|
||||
NotAuthorizedError: status.HTTP_403_FORBIDDEN,
|
||||
UserPaywalledError: status.HTTP_402_PAYMENT_REQUIRED,
|
||||
PreconditionFailed: status.HTTP_428_PRECONDITION_REQUIRED,
|
||||
FolderAlreadyExistsError: status.HTTP_409_CONFLICT,
|
||||
ConflictError: status.HTTP_409_CONFLICT,
|
||||
LibraryAgentInAnotherOrganizationError: status.HTTP_409_CONFLICT,
|
||||
FolderValidationError: status.HTTP_400_BAD_REQUEST,
|
||||
GraphActivationError: status.HTTP_400_BAD_REQUEST,
|
||||
WebhookRegistrationError: status.HTTP_400_BAD_REQUEST,
|
||||
ValueError: status.HTTP_400_BAD_REQUEST,
|
||||
MissingConfigError: status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
WebhookSetupUnavailableError: status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
PrismaError: status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
Exception: status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
}.items():
|
||||
app.add_exception_handler(exception, _handler_for(status_code))
|
||||
|
||||
app.add_exception_handler(StarletteHTTPException, _handle_http_exception)
|
||||
app.add_exception_handler(RequestValidationError, _handle_validation_error)
|
||||
app.add_exception_handler(pydantic.ValidationError, _handle_validation_error)
|
||||
|
||||
|
||||
def error_response(
|
||||
status_code: int,
|
||||
message: str,
|
||||
code: Optional[str] = None,
|
||||
details: Optional[dict[str, JsonValue]] = None,
|
||||
headers: Optional[Mapping[str, str]] = None,
|
||||
) -> fastapi.responses.JSONResponse:
|
||||
"""Build a v2 error body. Also used outside the handlers, by the middleware."""
|
||||
body = ErrorResponse(
|
||||
error=ErrorDetail(
|
||||
code=code or error_code_for(status_code),
|
||||
message=message,
|
||||
details=details,
|
||||
)
|
||||
)
|
||||
return fastapi.responses.JSONResponse(
|
||||
content=body.model_dump(mode="json"),
|
||||
status_code=status_code,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
def error_code_for(status_code: int) -> str:
|
||||
default = "internal_error" if status_code >= 500 else "request_failed"
|
||||
return _ERROR_CODES.get(status_code, default)
|
||||
|
||||
|
||||
_ERROR_CODES = {
|
||||
status.HTTP_400_BAD_REQUEST: "bad_request",
|
||||
status.HTTP_401_UNAUTHORIZED: "unauthorized",
|
||||
status.HTTP_402_PAYMENT_REQUIRED: "payment_required",
|
||||
status.HTTP_403_FORBIDDEN: "forbidden",
|
||||
status.HTTP_404_NOT_FOUND: "not_found",
|
||||
status.HTTP_405_METHOD_NOT_ALLOWED: "method_not_allowed",
|
||||
status.HTTP_409_CONFLICT: "conflict",
|
||||
status.HTTP_413_CONTENT_TOO_LARGE: "payload_too_large",
|
||||
status.HTTP_422_UNPROCESSABLE_CONTENT: "validation_error",
|
||||
status.HTTP_428_PRECONDITION_REQUIRED: "precondition_required",
|
||||
status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_exceeded",
|
||||
status.HTTP_500_INTERNAL_SERVER_ERROR: "internal_error",
|
||||
status.HTTP_502_BAD_GATEWAY: "upstream_error",
|
||||
status.HTTP_503_SERVICE_UNAVAILABLE: "service_unavailable",
|
||||
}
|
||||
|
||||
|
||||
def _handler_for(status_code: int):
|
||||
def handler(
|
||||
request: fastapi.Request, exc: Exception
|
||||
) -> fastapi.responses.JSONResponse:
|
||||
_log(request, status_code, exc)
|
||||
# A 5xx message describes our internals, not the caller's request.
|
||||
message = (
|
||||
f"Failed to process {request.method} {request.url.path}"
|
||||
if status_code >= 500
|
||||
else (str(exc) or type(exc).__name__)
|
||||
)
|
||||
return error_response(status_code, message)
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def _handle_http_exception(
|
||||
request: fastapi.Request, exc: Exception
|
||||
) -> fastapi.responses.JSONResponse:
|
||||
assert isinstance(exc, StarletteHTTPException)
|
||||
_log(request, exc.status_code, exc)
|
||||
# `detail` is a dict when a route raised HTTPException with structured context.
|
||||
if isinstance(exc.detail, dict):
|
||||
message = str(exc.detail.get("message", "Request failed"))
|
||||
details = {k: v for k, v in exc.detail.items() if k != "message"}
|
||||
else:
|
||||
message = str(exc.detail)
|
||||
details = None
|
||||
return error_response(
|
||||
exc.status_code, message, details=details, headers=exc.headers
|
||||
)
|
||||
|
||||
|
||||
def _handle_validation_error(
|
||||
request: fastapi.Request, exc: Exception
|
||||
) -> fastapi.responses.JSONResponse:
|
||||
_log(request, status.HTTP_422_UNPROCESSABLE_CONTENT, exc)
|
||||
errors = (
|
||||
exc.errors()
|
||||
if isinstance(exc, (RequestValidationError, pydantic.ValidationError))
|
||||
else []
|
||||
)
|
||||
return error_response(
|
||||
status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
f"Invalid data for {request.method} {request.url.path}",
|
||||
details={
|
||||
# A `@field_validator` leaves its live ValueError in `ctx`; without a
|
||||
# fallback the dump raises, and the ValueError -> 400 entry above
|
||||
# then swallows this handler's own failure.
|
||||
"errors": pydantic.TypeAdapter(JsonValue).dump_python(
|
||||
errors, mode="json", warnings=False, fallback=str
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# Routine client-side outcomes: a caller without the scope, without the credits,
|
||||
# or asking for something that is not there is not a server event worth logging.
|
||||
_UNLOGGED_STATUSES = (
|
||||
status.HTTP_401_UNAUTHORIZED,
|
||||
status.HTTP_402_PAYMENT_REQUIRED,
|
||||
status.HTTP_403_FORBIDDEN,
|
||||
status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
|
||||
def _log(request: fastapi.Request, status_code: int, exc: Exception) -> None:
|
||||
if status_code >= 500:
|
||||
logger.exception(
|
||||
f"{request.method} {request.url.path} failed: {exc}", exc_info=exc
|
||||
)
|
||||
elif status_code not in _UNLOGGED_STATUSES:
|
||||
logger.warning(f"{request.method} {request.url.path} -> {status_code}: {exc}")
|
||||
@@ -0,0 +1,214 @@
|
||||
"""
|
||||
V2 External API - Files Endpoints
|
||||
|
||||
Provides file upload, download, listing, metadata, and deletion functionality.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, Query, Security, UploadFile
|
||||
from fastapi.responses import Response
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.workspace.routes import create_file_download_response
|
||||
from backend.api.features.workspace.service import store_workspace_upload
|
||||
from backend.data.workspace import (
|
||||
count_workspace_files,
|
||||
get_workspace,
|
||||
get_workspace_file,
|
||||
list_workspace_files,
|
||||
)
|
||||
from backend.util.workspace import WorkspaceManager
|
||||
|
||||
from .models import UploadWorkspaceFileResponse, WorkspaceFileInfo
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .rate_limit import enforce, file_upload_limiter
|
||||
from .tenancy import TenantContext, require_paid_permission, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
file_workspace_router = APIRouter(tags=["files"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@file_workspace_router.get(
|
||||
path="",
|
||||
summary="List workspace files",
|
||||
operation_id="listWorkspaceFiles",
|
||||
)
|
||||
async def list_files(
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_FILES)),
|
||||
) -> Page[WorkspaceFileInfo]:
|
||||
"""List files in the user's workspace."""
|
||||
workspace = await get_workspace(auth.user_id)
|
||||
if workspace is None:
|
||||
return page.paged([], total_count=0)
|
||||
|
||||
# Otto's working files and skill packages stay out, as on the Files page.
|
||||
total_count = await count_workspace_files(workspace.id, include_internal=False)
|
||||
files = await list_workspace_files(
|
||||
workspace_id=workspace.id,
|
||||
limit=page.limit,
|
||||
offset=(page.page - 1) * page.limit,
|
||||
include_internal=False,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[
|
||||
WorkspaceFileInfo(
|
||||
id=f.id,
|
||||
name=f.name,
|
||||
path=f.path,
|
||||
mime_type=f.mime_type,
|
||||
size_bytes=f.size_bytes,
|
||||
created_at=f.created_at,
|
||||
updated_at=f.updated_at,
|
||||
)
|
||||
for f in files
|
||||
],
|
||||
total_count=total_count,
|
||||
)
|
||||
|
||||
|
||||
@file_workspace_router.get(
|
||||
path="/{file_id}",
|
||||
summary="Get workspace file metadata",
|
||||
operation_id="getWorkspaceFileInfo",
|
||||
)
|
||||
async def get_file(
|
||||
file_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_FILES)),
|
||||
) -> WorkspaceFileInfo:
|
||||
"""Get metadata for a specific file in the user's workspace."""
|
||||
workspace = await get_workspace(auth.user_id)
|
||||
if workspace is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Workspace not found",
|
||||
)
|
||||
|
||||
file = await get_workspace_file(file_id, workspace.id)
|
||||
if file is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"File #{file_id} not found",
|
||||
)
|
||||
|
||||
return WorkspaceFileInfo(
|
||||
id=file.id,
|
||||
name=file.name,
|
||||
path=file.path,
|
||||
mime_type=file.mime_type,
|
||||
size_bytes=file.size_bytes,
|
||||
created_at=file.created_at,
|
||||
updated_at=file.updated_at,
|
||||
)
|
||||
|
||||
|
||||
@file_workspace_router.delete(
|
||||
path="/{file_id}",
|
||||
summary="Delete file from workspace",
|
||||
operation_id="deleteWorkspaceFile",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_file(
|
||||
file_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_FILES)),
|
||||
) -> None:
|
||||
"""Delete a file from the user's workspace, with its stored content."""
|
||||
workspace = await get_workspace(auth.user_id)
|
||||
if workspace is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Workspace not found",
|
||||
)
|
||||
|
||||
# The manager also removes the stored blob and its search entry, as the
|
||||
# internal route does; soft-deleting the row alone freed no storage.
|
||||
deleted = await WorkspaceManager(auth.user_id, workspace.id).delete_file(file_id)
|
||||
if not deleted:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"File #{file_id} not found",
|
||||
)
|
||||
|
||||
|
||||
@file_workspace_router.post(
|
||||
path="/upload",
|
||||
summary="Upload file to workspace",
|
||||
operation_id="uploadWorkspaceFile",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def upload_file(
|
||||
response: Response,
|
||||
file: UploadFile = File(...),
|
||||
overwrite: bool = Query(
|
||||
default=False, description="Replace an existing file with the same path"
|
||||
),
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_FILES)
|
||||
),
|
||||
) -> UploadWorkspaceFileResponse:
|
||||
"""
|
||||
Upload a file to the user's workspace.
|
||||
|
||||
The file is listed by `GET /files` and can be passed to agent file inputs
|
||||
as the returned `file_uri`. Uploads are virus-scanned before storage and
|
||||
count against the account's storage quota.
|
||||
|
||||
**Rate limit:** 20 requests per 5 minutes per user.
|
||||
"""
|
||||
await enforce(file_upload_limiter, auth.user_id, response)
|
||||
|
||||
workspace_file = await store_workspace_upload(
|
||||
auth.user_id, file, overwrite=overwrite
|
||||
)
|
||||
|
||||
return UploadWorkspaceFileResponse(
|
||||
id=workspace_file.id,
|
||||
name=workspace_file.name,
|
||||
path=workspace_file.path,
|
||||
mime_type=workspace_file.mime_type,
|
||||
size_bytes=workspace_file.size_bytes,
|
||||
created_at=workspace_file.created_at,
|
||||
updated_at=workspace_file.updated_at,
|
||||
file_uri=f"workspace://{workspace_file.id}#{workspace_file.mime_type}",
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints - Download
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@file_workspace_router.get(
|
||||
path="/{file_id}/download",
|
||||
summary="Download file from workspace",
|
||||
operation_id="getWorkspaceFileDownload",
|
||||
)
|
||||
async def download_file(
|
||||
file_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_FILES)),
|
||||
) -> Response:
|
||||
"""Download a file from the user's workspace."""
|
||||
workspace = await get_workspace(auth.user_id)
|
||||
if workspace is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Workspace not found",
|
||||
)
|
||||
|
||||
file = await get_workspace_file(file_id, workspace.id)
|
||||
if file is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"File #{file_id} not found",
|
||||
)
|
||||
|
||||
return await create_file_download_response(file)
|
||||
@@ -0,0 +1,332 @@
|
||||
"""Uploads and the file listing must address the same storage.
|
||||
|
||||
Before this, `POST /files/upload` wrote to the temporary cloud bucket while
|
||||
every other file endpoint read the persistent workspace, so an uploaded file
|
||||
never appeared in the list and had no id to download or delete.
|
||||
"""
|
||||
|
||||
import io
|
||||
import mimetypes
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import FastAPI, HTTPException, Response, UploadFile
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
from backend.api.external.v2.files import (
|
||||
delete_file,
|
||||
file_upload_limiter,
|
||||
file_workspace_router,
|
||||
get_file,
|
||||
list_files,
|
||||
upload_file,
|
||||
)
|
||||
from backend.api.external.v2.global_rate_limit import GlobalRateLimitMiddleware
|
||||
from backend.api.external.v2.pagination import PageRequest
|
||||
from backend.api.external.v2.tenancy import TenantContext, require_auth
|
||||
from backend.api.utils.rate_limit import RateLimitState
|
||||
from backend.data.workspace import Workspace, WorkspaceFile
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG_ID = "org-1"
|
||||
WORKSPACE_ID = "ws-1"
|
||||
|
||||
|
||||
async def test_an_uploaded_file_is_listed_and_addressable(
|
||||
workspace: dict[str, WorkspaceFile],
|
||||
) -> None:
|
||||
"""The whole of item 1.9: upload, then find the same file by the id it returned."""
|
||||
uploaded = await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("report.csv", b"a,b\n1,2\n", "text/csv"),
|
||||
overwrite=False,
|
||||
auth=_auth(),
|
||||
)
|
||||
|
||||
listing = await list_files(page=PageRequest(limit=25), auth=_auth())
|
||||
assert [f.id for f in listing.items] == [uploaded.id]
|
||||
|
||||
found = await get_file(file_id=uploaded.id, auth=_auth())
|
||||
assert (found.name, found.path) == ("report.csv", "report.csv")
|
||||
|
||||
|
||||
async def test_the_upload_returns_a_uri_agent_file_inputs_accept(
|
||||
workspace: dict[str, WorkspaceFile],
|
||||
) -> None:
|
||||
"""`workspace://<id>#<mime>` is what `store_media_file` resolves for a run."""
|
||||
uploaded = await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("clip.mp4", b"\x00\x01", "video/mp4"),
|
||||
overwrite=False,
|
||||
auth=_auth(),
|
||||
)
|
||||
|
||||
assert uploaded.file_uri == f"workspace://{uploaded.id}#video/mp4"
|
||||
|
||||
|
||||
async def test_a_second_upload_of_the_same_name_conflicts_unless_overwritten(
|
||||
workspace: dict[str, WorkspaceFile],
|
||||
) -> None:
|
||||
await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("notes.txt", b"one"),
|
||||
overwrite=False,
|
||||
auth=_auth(),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("notes.txt", b"two"),
|
||||
overwrite=False,
|
||||
auth=_auth(),
|
||||
)
|
||||
assert raised.value.status_code == 409
|
||||
|
||||
replaced = await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("notes.txt", b"two"),
|
||||
overwrite=True,
|
||||
auth=_auth(),
|
||||
)
|
||||
assert list(workspace) == [replaced.id]
|
||||
|
||||
|
||||
async def test_deleting_a_file_removes_its_stored_content(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
"""A soft-deleted row alone left the blob in storage and off the quota."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(id=WORKSPACE_ID)),
|
||||
)
|
||||
manager_cls = mocker.patch("backend.api.external.v2.files.WorkspaceManager")
|
||||
manager_cls.return_value.delete_file = mock.AsyncMock(return_value=True)
|
||||
|
||||
await delete_file(file_id="file-1", auth=_auth())
|
||||
|
||||
manager_cls.assert_called_once_with(USER_ID, WORKSPACE_ID)
|
||||
manager_cls.return_value.delete_file.assert_awaited_once_with("file-1")
|
||||
|
||||
|
||||
async def test_deleting_a_missing_file_is_a_404(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(id=WORKSPACE_ID)),
|
||||
)
|
||||
manager_cls = mocker.patch("backend.api.external.v2.files.WorkspaceManager")
|
||||
manager_cls.return_value.delete_file = mock.AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await delete_file(file_id="file-1", auth=_auth())
|
||||
assert raised.value.status_code == 404
|
||||
|
||||
|
||||
async def test_an_upload_over_the_storage_quota_is_undone(
|
||||
workspace: dict[str, WorkspaceFile],
|
||||
quota_bytes: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""A write that breaches the quota must not survive as an unlisted file."""
|
||||
quota_bytes.return_value = 1
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await upload_file(
|
||||
response=Response(),
|
||||
file=_upload("big.bin", b"0" * 64),
|
||||
overwrite=False,
|
||||
auth=_auth(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 413
|
||||
assert workspace == {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def quota_bytes(mocker: pytest_mock.MockFixture) -> mock.AsyncMock:
|
||||
return mocker.patch(
|
||||
"backend.api.features.workspace.service.get_workspace_storage_limit_bytes",
|
||||
new=mock.AsyncMock(return_value=0),
|
||||
)
|
||||
|
||||
|
||||
async def test_the_listing_leaves_out_otto_s_working_files(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""Otto's scratch files and skill packages are not the user's files."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(id=WORKSPACE_ID)),
|
||||
)
|
||||
count = mocker.patch(
|
||||
"backend.api.external.v2.files.count_workspace_files",
|
||||
new=mock.AsyncMock(return_value=0),
|
||||
)
|
||||
listed = mocker.patch(
|
||||
"backend.api.external.v2.files.list_workspace_files",
|
||||
new=mock.AsyncMock(return_value=[]),
|
||||
)
|
||||
|
||||
await list_files(page=PageRequest(limit=25), auth=_auth())
|
||||
|
||||
assert count.await_args.kwargs["include_internal"] is False
|
||||
assert listed.await_args.kwargs["include_internal"] is False
|
||||
|
||||
|
||||
def test_a_successful_upload_reports_the_upload_cap_not_the_global_one(
|
||||
workspace: dict[str, WorkspaceFile], mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""Uploads have their own, narrower cap; a client backing off on the global
|
||||
window's numbers would walk into a 429 it was told was far away."""
|
||||
mocker.patch.object(
|
||||
file_upload_limiter,
|
||||
"check",
|
||||
new=mock.AsyncMock(
|
||||
return_value=RateLimitState(limit=20, remaining=19, reset_seconds=300)
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(user_id=USER_ID)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit._authenticated_limiter.check",
|
||||
new=mock.AsyncMock(
|
||||
return_value=RateLimitState(limit=200, remaining=199, reset_seconds=60)
|
||||
),
|
||||
)
|
||||
app = FastAPI()
|
||||
app.include_router(file_workspace_router, prefix="/files")
|
||||
app.add_middleware(GlobalRateLimitMiddleware)
|
||||
app.dependency_overrides[require_auth] = _auth
|
||||
|
||||
response = TestClient(app).post(
|
||||
"/files/upload", files={"file": ("notes.txt", b"hello")}
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.headers.get_list("x-ratelimit-limit") == ["20"]
|
||||
assert response.headers["x-ratelimit-remaining"] == "19"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def workspace(
|
||||
mocker: pytest_mock.MockFixture, quota_bytes: mock.AsyncMock
|
||||
) -> dict[str, WorkspaceFile]:
|
||||
"""One in-memory workspace behind both the write path and the read path.
|
||||
|
||||
The point of the fixture: the upload endpoint and the listing endpoints
|
||||
reach the *same* store, so a test can only pass if they do in production.
|
||||
"""
|
||||
files: dict[str, WorkspaceFile] = {}
|
||||
|
||||
async def write_file(
|
||||
content: bytes, filename: str, *, overwrite: bool = False, metadata=None
|
||||
) -> WorkspaceFile:
|
||||
existing = next((f for f in files.values() if f.path == filename), None)
|
||||
if existing and not overwrite:
|
||||
raise ValueError(f"File already exists at path: {filename}")
|
||||
if existing:
|
||||
del files[existing.id]
|
||||
file = _workspace_file(filename, len(content))
|
||||
files[file.id] = file
|
||||
return file
|
||||
|
||||
async def delete_file(file_id: str) -> None:
|
||||
files.pop(file_id, None)
|
||||
|
||||
manager = mocker.patch(
|
||||
"backend.api.features.workspace.service.WorkspaceManager"
|
||||
).return_value
|
||||
manager.write_file = mock.AsyncMock(side_effect=write_file)
|
||||
manager.delete_file = mock.AsyncMock(side_effect=delete_file)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.workspace.service.get_or_create_workspace",
|
||||
new=mock.AsyncMock(
|
||||
return_value=Workspace(
|
||||
id=WORKSPACE_ID,
|
||||
user_id=USER_ID,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.workspace.service.get_workspace_total_size",
|
||||
new=mock.AsyncMock(
|
||||
side_effect=lambda _: sum(f.size_bytes for f in files.values())
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.file_upload_limiter.check",
|
||||
new=mock.AsyncMock(return_value=None),
|
||||
)
|
||||
|
||||
async def get_workspace(user_id: str) -> Workspace:
|
||||
return Workspace(
|
||||
id=WORKSPACE_ID,
|
||||
user_id=user_id,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
updated_at=datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace",
|
||||
new=mock.AsyncMock(side_effect=get_workspace),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.count_workspace_files",
|
||||
new=mock.AsyncMock(side_effect=lambda _, **__: len(files)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.list_workspace_files",
|
||||
new=mock.AsyncMock(side_effect=lambda **_: list(files.values())),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.files.get_workspace_file",
|
||||
new=mock.AsyncMock(side_effect=lambda file_id, _: files.get(file_id)),
|
||||
)
|
||||
return files
|
||||
|
||||
|
||||
def _workspace_file(name: str, size_bytes: int) -> WorkspaceFile:
|
||||
now = datetime.now(timezone.utc)
|
||||
return WorkspaceFile(
|
||||
id=str(uuid.uuid4()),
|
||||
workspace_id=WORKSPACE_ID,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
name=name,
|
||||
path=name,
|
||||
storage_path=f"local://{name}",
|
||||
# Same derivation the real WorkspaceManager applies.
|
||||
mime_type=mimetypes.guess_type(name)[0] or "application/octet-stream",
|
||||
size_bytes=size_bytes,
|
||||
)
|
||||
|
||||
|
||||
def _upload(
|
||||
filename: str, content: bytes, content_type: Optional[str] = None
|
||||
) -> UploadFile:
|
||||
return UploadFile(
|
||||
file=io.BytesIO(content),
|
||||
filename=filename,
|
||||
headers=Headers({"content-type": content_type} if content_type else {}),
|
||||
)
|
||||
|
||||
|
||||
def _auth() -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=[APIKeyPermission.READ_FILES, APIKeyPermission.WRITE_FILES],
|
||||
type="api_key",
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
@@ -0,0 +1,232 @@
|
||||
"""
|
||||
V2 External API - Global Rate Limit Middleware
|
||||
|
||||
ASGI middleware that enforces per-user and per-IP request caps across all v2
|
||||
endpoints. Authenticated users get 200 req/min keyed by user ID; unauthenticated
|
||||
sessions get 5 req/min keyed by client IP.
|
||||
|
||||
Identifies the user through the auth middleware's `resolve_request_auth`.
|
||||
Checking a credential costs work the per-user cap can't bound, because it can
|
||||
only count a request once the check is done: a Scrypt hash for an API key, a
|
||||
database lookup for an OAuth access token. So each credential is counted
|
||||
first, per client IP: an API key by its head (the part its lookup matches
|
||||
on), a token by a digest of it. One presented too often, or failing too
|
||||
often, is refused unchecked until the window rolls over, as is an address
|
||||
failing too often across credentials. One client's bad key or flood doesn't
|
||||
lock out the other credentials behind the same address.
|
||||
|
||||
Every response carries the caller's `X-RateLimit-*` position, and a 429 adds
|
||||
`Retry-After`, so a client can back off on the numbers instead of guessing.
|
||||
|
||||
On auth-resolution failure or Redis errors the request passes through — the
|
||||
endpoint's own auth dependency handles 401, and the rate limiter fails open.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from autogpt_libs.api_key.keysmith import APIKeySmith
|
||||
from fastapi import HTTPException
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
from starlette.datastructures import Headers
|
||||
from starlette.responses import Response
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from backend.api.external.middleware import resolve_request_auth
|
||||
from backend.api.utils.rate_limit import RateLimiter, RateLimitState
|
||||
from backend.data.auth.oauth import is_access_token
|
||||
from backend.util.settings import Settings
|
||||
|
||||
from .errors import error_response
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
settings = Settings()
|
||||
|
||||
_authenticated_limiter = RateLimiter("v2:global", max_requests=200, window_seconds=60)
|
||||
_anonymous_limiter = RateLimiter("v2:global:anon", max_requests=5, window_seconds=60)
|
||||
# Above the per-user cap, so a client within its own limit never meets it.
|
||||
_presented_credential_limiter = RateLimiter(
|
||||
"v2:global:credential-presented", max_requests=300, window_seconds=60
|
||||
)
|
||||
_failed_credential_limiter = RateLimiter(
|
||||
"v2:global:credential-failures", max_requests=30, window_seconds=60
|
||||
)
|
||||
# Failures spread over many credentials, each of which costs a check.
|
||||
_failed_auth_limiter = RateLimiter(
|
||||
"v2:global:auth-failures", max_requests=300, window_seconds=60
|
||||
)
|
||||
# The docs tell agents to read the spec before they have a key; on the 5/min
|
||||
# bucket those reads would use up what the first real calls need.
|
||||
_anonymous_docs_limiter = RateLimiter(
|
||||
"v2:global:anon-docs", max_requests=60, window_seconds=60
|
||||
)
|
||||
_DOCS_PATHS = frozenset({"/docs", "/docs/oauth2-redirect", "/redoc", "/openapi.json"})
|
||||
|
||||
|
||||
class GlobalRateLimitMiddleware:
|
||||
def __init__(self, app: ASGIApp):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
# The first of a repeated header, as the route's own dependency reads
|
||||
# it, so both charge and authenticate the same credential.
|
||||
request_headers = Headers(scope=scope)
|
||||
|
||||
api_key = request_headers.get("x-api-key") or None
|
||||
auth_header = request_headers.get("authorization", "")
|
||||
bearer = None
|
||||
if auth_header.lower().startswith("bearer "):
|
||||
bearer = HTTPAuthorizationCredentials(
|
||||
scheme="Bearer", credentials=auth_header[7:]
|
||||
)
|
||||
|
||||
ip = client_ip(scope)
|
||||
bucket = _credential_bucket(ip, api_key, bearer)
|
||||
if bucket is not None and (
|
||||
refusal := await _refuse_before_verifying(ip, bucket)
|
||||
):
|
||||
await refusal(scope, receive, send)
|
||||
return
|
||||
|
||||
try:
|
||||
auth = await resolve_request_auth(scope, api_key=api_key, bearer=bearer)
|
||||
except HTTPException as rejection:
|
||||
auth = None
|
||||
if rejection.status_code == 401 and bucket is not None:
|
||||
with contextlib.suppress(HTTPException):
|
||||
await _failed_credential_limiter.check(bucket)
|
||||
with contextlib.suppress(HTTPException):
|
||||
await _failed_auth_limiter.check(ip)
|
||||
except Exception as exc:
|
||||
# Fail open on anything the auth backend throws that is not a
|
||||
# rejection; the route's own dependency will answer 401 or 500.
|
||||
logger.warning(f"Rate-limit auth resolution failed: {exc}")
|
||||
auth = None
|
||||
|
||||
try:
|
||||
if auth:
|
||||
state = await _authenticated_limiter.check(auth.user_id)
|
||||
elif app_path(scope) in _DOCS_PATHS:
|
||||
state = await _anonymous_docs_limiter.check(ip)
|
||||
else:
|
||||
state = await _anonymous_limiter.check(ip)
|
||||
except HTTPException as exc:
|
||||
# The middleware sits outside the app, so the v2 exception handlers
|
||||
# never see this — build the same envelope by hand.
|
||||
response = error_response(
|
||||
exc.status_code, str(exc.detail), headers=exc.headers
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
|
||||
await self.app(scope, receive, _with_rate_limit_headers(send, state))
|
||||
|
||||
|
||||
def _credential_bucket(
|
||||
ip: str, api_key: Optional[str], bearer: Optional[HTTPAuthorizationCredentials]
|
||||
) -> Optional[str]:
|
||||
"""The pre-verification bucket of a credential that costs work to check, or None.
|
||||
|
||||
An API key is bucketed by its head. `X-API-Key` is always tried as a key,
|
||||
so its value is too, even in the token format. A bearer in the access-token
|
||||
format is only ever looked up as a token (see `resolve_auth_info`), so it
|
||||
is bucketed by a digest of the token. Any other value is refused without a
|
||||
hash or a lookup, so it needs no bucket.
|
||||
"""
|
||||
if api_key is not None:
|
||||
credential = api_key
|
||||
elif bearer is not None:
|
||||
if is_access_token(bearer.credentials):
|
||||
return f"{ip}:token:{_token_digest(bearer.credentials)}"
|
||||
credential = bearer.credentials
|
||||
else:
|
||||
return None
|
||||
if not credential.startswith(APIKeySmith.PREFIX):
|
||||
return None
|
||||
return f"{ip}:{credential[: APIKeySmith.HEAD_LENGTH]}"
|
||||
|
||||
|
||||
def _token_digest(token: str) -> str:
|
||||
"""Names a token's bucket without putting the token, or the digest it is
|
||||
stored under, in Redis."""
|
||||
return hashlib.sha256(f"v2-rate-limit:{token}".encode()).hexdigest()[:32]
|
||||
|
||||
|
||||
async def _refuse_before_verifying(ip: str, bucket: str) -> Optional[Response]:
|
||||
"""A 429 for a credential that may not be checked now, or None to go on.
|
||||
|
||||
Counts the presentation itself, so a valid credential flooding the API is
|
||||
refused unchecked past the cap, as is one, or an address, that keeps
|
||||
failing.
|
||||
"""
|
||||
if await _failed_credential_limiter.exhausted(bucket) or (
|
||||
await _failed_auth_limiter.exhausted(ip)
|
||||
):
|
||||
return error_response(
|
||||
429,
|
||||
"Too many failed authentication attempts. Try again shortly.",
|
||||
headers={"Retry-After": str(_failed_credential_limiter.window_seconds)},
|
||||
)
|
||||
try:
|
||||
await _presented_credential_limiter.check(bucket)
|
||||
except HTTPException as exc:
|
||||
return error_response(exc.status_code, str(exc.detail), headers=exc.headers)
|
||||
return None
|
||||
|
||||
|
||||
def app_path(scope: Scope) -> str:
|
||||
"""The request path inside this app, past the prefix it is mounted at."""
|
||||
path: str = scope.get("path", "")
|
||||
root: str = scope.get("root_path", "")
|
||||
return path[len(root) :] if root and path.startswith(root) else path
|
||||
|
||||
|
||||
def client_ip(scope: Scope) -> str:
|
||||
"""The caller's address, trusting only the proxies in front of us.
|
||||
|
||||
Our proxies append `trusted_proxy_count` entries, so the client is that
|
||||
many from the right; anything further left the caller wrote itself and
|
||||
could use to spread its requests over an unlimited number of buckets.
|
||||
It counts entries, not proxies: Google's load balancer appends two.
|
||||
"""
|
||||
peer = (scope.get("client") or ("unknown",))[0]
|
||||
hops = settings.config.trusted_proxy_count
|
||||
if hops < 1:
|
||||
return peer
|
||||
|
||||
# Every line, in order, as one list: a repeated header means the same as
|
||||
# its lines joined, and a proxy may add a line rather than extend one.
|
||||
forwarded = [
|
||||
value.strip()
|
||||
for line in Headers(scope=scope).getlist("x-forwarded-for")
|
||||
for value in line.split(",")
|
||||
if value.strip()
|
||||
]
|
||||
return forwarded[-hops] if len(forwarded) >= hops else peer
|
||||
|
||||
|
||||
def _with_rate_limit_headers(send: Send, state: Optional[RateLimitState]) -> Send:
|
||||
"""Attach the caller's window position to the response headers."""
|
||||
if state is None:
|
||||
return send
|
||||
|
||||
encoded = [
|
||||
(name.lower().encode(), value.encode())
|
||||
for name, value in state.headers().items()
|
||||
]
|
||||
|
||||
async def send_with_headers(message: Message) -> None:
|
||||
if message["type"] == "http.response.start":
|
||||
headers = message.setdefault("headers", [])
|
||||
# An endpoint with its own, narrower limiter has already set these.
|
||||
present = {name.lower() for name, _ in headers}
|
||||
headers.extend(h for h in encoded if h[0] not in present)
|
||||
await send(message)
|
||||
|
||||
return send_with_headers
|
||||
@@ -0,0 +1,98 @@
|
||||
"""The anonymous bucket is keyed on an address the caller cannot choose."""
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
|
||||
from backend.api.external.v2 import global_rate_limit
|
||||
from backend.api.external.v2.global_rate_limit import (
|
||||
GlobalRateLimitMiddleware,
|
||||
client_ip,
|
||||
)
|
||||
from backend.util.settings import Config
|
||||
|
||||
PEER = "10.0.0.9"
|
||||
CLIENT = "203.0.113.7"
|
||||
# What the hosted platform's proxies append after the client: Cloudflare's
|
||||
# edge address, then Google's load balancer appends that and its own.
|
||||
OUR_PROXIES = "172.70.1.1, 34.1.1.1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spoofed", ["", "6.6.6.6, ", "6.6.6.6, 7.7.7.7, "])
|
||||
async def test_a_spoofed_hop_does_not_move_the_anonymous_bucket(
|
||||
mocker: pytest_mock.MockFixture, spoofed: str
|
||||
) -> None:
|
||||
"""A caller that can set the key picks its own bucket, so there is no cap."""
|
||||
mocker.patch.object(
|
||||
global_rate_limit.settings.config,
|
||||
"trusted_proxy_count",
|
||||
Config.model_fields["trusted_proxy_count"].default,
|
||||
)
|
||||
anonymous = mocker.patch.object(
|
||||
global_rate_limit._anonymous_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
await _call_middleware(f"{spoofed}{CLIENT}, {OUR_PROXIES}")
|
||||
|
||||
anonymous.assert_awaited_once_with(CLIENT)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"hops, forwarded, expected",
|
||||
[
|
||||
(3, f"{CLIENT}, {OUR_PROXIES}", CLIENT),
|
||||
(1, f"1.2.3.4, {CLIENT}", CLIENT),
|
||||
# Fewer entries than our proxies append: the request did not come
|
||||
# through them, so the socket peer is the only trustworthy value.
|
||||
(3, f"{CLIENT}, 34.1.1.1", PEER),
|
||||
(1, "", PEER),
|
||||
(0, f"{CLIENT}, {OUR_PROXIES}", PEER),
|
||||
],
|
||||
)
|
||||
def test_the_client_is_the_entry_our_first_proxy_appended(
|
||||
mocker: pytest_mock.MockFixture, hops: int, forwarded: str, expected: str
|
||||
) -> None:
|
||||
mocker.patch.object(global_rate_limit.settings.config, "trusted_proxy_count", hops)
|
||||
|
||||
headers = [(b"x-forwarded-for", forwarded.encode())] if forwarded else []
|
||||
scope = {"type": "http", "client": (PEER, 0), "headers": headers}
|
||||
assert client_ip(scope) == expected
|
||||
|
||||
|
||||
def test_a_forwarded_header_on_several_lines_is_read_as_one_list(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Each of our proxies may add a line of its own instead of extending one;
|
||||
reading only the last line found too few entries and fell back to the
|
||||
peer, every client in one bucket."""
|
||||
mocker.patch.object(global_rate_limit.settings.config, "trusted_proxy_count", 2)
|
||||
scope = {
|
||||
"type": "http",
|
||||
"client": (PEER, 0),
|
||||
"headers": [
|
||||
(b"x-forwarded-for", CLIENT.encode()),
|
||||
(b"x-forwarded-for", b"34.1.1.1"),
|
||||
],
|
||||
}
|
||||
|
||||
assert client_ip(scope) == CLIENT
|
||||
|
||||
|
||||
async def _call_middleware(forwarded: str) -> None:
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
await send({"type": "http.response.body", "body": b""})
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b"", "more_body": False}
|
||||
|
||||
async def send(message):
|
||||
pass
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "GET",
|
||||
"path": "/runs",
|
||||
"headers": [(b"x-forwarded-for", forwarded.encode())],
|
||||
"client": (PEER, 0),
|
||||
}
|
||||
await GlobalRateLimitMiddleware(app)(scope, receive, send)
|
||||
@@ -0,0 +1,569 @@
|
||||
"""
|
||||
V2 External API - Graphs Endpoints
|
||||
|
||||
Provides endpoints for managing agent graphs (CRUD operations).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.store import db as store_db
|
||||
from backend.data import graph as graph_db
|
||||
from backend.data.tenancy import writable_organization_ids
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import (
|
||||
before_graph_activate,
|
||||
clear_unowned_auto_credentials,
|
||||
on_graph_deactivate,
|
||||
)
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
from .integrations.helpers import get_credential_requirements
|
||||
from .models import (
|
||||
BlockInfo,
|
||||
CredentialRequirement,
|
||||
Graph,
|
||||
GraphCreateRequest,
|
||||
GraphMeta,
|
||||
GraphSetActiveVersionRequest,
|
||||
GraphSettings,
|
||||
LibraryAgent,
|
||||
MarketplaceAgentDetails,
|
||||
)
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .tenancy import (
|
||||
TenantContext,
|
||||
in_tenant,
|
||||
owned_by_caller,
|
||||
require_paid_permission,
|
||||
require_permission,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
graphs_router = APIRouter(tags=["graphs"])
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="",
|
||||
summary="List graphs",
|
||||
operation_id="listGraphs",
|
||||
)
|
||||
async def list_graphs(
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_GRAPH)),
|
||||
) -> Page[GraphMeta]:
|
||||
"""List all graphs owned by the authenticated user."""
|
||||
graphs, pagination_info = await graph_db.list_graphs_paginated(
|
||||
user_id=auth.user_id,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
filter_by="active",
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
return page.paged(
|
||||
[GraphMeta.from_internal(g) for g in graphs],
|
||||
total_count=pagination_info.total_items,
|
||||
)
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}",
|
||||
summary="Get graph details",
|
||||
operation_id="getGraphDetails",
|
||||
)
|
||||
async def get_graph(
|
||||
graph_id: str,
|
||||
version: Optional[int] = Query(
|
||||
default=None,
|
||||
description="Specific version to retrieve (default: active version)",
|
||||
),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_GRAPH)),
|
||||
) -> Graph:
|
||||
"""
|
||||
Get detailed information about a specific graph.
|
||||
|
||||
Returns the active version by default. Pass `version` to retrieve
|
||||
a specific version instead.
|
||||
"""
|
||||
graph = await _get_graph_version(
|
||||
graph_id, version, auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
if not graph:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph #{graph_id} not found.",
|
||||
)
|
||||
return Graph.from_internal(_shared_if_not_owner(graph, auth.user_id))
|
||||
|
||||
|
||||
@graphs_router.post(
|
||||
path="",
|
||||
summary="Create graph",
|
||||
operation_id="createGraph",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_graph(
|
||||
create_graph: GraphCreateRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_GRAPH)
|
||||
),
|
||||
) -> Graph:
|
||||
"""Create a new agent graph."""
|
||||
from backend.api.features.library import db as library_db
|
||||
|
||||
internal_graph = create_graph.to_internal(id=str(uuid4()), version=1)
|
||||
|
||||
graph = graph_db.make_graph_model(internal_graph, auth.user_id)
|
||||
graph.reassign_ids(user_id=auth.user_id, reassign_graph_id=True)
|
||||
graph.validate_graph(for_run=False)
|
||||
# Before the writes: a graph that fails activation must not be saved, and
|
||||
# the credential edits activation makes must be.
|
||||
if graph.is_active:
|
||||
graph = await before_graph_activate(graph, user_id=auth.user_id)
|
||||
else:
|
||||
await clear_unowned_auto_credentials(graph, auth.user_id)
|
||||
|
||||
await graph_db.create_graph(
|
||||
graph,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
await library_db.create_library_agent(
|
||||
graph,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
|
||||
return Graph.from_internal(graph)
|
||||
|
||||
|
||||
@graphs_router.put(
|
||||
path="/{graph_id}",
|
||||
summary="Update graph by creating a new version",
|
||||
operation_id="updateGraphCreateVersion",
|
||||
)
|
||||
async def update_graph(
|
||||
graph_id: str,
|
||||
update_graph: GraphCreateRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_GRAPH)
|
||||
),
|
||||
) -> Graph:
|
||||
"""
|
||||
Update a graph by creating a new version.
|
||||
|
||||
This does not modify existing versions; it creates a new version
|
||||
with the provided graph definition.
|
||||
"""
|
||||
from backend.api.features.library import db as library_db
|
||||
|
||||
existing_versions = await graph_db.get_graph_all_versions(
|
||||
graph_id,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
if not existing_versions:
|
||||
raise HTTPException(
|
||||
status.HTTP_404_NOT_FOUND, detail=f"Graph #{graph_id} not found"
|
||||
)
|
||||
|
||||
# Versions share the graph's id but not its owner: a colleague's append
|
||||
# would land under their own user id.
|
||||
latest_version = owned_by_caller(
|
||||
max(existing_versions, key=lambda g: g.version), auth, f"Graph #{graph_id}"
|
||||
)
|
||||
organization_ids = await _writable_organizations(graph_id, auth)
|
||||
|
||||
internal_graph = update_graph.to_internal(
|
||||
id=graph_id, version=latest_version.version + 1
|
||||
)
|
||||
|
||||
current_active_version = next((v for v in existing_versions if v.is_active), None)
|
||||
graph = graph_db.make_graph_model(internal_graph, auth.user_id)
|
||||
graph.reassign_ids(user_id=auth.user_id, reassign_graph_id=False)
|
||||
graph.validate_graph(for_run=False)
|
||||
# Before the write, for the same reason as in create_graph.
|
||||
if graph.is_active:
|
||||
graph = await before_graph_activate(graph, user_id=auth.user_id)
|
||||
else:
|
||||
await clear_unowned_auto_credentials(graph, auth.user_id)
|
||||
|
||||
new_graph_version = await graph_db.create_graph(
|
||||
graph,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
|
||||
if new_graph_version.is_active:
|
||||
await library_db.update_agent_version_in_library(
|
||||
auth.user_id,
|
||||
new_graph_version.id,
|
||||
new_graph_version.version,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
await graph_db.set_graph_active_version(
|
||||
graph_id=graph_id, version=new_graph_version.version, user_id=auth.user_id
|
||||
)
|
||||
if current_active_version:
|
||||
await on_graph_deactivate(current_active_version, user_id=auth.user_id)
|
||||
await _move_webhook_presets(new_graph_version, auth, organization_ids)
|
||||
|
||||
new_graph_version_with_subgraphs = await graph_db.get_graph(
|
||||
graph_id,
|
||||
new_graph_version.version,
|
||||
user_id=auth.user_id,
|
||||
include_subgraphs=True,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
assert new_graph_version_with_subgraphs
|
||||
return Graph.from_internal(new_graph_version_with_subgraphs)
|
||||
|
||||
|
||||
# NOTE: we don't expose graph deletion in the UI, so this is commented for now
|
||||
# @graphs_router.delete(
|
||||
# path="/{graph_id}",
|
||||
# summary="Delete graph permanently",
|
||||
# status_code=status.HTTP_204_NO_CONTENT,
|
||||
# )
|
||||
# async def delete_graph(
|
||||
# graph_id: str,
|
||||
# auth: TenantContext = Security(
|
||||
# require_permission(APIKeyPermission.WRITE_GRAPH)
|
||||
# ),
|
||||
# ) -> None:
|
||||
# """
|
||||
# Permanently delete a graph and all its versions.
|
||||
|
||||
# This action cannot be undone. All associated executions will remain
|
||||
# but will reference a deleted graph.
|
||||
# """
|
||||
# if active_version := await graph_db.get_graph(
|
||||
# graph_id=graph_id, version=None, user_id=auth.user_id
|
||||
# ):
|
||||
# await on_graph_deactivate(active_version, user_id=auth.user_id)
|
||||
|
||||
# # FIXME: maybe only expose delete for library agents?
|
||||
# deleted_count = await graph_db.delete_graph(graph_id, user_id=auth.user_id)
|
||||
# if deleted_count == 0:
|
||||
# raise HTTPException(
|
||||
# status_code=status.HTTP_404_NOT_FOUND, detail=f"Graph {graph_id} not found"
|
||||
# )
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}/versions",
|
||||
summary="List graph versions",
|
||||
operation_id="listGraphVersions",
|
||||
)
|
||||
async def list_graph_versions(
|
||||
graph_id: str,
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_GRAPH)),
|
||||
) -> Page[Graph]:
|
||||
"""Get all versions of a specific graph."""
|
||||
# Paged in the query: each version's sub-graphs cost their own round trips.
|
||||
total_count, graphs = await asyncio.gather(
|
||||
graph_db.count_graph_versions(
|
||||
graph_id, user_id=auth.user_id, organization_id=auth.organization_id
|
||||
),
|
||||
graph_db.get_graph_all_versions(
|
||||
graph_id,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
include_subgraphs=True,
|
||||
limit=page.limit,
|
||||
offset=(page.page - 1) * page.limit,
|
||||
),
|
||||
)
|
||||
if not total_count:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph #{graph_id} not found.",
|
||||
)
|
||||
return page.paged(
|
||||
[Graph.from_internal(_shared_if_not_owner(g, auth.user_id)) for g in graphs],
|
||||
total_count,
|
||||
)
|
||||
|
||||
|
||||
@graphs_router.put(
|
||||
path="/{graph_id}/versions/active",
|
||||
summary="Set active graph version",
|
||||
operation_id="updateGraphSetActiveVersion",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def set_active_version(
|
||||
graph_id: str,
|
||||
request_body: GraphSetActiveVersionRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_GRAPH)
|
||||
),
|
||||
) -> None:
|
||||
"""
|
||||
Set which version of a graph is the active version.
|
||||
|
||||
The active version is the one used when executing the graph
|
||||
and what is shown to users in the UI.
|
||||
"""
|
||||
from backend.api.features.library import db as library_db
|
||||
|
||||
new_active_version = request_body.active_graph_version
|
||||
new_active_graph = await graph_db.get_graph(
|
||||
graph_id,
|
||||
new_active_version,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
if not new_active_graph:
|
||||
raise HTTPException(
|
||||
status.HTTP_404_NOT_FOUND,
|
||||
f"Graph #{graph_id} v{new_active_version} not found",
|
||||
)
|
||||
owned_by_caller(new_active_graph, auth, f"Graph #{graph_id}")
|
||||
organization_ids = await _writable_organizations(graph_id, auth)
|
||||
|
||||
# Not `get_graph(version=None)`, which returns the newest version, a draft
|
||||
# included.
|
||||
versions = await graph_db.get_graph_all_versions(
|
||||
graph_id, user_id=auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
current_active_graph = next((v for v in versions if v.is_active), None)
|
||||
|
||||
new_active_graph = await before_graph_activate(
|
||||
new_active_graph, user_id=auth.user_id
|
||||
)
|
||||
# Before the flag flips, as in update_graph: this can refuse (409), and
|
||||
# would leave the graph active on a version its library entry isn't on.
|
||||
# before_graph_activate wrote nothing outside the graph object, so a
|
||||
# refusal here leaves nothing to undo.
|
||||
await library_db.update_agent_version_in_library(
|
||||
auth.user_id,
|
||||
new_active_graph.id,
|
||||
new_active_graph.version,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
await graph_db.set_graph_active_version(
|
||||
graph_id=graph_id,
|
||||
version=new_active_version,
|
||||
user_id=auth.user_id,
|
||||
)
|
||||
|
||||
if current_active_graph and current_active_graph.version != new_active_version:
|
||||
await on_graph_deactivate(current_active_graph, user_id=auth.user_id)
|
||||
|
||||
await _move_webhook_presets(new_active_graph, auth, organization_ids)
|
||||
|
||||
|
||||
@graphs_router.patch(
|
||||
path="/{graph_id}/settings",
|
||||
summary="Update graph settings",
|
||||
operation_id="updateGraphSettings",
|
||||
)
|
||||
async def update_graph_settings(
|
||||
graph_id: str,
|
||||
settings: GraphSettings,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_GRAPH)),
|
||||
) -> GraphSettings:
|
||||
"""Update settings for a graph."""
|
||||
from backend.api.features.library import db as library_db
|
||||
|
||||
library_agent = in_tenant(
|
||||
await library_db.get_library_agent_by_graph_id(
|
||||
graph_id=graph_id, user_id=auth.user_id
|
||||
),
|
||||
auth,
|
||||
f"Graph #{graph_id}",
|
||||
)
|
||||
|
||||
updated_agent = await library_db.update_library_agent(
|
||||
user_id=auth.user_id,
|
||||
library_agent_id=library_agent.id,
|
||||
settings=settings.to_internal(),
|
||||
)
|
||||
|
||||
return GraphSettings(
|
||||
human_in_the_loop_safe_mode=updated_agent.settings.human_in_the_loop_safe_mode
|
||||
)
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}/library-agent",
|
||||
summary="Get library agent for graph",
|
||||
operation_id="getLibraryAgentForGraph",
|
||||
)
|
||||
async def get_library_agent_by_graph(
|
||||
graph_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> LibraryAgent:
|
||||
"""Get the library agent associated with a specific graph."""
|
||||
agent = in_tenant(
|
||||
await library_db.get_library_agent_by_graph_id(
|
||||
graph_id=graph_id,
|
||||
user_id=auth.user_id,
|
||||
),
|
||||
auth,
|
||||
f"Library agent for graph #{graph_id}",
|
||||
)
|
||||
return LibraryAgent.from_internal(agent)
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}/blocks",
|
||||
summary="List blocks used in a graph",
|
||||
operation_id="listBlocksInGraph",
|
||||
)
|
||||
async def list_graph_blocks(
|
||||
graph_id: str,
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_GRAPH)),
|
||||
) -> Page[BlockInfo]:
|
||||
"""List the unique blocks used by a graph."""
|
||||
from backend.blocks import get_block
|
||||
|
||||
graph = await _get_graph_version(
|
||||
graph_id, None, auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
if not graph:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph #{graph_id} not found.",
|
||||
)
|
||||
|
||||
seen_block_ids: set[str] = set()
|
||||
blocks: list[BlockInfo] = []
|
||||
|
||||
for node in graph.nodes:
|
||||
if node.block_id in seen_block_ids:
|
||||
continue
|
||||
seen_block_ids.add(node.block_id)
|
||||
|
||||
block = get_block(node.block_id)
|
||||
if block and not block.disabled:
|
||||
blocks.append(BlockInfo.from_internal(block))
|
||||
|
||||
return page.slice(blocks)
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}/credentials",
|
||||
summary="Get graph credentials",
|
||||
operation_id="getCredentialRequirementsForGraph",
|
||||
)
|
||||
async def list_graph_credential_requirements(
|
||||
graph_id: str,
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_INTEGRATIONS)
|
||||
),
|
||||
) -> Page[CredentialRequirement]:
|
||||
"""List credential requirements for a graph and matching user credentials."""
|
||||
graph = await _get_graph_version(
|
||||
graph_id, None, auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
if not graph:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND, detail=f"Graph #{graph_id} not found"
|
||||
)
|
||||
|
||||
requirements = await get_credential_requirements(
|
||||
graph.credentials_input_schema, auth.user_id
|
||||
)
|
||||
return page.slice(requirements)
|
||||
|
||||
|
||||
@graphs_router.get(
|
||||
path="/{graph_id}/marketplace-listing",
|
||||
summary="Get marketplace listing for graph",
|
||||
operation_id="getMarketplaceListingForGraph",
|
||||
)
|
||||
async def get_marketplace_listing_for_graph(
|
||||
graph_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_STORE)),
|
||||
) -> MarketplaceAgentDetails:
|
||||
"""Get the marketplace listing for a given graph, if one exists."""
|
||||
try:
|
||||
agent = await store_db.get_store_agent_by_graph_id(graph_id)
|
||||
except NotFoundError:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"No marketplace listing found for graph #{graph_id}",
|
||||
)
|
||||
return MarketplaceAgentDetails.from_internal(agent)
|
||||
|
||||
|
||||
async def _move_webhook_presets(
|
||||
new_graph: graph_db.GraphModel,
|
||||
auth: TenantContext,
|
||||
organization_ids: frozenset[str | None],
|
||||
) -> None:
|
||||
"""Keep the user's webhook presets firing the active version, if the
|
||||
credential may start runs.
|
||||
|
||||
A preset's webhook runs the version it is pinned to, with the credentials
|
||||
stored on the preset, so moving it decides what the next delivery runs:
|
||||
a start, which `WRITE_GRAPH` doesn't grant. Without `RUN_AGENT` the presets
|
||||
stay on the version they run now. Only presets in *organization_ids* move.
|
||||
A preset the migration skips is only logged: unlike the internal route's,
|
||||
this response has no field for it.
|
||||
"""
|
||||
if APIKeyPermission.RUN_AGENT not in auth.scopes:
|
||||
logger.info(
|
||||
f"Not moving graph #{new_graph.id}'s webhook presets to "
|
||||
f"v{new_graph.version}: the credential lacks RUN_AGENT"
|
||||
)
|
||||
return
|
||||
await library_db.migrate_webhook_presets_to_new_version(
|
||||
user_id=auth.user_id, new_graph=new_graph, organization_ids=organization_ids
|
||||
)
|
||||
|
||||
|
||||
async def _writable_organizations(
|
||||
graph_id: str, auth: TenantContext
|
||||
) -> frozenset[str | None]:
|
||||
"""The organizations a version saved or activated here may change rows in,
|
||||
once every version of the graph is checked to be the caller's, within them.
|
||||
|
||||
The version becomes the graph's active one and moves the library entry and
|
||||
webhook presets that follow it; on a graph that reaches another
|
||||
organization it would change what that organization runs.
|
||||
"""
|
||||
organization_ids = await writable_organization_ids(
|
||||
auth.user_id, auth.organization_id
|
||||
)
|
||||
await graph_db.assert_graph_within(graph_id, auth.user_id, organization_ids)
|
||||
return organization_ids
|
||||
|
||||
|
||||
async def _get_graph_version(
|
||||
graph_id: str, version: Optional[int], user_id: str, organization_id: str
|
||||
) -> Optional[graph_db.GraphModel]:
|
||||
"""`version`, or else the active one, with its sub-graphs."""
|
||||
if version is None:
|
||||
version = await graph_db.get_active_graph_version(
|
||||
graph_id, user_id, organization_id=organization_id
|
||||
)
|
||||
return await graph_db.get_graph(
|
||||
graph_id,
|
||||
version,
|
||||
user_id=user_id,
|
||||
include_subgraphs=True,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
|
||||
def _shared_if_not_owner(
|
||||
graph: graph_db.GraphModel, user_id: str
|
||||
) -> graph_db.GraphModel:
|
||||
"""As the internal graph routes answer: anyone but the owner, a colleague
|
||||
included, gets the copy without the owner's ids, credentials and secrets."""
|
||||
return graph if graph.user_id == user_id else graph.shared_copy()
|
||||
@@ -0,0 +1,740 @@
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
from uuid import uuid4
|
||||
|
||||
import prisma.models
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission, SubmissionStatus
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.library.exceptions import (
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
from backend.api.features.orgs.db import get_user_default_team
|
||||
from backend.blocks import _base
|
||||
from backend.blocks.agent import AgentExecutorBlock
|
||||
from backend.blocks.code_executor import ExecuteCodeBlock
|
||||
from backend.blocks.generic_webhook.triggers import GenericWebhookTriggerBlock
|
||||
from backend.data import graph as graph_db
|
||||
from backend.data.user import get_or_create_user
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import GraphActivationError
|
||||
from backend.util.exceptions import NotAuthorizedError
|
||||
from backend.util.test import SpinTestServer
|
||||
|
||||
from .graphs import (
|
||||
create_graph,
|
||||
get_graph,
|
||||
list_graph_blocks,
|
||||
list_graph_credential_requirements,
|
||||
list_graph_versions,
|
||||
set_active_version,
|
||||
update_graph,
|
||||
)
|
||||
from .library.agents import update_library_agent as patch_library_agent
|
||||
from .models import (
|
||||
GraphCreateRequest,
|
||||
GraphNode,
|
||||
GraphSetActiveVersionRequest,
|
||||
LibraryAgentUpdateRequest,
|
||||
)
|
||||
from .pagination import PageRequest
|
||||
from .tenancy import TenantContext
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
_ALL_BUT_RUN_AGENT = [p for p in APIKeyPermission if p != APIKeyPermission.RUN_AGENT]
|
||||
_REQUEST = GraphCreateRequest(name="Agent", nodes=[], links=[])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def graph_in_org_1(mocker: pytest_mock.MockFixture) -> AsyncMock:
|
||||
"""The credential acts in org-1, where every version of the graph is."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs.writable_organization_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=frozenset({"org-1"}),
|
||||
)
|
||||
return mocker.patch(
|
||||
"backend.data.graph.get_graph_version_owners",
|
||||
new_callable=AsyncMock,
|
||||
return_value={(_AUTH.user_id, "org-1")},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def writes(
|
||||
mocker: pytest_mock.MockFixture, graph_in_org_1: AsyncMock
|
||||
) -> dict[str, AsyncMock]:
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph_all_versions",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[Mock(version=1, is_active=True, user_id=_AUTH.user_id)],
|
||||
)
|
||||
return {
|
||||
name: mocker.patch(target, new_callable=AsyncMock)
|
||||
for name, target in {
|
||||
"create_graph": "backend.data.graph.create_graph",
|
||||
"create_library_agent": "backend.api.features.library.db.create_library_agent",
|
||||
"update_agent_version_in_library": "backend.api.features.library.db.update_agent_version_in_library",
|
||||
}.items()
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scopes, moves", [(list(APIKeyPermission), True), (_ALL_BUT_RUN_AGENT, False)]
|
||||
)
|
||||
async def test_activating_moves_webhook_presets_only_for_a_credential_that_may_run(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
graph_in_org_1: AsyncMock,
|
||||
scopes: list[APIKeyPermission],
|
||||
moves: bool,
|
||||
) -> None:
|
||||
graph = Mock(
|
||||
user_id=_AUTH.user_id, organization_id="org-1", id="graph-1", version=2
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph", new_callable=AsyncMock, return_value=graph
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph_all_versions",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[graph],
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate",
|
||||
new_callable=AsyncMock,
|
||||
return_value=graph,
|
||||
)
|
||||
mocker.patch("backend.data.graph.set_graph_active_version", new_callable=AsyncMock)
|
||||
move_entry = mocker.patch(
|
||||
"backend.api.features.library.db.update_agent_version_in_library",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
migrate = mocker.patch(
|
||||
"backend.api.features.library.db.migrate_webhook_presets_to_new_version",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
await set_active_version(
|
||||
"graph-1",
|
||||
GraphSetActiveVersionRequest(active_graph_version=2),
|
||||
auth=_AUTH.model_copy(update={"scopes": scopes}),
|
||||
)
|
||||
|
||||
# Only the credential's organization's entry and presets follow it.
|
||||
assert move_entry.await_args.kwargs["organization_ids"] == {"org-1"}
|
||||
assert migrate.await_count == (1 if moves else 0)
|
||||
if moves:
|
||||
assert migrate.await_args.kwargs["organization_ids"] == {"org-1"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"change",
|
||||
[
|
||||
lambda: update_graph("graph-1", _REQUEST, auth=_AUTH),
|
||||
lambda: set_active_version(
|
||||
"graph-1", GraphSetActiveVersionRequest(active_graph_version=1), auth=_AUTH
|
||||
),
|
||||
],
|
||||
ids=["update", "set-active-version"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"other_version",
|
||||
[
|
||||
# An untagged version is the personal org's, and org-1 isn't that.
|
||||
(_AUTH.user_id, None),
|
||||
# A colleague's version would be deactivated with the caller's
|
||||
# credentials, and stay active beside the new one.
|
||||
("user-2", "org-1"),
|
||||
],
|
||||
ids=["another-organization", "a-colleagues-version"],
|
||||
)
|
||||
async def test_a_graph_with_a_version_beyond_the_caller_is_not_changed(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
writes: dict[str, AsyncMock],
|
||||
graph_in_org_1: AsyncMock,
|
||||
change,
|
||||
other_version: tuple[str, str | None],
|
||||
) -> None:
|
||||
"""Its new active version would move another organization's library entry
|
||||
and presets, and be the active version there too."""
|
||||
graph_in_org_1.return_value = {(_AUTH.user_id, "org-1"), other_version}
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(user_id=_AUTH.user_id, organization_id=None, version=1),
|
||||
)
|
||||
activate = mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await change()
|
||||
|
||||
activate.assert_not_awaited()
|
||||
writes["create_graph"].assert_not_awaited()
|
||||
writes["update_agent_version_in_library"].assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_refused_library_move_leaves_the_active_version_alone(
|
||||
mocker: pytest_mock.MockFixture, graph_in_org_1: AsyncMock
|
||||
) -> None:
|
||||
"""The move can answer 409; after it, the graph would be active on a
|
||||
version its library entry isn't on."""
|
||||
graph = Mock(
|
||||
user_id=_AUTH.user_id, organization_id="org-1", id="graph-1", version=2
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph", new_callable=AsyncMock, return_value=graph
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph_all_versions",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[graph],
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate",
|
||||
new_callable=AsyncMock,
|
||||
return_value=graph,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.update_agent_version_in_library",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=LibraryAgentInAnotherOrganizationError("elsewhere"),
|
||||
)
|
||||
flip = mocker.patch(
|
||||
"backend.data.graph.set_graph_active_version", new_callable=AsyncMock
|
||||
)
|
||||
deactivate = mocker.patch(
|
||||
"backend.api.external.v2.graphs.on_graph_deactivate", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(LibraryAgentInAnotherOrganizationError):
|
||||
await set_active_version(
|
||||
"graph-1", GraphSetActiveVersionRequest(active_graph_version=2), auth=_AUTH
|
||||
)
|
||||
|
||||
flip.assert_not_awaited()
|
||||
# The activation hooks wrote nothing to undo, and the version may be the
|
||||
# one already active: deactivating it would detach its live webhooks.
|
||||
deactivate.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"save",
|
||||
[
|
||||
lambda: create_graph(_REQUEST, auth=_AUTH),
|
||||
lambda: update_graph("graph-1", _REQUEST, auth=_AUTH),
|
||||
],
|
||||
ids=["create", "update"],
|
||||
)
|
||||
async def test_a_graph_that_fails_activation_is_not_saved(
|
||||
mocker: pytest_mock.MockFixture, writes: dict[str, AsyncMock], save
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=GraphActivationError("missing credential"),
|
||||
)
|
||||
|
||||
with pytest.raises(GraphActivationError):
|
||||
await save()
|
||||
|
||||
for write in writes.values():
|
||||
write.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_created_graph_is_saved_with_its_activation_edits(
|
||||
mocker: pytest_mock.MockFixture, writes: dict[str, AsyncMock]
|
||||
) -> None:
|
||||
activated = Mock()
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate",
|
||||
new_callable=AsyncMock,
|
||||
return_value=activated,
|
||||
)
|
||||
mocker.patch("backend.api.external.v2.graphs.Graph.from_internal")
|
||||
|
||||
await create_graph(_REQUEST, auth=_AUTH)
|
||||
|
||||
assert writes["create_graph"].await_args.args[0] is activated
|
||||
assert writes["create_library_agent"].await_args.args[0] is activated
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"save",
|
||||
[
|
||||
lambda request: create_graph(request, auth=_AUTH),
|
||||
lambda request: update_graph("graph-1", request, auth=_AUTH),
|
||||
],
|
||||
ids=["create", "update"],
|
||||
)
|
||||
async def test_an_inactive_version_is_saved_without_another_users_credential_refs(
|
||||
mocker: pytest_mock.MockFixture, writes: dict[str, AsyncMock], save
|
||||
) -> None:
|
||||
activate = mocker.patch(
|
||||
"backend.api.external.v2.graphs.before_graph_activate",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=GraphActivationError("missing credential"),
|
||||
)
|
||||
clear = mocker.patch(
|
||||
"backend.api.external.v2.graphs.clear_unowned_auto_credentials",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
mocker.patch("backend.data.graph.get_graph", new_callable=AsyncMock)
|
||||
mocker.patch("backend.api.external.v2.graphs.Graph.from_internal")
|
||||
writes["create_graph"].return_value = Mock(is_active=False)
|
||||
inactive = _REQUEST.model_copy(update={"is_active": False})
|
||||
|
||||
await save(inactive)
|
||||
|
||||
clear.assert_awaited_once()
|
||||
assert clear.await_args.args[0] is writes["create_graph"].await_args.args[0]
|
||||
activate.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_listed_versions_carry_the_credentials_their_sub_graphs_need(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
parent = await _parent_of_a_sub_graph(user_id)
|
||||
auth = _AUTH.model_copy(update={"user_id": user_id})
|
||||
|
||||
single = await get_graph(parent.id, version=None, auth=auth)
|
||||
listed = await list_graph_versions(parent.id, page=PageRequest(limit=10), auth=auth)
|
||||
|
||||
assert single.credentials_input_schema["properties"]
|
||||
assert listed.items[0].credentials_input_schema == single.credentials_input_schema
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_version_change_returns_the_library_agent_a_fresh_read_does(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
"""A sub-graph's credentials included, which the active-version routes rely on."""
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
parent = await _parent_of_a_sub_graph(user_id)
|
||||
await library_db.create_library_agent(parent, user_id)
|
||||
|
||||
changed = await library_db.update_agent_version_in_library(
|
||||
user_id, parent.id, parent.version
|
||||
)
|
||||
fresh = await library_db.get_library_agent_by_graph_id(user_id, parent.id)
|
||||
|
||||
assert fresh and fresh.credentials_input_schema["properties"]
|
||||
assert changed.credentials_input_schema == fresh.credentials_input_schema
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_page_of_versions_resolves_only_its_own_sub_graphs(
|
||||
server: SpinTestServer, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""Every version's sub-graphs cost round trips, so a page pays only for its own."""
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
first = await graph_db.create_graph(
|
||||
graph_db.Graph(name="Versions", description="", nodes=[]), user_id
|
||||
)
|
||||
for version in (2, 3):
|
||||
later = graph_db.make_graph_model(
|
||||
graph_db.Graph(
|
||||
id=first.id, version=version, name="Versions", description=""
|
||||
),
|
||||
user_id,
|
||||
)
|
||||
later.reassign_ids(user_id=user_id, reassign_graph_id=False)
|
||||
await graph_db.create_graph(later, user_id)
|
||||
resolve = mocker.patch.object(
|
||||
graph_db, "get_sub_graphs_of_each", wraps=graph_db.get_sub_graphs_of_each
|
||||
)
|
||||
auth = _AUTH.model_copy(update={"user_id": user_id})
|
||||
|
||||
newest = await list_graph_versions(first.id, page=PageRequest(limit=1), auth=auth)
|
||||
older = await list_graph_versions(
|
||||
first.id, page=PageRequest(limit=1, cursor=newest.next_cursor), auth=auth
|
||||
)
|
||||
|
||||
assert [v.version for v in newest.items + older.items] == [3, 2]
|
||||
assert newest.total_count == 3
|
||||
assert [[g.version for g in call.args[0]] for call in resolve.await_args_list] == [
|
||||
[3],
|
||||
[2],
|
||||
]
|
||||
|
||||
|
||||
_ACTIVATING_V2 = pytest.mark.parametrize(
|
||||
"activate_v2",
|
||||
[
|
||||
lambda graph_id, request, auth: update_graph(graph_id, request, auth=auth),
|
||||
lambda graph_id, request, auth: _save_inactive_then_activate(
|
||||
graph_id, request, auth
|
||||
),
|
||||
],
|
||||
ids=["update", "set-active-version"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
@_ACTIVATING_V2
|
||||
async def test_activating_a_version_moves_its_webhook_presets_onto_it(
|
||||
server: SpinTestServer, monkeypatch: pytest.MonkeyPatch, activate_v2
|
||||
) -> None:
|
||||
"""Left on v1, a preset's webhook URL keeps running the deactivated version."""
|
||||
assert await _preset_version_once_v2_is_active(monkeypatch, activate_v2, _AUTH) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
@_ACTIVATING_V2
|
||||
async def test_a_credential_that_cannot_start_runs_leaves_webhook_presets_in_place(
|
||||
server: SpinTestServer, monkeypatch: pytest.MonkeyPatch, activate_v2
|
||||
) -> None:
|
||||
"""Moving a preset decides what its next delivery runs, with the
|
||||
credentials stored on it: a start, which WRITE_GRAPH alone doesn't grant."""
|
||||
auth = _AUTH.model_copy(update={"scopes": _ALL_BUT_RUN_AGENT})
|
||||
|
||||
assert await _preset_version_once_v2_is_active(monkeypatch, activate_v2, auth) == 1
|
||||
|
||||
|
||||
async def _preset_version_once_v2_is_active(
|
||||
monkeypatch: pytest.MonkeyPatch, activate_v2, auth: TenantContext
|
||||
) -> int:
|
||||
"""Pin a webhook preset to a triggered graph's v1, activate v2 with `auth`,
|
||||
and read which version the preset runs."""
|
||||
# Trigger blocks are disabled without a platform URL, and CI sets none.
|
||||
monkeypatch.setattr(_base.app_config, "platform_base_url", "https://example.com")
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
auth = auth.model_copy(update={"user_id": user_id})
|
||||
triggered = GraphCreateRequest(
|
||||
name="Triggered",
|
||||
nodes=[GraphNode(id="trigger", block_id=GenericWebhookTriggerBlock().id)],
|
||||
links=[],
|
||||
)
|
||||
v1 = await create_graph(triggered, auth=auth)
|
||||
webhook = await prisma.models.IntegrationWebhook.prisma().create(
|
||||
data={
|
||||
"userId": user_id,
|
||||
"provider": "generic_webhook",
|
||||
"credentialsId": "",
|
||||
"webhookType": "plain",
|
||||
"resource": "",
|
||||
"events": [],
|
||||
"config": "{}",
|
||||
"secret": "",
|
||||
"providerWebhookId": "",
|
||||
}
|
||||
)
|
||||
preset = await prisma.models.AgentPreset.prisma().create(
|
||||
data={
|
||||
"userId": user_id,
|
||||
"name": "On webhook",
|
||||
"description": "",
|
||||
"agentGraphId": v1.id,
|
||||
"agentGraphVersion": v1.version,
|
||||
"webhookId": webhook.id,
|
||||
# Presets take their graph's organization.
|
||||
"organizationId": auth.organization_id,
|
||||
}
|
||||
)
|
||||
|
||||
await activate_v2(v1.id, triggered, auth)
|
||||
|
||||
preset = await prisma.models.AgentPreset.prisma().find_unique_or_raise(
|
||||
where={"id": preset.id}
|
||||
)
|
||||
return preset.agentGraphVersion
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_graph_read_without_a_version_is_of_the_active_one(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
auth = await _new_user_auth()
|
||||
graph_id = await _active_v1_and_a_newer_draft(auth)
|
||||
|
||||
graph = await get_graph(graph_id, version=None, auth=auth)
|
||||
|
||||
assert graph.version == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"read", [list_graph_blocks, list_graph_credential_requirements]
|
||||
)
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_page_read_of_a_graph_is_of_its_active_version(
|
||||
server: SpinTestServer, read
|
||||
) -> None:
|
||||
auth = await _new_user_auth()
|
||||
graph_id = await _active_v1_and_a_newer_draft(auth)
|
||||
|
||||
page = await read(graph_id, page=PageRequest(limit=10), auth=auth)
|
||||
|
||||
# The active v1 has no nodes; the draft's code block needs credentials.
|
||||
assert page.items == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_library_agent_saved_past_its_listed_version_lists_as_published(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
"""As its single read reports it: a later save is still the listed agent."""
|
||||
auth = await _new_user_auth()
|
||||
v1 = await create_graph(_REQUEST, auth=auth)
|
||||
listing = await prisma.models.StoreListing.prisma().create(
|
||||
data={
|
||||
"slug": f"listed-{v1.id}",
|
||||
"agentGraphId": v1.id,
|
||||
"owningUserId": auth.user_id,
|
||||
"hasApprovedVersion": True,
|
||||
}
|
||||
)
|
||||
await prisma.models.StoreListingVersion.prisma().create(
|
||||
data={
|
||||
"agentGraphId": v1.id,
|
||||
"agentGraphVersion": 1,
|
||||
"name": "Agent",
|
||||
"subHeading": "",
|
||||
"description": "",
|
||||
"storeListingId": listing.id,
|
||||
"submissionStatus": SubmissionStatus.APPROVED,
|
||||
"isAvailable": True,
|
||||
"imageUrls": [],
|
||||
"categories": [],
|
||||
}
|
||||
)
|
||||
await update_graph(v1.id, _REQUEST, auth=auth)
|
||||
|
||||
listed = await library_db.list_library_agents(auth.user_id, published=True)
|
||||
unlisted = await library_db.list_library_agents(auth.user_id, published=False)
|
||||
|
||||
assert [agent.graph_version for agent in listed.agents] == [2]
|
||||
assert unlisted.agents == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_activating_a_draft_deactivates_the_version_that_was_active(
|
||||
server: SpinTestServer, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""Left active, the old version's webhook triggers keep firing."""
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
auth = _AUTH.model_copy(update={"user_id": user_id})
|
||||
v1 = await create_graph(_REQUEST, auth=auth)
|
||||
deactivate = mocker.patch(
|
||||
"backend.api.external.v2.graphs.on_graph_deactivate", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
await _save_inactive_then_activate(v1.id, _REQUEST, auth)
|
||||
|
||||
assert [call.args[0].version for call in deactivate.await_args_list] == [1]
|
||||
|
||||
|
||||
async def _new_user_auth() -> TenantContext:
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
return _AUTH.model_copy(update={"user_id": user_id})
|
||||
|
||||
|
||||
async def _active_v1_and_a_newer_draft(auth: TenantContext) -> str:
|
||||
v1 = await create_graph(_REQUEST, auth=auth)
|
||||
draft = GraphCreateRequest(
|
||||
name="Agent",
|
||||
nodes=[GraphNode(id="code", block_id=ExecuteCodeBlock().id)],
|
||||
links=[],
|
||||
is_active=False,
|
||||
)
|
||||
await update_graph(v1.id, draft, auth=auth)
|
||||
return v1.id
|
||||
|
||||
|
||||
async def _parent_of_a_sub_graph(user_id: str) -> graph_db.GraphModel:
|
||||
sub = await graph_db.create_graph(
|
||||
graph_db.Graph(
|
||||
name="Sub",
|
||||
description="",
|
||||
nodes=[graph_db.Node(block_id=ExecuteCodeBlock().id)],
|
||||
),
|
||||
user_id,
|
||||
)
|
||||
return await graph_db.create_graph(
|
||||
graph_db.Graph(
|
||||
name="Parent",
|
||||
description="",
|
||||
nodes=[
|
||||
graph_db.Node(
|
||||
block_id=AgentExecutorBlock().id,
|
||||
input_default={
|
||||
"user_id": user_id,
|
||||
"graph_id": sub.id,
|
||||
"graph_version": sub.version,
|
||||
"inputs": {},
|
||||
"input_schema": {},
|
||||
"output_schema": {},
|
||||
},
|
||||
)
|
||||
],
|
||||
),
|
||||
user_id,
|
||||
)
|
||||
|
||||
|
||||
async def _save_inactive_then_activate(
|
||||
graph_id: str, request: GraphCreateRequest, auth: TenantContext
|
||||
) -> None:
|
||||
await update_graph(
|
||||
graph_id, request.model_copy(update={"is_active": False}), auth=auth
|
||||
)
|
||||
await set_active_version(
|
||||
graph_id, GraphSetActiveVersionRequest(active_graph_version=2), auth=auth
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_another_organizations_key_cannot_version_a_personal_graph(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
"""A graph saved without an organization is the personal org's. A key in
|
||||
another organization can read it, but a version it saved or activated would
|
||||
move the personal org's library entry and the graph's active version."""
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
personal_org, _ = await get_user_default_team(user_id)
|
||||
assert personal_org
|
||||
v1 = await _untagged_version(user_id, str(uuid4()), version=1)
|
||||
await library_db.create_library_agent(v1, user_id=user_id)
|
||||
await _untagged_version(user_id, v1.id, version=2, is_active=False)
|
||||
activate_v2 = GraphSetActiveVersionRequest(active_graph_version=2)
|
||||
elsewhere = _AUTH.model_copy(
|
||||
update={"user_id": user_id, "organization_id": str(uuid4())}
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await update_graph(v1.id, _REQUEST, auth=elsewhere)
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await set_active_version(v1.id, activate_v2, auth=elsewhere)
|
||||
|
||||
versions = await graph_db.get_graph_all_versions(v1.id, user_id=user_id)
|
||||
assert sorted(v.version for v in versions) == [1, 2]
|
||||
assert await _library_version(user_id, v1.id) == 1
|
||||
|
||||
personal = _AUTH.model_copy(
|
||||
update={"user_id": user_id, "organization_id": personal_org}
|
||||
)
|
||||
await set_active_version(v1.id, activate_v2, auth=personal)
|
||||
|
||||
assert await _library_version(user_id, v1.id) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_library_agent_is_not_pointed_at_another_organizations_version(
|
||||
server: SpinTestServer,
|
||||
) -> None:
|
||||
"""The web app saves a version in whichever organization is selected, so a
|
||||
graph can span two. A library agent pointed at the other one's version
|
||||
would run it, billed to the first."""
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
personal_org, _ = await get_user_default_team(user_id)
|
||||
assert personal_org
|
||||
v1 = await _untagged_version(user_id, str(uuid4()), version=1)
|
||||
[entry] = await library_db.create_library_agent(v1, user_id=user_id)
|
||||
await _untagged_version(
|
||||
user_id, v1.id, version=2, is_active=False, organization_id=str(uuid4())
|
||||
)
|
||||
await _untagged_version(user_id, v1.id, version=3, is_active=False)
|
||||
personal = _AUTH.model_copy(
|
||||
update={"user_id": user_id, "organization_id": personal_org}
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await patch_library_agent(
|
||||
LibraryAgentUpdateRequest(graph_version=2), entry.id, auth=personal
|
||||
)
|
||||
assert await _library_version(user_id, v1.id) == 1
|
||||
|
||||
await patch_library_agent(
|
||||
LibraryAgentUpdateRequest(graph_version=3), entry.id, auth=personal
|
||||
)
|
||||
assert await _library_version(user_id, v1.id) == 3
|
||||
|
||||
|
||||
async def _untagged_version(
|
||||
user_id: str,
|
||||
graph_id: str,
|
||||
*,
|
||||
version: int,
|
||||
is_active: bool = True,
|
||||
organization_id: str | None = None,
|
||||
) -> graph_db.GraphModel:
|
||||
"""A version saved with no organization, as the app's chat saves them, or
|
||||
in *organization_id*."""
|
||||
internal = _REQUEST.model_copy(update={"is_active": is_active}).to_internal(
|
||||
id=graph_id, version=version
|
||||
)
|
||||
graph = graph_db.make_graph_model(internal, user_id)
|
||||
graph.reassign_ids(user_id=user_id, reassign_graph_id=False)
|
||||
return await graph_db.create_graph(
|
||||
graph, user_id=user_id, organization_id=organization_id
|
||||
)
|
||||
|
||||
|
||||
async def _library_version(user_id: str, graph_id: str) -> int:
|
||||
entry = await library_db.get_library_agent_by_graph_id(user_id, graph_id)
|
||||
assert entry
|
||||
return entry.graph_version
|
||||
|
||||
|
||||
_OWNERS_SECRET = "owner-secret-do-not-leak"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reader", ["get_graph", "list_graph_versions"])
|
||||
@pytest.mark.parametrize("owner, shared", [("user-1", False), ("user-2", True)])
|
||||
async def test_a_graph_someone_else_owns_is_read_as_its_shared_copy(
|
||||
mocker: pytest_mock.MockFixture, reader: str, owner: str, shared: bool
|
||||
) -> None:
|
||||
"""As on the internal routes: no owner's ids, credentials or secret inputs."""
|
||||
graph = graph_db.make_graph_model(
|
||||
GraphCreateRequest(
|
||||
name="Agent",
|
||||
nodes=[
|
||||
GraphNode(
|
||||
id="trigger",
|
||||
block_id=GenericWebhookTriggerBlock().id,
|
||||
input_default={"secret_token": _OWNERS_SECRET},
|
||||
)
|
||||
],
|
||||
links=[],
|
||||
).to_internal(id="graph-1", version=1),
|
||||
owner,
|
||||
).model_copy(update={"organization_id": "org-1"})
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.graphs._get_graph_version",
|
||||
new_callable=AsyncMock,
|
||||
return_value=graph,
|
||||
)
|
||||
mocker.patch.object(
|
||||
graph_db, "count_graph_versions", new_callable=AsyncMock, return_value=1
|
||||
)
|
||||
mocker.patch.object(
|
||||
graph_db, "get_graph_all_versions", new_callable=AsyncMock, return_value=[graph]
|
||||
)
|
||||
|
||||
if reader == "get_graph":
|
||||
body = (
|
||||
await get_graph(graph_id="graph-1", version=None, auth=_AUTH)
|
||||
).model_dump_json()
|
||||
else:
|
||||
page = await list_graph_versions(
|
||||
graph_id="graph-1", page=PageRequest(limit=20), auth=_AUTH
|
||||
)
|
||||
body = page.model_dump_json()
|
||||
|
||||
assert (_OWNERS_SECRET not in body) is shared
|
||||
assert ('"organization_id":"org-1"' not in body) is shared
|
||||
@@ -0,0 +1,220 @@
|
||||
"""
|
||||
V2 External API - Idempotent run creation
|
||||
|
||||
Starting a run costs the caller money, and a timed-out request says nothing
|
||||
about whether it started. Without a key the only safe retry is no retry.
|
||||
|
||||
The caller sends `Idempotency-Key: <a value they choose>`. The first request to
|
||||
claim it runs; a later request with the same key gets back the run the first one
|
||||
created. A request that arrives while the first is still in flight is a 409 —
|
||||
there is no run to hand back yet. A key is bound to the request it first came
|
||||
with: reusing it for another agent or other inputs is a 422, never a replay of
|
||||
a run the caller didn't ask for.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncIterator, Optional
|
||||
|
||||
from fastapi import Header, HTTPException
|
||||
from starlette import status
|
||||
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data.redis_client import get_redis_async
|
||||
|
||||
from .models import AgentGraphRun
|
||||
from .tenancy import TenantContext, in_tenant
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
IDEMPOTENCY_HEADER = "Idempotency-Key"
|
||||
|
||||
# Long enough to cover any sane retry window, short enough that a key is not a
|
||||
# permanent reservation on a name the caller picked.
|
||||
_TTL_SECONDS = 24 * 60 * 60
|
||||
|
||||
|
||||
def idempotency_key(
|
||||
key: Optional[str] = Header(
|
||||
default=None,
|
||||
alias=IDEMPOTENCY_HEADER,
|
||||
max_length=255,
|
||||
description=(
|
||||
"Retry-safety token. Repeating a request with the same value returns "
|
||||
"the run the first one started instead of starting another. Scoped to "
|
||||
"the caller and its organization, and bound to the request it first "
|
||||
"came with: reusing it for another agent or other inputs is a 422. "
|
||||
"Expires after 24 hours."
|
||||
),
|
||||
),
|
||||
) -> Optional[str]:
|
||||
return (key or "").strip() or None
|
||||
|
||||
|
||||
def request_fingerprint(**request: Any) -> str:
|
||||
"""A digest of what a request asks for, to bind its idempotency key to."""
|
||||
canonical = json.dumps(request, sort_keys=True, default=str)
|
||||
return hashlib.sha256(canonical.encode()).hexdigest()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def idempotent_run(
|
||||
key: Optional[str], auth: TenantContext, fingerprint: str
|
||||
) -> AsyncIterator["RunClaim"]:
|
||||
"""Claim `key` for this request, or report the run it already produced.
|
||||
|
||||
A key the store can't claim is a retryable 503: running anyway would make
|
||||
the retry the caller sent the key for a second run and a second charge.
|
||||
"""
|
||||
claim = RunClaim(
|
||||
key=key,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
if key is None:
|
||||
yield claim
|
||||
return
|
||||
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
claimed = await redis.set(
|
||||
claim.redis_key, claim.stored_value(), nx=True, ex=_TTL_SECONDS
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Idempotency store unavailable: {e}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail=f"Cannot claim the {IDEMPOTENCY_HEADER} right now; retry shortly.",
|
||||
headers={"Retry-After": "5"},
|
||||
)
|
||||
|
||||
if not claimed:
|
||||
claim.existing_run_id = await claim.resolve_existing()
|
||||
yield claim
|
||||
return
|
||||
|
||||
claim.holds_key = True
|
||||
try:
|
||||
yield claim
|
||||
except BaseException:
|
||||
# No run came of it (a cancelled request included), so free the key for a
|
||||
# retry. A run that started keeps it, so a retry finds it, not a second run.
|
||||
await claim.release()
|
||||
raise
|
||||
|
||||
|
||||
async def replayed_run(claim: "RunClaim", auth: TenantContext) -> AgentGraphRun:
|
||||
"""The run a previous request with this key started."""
|
||||
return AgentGraphRun.from_internal(
|
||||
in_tenant(
|
||||
await execution_db.get_graph_execution(
|
||||
user_id=auth.user_id,
|
||||
execution_id=claim.existing_run_id or "",
|
||||
organization_id=auth.organization_id,
|
||||
),
|
||||
auth,
|
||||
f"Run #{claim.existing_run_id}",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class RunClaim:
|
||||
"""One request's hold on an idempotency key."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
key: Optional[str],
|
||||
user_id: str,
|
||||
organization_id: str,
|
||||
fingerprint: str,
|
||||
team_id: Optional[str] = None,
|
||||
) -> None:
|
||||
self.key = key
|
||||
self.user_id = user_id
|
||||
self.organization_id = organization_id
|
||||
self.team_id = team_id
|
||||
self.fingerprint = fingerprint
|
||||
self.holds_key = False
|
||||
self.run_id: Optional[str] = None
|
||||
self.existing_run_id: Optional[str] = None
|
||||
|
||||
@property
|
||||
def redis_key(self) -> str:
|
||||
# Per organization and team, where a run lands: a shared key there must
|
||||
# start its own run rather than replay one from elsewhere.
|
||||
team = self.team_id or "-"
|
||||
return f"v2:idem:{self.user_id}:{self.organization_id}:{team}:{self.key}"
|
||||
|
||||
def stored_value(self, run_id: Optional[str] = None) -> str:
|
||||
"""What the key holds: the request it was claimed for, and its run."""
|
||||
return json.dumps({"request": self.fingerprint, "run": run_id})
|
||||
|
||||
async def resolve_existing(self) -> str:
|
||||
"""The run the first request recorded; 409 while it is still running,
|
||||
422 if the key was first sent with a different request."""
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
stored = await redis.get(self.redis_key)
|
||||
except Exception as e:
|
||||
logger.warning(f"Idempotency store unavailable on replay: {e}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="Cannot tell whether this request already ran; retry shortly.",
|
||||
headers={"Retry-After": "5"},
|
||||
)
|
||||
claimed = _parse(stored)
|
||||
if claimed is not None and claimed.get("request") != self.fingerprint:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail=f"{IDEMPOTENCY_HEADER} '{self.key}' was already used for a "
|
||||
"different request. Use a new key for a new request.",
|
||||
)
|
||||
run_id = claimed.get("run") if claimed is not None else None
|
||||
if not isinstance(run_id, str) or not run_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"A request with {IDEMPOTENCY_HEADER} '{self.key}' is still "
|
||||
"in flight. Retry once it completes.",
|
||||
)
|
||||
return run_id
|
||||
|
||||
async def record(self, run_id: str) -> None:
|
||||
"""Point the key at the run, so a later retry gets this one back."""
|
||||
self.run_id = run_id
|
||||
if not self.holds_key:
|
||||
return
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
await redis.set(self.redis_key, self.stored_value(run_id), ex=_TTL_SECONDS)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not record idempotent run {run_id}: {e}")
|
||||
|
||||
async def release(self) -> None:
|
||||
# Once a run exists the key stays, even unrecorded: a retry then gets a
|
||||
# 409 until the key expires, never a second run.
|
||||
if not self.holds_key or self.run_id is not None:
|
||||
return
|
||||
try:
|
||||
redis = await get_redis_async()
|
||||
await redis.delete(self.redis_key)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not release idempotency key: {e}")
|
||||
|
||||
|
||||
def _parse(stored: Any) -> Optional[dict[str, Any]]:
|
||||
"""What a key holds, or None while its first request hasn't written it.
|
||||
|
||||
The client decodes responses; the redis stubs still admit bytes.
|
||||
"""
|
||||
if stored is None:
|
||||
return None
|
||||
text = stored if isinstance(stored, str) else stored.decode()
|
||||
try:
|
||||
value = json.loads(text)
|
||||
except ValueError:
|
||||
return None
|
||||
return value if isinstance(value, dict) else None
|
||||
@@ -0,0 +1,333 @@
|
||||
"""
|
||||
Tests for retry safety on the endpoints that spend money.
|
||||
|
||||
A run costs the caller credits, and a request that times out says nothing about
|
||||
whether it started. These pin the three outcomes a client needs to be able to
|
||||
tell apart: the run started, the run already started, and it is still starting.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Optional
|
||||
from unittest import mock
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.api.external.v2.idempotency import (
|
||||
IDEMPOTENCY_HEADER,
|
||||
idempotency_key,
|
||||
idempotent_run,
|
||||
replayed_run,
|
||||
request_fingerprint,
|
||||
)
|
||||
from backend.api.external.v2.tenancy import TenantContext
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG_ID = "org-1"
|
||||
KEY = "client-chosen-key"
|
||||
REQUEST = request_fingerprint(agent_id="agent-1", request={"inputs": {"q": "a"}})
|
||||
|
||||
|
||||
def _held(run_id: Optional[str] = None, request: str = REQUEST) -> str:
|
||||
"""What the store holds for a key claimed by `request`."""
|
||||
return json.dumps({"request": request, "run": run_id})
|
||||
|
||||
|
||||
async def test_a_first_request_claims_the_key_and_records_its_run(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
assert claim.existing_run_id is None
|
||||
await claim.record("run-1")
|
||||
|
||||
redis.set.assert_any_await(
|
||||
f"v2:idem:{USER_ID}:{ORG_ID}:-:{KEY}", _held(), nx=True, ex=86400
|
||||
)
|
||||
redis.set.assert_awaited_with(
|
||||
f"v2:idem:{USER_ID}:{ORG_ID}:-:{KEY}", _held("run-1"), ex=86400
|
||||
)
|
||||
|
||||
|
||||
async def test_a_retry_with_the_same_key_reports_the_first_run(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""Without this the retry is a second execution and a second charge."""
|
||||
redis.set.return_value = None # the key is already claimed
|
||||
redis.get.return_value = _held("run-1")
|
||||
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
assert claim.existing_run_id == "run-1"
|
||||
|
||||
|
||||
async def test_a_retry_while_the_first_is_still_running_is_a_conflict(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
redis.set.return_value = None
|
||||
redis.get.return_value = _held()
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
pass
|
||||
|
||||
assert raised.value.status_code == 409
|
||||
|
||||
|
||||
async def test_a_failed_run_releases_its_key_for_a_retry(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""A key held by a request that never produced a run would lock the caller
|
||||
out of retrying with the value they chose, for 24 hours."""
|
||||
with pytest.raises(RuntimeError):
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
raise RuntimeError("enqueue failed")
|
||||
|
||||
redis.delete.assert_awaited_once_with(f"v2:idem:{USER_ID}:{ORG_ID}:-:{KEY}")
|
||||
|
||||
|
||||
async def test_a_run_that_started_keeps_its_key_even_if_the_response_fails(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""Releasing here would let the retry start a second run and charge again —
|
||||
the first one is already enqueued."""
|
||||
with pytest.raises(RuntimeError):
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
await claim.record("run-1")
|
||||
raise RuntimeError("serialising the response failed")
|
||||
|
||||
redis.delete.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_run_whose_key_could_not_be_recorded_is_never_run_twice(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Recording can fail; if the response then fails too, releasing the key
|
||||
would let the retry start, and charge for, a second run."""
|
||||
store: dict[str, str] = {}
|
||||
|
||||
async def set_(name: str, value: str, nx: bool = False, ex: int = 0) -> bool:
|
||||
if not nx:
|
||||
raise ConnectionError("redis dropped the write")
|
||||
if name in store:
|
||||
return False
|
||||
store[name] = value
|
||||
return True
|
||||
|
||||
async def delete(name: str) -> None:
|
||||
store.pop(name, None)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.idempotency.get_redis_async",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(
|
||||
set=mock.AsyncMock(side_effect=set_),
|
||||
get=mock.AsyncMock(side_effect=store.get),
|
||||
delete=mock.AsyncMock(side_effect=delete),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
await claim.record("run-1")
|
||||
raise RuntimeError("serialising the response failed")
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as retry:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
pass
|
||||
assert retry.value.status_code == 409
|
||||
|
||||
|
||||
async def test_no_key_means_no_claim(redis: mock.AsyncMock) -> None:
|
||||
async with idempotent_run(None, _tenant(), REQUEST) as claim:
|
||||
assert claim.existing_run_id is None
|
||||
await claim.record("run-1")
|
||||
|
||||
redis.set.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_key_that_cannot_be_claimed_is_a_retryable_refusal(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The caller sent the key so a retry can't run twice; running without the
|
||||
claim would make the retry exactly that."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.idempotency.get_redis_async",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=RuntimeError("redis is down"),
|
||||
)
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
pass
|
||||
|
||||
assert raised.value.status_code == 503
|
||||
assert (raised.value.headers or {}).get("Retry-After")
|
||||
|
||||
|
||||
async def test_a_key_reused_for_a_different_request_is_refused(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""Replaying would hand back a run of another agent, or other inputs."""
|
||||
redis.set.return_value = None
|
||||
redis.get.return_value = _held("run-1", request="another-request")
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
pass
|
||||
|
||||
assert raised.value.status_code == 422
|
||||
|
||||
|
||||
def test_the_fingerprint_covers_the_agent_and_the_body() -> None:
|
||||
same = request_fingerprint(agent_id="agent-1", request={"inputs": {"q": "a"}})
|
||||
assert same == REQUEST
|
||||
assert REQUEST != request_fingerprint(
|
||||
agent_id="agent-2", request={"inputs": {"q": "a"}}
|
||||
)
|
||||
assert REQUEST != request_fingerprint(
|
||||
agent_id="agent-1", request={"inputs": {"q": "b"}}
|
||||
)
|
||||
|
||||
|
||||
async def test_a_replay_is_scoped_to_the_key_s_tenant(
|
||||
redis: mock.AsyncMock, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""The recorded run id is a bare string; it must still pass the tenant check."""
|
||||
redis.set.return_value = None
|
||||
redis.get.return_value = _held("run-in-another-org")
|
||||
mocker.patch(
|
||||
"backend.data.execution.get_graph_execution",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(organization_id="another-org"),
|
||||
)
|
||||
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
with pytest.raises(NotFoundError):
|
||||
await replayed_run(claim, _tenant())
|
||||
|
||||
|
||||
async def test_the_same_key_in_two_organizations_claims_twice(
|
||||
key_store: dict[str, str],
|
||||
) -> None:
|
||||
"""A run in one org is a 404 to the same user's other org, so a shared key
|
||||
must not replay it there."""
|
||||
other_org = _tenant().model_copy(update={"organization_id": "org-2"})
|
||||
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as claim:
|
||||
await claim.record("run-1")
|
||||
async with idempotent_run(KEY, other_org, REQUEST) as claim:
|
||||
assert claim.existing_run_id is None
|
||||
|
||||
|
||||
async def test_the_same_key_in_two_teams_claims_twice(
|
||||
key_store: dict[str, str],
|
||||
) -> None:
|
||||
"""A run lands in the caller's team, so a key reused under another
|
||||
`X-Team-Id` must start a run there rather than replay the first."""
|
||||
one_team = _tenant().model_copy(update={"team_id": "team-1"})
|
||||
other_team = _tenant().model_copy(update={"team_id": "team-2"})
|
||||
|
||||
async with idempotent_run(KEY, one_team, REQUEST) as claim:
|
||||
await claim.record("run-1")
|
||||
async with idempotent_run(KEY, other_team, REQUEST) as claim:
|
||||
assert claim.existing_run_id is None
|
||||
async with idempotent_run(KEY, one_team, REQUEST) as claim:
|
||||
assert claim.existing_run_id == "run-1"
|
||||
|
||||
|
||||
async def test_a_cancelled_request_releases_its_key(
|
||||
key_store: dict[str, str],
|
||||
) -> None:
|
||||
"""A client disconnect cancels the request before any run exists; a key left
|
||||
in flight would answer every retry with a 409 for 24 hours."""
|
||||
claimed = asyncio.Event()
|
||||
|
||||
async def request() -> None:
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST):
|
||||
claimed.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
task = asyncio.create_task(request())
|
||||
await claimed.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
async with idempotent_run(KEY, _tenant(), REQUEST) as retry:
|
||||
assert retry.holds_key
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sent,expected", [(None, None), ("", None), (" ", None), (" k ", "k")]
|
||||
)
|
||||
def test_a_blank_header_is_no_key(sent: Optional[str], expected: Optional[str]) -> None:
|
||||
assert idempotency_key(sent) == expected
|
||||
|
||||
|
||||
def test_the_run_endpoint_accepts_the_header() -> None:
|
||||
from .app import v2_app
|
||||
|
||||
post = v2_app.openapi()["paths"]["/library/agents/{agent_id}/runs"]["post"]
|
||||
assert IDEMPOTENCY_HEADER in {p["name"] for p in post.get("parameters", [])}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def key_store(mocker: pytest_mock.MockFixture) -> dict[str, str]:
|
||||
"""A key store that keeps what it is given."""
|
||||
store: dict[str, str] = {}
|
||||
|
||||
async def set_(name: str, value: str, nx: bool = False, ex: int = 0) -> bool:
|
||||
if nx and name in store:
|
||||
return False
|
||||
store[name] = value
|
||||
return True
|
||||
|
||||
async def delete(name: str) -> None:
|
||||
store.pop(name, None)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.idempotency.get_redis_async",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(
|
||||
set=mock.AsyncMock(side_effect=set_),
|
||||
get=mock.AsyncMock(side_effect=store.get),
|
||||
delete=mock.AsyncMock(side_effect=delete),
|
||||
),
|
||||
)
|
||||
return store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis(mocker: pytest_mock.MockFixture) -> mock.AsyncMock:
|
||||
client = mock.AsyncMock()
|
||||
client.set.return_value = True
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.idempotency.get_redis_async",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=client,
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
def _tenant() -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
|
||||
def test_the_header_parameter_defaults_to_no_key_when_called_directly() -> None:
|
||||
"""`= Depends(...)` leaves that sentinel as the Python default, so a direct
|
||||
call — which is how this suite exercises handlers — arrived at the key store
|
||||
with a `Depends` object for a key and opened a Redis client for it."""
|
||||
import inspect
|
||||
|
||||
from .library.agents import execute_agent
|
||||
|
||||
default = inspect.signature(execute_agent).parameters["idempotency"].default
|
||||
assert default is None, f"execute_agent defaults to {default!r}"
|
||||
@@ -0,0 +1,53 @@
|
||||
"""V2 External API - Identity Endpoint"""
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.api.features.orgs.db import get_org
|
||||
from backend.data.db import prisma
|
||||
from backend.data.user import get_user_by_id
|
||||
|
||||
from .models import Identity, TenantOrganization, TenantTeam
|
||||
from .tenancy import TenantContext, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
identity_router = APIRouter(tags=["identity"])
|
||||
|
||||
|
||||
@identity_router.get(
|
||||
path="/me",
|
||||
summary="Get the identity and tenancy of the current credentials",
|
||||
operation_id="getIdentity",
|
||||
)
|
||||
async def get_identity(
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.IDENTITY)),
|
||||
) -> Identity:
|
||||
"""Get the user, organization and team these credentials act as.
|
||||
|
||||
The organization is fixed when the credentials are created; `team` reflects
|
||||
the team they are pinned to, or the one selected with `X-Team-Id`.
|
||||
"""
|
||||
user = await get_user_by_id(auth.user_id)
|
||||
org = await get_org(auth.organization_id)
|
||||
|
||||
team = None
|
||||
if auth.team_id and (
|
||||
row := await prisma.team.find_unique(where={"id": auth.team_id})
|
||||
):
|
||||
team = TenantTeam(id=row.id, name=row.name)
|
||||
|
||||
return Identity(
|
||||
user_id=user.id,
|
||||
email=user.email,
|
||||
name=user.name,
|
||||
timezone=user.timezone,
|
||||
organization=TenantOrganization(
|
||||
id=org.id, name=org.name, is_personal=org.is_personal
|
||||
),
|
||||
team=team,
|
||||
scopes=[scope.value for scope in auth.scopes],
|
||||
credential_type=auth.type,
|
||||
)
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
V2 External API - Integrations Package
|
||||
|
||||
Aggregates all integration-related sub-routers.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from .credentials import credentials_router
|
||||
|
||||
integrations_router = APIRouter(tags=["integrations"])
|
||||
|
||||
integrations_router.include_router(credentials_router)
|
||||
+175
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
V2 External API - Credential CRUD Endpoints
|
||||
|
||||
Provides endpoints for managing integration credentials.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Annotated, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, HTTPException, Query, Request, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import SecretStr
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.integrations.router import delete_credentials_and_dependents
|
||||
from backend.data.model import (
|
||||
APIKeyCredentials,
|
||||
HostScopedCredentials,
|
||||
UserPasswordCredentials,
|
||||
is_sdk_default,
|
||||
)
|
||||
from backend.integrations.credentials_store import is_system_credential
|
||||
from backend.integrations.providers import ProviderName
|
||||
from backend.util.exceptions import NeedConfirmation
|
||||
|
||||
from ..models import CredentialCreateRequest, CredentialInfo
|
||||
from ..pagination import Page, PageRequest, page_request
|
||||
from ..tenancy import TenantContext, require_paid_permission, require_permission
|
||||
from .helpers import creds_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
credentials_router = APIRouter()
|
||||
|
||||
|
||||
@credentials_router.get(
|
||||
path="/credentials",
|
||||
summary="List integration credentials",
|
||||
operation_id="listIntegrationCredentials",
|
||||
)
|
||||
async def list_credentials(
|
||||
provider: Optional[str] = Query(
|
||||
default=None,
|
||||
description="Filter by provider name (e.g., 'github', 'google')",
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_INTEGRATIONS)
|
||||
),
|
||||
) -> Page[CredentialInfo]:
|
||||
"""List integration credentials for the authenticated user."""
|
||||
credentials = await creds_manager.store.get_all_creds(auth.user_id)
|
||||
|
||||
if provider:
|
||||
credentials = [c for c in credentials if c.provider.lower() == provider.lower()]
|
||||
|
||||
return page.slice([CredentialInfo.from_internal(c) for c in credentials])
|
||||
|
||||
|
||||
@credentials_router.post(
|
||||
path="/credentials",
|
||||
summary="Create integration credential",
|
||||
operation_id="createIntegrationCredential",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_credential(
|
||||
request: Annotated[CredentialCreateRequest, Body(discriminator="type")],
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.MANAGE_INTEGRATIONS)
|
||||
),
|
||||
) -> CredentialInfo:
|
||||
"""
|
||||
Create a new integration credential.
|
||||
|
||||
Supports `api_key`, `user_password`, and `host_scoped` credential types.
|
||||
OAuth credentials must be set up through the web UI.
|
||||
"""
|
||||
cred_id = str(uuid4())
|
||||
|
||||
if request.type == "api_key":
|
||||
credentials = APIKeyCredentials(
|
||||
id=cred_id,
|
||||
provider=request.provider,
|
||||
title=request.title,
|
||||
api_key=SecretStr(request.api_key),
|
||||
)
|
||||
elif request.type == "user_password":
|
||||
credentials = UserPasswordCredentials(
|
||||
id=cred_id,
|
||||
provider=request.provider,
|
||||
title=request.title,
|
||||
username=SecretStr(request.username),
|
||||
password=SecretStr(request.password),
|
||||
)
|
||||
else:
|
||||
credentials = HostScopedCredentials(
|
||||
id=cred_id,
|
||||
provider=request.provider,
|
||||
title=request.title,
|
||||
host=request.host,
|
||||
headers={k: SecretStr(v) for k, v in request.headers.items()},
|
||||
)
|
||||
|
||||
await creds_manager.create(auth.user_id, credentials)
|
||||
return CredentialInfo.from_internal(credentials)
|
||||
|
||||
|
||||
@credentials_router.delete(
|
||||
path="/credentials/{credential_id}",
|
||||
summary="Delete integration credential",
|
||||
operation_id="deleteIntegrationCredential",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_credential(
|
||||
request: Request,
|
||||
credential_id: str,
|
||||
force: Annotated[
|
||||
bool,
|
||||
Query(
|
||||
description=(
|
||||
"Delete even if a trigger registered with this credential is "
|
||||
"still used by an agent; the trigger stops firing"
|
||||
),
|
||||
),
|
||||
] = False,
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.DELETE_INTEGRATIONS)
|
||||
),
|
||||
) -> None:
|
||||
"""
|
||||
Delete an integration credential.
|
||||
|
||||
Webhooks registered with it are removed, and OAuth tokens are revoked at
|
||||
the provider. While an agent still uses one of those webhooks the request
|
||||
is refused with 409 unless `force=true`.
|
||||
|
||||
Platform-provided credentials cannot be deleted. Any agents using this
|
||||
credential will fail on their next run.
|
||||
"""
|
||||
# An SDK default is not the caller's credential at all, so it is not there
|
||||
# to find; a system or AutoGPT-managed one is, and refusing says so.
|
||||
if is_sdk_default(credential_id):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Credential #{credential_id} not found",
|
||||
)
|
||||
|
||||
existing = await creds_manager.store.get_creds_by_id(
|
||||
user_id=auth.user_id, credentials_id=credential_id
|
||||
)
|
||||
if not existing:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Credential #{credential_id} not found",
|
||||
)
|
||||
if is_system_credential(credential_id) or existing.is_managed:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Platform-managed credentials cannot be deleted",
|
||||
)
|
||||
|
||||
try:
|
||||
await delete_credentials_and_dependents(
|
||||
request,
|
||||
auth.user_id,
|
||||
existing,
|
||||
ProviderName(existing.provider),
|
||||
force,
|
||||
)
|
||||
except NeedConfirmation as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"{e}. Pass force=true to delete it anyway.",
|
||||
)
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
V2 External API - Integration Helpers
|
||||
|
||||
Shared logic for credential-related operations.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from backend.integrations.creds_manager import IntegrationCredentialsManager
|
||||
|
||||
from ..models import CredentialInfo, CredentialRequirement
|
||||
|
||||
creds_manager = IntegrationCredentialsManager()
|
||||
|
||||
|
||||
async def get_credential_requirements(
|
||||
creds_schema: dict[str, Any],
|
||||
user_id: str,
|
||||
) -> list[CredentialRequirement]:
|
||||
"""
|
||||
Extract credential requirements from a graph's credentials input schema
|
||||
and match them against the user's existing credentials.
|
||||
"""
|
||||
all_credentials = await creds_manager.store.get_all_creds(user_id)
|
||||
|
||||
return [
|
||||
CredentialRequirement(
|
||||
field_name=field_name,
|
||||
provider=provider,
|
||||
supported_types=list(field_schema.get("credentials_types") or []),
|
||||
required_scopes=list(field_schema.get("credentials_scopes") or []),
|
||||
matching_credentials=[
|
||||
CredentialInfo.from_internal(c)
|
||||
for c in all_credentials
|
||||
if c.provider.lower() == provider.lower()
|
||||
],
|
||||
)
|
||||
for field_name, field_schema in (creds_schema.get("properties") or {}).items()
|
||||
for provider in field_schema.get("credentials_provider") or []
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
V2 External API - Library Package
|
||||
|
||||
Aggregates all library-related sub-routers (agents, folders, presets).
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from .agents import agents_router
|
||||
from .folders import folders_router
|
||||
from .presets import presets_router
|
||||
|
||||
library_router = APIRouter()
|
||||
|
||||
library_router.include_router(agents_router)
|
||||
library_router.include_router(folders_router)
|
||||
library_router.include_router(presets_router)
|
||||
@@ -0,0 +1,309 @@
|
||||
"""V2 External API - Library Agent Endpoints"""
|
||||
|
||||
import logging
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.library.model import LibraryAgentSort
|
||||
from backend.data import graph as graph_db
|
||||
from backend.data.tenancy import writable_organization_ids
|
||||
from backend.executor import utils as execution_utils
|
||||
|
||||
from ..errors import ErrorResponse
|
||||
from ..idempotency import (
|
||||
idempotency_key,
|
||||
idempotent_run,
|
||||
replayed_run,
|
||||
request_fingerprint,
|
||||
)
|
||||
from ..integrations.helpers import get_credential_requirements
|
||||
from ..models import (
|
||||
AgentGraphRun,
|
||||
AgentRunRequest,
|
||||
CredentialRequirement,
|
||||
LibraryAgent,
|
||||
LibraryAgentUpdateRequest,
|
||||
)
|
||||
from ..pagination import Page, PageRequest, page_request
|
||||
from ..rate_limit import enforce, graph_exec_limiter
|
||||
from ..tenancy import (
|
||||
TenantContext,
|
||||
in_tenant,
|
||||
require_paid_permission,
|
||||
require_permission,
|
||||
)
|
||||
from .helpers import assert_can_pay, assert_inputs_match
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
agents_router = APIRouter(tags=["library"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@agents_router.get(
|
||||
path="/agents",
|
||||
summary="List library agents",
|
||||
operation_id="listLibraryAgents",
|
||||
)
|
||||
async def list_library_agents(
|
||||
published: Optional[bool] = Query(
|
||||
default=None,
|
||||
description="Filter by marketplace publish status",
|
||||
),
|
||||
favorite: Optional[bool] = Query(
|
||||
default=None,
|
||||
description="Filter by `isFavorite` attribute",
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> Page[LibraryAgent]:
|
||||
"""List agents in the user's library, oldest first."""
|
||||
# Creation order, not the default last-updated: every run bumps an agent's
|
||||
# updatedAt, which reshuffled the pages while a client walked them.
|
||||
result = await library_db.list_library_agents(
|
||||
user_id=auth.user_id,
|
||||
sort_by=LibraryAgentSort.CREATED_AT,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
published=published,
|
||||
favorite=favorite,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[LibraryAgent.from_internal(a) for a in result.agents],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@agents_router.get(
|
||||
path="/agents/{agent_id}",
|
||||
summary="Get library agent",
|
||||
operation_id="getLibraryAgent",
|
||||
)
|
||||
async def get_library_agent(
|
||||
agent_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> LibraryAgent:
|
||||
"""Get detailed information about a specific agent in the user's library."""
|
||||
agent = in_tenant(
|
||||
await library_db.get_library_agent(id=agent_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Agent #{agent_id}",
|
||||
)
|
||||
return LibraryAgent.from_internal(agent)
|
||||
|
||||
|
||||
@agents_router.patch(
|
||||
path="/agents/{agent_id}",
|
||||
summary="Update library agent",
|
||||
operation_id="updateLibraryAgent",
|
||||
)
|
||||
async def update_library_agent(
|
||||
request: LibraryAgentUpdateRequest,
|
||||
agent_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> LibraryAgent:
|
||||
"""Update properties of a library agent."""
|
||||
await _assert_agent_in_tenant(agent_id, auth)
|
||||
if request.folder_id:
|
||||
in_tenant(
|
||||
await library_db.get_folder(
|
||||
folder_id=request.folder_id, user_id=auth.user_id
|
||||
),
|
||||
auth,
|
||||
f"Folder #{request.folder_id}",
|
||||
)
|
||||
|
||||
# Only this organization's entries may give way to the chosen version.
|
||||
organization_ids = (
|
||||
await writable_organization_ids(auth.user_id, auth.organization_id)
|
||||
if request.graph_version is not None
|
||||
else None
|
||||
)
|
||||
updated = await library_db.update_library_agent(
|
||||
library_agent_id=agent_id,
|
||||
user_id=auth.user_id,
|
||||
auto_update_version=request.auto_update_version,
|
||||
graph_version=request.graph_version,
|
||||
is_favorite=request.is_favorite,
|
||||
is_archived=request.is_archived,
|
||||
folder_id=request.folder_id,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
return LibraryAgent.from_internal(updated)
|
||||
|
||||
|
||||
@agents_router.delete(
|
||||
path="/agents/{agent_id}",
|
||||
summary="Delete library agent",
|
||||
operation_id="deleteLibraryAgent",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_library_agent(
|
||||
agent_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> None:
|
||||
"""Remove an agent from the user's library."""
|
||||
await _assert_agent_in_tenant(agent_id, auth)
|
||||
|
||||
await library_db.delete_library_agent(
|
||||
library_agent_id=agent_id,
|
||||
user_id=auth.user_id,
|
||||
)
|
||||
|
||||
|
||||
@agents_router.post(
|
||||
path="/agents/{agent_id}/fork",
|
||||
summary="Fork library agent",
|
||||
operation_id="forkLibraryAgent",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def fork_library_agent(
|
||||
agent_id: str,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_LIBRARY)
|
||||
),
|
||||
) -> LibraryAgent:
|
||||
"""Fork (clone) a library agent.
|
||||
|
||||
Creates a deep copy of the agent's underlying graph and all its nodes,
|
||||
assigning new IDs. The cloned graph is added to the user's library as
|
||||
an independent agent that can be modified without affecting the original.
|
||||
"""
|
||||
await _assert_agent_in_tenant(agent_id, auth)
|
||||
|
||||
forked = await library_db.fork_library_agent(
|
||||
library_agent_id=agent_id,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
return LibraryAgent.from_internal(forked)
|
||||
|
||||
|
||||
@agents_router.post(
|
||||
path="/agents/{agent_id}/runs",
|
||||
summary="Execute library agent",
|
||||
operation_id="executeLibraryAgent",
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
responses={
|
||||
status.HTTP_402_PAYMENT_REQUIRED: {
|
||||
"model": ErrorResponse,
|
||||
"description": "No active plan, or a credit balance of zero",
|
||||
},
|
||||
status.HTTP_409_CONFLICT: {
|
||||
"model": ErrorResponse,
|
||||
"description": "The `Idempotency-Key`'s run has not been recorded yet",
|
||||
},
|
||||
status.HTTP_422_UNPROCESSABLE_CONTENT: {
|
||||
"model": ErrorResponse,
|
||||
"description": "Invalid inputs, or an `Idempotency-Key` already used "
|
||||
"for a different request",
|
||||
},
|
||||
status.HTTP_503_SERVICE_UNAVAILABLE: {
|
||||
"model": ErrorResponse,
|
||||
"description": "The `Idempotency-Key` cannot be checked right now",
|
||||
},
|
||||
},
|
||||
)
|
||||
async def execute_agent(
|
||||
response: Response,
|
||||
request: AgentRunRequest,
|
||||
agent_id: str,
|
||||
idempotency: Annotated[Optional[str], Depends(idempotency_key)] = None,
|
||||
auth: TenantContext = Security(require_paid_permission(APIKeyPermission.RUN_AGENT)),
|
||||
) -> AgentGraphRun:
|
||||
"""
|
||||
Execute an agent from the library.
|
||||
|
||||
Send an `Idempotency-Key` to make a retry safe: a second request carrying the
|
||||
same key returns the run the first one started rather than starting another.
|
||||
Reusing a key for a different agent or body is a 422.
|
||||
|
||||
**Rate limit:** 60 requests per minute per user.
|
||||
"""
|
||||
await enforce(graph_exec_limiter, auth.user_id, response)
|
||||
|
||||
fingerprint = request_fingerprint(
|
||||
agent_id=agent_id, request=request.model_dump(mode="json")
|
||||
)
|
||||
async with idempotent_run(idempotency, auth, fingerprint) as claim:
|
||||
if claim.existing_run_id:
|
||||
return await replayed_run(claim, auth)
|
||||
|
||||
await assert_can_pay(auth)
|
||||
|
||||
library_agent = in_tenant(
|
||||
await library_db.get_library_agent(id=agent_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Agent #{agent_id}",
|
||||
)
|
||||
assert_inputs_match(library_agent.input_schema, request.inputs)
|
||||
|
||||
result = await execution_utils.add_graph_execution(
|
||||
graph_id=library_agent.graph_id,
|
||||
user_id=auth.user_id,
|
||||
inputs=request.inputs,
|
||||
graph_version=library_agent.graph_version,
|
||||
graph_credentials_inputs=request.credentials_inputs,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
await claim.record(result.id)
|
||||
return AgentGraphRun.from_internal(result)
|
||||
|
||||
|
||||
@agents_router.get(
|
||||
path="/agents/{agent_id}/credentials",
|
||||
summary="Get library agent credential requirements",
|
||||
operation_id="getCredentialRequirementsForLibraryAgent",
|
||||
)
|
||||
async def list_agent_credential_requirements(
|
||||
agent_id: str,
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_INTEGRATIONS)
|
||||
),
|
||||
) -> Page[CredentialRequirement]:
|
||||
"""List credential requirements and matching user credentials for a library agent."""
|
||||
library_agent = in_tenant(
|
||||
await library_db.get_library_agent(agent_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Agent #{agent_id}",
|
||||
)
|
||||
|
||||
graph = await graph_db.get_graph(
|
||||
graph_id=library_agent.graph_id,
|
||||
version=library_agent.graph_version,
|
||||
user_id=auth.user_id,
|
||||
include_subgraphs=True,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
if not graph:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph for agent #{agent_id} not found",
|
||||
)
|
||||
|
||||
requirements = await get_credential_requirements(
|
||||
graph.credentials_input_schema, auth.user_id
|
||||
)
|
||||
return page.slice(requirements)
|
||||
|
||||
|
||||
async def _assert_agent_in_tenant(agent_id: str, auth: TenantContext) -> None:
|
||||
"""404 before mutating an agent the credentials cannot reach."""
|
||||
in_tenant(
|
||||
await library_db.get_library_agent(id=agent_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Agent #{agent_id}",
|
||||
)
|
||||
@@ -0,0 +1,186 @@
|
||||
"""V2 External API - Library Folder Endpoints"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Query, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
|
||||
from ..models import (
|
||||
LibraryFolder,
|
||||
LibraryFolderCreateRequest,
|
||||
LibraryFolderMoveRequest,
|
||||
LibraryFolderTree,
|
||||
LibraryFolderUpdateRequest,
|
||||
)
|
||||
from ..pagination import Page, PageRequest, page_request
|
||||
from ..tenancy import TenantContext, in_tenant, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
folders_router = APIRouter(tags=["library"])
|
||||
|
||||
|
||||
@folders_router.get(
|
||||
path="/folders",
|
||||
summary="List folders in library",
|
||||
operation_id="listLibraryFolders",
|
||||
)
|
||||
async def list_folders(
|
||||
parent_id: Optional[str] = Query(
|
||||
default=None, description="Filter by parent folder ID. Omit for root folders."
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> Page[LibraryFolder]:
|
||||
"""List folders in the user's library."""
|
||||
folders = await library_db.list_folders(
|
||||
user_id=auth.user_id,
|
||||
parent_id=parent_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.slice([LibraryFolder.from_internal(f) for f in folders])
|
||||
|
||||
|
||||
@folders_router.get(
|
||||
path="/folders/tree",
|
||||
summary="Get library folder tree",
|
||||
operation_id="getLibraryFolderTree",
|
||||
)
|
||||
async def get_folder_tree(
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> Page[LibraryFolderTree]:
|
||||
"""Get the full folder tree for the user's library.
|
||||
|
||||
Pagination applies to the root folders; each item carries its whole subtree.
|
||||
"""
|
||||
tree = await library_db.get_folder_tree(
|
||||
user_id=auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
|
||||
return page.slice([LibraryFolderTree.from_internal(f) for f in tree])
|
||||
|
||||
|
||||
@folders_router.get(
|
||||
path="/folders/{folder_id}",
|
||||
summary="Get folder in library",
|
||||
operation_id="getLibraryFolder",
|
||||
)
|
||||
async def get_folder(
|
||||
folder_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> LibraryFolder:
|
||||
"""Get details of a specific folder."""
|
||||
folder = in_tenant(
|
||||
await library_db.get_folder(folder_id=folder_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Folder #{folder_id}",
|
||||
)
|
||||
return LibraryFolder.from_internal(folder)
|
||||
|
||||
|
||||
@folders_router.post(
|
||||
path="/folders",
|
||||
summary="Create folder in library",
|
||||
operation_id="createLibraryFolder",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_folder(
|
||||
request: LibraryFolderCreateRequest,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> LibraryFolder:
|
||||
"""Create a new folder in the user's library."""
|
||||
if request.parent_id:
|
||||
await _assert_folder_in_tenant(request.parent_id, auth)
|
||||
|
||||
folder = await library_db.create_folder(
|
||||
user_id=auth.user_id,
|
||||
name=request.name,
|
||||
parent_id=request.parent_id,
|
||||
icon=request.icon,
|
||||
color=request.color,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
return LibraryFolder.from_internal(folder)
|
||||
|
||||
|
||||
@folders_router.patch(
|
||||
path="/folders/{folder_id}",
|
||||
summary="Update folder in library",
|
||||
operation_id="updateLibraryFolder",
|
||||
)
|
||||
async def update_folder(
|
||||
request: LibraryFolderUpdateRequest,
|
||||
folder_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> LibraryFolder:
|
||||
"""Update properties of a folder."""
|
||||
await _assert_folder_in_tenant(folder_id, auth)
|
||||
|
||||
folder = await library_db.update_folder(
|
||||
folder_id=folder_id,
|
||||
user_id=auth.user_id,
|
||||
name=request.name,
|
||||
icon=request.icon,
|
||||
color=request.color,
|
||||
)
|
||||
return LibraryFolder.from_internal(folder)
|
||||
|
||||
|
||||
@folders_router.post(
|
||||
path="/folders/{folder_id}/move",
|
||||
summary="Move folder in library",
|
||||
operation_id="moveLibraryFolder",
|
||||
)
|
||||
async def move_folder(
|
||||
request: LibraryFolderMoveRequest,
|
||||
folder_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> LibraryFolder:
|
||||
"""Move a folder to a new parent. Set target_parent_id to null to move to root."""
|
||||
await _assert_folder_in_tenant(folder_id, auth)
|
||||
if request.target_parent_id:
|
||||
await _assert_folder_in_tenant(request.target_parent_id, auth)
|
||||
|
||||
folder = await library_db.move_folder(
|
||||
folder_id=folder_id,
|
||||
user_id=auth.user_id,
|
||||
target_parent_id=request.target_parent_id,
|
||||
)
|
||||
return LibraryFolder.from_internal(folder)
|
||||
|
||||
|
||||
@folders_router.delete(
|
||||
path="/folders/{folder_id}",
|
||||
summary="Delete folder in library",
|
||||
operation_id="deleteLibraryFolder",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_folder(
|
||||
folder_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> None:
|
||||
"""
|
||||
Delete a folder and its subfolders. Agents in this folder will be moved to root.
|
||||
"""
|
||||
await _assert_folder_in_tenant(folder_id, auth)
|
||||
|
||||
await library_db.delete_folder(
|
||||
folder_id=folder_id,
|
||||
user_id=auth.user_id,
|
||||
)
|
||||
|
||||
|
||||
async def _assert_folder_in_tenant(folder_id: str, auth: TenantContext) -> None:
|
||||
"""404 before mutating a folder the credentials cannot reach."""
|
||||
in_tenant(
|
||||
await library_db.get_folder(folder_id=folder_id, user_id=auth.user_id),
|
||||
auth,
|
||||
f"Folder #{folder_id}",
|
||||
)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
V2 External API - Library route helpers
|
||||
|
||||
Shared logic for the endpoints that start runs.
|
||||
"""
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from fastapi import HTTPException
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from starlette import status
|
||||
|
||||
from backend.data.credit import get_credit_model
|
||||
from backend.data.graph import missing_inputs, unknown_inputs
|
||||
|
||||
from ..tenancy import TenantContext
|
||||
|
||||
|
||||
async def assert_can_pay(auth: TenantContext) -> None:
|
||||
"""Refuse a run on a zero balance, as the internal run route does.
|
||||
|
||||
Not the same guard as the paywall inside `add_graph_execution`, which asks
|
||||
whether the user has a subscription at all.
|
||||
"""
|
||||
credit_model = await get_credit_model(auth.user_id, auth.organization_id)
|
||||
if await credit_model.get_credits(auth.user_id) <= 0:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
||||
detail="Insufficient balance to execute the agent. "
|
||||
"Please top up your account.",
|
||||
)
|
||||
|
||||
|
||||
def assert_inputs_match(
|
||||
input_schema: dict[str, Any], inputs: Mapping[str, Any], *, partial: bool = False
|
||||
) -> None:
|
||||
"""422 on an input the graph does not have, or a required one omitted or null:
|
||||
the executor runs either and finishes with no outputs and no error.
|
||||
|
||||
`partial` allows omissions, for inputs a later request completes."""
|
||||
missing = [] if partial else missing_inputs(input_schema, inputs)
|
||||
errors = [
|
||||
{
|
||||
"type": "extra_forbidden",
|
||||
"loc": ["body", "inputs", name],
|
||||
"msg": "Unknown input",
|
||||
}
|
||||
for name in unknown_inputs(input_schema, inputs)
|
||||
] + [
|
||||
{"type": "missing", "loc": ["body", "inputs", name], "msg": "Input required"}
|
||||
for name in missing
|
||||
]
|
||||
if errors:
|
||||
raise RequestValidationError(errors)
|
||||
@@ -0,0 +1,398 @@
|
||||
"""
|
||||
V2 External API - Library Preset Endpoints
|
||||
|
||||
Provides endpoints for managing agent presets (saved run configurations) and
|
||||
the webhook triggers built on them.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.experts import experts_db
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.library import triggers
|
||||
from backend.api.features.library.model import (
|
||||
LibraryAgentPreset,
|
||||
LibraryAgentPresetCreatable,
|
||||
)
|
||||
from backend.blocks import get_blocks
|
||||
from backend.data import graph as graph_db
|
||||
from backend.executor import utils as execution_utils
|
||||
from backend.util.exceptions import ConflictError, ExpertRunPausedError, NotFoundError
|
||||
|
||||
from ..errors import ErrorResponse
|
||||
from ..idempotency import (
|
||||
idempotency_key,
|
||||
idempotent_run,
|
||||
replayed_run,
|
||||
request_fingerprint,
|
||||
)
|
||||
from ..models import (
|
||||
AgentGraphRun,
|
||||
AgentPreset,
|
||||
AgentPresetCreateRequest,
|
||||
AgentPresetRunRequest,
|
||||
AgentPresetUpdateRequest,
|
||||
AgentTriggerSetupRequest,
|
||||
)
|
||||
from ..pagination import Page, PageRequest, page_request
|
||||
from ..rate_limit import enforce, graph_exec_limiter
|
||||
from ..tenancy import (
|
||||
TenantContext,
|
||||
in_tenant,
|
||||
require_paid_permission,
|
||||
require_permission,
|
||||
)
|
||||
from .helpers import assert_can_pay, assert_inputs_match
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
presets_router = APIRouter(tags=["presets"])
|
||||
|
||||
|
||||
@presets_router.get(
|
||||
path="/presets",
|
||||
summary="List agent execution presets",
|
||||
operation_id="listAgentRunPresets",
|
||||
)
|
||||
async def list_presets(
|
||||
graph_id: Optional[str] = Query(default=None, description="Filter by graph ID"),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> Page[AgentPreset]:
|
||||
"""List presets in the user's library, optionally filtered by graph ID."""
|
||||
result = await library_db.list_presets(
|
||||
user_id=auth.user_id,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
graph_id=graph_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[AgentPreset.from_internal(p, secret_trigger_inputs()) for p in result.presets],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@presets_router.get(
|
||||
path="/presets/{preset_id}",
|
||||
summary="Get agent execution preset",
|
||||
operation_id="getAgentRunPreset",
|
||||
)
|
||||
async def get_preset(
|
||||
preset_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_LIBRARY)),
|
||||
) -> AgentPreset:
|
||||
"""Get details of a specific preset."""
|
||||
return AgentPreset.from_internal(
|
||||
await _preset_in_tenant(preset_id, auth), secret_trigger_inputs()
|
||||
)
|
||||
|
||||
|
||||
@presets_router.post(
|
||||
path="/presets",
|
||||
summary="Create agent execution preset",
|
||||
operation_id="createAgentRunPreset",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_preset(
|
||||
request: AgentPresetCreateRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_LIBRARY)
|
||||
),
|
||||
) -> AgentPreset:
|
||||
"""Create a new preset with saved inputs and credentials for an agent.
|
||||
|
||||
The inputs may leave out required ones, which a run then supplies. An
|
||||
active preset of an agent with a webhook trigger also needs `RUN_AGENT`.
|
||||
"""
|
||||
graph = await _graph_in_tenant(request.graph_id, request.graph_version, auth)
|
||||
if (
|
||||
request.is_active
|
||||
and graph.webhook_input_node
|
||||
and APIKeyPermission.RUN_AGENT not in auth.scopes
|
||||
):
|
||||
raise _run_permission_missing()
|
||||
assert_inputs_match(graph.input_schema, request.inputs, partial=True)
|
||||
|
||||
preset = await library_db.create_preset(
|
||||
user_id=auth.user_id,
|
||||
preset=LibraryAgentPresetCreatable(
|
||||
graph_id=request.graph_id,
|
||||
graph_version=request.graph_version,
|
||||
name=request.name,
|
||||
description=request.description,
|
||||
inputs=request.inputs,
|
||||
credentials=request.credentials_inputs,
|
||||
is_active=request.is_active,
|
||||
),
|
||||
# As on the internal routes: a graph one hired expert installed runs as
|
||||
# that expert's, under its budget and pause.
|
||||
expert_id=await experts_db.resolve_expert_for_graph(
|
||||
auth.user_id, request.graph_id
|
||||
),
|
||||
)
|
||||
return AgentPreset.from_internal(preset)
|
||||
|
||||
|
||||
@presets_router.post(
|
||||
path="/presets/setup-trigger",
|
||||
summary="Setup triggered preset",
|
||||
operation_id="setupAgentRunTrigger",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def setup_trigger(
|
||||
request: AgentTriggerSetupRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(
|
||||
APIKeyPermission.WRITE_LIBRARY, APIKeyPermission.RUN_AGENT
|
||||
)
|
||||
),
|
||||
) -> AgentPreset:
|
||||
"""
|
||||
Create a preset with a webhook trigger for automatic execution.
|
||||
|
||||
The agent's `trigger_setup_info` describes the required trigger configuration
|
||||
schema and credentials. Use it to populate `trigger_config` and
|
||||
`credentials_inputs`. For a trigger you set up by hand in the other service,
|
||||
the response's `webhook_url` is where that service must send its events.
|
||||
This response is the only one that returns the trigger's secret settings.
|
||||
"""
|
||||
await _graph_in_tenant(request.graph_id, request.graph_version, auth)
|
||||
|
||||
preset = await triggers.setup_triggered_preset(
|
||||
user_id=auth.user_id,
|
||||
graph_id=request.graph_id,
|
||||
graph_version=request.graph_version,
|
||||
name=request.name,
|
||||
description=request.description,
|
||||
trigger_config=request.trigger_config,
|
||||
agent_credentials=request.credentials_inputs,
|
||||
expert_id=await experts_db.resolve_expert_for_graph(
|
||||
auth.user_id, request.graph_id
|
||||
),
|
||||
)
|
||||
return AgentPreset.from_internal(preset)
|
||||
|
||||
|
||||
@presets_router.patch(
|
||||
path="/presets/{preset_id}",
|
||||
operation_id="updateAgentRunPreset",
|
||||
summary="Update agent execution preset",
|
||||
)
|
||||
async def update_preset(
|
||||
request: AgentPresetUpdateRequest,
|
||||
preset_id: str,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_LIBRARY)
|
||||
),
|
||||
) -> AgentPreset:
|
||||
"""Update properties of a preset. Only provided fields will be updated.
|
||||
|
||||
New `inputs` or `credentials_inputs` on a webhook trigger re-register its
|
||||
webhook; whichever the request leaves out keeps its saved value, and new
|
||||
inputs keep the secret settings in its `hidden_inputs` unless they set them.
|
||||
Setting `is_active` to true, or reconfiguring a trigger that stays active,
|
||||
also needs `RUN_AGENT`: an active trigger starts runs.
|
||||
"""
|
||||
current = await _preset_in_tenant(preset_id, auth)
|
||||
if APIKeyPermission.RUN_AGENT not in auth.scopes and await _leaves_a_live_trigger(
|
||||
request, current, auth.user_id
|
||||
):
|
||||
raise _run_permission_missing()
|
||||
|
||||
inputs = request.inputs
|
||||
credentials = request.credentials_inputs
|
||||
if inputs is not None and current.webhook_id:
|
||||
inputs = {**_kept_secrets(current.inputs), **inputs}
|
||||
if inputs is not None and credentials is None:
|
||||
# Without credentials the trigger's webhook would keep the old config.
|
||||
credentials = current.credentials
|
||||
elif credentials is not None and inputs is None:
|
||||
# Saved whole, so a trigger's secret settings survive the rewrite.
|
||||
inputs = current.inputs
|
||||
|
||||
preset = await triggers.update_triggered_preset(
|
||||
user_id=auth.user_id,
|
||||
preset_id=preset_id,
|
||||
inputs=inputs,
|
||||
credentials=credentials,
|
||||
name=request.name,
|
||||
description=request.description,
|
||||
is_active=request.is_active,
|
||||
)
|
||||
return AgentPreset.from_internal(preset, secret_trigger_inputs())
|
||||
|
||||
|
||||
@presets_router.delete(
|
||||
path="/presets/{preset_id}",
|
||||
summary="Delete agent execution preset",
|
||||
operation_id="deleteAgentRunPreset",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_preset(
|
||||
preset_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_LIBRARY)),
|
||||
) -> None:
|
||||
"""Delete a preset. A webhook trigger's webhook is deregistered with it."""
|
||||
await _preset_in_tenant(preset_id, auth)
|
||||
|
||||
await triggers.delete_preset_with_webhook_cleanup(
|
||||
user_id=auth.user_id, preset_id=preset_id
|
||||
)
|
||||
|
||||
|
||||
@presets_router.post(
|
||||
path="/presets/{preset_id}/execute",
|
||||
summary="Execute agent preset",
|
||||
operation_id="executeAgentRunPreset",
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
responses={
|
||||
status.HTTP_402_PAYMENT_REQUIRED: {
|
||||
"model": ErrorResponse,
|
||||
"description": "No active plan, or a credit balance of zero",
|
||||
},
|
||||
status.HTTP_409_CONFLICT: {
|
||||
"model": ErrorResponse,
|
||||
"description": "The `Idempotency-Key`'s run has not been recorded yet, "
|
||||
"or the expert this preset runs for is paused",
|
||||
},
|
||||
status.HTTP_422_UNPROCESSABLE_CONTENT: {
|
||||
"model": ErrorResponse,
|
||||
"description": "Invalid inputs, or an `Idempotency-Key` already used "
|
||||
"for a different request",
|
||||
},
|
||||
status.HTTP_503_SERVICE_UNAVAILABLE: {
|
||||
"model": ErrorResponse,
|
||||
"description": "The `Idempotency-Key` cannot be checked right now",
|
||||
},
|
||||
},
|
||||
)
|
||||
async def execute_preset(
|
||||
response: Response,
|
||||
preset_id: str,
|
||||
request: AgentPresetRunRequest = AgentPresetRunRequest(),
|
||||
idempotency: Annotated[Optional[str], Depends(idempotency_key)] = None,
|
||||
auth: TenantContext = Security(require_paid_permission(APIKeyPermission.RUN_AGENT)),
|
||||
) -> AgentGraphRun:
|
||||
"""
|
||||
Execute a preset, optionally overriding saved inputs and credentials.
|
||||
|
||||
Send an `Idempotency-Key` to make a retry safe, as with
|
||||
`POST /library/agents/{agent_id}/runs`.
|
||||
|
||||
**Rate limit:** 60 requests per minute per user, shared with that endpoint.
|
||||
"""
|
||||
await enforce(graph_exec_limiter, auth.user_id, response)
|
||||
|
||||
fingerprint = request_fingerprint(
|
||||
preset_id=preset_id, request=request.model_dump(mode="json")
|
||||
)
|
||||
async with idempotent_run(idempotency, auth, fingerprint) as claim:
|
||||
if claim.existing_run_id:
|
||||
return await replayed_run(claim, auth)
|
||||
|
||||
await assert_can_pay(auth)
|
||||
|
||||
preset = await _preset_in_tenant(preset_id, auth)
|
||||
merged_inputs = {**preset.inputs, **request.inputs}
|
||||
merged_credentials = {**preset.credentials, **request.credentials_inputs}
|
||||
if not preset.webhook_id:
|
||||
# A trigger's inputs are its block's config, not the graph's inputs.
|
||||
graph = await _graph_in_tenant(preset.graph_id, preset.graph_version, auth)
|
||||
assert_inputs_match(graph.input_schema, merged_inputs)
|
||||
|
||||
# The run belongs where the preset does, as on the internal API.
|
||||
organization_id, team_id = (
|
||||
(preset.organization_id, preset.team_id)
|
||||
if preset.organization_id
|
||||
else (auth.organization_id, auth.team_id)
|
||||
)
|
||||
try:
|
||||
result = await execution_utils.add_graph_execution(
|
||||
graph_id=preset.graph_id,
|
||||
user_id=auth.user_id,
|
||||
inputs=merged_inputs,
|
||||
graph_version=preset.graph_version,
|
||||
graph_credentials_inputs=merged_credentials,
|
||||
preset_id=preset_id,
|
||||
expert_id=preset.expert_id,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
except ExpertRunPausedError as e:
|
||||
# A ValueError, which the global handlers would answer with a 400.
|
||||
raise ConflictError(str(e)) from e
|
||||
await claim.record(result.id)
|
||||
return AgentGraphRun.from_internal(result)
|
||||
|
||||
|
||||
async def _leaves_a_live_trigger(
|
||||
request: AgentPresetUpdateRequest, current: LibraryAgentPreset, user_id: str
|
||||
) -> bool:
|
||||
"""Whether the update turns the preset on, or re-registers the webhook of a
|
||||
trigger that stays on."""
|
||||
if request.is_active:
|
||||
return True
|
||||
if request.is_active is False or not current.is_active:
|
||||
return False
|
||||
if request.inputs is None and request.credentials_inputs is None:
|
||||
return False
|
||||
graph = await graph_db.get_graph(
|
||||
current.graph_id, current.graph_version, user_id=user_id
|
||||
)
|
||||
return bool(graph and graph.webhook_input_node)
|
||||
|
||||
|
||||
def _run_permission_missing() -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Missing required permission(s): RUN_AGENT",
|
||||
)
|
||||
|
||||
|
||||
async def _preset_in_tenant(preset_id: str, auth: TenantContext) -> LibraryAgentPreset:
|
||||
"""The caller's preset, or 404 if it is missing or in another organization."""
|
||||
return in_tenant(
|
||||
await library_db.get_preset(user_id=auth.user_id, preset_id=preset_id),
|
||||
auth,
|
||||
f"Preset #{preset_id}",
|
||||
)
|
||||
|
||||
|
||||
async def _graph_in_tenant(
|
||||
graph_id: str, graph_version: int, auth: TenantContext
|
||||
) -> graph_db.GraphModel:
|
||||
"""404 for a graph the credentials' organization can't see: a preset takes
|
||||
its graph's organization, so this keeps it from landing in another one."""
|
||||
graph = await graph_db.get_graph(
|
||||
graph_id=graph_id,
|
||||
version=graph_version,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
if not graph:
|
||||
raise NotFoundError(f"Graph #{graph_id} v{graph_version} not found")
|
||||
return graph
|
||||
|
||||
|
||||
@functools.cache
|
||||
def secret_trigger_inputs() -> frozenset[str]:
|
||||
"""The inputs any webhook trigger block marks secret, which reads leave out."""
|
||||
return frozenset(
|
||||
name
|
||||
for block_cls in get_blocks().values()
|
||||
if (block := block_cls()).webhook_config
|
||||
for name, spec in block.input_schema.jsonschema().get("properties", {}).items()
|
||||
if isinstance(spec, dict) and spec.get("secret") is True
|
||||
)
|
||||
|
||||
|
||||
def _kept_secrets(saved: dict) -> dict:
|
||||
return {name: saved[name] for name in secret_trigger_inputs() if name in saved}
|
||||
@@ -0,0 +1,729 @@
|
||||
"""What the v2 preset routes do beyond the tenant check: webhook upkeep, the
|
||||
expert a preset runs for, and what a response leaves out."""
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import prisma.models
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException, Response
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.api.features.library.model import LibraryAgentPreset
|
||||
from backend.copilot.tools.manage_presets import UpdatePresetTool
|
||||
from backend.data.execution import ExecutionStatus, GraphExecutionMeta
|
||||
from backend.data.model import CredentialsMetaInput
|
||||
from backend.util.exceptions import ConflictError, ExpertRunPausedError
|
||||
|
||||
from ..models import (
|
||||
AgentPreset,
|
||||
AgentPresetCreateRequest,
|
||||
AgentPresetRunRequest,
|
||||
AgentPresetUpdateRequest,
|
||||
AgentTriggerSetupRequest,
|
||||
)
|
||||
from ..pagination import PageRequest
|
||||
from ..tenancy import TenantContext
|
||||
from . import presets
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG = "org-1"
|
||||
SECRET = "whsec-do-not-leak"
|
||||
PROVIDER_HOOK_ID = "provider-hook-do-not-leak"
|
||||
TRIGGER_SECRET = "delivery-secret-do-not-leak"
|
||||
RUNNER = [APIKeyPermission.WRITE_LIBRARY, APIKeyPermission.RUN_AGENT]
|
||||
NEW_CREDENTIALS = {
|
||||
"github_credentials": CredentialsMetaInput(
|
||||
id="cred-2", provider="github", type="api_key", title="GitHub"
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_rate_limit(mocker: pytest_mock.MockFixture) -> None:
|
||||
mocker.patch.object(
|
||||
presets.graph_exec_limiter, "check", AsyncMock(return_value=None)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def can_pay(mocker: pytest_mock.MockFixture) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=100)),
|
||||
)
|
||||
|
||||
|
||||
def test_a_preset_response_carries_no_webhook_secret():
|
||||
body = AgentPreset.from_internal(_preset(webhook=True)).model_dump_json()
|
||||
|
||||
assert SECRET not in body
|
||||
assert PROVIDER_HOOK_ID not in body
|
||||
assert "/ingress" in body
|
||||
|
||||
|
||||
async def test_new_trigger_inputs_re_register_the_webhook_with_saved_credentials(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
current = _preset(webhook=True)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=current,
|
||||
)
|
||||
update = mocker.patch(
|
||||
"backend.api.features.library.triggers.update_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=current,
|
||||
)
|
||||
|
||||
await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(inputs={"repo": "a/b"}),
|
||||
preset_id=current.id,
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
assert update.await_args.kwargs["inputs"] == {"repo": "a/b"}
|
||||
assert update.await_args.kwargs["credentials"] == current.credentials
|
||||
|
||||
|
||||
async def test_deleting_a_trigger_cleans_up_its_webhook(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(webhook=True),
|
||||
)
|
||||
cleanup = mocker.patch(
|
||||
"backend.api.features.library.triggers.delete_preset_with_webhook_cleanup",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
await presets.delete_preset(preset_id="p-1", auth=_key())
|
||||
|
||||
cleanup.assert_awaited_once_with(user_id=USER_ID, preset_id="p-1")
|
||||
|
||||
|
||||
async def test_a_run_from_an_expert_preset_counts_against_that_expert(
|
||||
mocker: pytest_mock.MockFixture, can_pay: None
|
||||
):
|
||||
_stub_preset_and_graph(mocker, _preset(expert_id="expert-1"))
|
||||
add = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
|
||||
await _execute()
|
||||
|
||||
assert add.await_args.kwargs["expert_id"] == "expert-1"
|
||||
assert add.await_args.kwargs["preset_id"] == "p-1"
|
||||
|
||||
|
||||
async def test_a_paused_expert_s_preset_is_a_conflict(
|
||||
mocker: pytest_mock.MockFixture, can_pay: None
|
||||
):
|
||||
_stub_preset_and_graph(mocker, _preset(expert_id="expert-1"))
|
||||
mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=ExpertRunPausedError("paused", "expert-1"),
|
||||
)
|
||||
|
||||
with pytest.raises(ConflictError):
|
||||
await _execute()
|
||||
|
||||
|
||||
async def test_a_run_missing_a_required_input_is_refused(
|
||||
mocker: pytest_mock.MockFixture, can_pay: None
|
||||
):
|
||||
_stub_preset_and_graph(
|
||||
mocker,
|
||||
_preset(),
|
||||
input_schema={"properties": {"topic": {}}, "required": ["topic"]},
|
||||
)
|
||||
add = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(RequestValidationError):
|
||||
await _execute()
|
||||
add.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_trigger_s_settings_are_not_checked_as_graph_inputs(
|
||||
mocker: pytest_mock.MockFixture, can_pay: None
|
||||
):
|
||||
"""A trigger's inputs configure its block; the graph's inputs come from
|
||||
each event."""
|
||||
_stub_preset_and_graph(
|
||||
mocker,
|
||||
_trigger(),
|
||||
input_schema={"properties": {"topic": {}}, "required": ["topic"]},
|
||||
)
|
||||
add = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
|
||||
await _execute()
|
||||
|
||||
assert add.await_args.kwargs["inputs"]["secret_token"] == TRIGGER_SECRET
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["create", "setup-trigger"])
|
||||
async def test_a_preset_for_a_graph_a_hired_expert_installed_is_that_expert_s(
|
||||
mocker: pytest_mock.MockFixture, route: str
|
||||
):
|
||||
"""As on the internal routes, so its runs meet the expert's budget and pause."""
|
||||
mocker.patch(
|
||||
"backend.api.features.experts.experts_db.resolve_expert_for_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value="expert-1",
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(input_schema={"properties": {}}, webhook_input_node=None),
|
||||
)
|
||||
save = mocker.patch(
|
||||
(
|
||||
"backend.api.features.library.db.create_preset"
|
||||
if route == "create"
|
||||
else "backend.api.features.library.triggers.setup_triggered_preset"
|
||||
),
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(),
|
||||
)
|
||||
|
||||
if route == "create":
|
||||
await presets.create_preset(
|
||||
request=AgentPresetCreateRequest(
|
||||
name="Preset", graph_id="graph-1", graph_version=1
|
||||
),
|
||||
auth=_key(),
|
||||
)
|
||||
else:
|
||||
await presets.setup_trigger(
|
||||
request=AgentTriggerSetupRequest(
|
||||
name="On delivery",
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
trigger_config={},
|
||||
),
|
||||
auth=_key(),
|
||||
)
|
||||
|
||||
assert save.await_args.kwargs["expert_id"] == "expert-1"
|
||||
|
||||
|
||||
def test_the_hidden_settings_are_every_trigger_block_s_secret_inputs():
|
||||
"""Read from the blocks, so a new trigger's secret is hidden without a list."""
|
||||
assert {"secret_token", "signing_secret"} <= presets.secret_trigger_inputs()
|
||||
|
||||
|
||||
async def test_a_read_withholds_a_trigger_s_secret_settings(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
|
||||
read = await presets.get_preset(preset_id="p-1", auth=_key())
|
||||
|
||||
assert TRIGGER_SECRET not in read.model_dump_json()
|
||||
assert read.inputs == {"constants": {"a": 1}}
|
||||
assert read.hidden_inputs == ["secret_token"]
|
||||
|
||||
|
||||
async def test_a_listing_withholds_a_trigger_s_secret_settings(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.list_presets",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(presets=[_trigger()], pagination=Mock(total_items=1)),
|
||||
)
|
||||
|
||||
listed = await presets.list_presets(
|
||||
graph_id=None, page=PageRequest(limit=20), auth=_key()
|
||||
)
|
||||
|
||||
assert TRIGGER_SECRET not in listed.model_dump_json()
|
||||
assert listed.items[0].hidden_inputs == ["secret_token"]
|
||||
|
||||
|
||||
async def test_setting_up_a_trigger_returns_its_secret_settings_once(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph", new_callable=AsyncMock, return_value=Mock()
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.experts.experts_db.resolve_expert_for_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.setup_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
|
||||
created = await presets.setup_trigger(
|
||||
request=AgentTriggerSetupRequest(
|
||||
name="On delivery", graph_id="graph-1", graph_version=1, trigger_config={}
|
||||
),
|
||||
auth=_key(),
|
||||
)
|
||||
|
||||
assert created.inputs["secret_token"] == TRIGGER_SECRET
|
||||
assert created.hidden_inputs == []
|
||||
|
||||
|
||||
async def test_an_update_that_leaves_out_a_hidden_setting_keeps_it(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
):
|
||||
"""A client sends back what it read; that must not switch off the check."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
update = mocker.patch(
|
||||
"backend.api.features.library.triggers.update_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
|
||||
updated = await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(inputs={"constants": {"a": 2}}),
|
||||
preset_id="p-1",
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
assert update.await_args.kwargs["inputs"] == {
|
||||
"constants": {"a": 2},
|
||||
"secret_token": TRIGGER_SECRET,
|
||||
}
|
||||
assert TRIGGER_SECRET not in updated.model_dump_json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scopes, allowed", [([], False), (["RUN_AGENT"], True)])
|
||||
async def test_turning_a_preset_on_needs_the_run_permission(
|
||||
mocker: pytest_mock.MockFixture, scopes: list[str], allowed: bool
|
||||
):
|
||||
"""An active trigger starts runs, which spend credits."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
update = mocker.patch(
|
||||
"backend.api.features.library.triggers.update_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
request = AgentPresetUpdateRequest(is_active=True)
|
||||
auth = _key([APIKeyPermission.WRITE_LIBRARY, *map(APIKeyPermission, scopes)])
|
||||
|
||||
if allowed:
|
||||
await presets.update_preset(request=request, preset_id="p-1", auth=auth)
|
||||
update.assert_awaited_once()
|
||||
else:
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await presets.update_preset(request=request, preset_id="p-1", auth=auth)
|
||||
assert refused.value.status_code == 403
|
||||
update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"is_active, trigger_graph, request_active, refused",
|
||||
[
|
||||
(True, True, None, True),
|
||||
(True, True, False, False),
|
||||
(True, False, None, False),
|
||||
(False, True, None, False),
|
||||
],
|
||||
ids=["arms-a-trigger", "turns-it-off", "not-a-trigger", "stays-off"],
|
||||
)
|
||||
async def test_new_inputs_on_a_trigger_that_stays_on_need_the_run_permission(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
is_active: bool,
|
||||
trigger_graph: bool,
|
||||
request_active: bool | None,
|
||||
refused: bool,
|
||||
):
|
||||
"""New inputs register its webhook, and an active trigger starts runs."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset().model_copy(update={"is_active": is_active}),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(webhook_input_node=Mock() if trigger_graph else None),
|
||||
)
|
||||
update = mocker.patch(
|
||||
"backend.api.features.library.triggers.update_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(),
|
||||
)
|
||||
request = AgentPresetUpdateRequest(inputs={"repo": "a/b"}, is_active=request_active)
|
||||
writer = _key([APIKeyPermission.WRITE_LIBRARY])
|
||||
|
||||
if not refused:
|
||||
await presets.update_preset(request=request, preset_id="p-1", auth=writer)
|
||||
update.assert_awaited_once()
|
||||
return
|
||||
with pytest.raises(HTTPException) as refusal:
|
||||
await presets.update_preset(request=request, preset_id="p-1", auth=writer)
|
||||
assert refusal.value.status_code == 403
|
||||
update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"trigger_graph, is_active, refused",
|
||||
[(True, True, True), (True, False, False), (False, True, False)],
|
||||
ids=["active-trigger", "inactive-trigger", "not-a-trigger"],
|
||||
)
|
||||
async def test_creating_an_active_trigger_preset_needs_the_run_permission(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
trigger_graph: bool,
|
||||
is_active: bool,
|
||||
refused: bool,
|
||||
):
|
||||
"""New inputs would then arm it, and turning it on needs the permission."""
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(
|
||||
input_schema={"properties": {}},
|
||||
webhook_input_node=Mock() if trigger_graph else None,
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.experts.experts_db.resolve_expert_for_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
create = mocker.patch(
|
||||
"backend.api.features.library.db.create_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(),
|
||||
)
|
||||
request = AgentPresetCreateRequest(
|
||||
name="Preset", graph_id="graph-1", graph_version=1, is_active=is_active
|
||||
)
|
||||
writer = _key([APIKeyPermission.WRITE_LIBRARY])
|
||||
|
||||
if not refused:
|
||||
await presets.create_preset(request=request, auth=writer)
|
||||
create.assert_awaited_once()
|
||||
return
|
||||
with pytest.raises(HTTPException) as refusal:
|
||||
await presets.create_preset(request=request, auth=writer)
|
||||
assert refusal.value.status_code == 403
|
||||
create.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args, needed",
|
||||
[
|
||||
({"preset_id": "p-1", "is_active": True}, [APIKeyPermission.RUN_AGENT]),
|
||||
({"preset_id": "p-1", "is_active": False}, []),
|
||||
({"preset_id": "p-1", "name": "x"}, []),
|
||||
# New inputs re-register a trigger's webhook, and the tool can't see
|
||||
# the graph before it runs.
|
||||
({"preset_id": "p-1", "inputs": {"a": 1}}, [APIKeyPermission.RUN_AGENT]),
|
||||
({"preset_id": "p-1", "inputs": {"a": 1}, "is_active": False}, []),
|
||||
],
|
||||
)
|
||||
def test_turning_a_preset_on_over_mcp_needs_the_run_permission(
|
||||
args: dict, needed: list[APIKeyPermission]
|
||||
):
|
||||
assert UpdatePresetTool().external_permissions(args) == needed
|
||||
|
||||
|
||||
async def test_new_inputs_on_a_preset_without_credentials_are_saved(
|
||||
mocker: pytest_mock.MockFixture, stored: AsyncMock
|
||||
):
|
||||
"""Its saved credentials, `{}`, go along with the new inputs."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(webhook_input_node=None),
|
||||
)
|
||||
|
||||
await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(inputs={"topic": "x"}),
|
||||
preset_id="p-1",
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
rows = stored.update.await_args.kwargs["data"]["InputPresets"]["create"]
|
||||
assert [row["name"] for row in rows] == ["topic"]
|
||||
|
||||
|
||||
async def test_new_inputs_on_a_trigger_without_credentials_move_it_to_its_new_webhook(
|
||||
mocker: pytest_mock.MockFixture, stored: AsyncMock
|
||||
):
|
||||
"""The new webhook is registered before the inputs are saved, so a refused
|
||||
save would leave it attached to nothing."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(
|
||||
webhook_input_node=Mock(id="node-1"), organization_id=ORG, team_id=None
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.resolve_member_tenancy",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(ORG, None),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.make_node_credentials_input_map",
|
||||
return_value={},
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.setup_webhook_for_block",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(Mock(id="hook-2"), None),
|
||||
)
|
||||
attach = mocker.patch(
|
||||
"backend.api.features.library.db.set_preset_webhook",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
prune = mocker.patch(
|
||||
"backend.api.features.library.triggers._prune_dangling_webhook",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(inputs={"constants": {"a": 2}}),
|
||||
preset_id="p-1",
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
attach.assert_awaited_once_with(USER_ID, "p-1", "hook-2")
|
||||
prune.assert_awaited_once_with(USER_ID, "hook-1")
|
||||
|
||||
|
||||
async def test_new_credentials_alone_are_saved_with_the_saved_inputs(
|
||||
mocker: pytest_mock.MockFixture, stored: AsyncMock
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_preset(inputs={"topic": "x"}),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(webhook_input_node=None),
|
||||
)
|
||||
|
||||
await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(credentials_inputs=NEW_CREDENTIALS),
|
||||
preset_id="p-1",
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
rows = stored.update.await_args.kwargs["data"]["InputPresets"]["create"]
|
||||
assert sorted(row["name"] for row in rows) == ["github_credentials", "topic"]
|
||||
|
||||
|
||||
async def test_new_credentials_alone_move_a_trigger_keeping_its_secret_settings(
|
||||
mocker: pytest_mock.MockFixture, stored: AsyncMock
|
||||
):
|
||||
"""Reads leave the secret settings out, so they come from the saved preset."""
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(
|
||||
webhook_input_node=Mock(id="node-1"), organization_id=ORG, team_id=None
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.resolve_member_tenancy",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(ORG, None),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers.make_node_credentials_input_map",
|
||||
return_value={},
|
||||
)
|
||||
register = mocker.patch(
|
||||
"backend.api.features.library.triggers.setup_webhook_for_block",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(Mock(id="hook-2"), None),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.set_preset_webhook",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_trigger(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.triggers._prune_dangling_webhook",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
await presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(credentials_inputs=NEW_CREDENTIALS),
|
||||
preset_id="p-1",
|
||||
auth=_key(RUNNER),
|
||||
)
|
||||
|
||||
assert register.await_args.kwargs["trigger_config"]["secret_token"] == (
|
||||
TRIGGER_SECRET
|
||||
)
|
||||
rows = stored.update.await_args.kwargs["data"]["InputPresets"]["create"]
|
||||
assert sorted(row["name"] for row in rows) == [
|
||||
"constants",
|
||||
"github_credentials",
|
||||
"secret_token",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stored(mocker: pytest_mock.MockFixture) -> AsyncMock:
|
||||
"""`db.update_preset` runs for real; only its Prisma calls are stubbed."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def no_transaction():
|
||||
yield None
|
||||
|
||||
client = AsyncMock()
|
||||
mocker.patch("backend.api.features.library.db.transaction", no_transaction)
|
||||
mocker.patch.object(prisma.models.AgentPreset, "prisma", return_value=client)
|
||||
mocker.patch.object(
|
||||
prisma.models.AgentNodeExecutionInputOutput, "prisma", return_value=AsyncMock()
|
||||
)
|
||||
mocker.patch.object(LibraryAgentPreset, "from_db", return_value=_preset())
|
||||
return client
|
||||
|
||||
|
||||
def _trigger() -> LibraryAgentPreset:
|
||||
return _preset(
|
||||
webhook=True,
|
||||
inputs={"secret_token": TRIGGER_SECRET, "constants": {"a": 1}},
|
||||
)
|
||||
|
||||
|
||||
def _stub_preset_and_graph(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
preset: LibraryAgentPreset,
|
||||
input_schema: dict | None = None,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=preset,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(input_schema=input_schema or {"properties": {}}),
|
||||
)
|
||||
|
||||
|
||||
async def _execute():
|
||||
return await presets.execute_preset(
|
||||
response=Response(),
|
||||
preset_id="p-1",
|
||||
request=AgentPresetRunRequest(),
|
||||
auth=_key(),
|
||||
)
|
||||
|
||||
|
||||
def _preset(
|
||||
webhook: bool = False, expert_id: str | None = None, inputs: dict | None = None
|
||||
) -> LibraryAgentPreset:
|
||||
now = datetime.now(tz=timezone.utc)
|
||||
return LibraryAgentPreset.model_validate(
|
||||
{
|
||||
"id": "p-1",
|
||||
"user_id": USER_ID,
|
||||
"graph_id": "graph-1",
|
||||
"graph_version": 1,
|
||||
"name": "Preset",
|
||||
"description": "",
|
||||
"inputs": inputs or {},
|
||||
"credentials": {},
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"organization_id": ORG,
|
||||
"expert_id": expert_id,
|
||||
"webhook_id": "hook-1" if webhook else None,
|
||||
"webhook": (
|
||||
{
|
||||
"id": "hook-1",
|
||||
"user_id": USER_ID,
|
||||
"provider": "generic_webhook",
|
||||
"credentials_id": "",
|
||||
"webhook_type": "plain",
|
||||
"resource": "",
|
||||
"events": [],
|
||||
"secret": SECRET,
|
||||
"provider_webhook_id": PROVIDER_HOOK_ID,
|
||||
}
|
||||
if webhook
|
||||
else None
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _run() -> GraphExecutionMeta:
|
||||
return GraphExecutionMeta.model_construct(
|
||||
id="run-1",
|
||||
user_id=USER_ID,
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
preset_id="p-1",
|
||||
status=ExecutionStatus.QUEUED,
|
||||
started_at=None,
|
||||
ended_at=None,
|
||||
inputs={},
|
||||
credential_inputs=None,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
stats=None,
|
||||
organization_id=ORG,
|
||||
team_id=None,
|
||||
)
|
||||
|
||||
|
||||
def _key(scopes: list[APIKeyPermission] | None = None) -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=scopes or [],
|
||||
type="api_key",
|
||||
organization_id=ORG,
|
||||
)
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Library writes stay inside the organization the credential acts in.
|
||||
|
||||
Library rows are the user's, and a lookup by id finds the user's row in any of
|
||||
their organizations; every folder a write names, and every row it creates,
|
||||
must be the credential's organization's.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
from ..models import (
|
||||
LibraryAgentUpdateRequest,
|
||||
LibraryFolderCreateRequest,
|
||||
LibraryFolderMoveRequest,
|
||||
)
|
||||
from ..tenancy import TenantContext
|
||||
from .agents import fork_library_agent, update_library_agent
|
||||
from .folders import create_folder, move_folder
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-a",
|
||||
team_id="team-a",
|
||||
)
|
||||
|
||||
|
||||
def _folder(organization_id: str | None) -> Mock:
|
||||
return Mock(organization_id=organization_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def folders(mocker: pytest_mock.MockFixture) -> dict[str, AsyncMock]:
|
||||
"""Folder f-a is org A's, f-b org B's; the writes are recorded."""
|
||||
by_id = {"f-a": _folder("org-a"), "f-b": _folder("org-b")}
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_folder",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=lambda folder_id, user_id: by_id[folder_id],
|
||||
)
|
||||
mocker.patch("backend.api.external.v2.library.folders.LibraryFolder.from_internal")
|
||||
mocker.patch("backend.api.external.v2.library.agents.LibraryAgent.from_internal")
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id="org-a"),
|
||||
)
|
||||
return {
|
||||
name: mocker.patch(
|
||||
f"backend.api.features.library.db.{name}", new_callable=AsyncMock
|
||||
)
|
||||
for name in (
|
||||
"create_folder",
|
||||
"move_folder",
|
||||
"update_library_agent",
|
||||
"fork_library_agent",
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
async def test_a_new_folder_is_tagged_with_the_credentials_tenant(
|
||||
folders: dict[str, AsyncMock],
|
||||
) -> None:
|
||||
await create_folder(
|
||||
LibraryFolderCreateRequest(name="Reports", parent_id="f-a"), auth=_AUTH
|
||||
)
|
||||
|
||||
kwargs = folders["create_folder"].await_args.kwargs
|
||||
assert (kwargs["organization_id"], kwargs["team_id"]) == ("org-a", "team-a")
|
||||
|
||||
|
||||
async def test_a_folder_cannot_be_created_under_another_organizations(
|
||||
folders: dict[str, AsyncMock],
|
||||
) -> None:
|
||||
with pytest.raises(NotFoundError):
|
||||
await create_folder(
|
||||
LibraryFolderCreateRequest(name="Reports", parent_id="f-b"), auth=_AUTH
|
||||
)
|
||||
folders["create_folder"].assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_folder_cannot_be_moved_under_another_organizations(
|
||||
folders: dict[str, AsyncMock],
|
||||
) -> None:
|
||||
with pytest.raises(NotFoundError):
|
||||
await move_folder(
|
||||
LibraryFolderMoveRequest(target_parent_id="f-b"),
|
||||
folder_id="f-a",
|
||||
auth=_AUTH,
|
||||
)
|
||||
folders["move_folder"].assert_not_awaited()
|
||||
|
||||
|
||||
async def test_an_agent_cannot_be_filed_in_another_organizations_folder(
|
||||
folders: dict[str, AsyncMock],
|
||||
) -> None:
|
||||
with pytest.raises(NotFoundError):
|
||||
await update_library_agent(
|
||||
LibraryAgentUpdateRequest(folder_id="f-b"), agent_id="agent-1", auth=_AUTH
|
||||
)
|
||||
folders["update_library_agent"].assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_fork_lands_in_the_credentials_tenant(
|
||||
folders: dict[str, AsyncMock],
|
||||
) -> None:
|
||||
await fork_library_agent(agent_id="agent-1", auth=_AUTH)
|
||||
|
||||
kwargs = folders["fork_library_agent"].await_args.kwargs
|
||||
assert (kwargs["organization_id"], kwargs["team_id"]) == ("org-a", "team-a")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"graph_version, organization_ids", [(2, frozenset({"org-a"})), (None, None)]
|
||||
)
|
||||
async def test_only_this_organizations_entries_give_way_to_a_chosen_version(
|
||||
folders: dict[str, AsyncMock],
|
||||
mocker: pytest_mock.MockFixture,
|
||||
graph_version: int | None,
|
||||
organization_ids: frozenset[str | None] | None,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.agents.writable_organization_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=frozenset({"org-a"}),
|
||||
)
|
||||
|
||||
await update_library_agent(
|
||||
LibraryAgentUpdateRequest(graph_version=graph_version),
|
||||
agent_id="agent-1",
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
kwargs = folders["update_library_agent"].await_args.kwargs
|
||||
assert kwargs["organization_ids"] == organization_ids
|
||||
@@ -0,0 +1,445 @@
|
||||
"""
|
||||
V2 External API - Marketplace Endpoints
|
||||
|
||||
Provides access to the agent marketplace (store).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import urllib.parse
|
||||
from typing import Literal, Optional
|
||||
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
Depends,
|
||||
File,
|
||||
HTTPException,
|
||||
Path,
|
||||
Query,
|
||||
Response,
|
||||
Security,
|
||||
UploadFile,
|
||||
)
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.library.exceptions import (
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
from backend.api.features.store import cache as store_cache
|
||||
from backend.api.features.store import db as store_db
|
||||
from backend.api.features.store import media as store_media
|
||||
from backend.api.features.store.db import (
|
||||
StoreAgentsSortOptions,
|
||||
StoreCreatorsSortOptions,
|
||||
)
|
||||
from backend.api.features.store.model import ProfileUpdateRequest
|
||||
from backend.api.features.workspace.service import read_within_size_limit
|
||||
from backend.util.virus_scanner import scan_content_safe
|
||||
|
||||
from .models import (
|
||||
LibraryAgent,
|
||||
MarketplaceAgent,
|
||||
MarketplaceAgentDetails,
|
||||
MarketplaceAgentSubmission,
|
||||
MarketplaceAgentSubmissionCreateRequest,
|
||||
MarketplaceAgentSubmissionEditRequest,
|
||||
MarketplaceCreatorDetails,
|
||||
MarketplaceMediaUploadResponse,
|
||||
MarketplaceUserProfile,
|
||||
MarketplaceUserProfileUpdateRequest,
|
||||
)
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .rate_limit import enforce, media_upload_limiter
|
||||
from .tenancy import (
|
||||
TenantContext,
|
||||
require_auth,
|
||||
require_paid_permission,
|
||||
require_permission,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
marketplace_router = APIRouter(tags=["marketplace"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Agents
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/agents",
|
||||
summary="List or search marketplace agents",
|
||||
operation_id="listMarketplaceAgents",
|
||||
)
|
||||
async def list_agents(
|
||||
featured: bool = Query(
|
||||
default=False, description="Filter to only show featured agents"
|
||||
),
|
||||
creator: Optional[str] = Query(
|
||||
default=None, description="Filter by creator username"
|
||||
),
|
||||
category: Optional[str] = Query(default=None, description="Filter by category"),
|
||||
search_query: Optional[str] = Query(
|
||||
default=None, description="Literal + semantic search on names and descriptions"
|
||||
),
|
||||
sorted_by: Optional[Literal["rating", "runs", "name", "updated_at"]] = Query(
|
||||
default=None,
|
||||
description="Property to sort results by. Ignored if search_query is provided.",
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
# This data is public, but we still require auth for access tracking and rate limits
|
||||
auth: TenantContext = Security(require_auth),
|
||||
) -> Page[MarketplaceAgent]:
|
||||
"""List agents available in the marketplace, with optional filtering and sorting."""
|
||||
result = await store_cache.get_cached_store_agents(
|
||||
featured=featured,
|
||||
creator=creator,
|
||||
sorted_by=StoreAgentsSortOptions(sorted_by) if sorted_by else None,
|
||||
search_query=search_query,
|
||||
category=category,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[MarketplaceAgent.from_internal(a) for a in result.agents],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/agents/by-version/{version_id}",
|
||||
summary="Get marketplace agent by version ID",
|
||||
operation_id="getMarketplaceAgentByListingVersion",
|
||||
)
|
||||
async def get_agent_by_version(
|
||||
version_id: str,
|
||||
# This data is public, but we still require auth for access tracking and rate limits
|
||||
auth: TenantContext = Security(require_auth),
|
||||
) -> MarketplaceAgentDetails:
|
||||
"""Get details of a marketplace agent by its store listing version ID."""
|
||||
agent = await store_db.get_store_agent_by_version_id(version_id)
|
||||
return MarketplaceAgentDetails.from_internal(agent)
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/agents/{username}/{agent_name}",
|
||||
summary="Get marketplace agent details",
|
||||
operation_id="getMarketplaceAgent",
|
||||
)
|
||||
async def get_agent_details(
|
||||
username: str,
|
||||
agent_name: str,
|
||||
# This data is public, but we still require auth for access tracking and rate limits
|
||||
auth: TenantContext = Security(require_auth),
|
||||
) -> MarketplaceAgentDetails:
|
||||
"""Get details of a specific marketplace agent."""
|
||||
username = urllib.parse.unquote(username).lower()
|
||||
agent_name = urllib.parse.unquote(agent_name).lower()
|
||||
|
||||
agent = await store_cache.get_cached_agent_details(
|
||||
username=username, agent_name=agent_name
|
||||
)
|
||||
|
||||
return MarketplaceAgentDetails.from_internal(agent)
|
||||
|
||||
|
||||
@marketplace_router.post(
|
||||
path="/agents/{username}/{agent_name}/add-to-library",
|
||||
summary="Add marketplace agent to library",
|
||||
operation_id="addMarketplaceAgentToLibrary",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def add_agent_to_library(
|
||||
username: str,
|
||||
agent_name: str,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_LIBRARY)
|
||||
),
|
||||
) -> LibraryAgent:
|
||||
"""Add a marketplace agent to the authenticated user's library."""
|
||||
username = urllib.parse.unquote(username).lower()
|
||||
agent_name = urllib.parse.unquote(agent_name).lower()
|
||||
|
||||
agent_details = await store_cache.get_cached_agent_details(
|
||||
username=username, agent_name=agent_name
|
||||
)
|
||||
|
||||
try:
|
||||
agent = await library_db.add_store_agent_to_library(
|
||||
store_listing_version_id=agent_details.store_listing_version_id,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
except LibraryAgentInAnotherOrganizationError as e:
|
||||
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e))
|
||||
|
||||
return LibraryAgent.from_internal(agent)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Creators
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/creators",
|
||||
summary="List marketplace creators",
|
||||
operation_id="listMarketplaceCreators",
|
||||
)
|
||||
async def list_creators(
|
||||
featured: bool = Query(
|
||||
default=False, description="Filter to featured creators only"
|
||||
),
|
||||
search_query: Optional[str] = Query(
|
||||
default=None,
|
||||
max_length=100,
|
||||
description="Literal + semantic search on names and descriptions",
|
||||
),
|
||||
sorted_by: Optional[Literal["agent_rating", "agent_runs", "num_agents"]] = Query(
|
||||
default=None, description="Sort field"
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
# This data is public, but we still require auth for access tracking and rate limits
|
||||
auth: TenantContext = Security(require_auth),
|
||||
) -> Page[MarketplaceCreatorDetails]:
|
||||
"""List or search marketplace creators."""
|
||||
result = await store_cache.get_cached_store_creators(
|
||||
featured=featured,
|
||||
# A blank search lists everyone, rather than failing in the query.
|
||||
search_query=(search_query or "").strip() or None,
|
||||
sorted_by=StoreCreatorsSortOptions(sorted_by) if sorted_by else None,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[MarketplaceCreatorDetails.from_internal(c) for c in result.creators],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/creators/{username}",
|
||||
summary="Get marketplace creator details",
|
||||
operation_id="getMarketplaceCreator",
|
||||
)
|
||||
async def get_creator_details(
|
||||
username: str,
|
||||
# This data is public, but we still require auth for access tracking and rate limits
|
||||
auth: TenantContext = Security(require_auth),
|
||||
) -> MarketplaceCreatorDetails:
|
||||
"""Get a marketplace creator's profile w/ stats."""
|
||||
username = urllib.parse.unquote(username).lower()
|
||||
creator = await store_cache.get_cached_creator_details(username=username)
|
||||
return MarketplaceCreatorDetails.from_internal(creator)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Profile
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/profile",
|
||||
summary="Get my marketplace profile",
|
||||
operation_id="getMarketplaceMyProfile",
|
||||
)
|
||||
async def get_profile(
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_STORE)),
|
||||
) -> MarketplaceCreatorDetails:
|
||||
"""Get the authenticated user's marketplace profile w/ creator stats."""
|
||||
profile = await store_db.get_user_profile(auth.user_id)
|
||||
if not profile:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Profile not found",
|
||||
)
|
||||
|
||||
creator = await store_cache.get_cached_creator_details(username=profile.username)
|
||||
return MarketplaceCreatorDetails.from_internal(creator)
|
||||
|
||||
|
||||
@marketplace_router.patch(
|
||||
path="/profile",
|
||||
summary="Update my marketplace profile",
|
||||
operation_id="updateMarketplaceMyProfile",
|
||||
)
|
||||
async def update_profile(
|
||||
request: MarketplaceUserProfileUpdateRequest,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_STORE)),
|
||||
) -> MarketplaceUserProfile:
|
||||
"""Update the authenticated user's marketplace profile.
|
||||
|
||||
Only the fields present in the request body are changed. Set `avatar_url`
|
||||
to `null` to remove the avatar. A user without a profile yet can create one
|
||||
here by sending at least `username`, `name` and `description`.
|
||||
"""
|
||||
# exclude_unset keeps "omitted" distinct from "explicitly null" downstream
|
||||
profile = ProfileUpdateRequest.model_validate(
|
||||
request.model_dump(exclude_unset=True)
|
||||
)
|
||||
|
||||
updated_profile = await store_db.update_profile(auth.user_id, profile)
|
||||
return MarketplaceUserProfile.from_internal(updated_profile)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Submissions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@marketplace_router.get(
|
||||
path="/submissions",
|
||||
summary="List my marketplace submissions",
|
||||
operation_id="listMarketplaceSubmissions",
|
||||
)
|
||||
async def list_submissions(
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_STORE)),
|
||||
) -> Page[MarketplaceAgentSubmission]:
|
||||
"""List the authenticated user's marketplace listing submissions."""
|
||||
result = await store_db.get_store_submissions(
|
||||
user_id=auth.user_id,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[MarketplaceAgentSubmission.from_internal(s) for s in result.submissions],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@marketplace_router.post(
|
||||
path="/submissions",
|
||||
summary="Create marketplace submission",
|
||||
operation_id="createMarketplaceSubmission",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_submission(
|
||||
request: MarketplaceAgentSubmissionCreateRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_STORE)
|
||||
),
|
||||
) -> MarketplaceAgentSubmission:
|
||||
"""Submit a new marketplace listing for review."""
|
||||
submission = await store_db.create_store_submission(
|
||||
user_id=auth.user_id,
|
||||
graph_id=request.graph_id,
|
||||
graph_version=request.graph_version,
|
||||
slug=request.slug,
|
||||
name=request.name,
|
||||
sub_heading=request.sub_heading,
|
||||
description=request.description,
|
||||
instructions=request.instructions,
|
||||
categories=request.categories,
|
||||
image_urls=request.image_urls,
|
||||
video_url=request.video_url,
|
||||
agent_output_demo_url=request.agent_output_demo_url,
|
||||
changes_summary=request.changes_summary or "Initial Submission",
|
||||
recommended_schedule_cron=request.recommended_schedule_cron,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return MarketplaceAgentSubmission.from_internal(submission)
|
||||
|
||||
|
||||
@marketplace_router.put(
|
||||
path="/submissions/{version_id}",
|
||||
summary="Edit marketplace submission",
|
||||
operation_id="updateMarketplaceSubmission",
|
||||
)
|
||||
async def edit_submission(
|
||||
request: MarketplaceAgentSubmissionEditRequest,
|
||||
version_id: str = Path(description="Store listing version ID"),
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_STORE)
|
||||
),
|
||||
) -> MarketplaceAgentSubmission:
|
||||
"""Replace a pending marketplace listing submission.
|
||||
|
||||
Every field is written, so send the whole listing; omitting one clears it.
|
||||
"""
|
||||
submission = await store_db.edit_store_submission(
|
||||
user_id=auth.user_id,
|
||||
store_listing_version_id=version_id,
|
||||
name=request.name,
|
||||
sub_heading=request.sub_heading,
|
||||
description=request.description,
|
||||
image_urls=request.image_urls,
|
||||
video_url=request.video_url,
|
||||
agent_output_demo_url=request.agent_output_demo_url,
|
||||
categories=request.categories,
|
||||
changes_summary=request.changes_summary,
|
||||
recommended_schedule_cron=request.recommended_schedule_cron,
|
||||
instructions=request.instructions,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
return MarketplaceAgentSubmission.from_internal(submission)
|
||||
|
||||
|
||||
@marketplace_router.delete(
|
||||
path="/submissions/{version_id}",
|
||||
summary="Delete marketplace submission",
|
||||
operation_id="deleteMarketplaceSubmission",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_submission(
|
||||
version_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_STORE)),
|
||||
) -> None:
|
||||
"""Delete a marketplace listing submission. Approved listings can not be deleted."""
|
||||
# Raising, so an approved listing is a 400 and a failure a 500, not a 404.
|
||||
await store_db.delete_store_submission_or_raise(
|
||||
user_id=auth.user_id,
|
||||
store_listing_version_id=version_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Submission Media
|
||||
# ============================================================================
|
||||
|
||||
MEDIA_MAX_SIZE_MB = 10
|
||||
|
||||
|
||||
@marketplace_router.post(
|
||||
path="/submissions/media",
|
||||
summary="Upload marketplace submission media",
|
||||
operation_id="uploadMarketplaceSubmissionMedia",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def upload_submission_media(
|
||||
response: Response,
|
||||
file: UploadFile = File(...),
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_STORE)
|
||||
),
|
||||
) -> MarketplaceMediaUploadResponse:
|
||||
"""
|
||||
Upload an image or video for a marketplace submission. Max size: 10 MB.
|
||||
|
||||
**Rate limit:** 10 requests per 5 minutes per user.
|
||||
"""
|
||||
await enforce(media_upload_limiter, auth.user_id, response)
|
||||
|
||||
content = await read_within_size_limit(file, max_size_mb=MEDIA_MAX_SIZE_MB)
|
||||
|
||||
# Virus scan
|
||||
await scan_content_safe(content, filename=file.filename or "upload")
|
||||
|
||||
# Reset file position for store_media to read
|
||||
await file.seek(0)
|
||||
|
||||
url = await store_media.upload_media(
|
||||
user_id=auth.user_id,
|
||||
file=file,
|
||||
)
|
||||
|
||||
return MarketplaceMediaUploadResponse(url=url)
|
||||
@@ -0,0 +1,202 @@
|
||||
"""The v2 marketplace routes: searches, library adds and media uploads."""
|
||||
|
||||
import io
|
||||
from typing import Optional
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import prisma.models
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException, Response, UploadFile
|
||||
from prisma.enums import APIKeyPermission, SubmissionStatus
|
||||
|
||||
from backend.api.features.library.exceptions import (
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
|
||||
from .errors import add_v2_exception_handlers
|
||||
from .marketplace import (
|
||||
add_agent_to_library,
|
||||
marketplace_router,
|
||||
upload_submission_media,
|
||||
)
|
||||
from .routes import v2_router
|
||||
from .tenancy import TenantContext, require_auth
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
_LIMIT = 10 * 1024 * 1024
|
||||
_ELEVEN_MB = 11 * 1024 * 1024
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def creators(mocker: pytest_mock.MockFixture) -> AsyncMock:
|
||||
return mocker.patch(
|
||||
"backend.api.external.v2.marketplace.store_cache.get_cached_store_creators",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(creators=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> fastapi.testclient.TestClient:
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(marketplace_router, prefix="/marketplace")
|
||||
app.dependency_overrides[require_auth] = lambda: _AUTH
|
||||
add_v2_exception_handlers(app)
|
||||
return fastapi.testclient.TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
def test_a_blank_creator_search_lists_every_creator(
|
||||
client: fastapi.testclient.TestClient, creators: AsyncMock
|
||||
) -> None:
|
||||
response = client.get("/marketplace/creators", params={"search_query": " "})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert creators.await_args.kwargs["search_query"] is None
|
||||
|
||||
|
||||
def test_an_overlong_creator_search_is_refused(
|
||||
client: fastapi.testclient.TestClient, creators: AsyncMock
|
||||
) -> None:
|
||||
response = client.get("/marketplace/creators", params={"search_query": "x" * 101})
|
||||
|
||||
assert response.status_code == 422
|
||||
creators.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_adding_a_listing_lands_in_the_credentials_organization(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.marketplace.store_cache.get_cached_agent_details",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(store_listing_version_id="slv-1"),
|
||||
)
|
||||
add = mocker.patch(
|
||||
"backend.api.external.v2.marketplace.library_db.add_store_agent_to_library",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
mocker.patch("backend.api.external.v2.marketplace.LibraryAgent.from_internal")
|
||||
|
||||
await add_agent_to_library(username="someone", agent_name="agent", auth=_AUTH)
|
||||
|
||||
assert add.await_args.kwargs["organization_id"] == "org-1"
|
||||
|
||||
|
||||
def test_a_listing_already_in_another_organizations_library_is_a_conflict(
|
||||
client: fastapi.testclient.TestClient, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""The entry is unique per user, so it can be neither duplicated nor taken."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.marketplace.store_cache.get_cached_agent_details",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(store_listing_version_id="slv-1"),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.marketplace.library_db.add_store_agent_to_library",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=LibraryAgentInAnotherOrganizationError("elsewhere"),
|
||||
)
|
||||
|
||||
response = client.post("/marketplace/agents/someone/agent/add-to-library")
|
||||
|
||||
assert response.status_code == 409
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"declared_size", [_ELEVEN_MB, None], ids=["size-known", "size-unknown"]
|
||||
)
|
||||
async def test_oversized_media_is_refused_without_being_read_whole(
|
||||
mocker: pytest_mock.MockFixture, declared_size: Optional[int]
|
||||
) -> None:
|
||||
"""The whole body in memory is what the cap exists to prevent."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.marketplace.media_upload_limiter.check",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
scan = mocker.patch(
|
||||
"backend.api.external.v2.marketplace.scan_content_safe",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
store = mocker.patch(
|
||||
"backend.api.features.store.media.upload_media", new_callable=AsyncMock
|
||||
)
|
||||
body = _CountingBody(b"\0" * _ELEVEN_MB)
|
||||
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await upload_submission_media(
|
||||
response=Response(),
|
||||
file=UploadFile(file=body, size=declared_size, filename="big.png"),
|
||||
auth=TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=[APIKeyPermission.WRITE_STORE],
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
),
|
||||
)
|
||||
|
||||
assert refused.value.status_code == 413
|
||||
assert body.bytes_read <= (0 if declared_size else _LIMIT + 64 * 1024)
|
||||
scan.assert_not_awaited()
|
||||
store.assert_not_awaited()
|
||||
|
||||
|
||||
class _CountingBody(io.BytesIO):
|
||||
bytes_read = 0
|
||||
|
||||
def read(self, size: Optional[int] = -1) -> bytes:
|
||||
chunk = super().read(size)
|
||||
self.bytes_read += len(chunk)
|
||||
return chunk
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"owning_org_id, status, expected",
|
||||
[
|
||||
("org-1", SubmissionStatus.APPROVED, 400),
|
||||
("org-2", SubmissionStatus.PENDING, 404),
|
||||
("org-1", SubmissionStatus.PENDING, 204),
|
||||
],
|
||||
ids=["approved", "another-org", "deleted"],
|
||||
)
|
||||
def test_a_submission_that_cannot_be_deleted_says_why(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
owning_org_id: str,
|
||||
status: SubmissionStatus,
|
||||
expected: int,
|
||||
):
|
||||
version = Mock(
|
||||
id="slv-1",
|
||||
storeListingId="sl-1",
|
||||
submissionStatus=status,
|
||||
StoreListing=Mock(owningOrgId=owning_org_id, owningUserId="user-1"),
|
||||
)
|
||||
client = AsyncMock()
|
||||
client.find_unique.return_value = version
|
||||
client.count.return_value = 1
|
||||
mocker.patch.object(
|
||||
prisma.models.StoreListingVersion, "prisma", return_value=client
|
||||
)
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(v2_router)
|
||||
add_v2_exception_handlers(app)
|
||||
app.dependency_overrides[require_auth] = lambda: TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
|
||||
response = fastapi.testclient.TestClient(app).delete(
|
||||
"/marketplace/submissions/slv-1"
|
||||
)
|
||||
|
||||
assert response.status_code == expected, response.text
|
||||
@@ -0,0 +1,172 @@
|
||||
"""The checks an MCP tool call passes before the tool runs.
|
||||
|
||||
A Copilot tool trusts its caller: in a chat, that is AutoPilot's own model,
|
||||
working inside the chat's session. Over MCP the caller is anything holding a
|
||||
v2 credential, so the call is held to the credential's permissions first,
|
||||
the way each REST route is.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Sequence
|
||||
|
||||
from fastapi import HTTPException
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import ValidationError
|
||||
|
||||
from backend.api.utils.rate_limit import RateLimiter
|
||||
from backend.copilot.config import ChatConfig
|
||||
from backend.copilot.rate_limit import (
|
||||
RateLimitExceeded,
|
||||
RateLimitUnavailable,
|
||||
check_rate_limit,
|
||||
get_global_rate_limits,
|
||||
is_user_paywalled,
|
||||
)
|
||||
from backend.copilot.tools.base import BaseTool, without_card_hints
|
||||
from backend.copilot.tools.run_agent import RunAgentInput
|
||||
from backend.data.credit import get_credit_model
|
||||
|
||||
from .rate_limit import file_upload_limiter, graph_exec_limiter, search_limiter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
config = ChatConfig()
|
||||
|
||||
|
||||
def input_schema(tool: BaseTool) -> dict[str, Any]:
|
||||
"""The JSON Schema of the arguments an external caller may pass.
|
||||
|
||||
The chat's own schema leaves ``required`` out for the sake of a model whose
|
||||
tool call was cut short; an MCP client has no such excuse, so it is kept.
|
||||
The approval card's hints are left out, as they are from what the chat's
|
||||
model is sent: they aren't JSON Schema, and a client may refuse them.
|
||||
"""
|
||||
parameters = without_card_hints(tool.external_parameters)
|
||||
schema: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": parameters.get("properties", {}),
|
||||
"additionalProperties": False,
|
||||
}
|
||||
if required := parameters.get("required"):
|
||||
schema["required"] = list(required)
|
||||
return schema
|
||||
|
||||
|
||||
def check_arguments(tool: BaseTool, args: dict[str, Any]) -> None:
|
||||
"""Refuse an argument the tool doesn't advertise to external callers.
|
||||
|
||||
The tools take ``**kwargs``, so an unadvertised argument would reach code
|
||||
the advertised schema keeps external callers away from.
|
||||
"""
|
||||
allowed = tool.external_parameters.get("properties", {})
|
||||
if unknown := sorted(set(args) - set(allowed)):
|
||||
raise ToolError(f"Unknown argument(s): {', '.join(unknown)}")
|
||||
|
||||
|
||||
def missing_scopes(
|
||||
tool: BaseTool,
|
||||
required_scopes: Sequence[str],
|
||||
args: dict[str, Any],
|
||||
granted: Sequence[str],
|
||||
) -> list[str]:
|
||||
"""The permissions this call needs that the credential doesn't carry.
|
||||
|
||||
The tool's own permissions, plus those the arguments add: a branch that
|
||||
writes what those don't cover (``BaseTool.external_permissions``).
|
||||
"""
|
||||
needed = [*required_scopes, *(p.value for p in tool.external_permissions(args))]
|
||||
return [scope for scope in dict.fromkeys(needed) if scope not in granted]
|
||||
|
||||
|
||||
# The REST routes these tools mirror cap callers tighter than the global limit
|
||||
# (`rate_limit.py`); a call over MCP is held to the same caps, in the same
|
||||
# windows, so neither surface is a way around the other's.
|
||||
TOOL_LIMITERS: dict[str, RateLimiter] = {
|
||||
"run_agent": graph_exec_limiter,
|
||||
"write_workspace_file": file_upload_limiter,
|
||||
"find_agent": search_limiter,
|
||||
"find_library_agent": search_limiter,
|
||||
"search_docs": search_limiter,
|
||||
}
|
||||
|
||||
|
||||
async def check_tool_limits(
|
||||
tool: BaseTool, args: dict[str, Any], user_id: str, organization_id: str
|
||||
) -> None:
|
||||
"""The per-endpoint caps, and for a real run, the zero-balance refusal."""
|
||||
for limiter in _tool_limiters(tool, args):
|
||||
try:
|
||||
await limiter.check(user_id)
|
||||
except HTTPException as exc:
|
||||
raise ToolError(str(exc.detail)) from exc
|
||||
if tool.name == "run_agent" and _starts_a_paid_run(args):
|
||||
try:
|
||||
credit_model = await get_credit_model(user_id, organization_id)
|
||||
balance = await credit_model.get_credits(user_id)
|
||||
except Exception as exc:
|
||||
logger.warning(f"MCP balance check failed for {user_id}: {exc}")
|
||||
raise ToolError("Couldn't check your balance; retry shortly") from exc
|
||||
if balance <= 0:
|
||||
raise ToolError(
|
||||
"Insufficient balance to run the agent. Please top up your account."
|
||||
)
|
||||
|
||||
|
||||
def _tool_limiters(tool: BaseTool, args: dict[str, Any]) -> list[RateLimiter]:
|
||||
"""The tool's own cap, plus the upload cap when an argument has it write a
|
||||
workspace file, as `write_to` does: the write the upload route caps."""
|
||||
limiters = [TOOL_LIMITERS[tool.name]] if tool.name in TOOL_LIMITERS else []
|
||||
if APIKeyPermission.WRITE_FILES in tool.external_permissions(args):
|
||||
limiters.append(file_upload_limiter)
|
||||
return list(dict.fromkeys(limiters))
|
||||
|
||||
|
||||
def _starts_a_paid_run(args: dict[str, Any]) -> bool:
|
||||
"""A run now, not a dry run and not a schedule, as REST's run route refuses.
|
||||
|
||||
Read with the tool's own model, so this check and the tool agree on what
|
||||
the arguments ask for: a blank `cron` is no schedule to either. Arguments
|
||||
the model refuses start nothing.
|
||||
"""
|
||||
try:
|
||||
params = RunAgentInput(**args)
|
||||
except ValidationError:
|
||||
return True
|
||||
return not (params.dry_run or params.schedule_name or params.cron)
|
||||
|
||||
|
||||
async def check_plan(user_id: str) -> None:
|
||||
"""Every tool call needs an active plan, as every chat turn does."""
|
||||
try:
|
||||
paywalled = await is_user_paywalled(user_id)
|
||||
except Exception as exc:
|
||||
logger.warning(f"MCP paywall check failed for {user_id}: {exc}")
|
||||
raise ToolError("Couldn't check your plan; retry shortly") from exc
|
||||
if paywalled:
|
||||
raise ToolError("This tool needs an active subscription")
|
||||
|
||||
|
||||
async def check_spend_allowance(user_id: str) -> None:
|
||||
"""The usage caps a chat turn checks, for a tool that spends what a turn does.
|
||||
|
||||
Metering after the call only slows the next one down; without this a key
|
||||
past its usage cap keeps spending.
|
||||
"""
|
||||
try:
|
||||
daily_limit, weekly_limit, _ = await get_global_rate_limits(
|
||||
user_id,
|
||||
config.daily_cost_limit_microdollars,
|
||||
config.weekly_cost_limit_microdollars,
|
||||
)
|
||||
await check_rate_limit(
|
||||
user_id=user_id,
|
||||
daily_cost_limit=daily_limit,
|
||||
weekly_cost_limit=weekly_limit,
|
||||
)
|
||||
except RateLimitExceeded as exc:
|
||||
raise ToolError(str(exc)) from exc
|
||||
except RateLimitUnavailable as exc:
|
||||
# Fail closed, as the chat route does: an outage can't prove the
|
||||
# caller is under their cap.
|
||||
raise ToolError("Usage limits are unavailable; retry shortly") from exc
|
||||
@@ -0,0 +1,196 @@
|
||||
"""The `/mcp` endpoint answers over HTTP once the host app's lifespan has run.
|
||||
|
||||
Starlette never runs a mounted app's lifespan, and FastMCP starts the session
|
||||
manager that every request goes through in its lifespan. Mounted as-is, every
|
||||
request to `/api/v2/mcp/` was a 500: "Task group is not initialized".
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from backend.api.external.middleware import VerifiedCredential
|
||||
from backend.api.external.v2.mcp_server import (
|
||||
ExternalAPITokenVerifier,
|
||||
MCPMount,
|
||||
TenantedAccessToken,
|
||||
)
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import ACCESS_TOKEN_PREFIX, InvalidTokenError
|
||||
|
||||
INITIALIZE = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2025-06-18",
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": "test", "version": "1"},
|
||||
},
|
||||
}
|
||||
HEADERS = {
|
||||
"Authorization": "Bearer agpt_test",
|
||||
"Accept": "application/json, text/event-stream",
|
||||
}
|
||||
|
||||
|
||||
def _host(mount: MCPMount) -> fastapi.FastAPI:
|
||||
"""A host app wired the way the platform's own app wires the mount."""
|
||||
host = fastapi.FastAPI(lifespan=_lifespan(mount))
|
||||
host.mount("/mcp", mount)
|
||||
return host
|
||||
|
||||
|
||||
def _lifespan(mount: MCPMount):
|
||||
@contextlib.asynccontextmanager
|
||||
async def lifespan(_: fastapi.FastAPI):
|
||||
async with mount.running():
|
||||
yield
|
||||
|
||||
return lifespan
|
||||
|
||||
|
||||
def test_an_authenticated_initialize_succeeds_through_the_mount(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.ExternalAPITokenVerifier.verify_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=TenantedAccessToken(
|
||||
token="agpt_test",
|
||||
client_id="user-1",
|
||||
scopes=["IDENTITY"],
|
||||
organization_id="org-1",
|
||||
),
|
||||
)
|
||||
|
||||
# A public hostname: FastMCP's default host check admitted only localhost.
|
||||
with TestClient(_host(MCPMount()), base_url="https://backend.agpt.co") as client:
|
||||
response = client.post("/mcp/", headers=HEADERS, json=INITIALIZE)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert "autogpt-platform" in response.text
|
||||
|
||||
|
||||
def test_the_host_lifespan_can_run_more_than_once(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""A session manager runs once, so each lifespan needs its own server."""
|
||||
mount = MCPMount()
|
||||
for _ in range(2):
|
||||
with TestClient(_host(mount)):
|
||||
pass
|
||||
|
||||
|
||||
def test_a_server_that_fails_to_start_leaves_the_host_up_answering_503(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The REST API must not go down with its MCP endpoint."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.create_mcp_app",
|
||||
side_effect=RuntimeError("no tools today"),
|
||||
)
|
||||
|
||||
with TestClient(_host(MCPMount())) as client:
|
||||
response = client.post("/mcp/", headers=HEADERS, json=INITIALIZE)
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.json()["error"]["code"] == "service_unavailable"
|
||||
|
||||
|
||||
def test_before_the_lifespan_runs_the_mount_answers_in_the_v2_envelope() -> None:
|
||||
host = fastapi.FastAPI()
|
||||
host.mount("/mcp", MCPMount())
|
||||
|
||||
response = TestClient(host).post("/mcp/", headers=HEADERS, json=INITIALIZE)
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.json()["error"]["code"] == "service_unavailable"
|
||||
|
||||
|
||||
def test_a_credential_the_rate_limiter_verified_is_not_verified_again(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""An API key costs a Scrypt hash; the request pays for one, not two."""
|
||||
validate = mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.validate_api_key", new_callable=AsyncMock
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.resolve_credential_tenancy",
|
||||
new_callable=AsyncMock,
|
||||
return_value=("org-1", None),
|
||||
)
|
||||
auth = APIAuthorizationInfo(
|
||||
user_id="user-1", scopes=[], type="api_key", created_at=datetime.now(UTC)
|
||||
)
|
||||
|
||||
async def verified_by_the_limiter(scope, receive, send):
|
||||
scope.setdefault("state", {})["external_api_auth"] = VerifiedCredential(
|
||||
credential="agpt_test", auth=auth
|
||||
)
|
||||
await mount(scope, receive, send)
|
||||
|
||||
mount = MCPMount()
|
||||
host = fastapi.FastAPI(lifespan=_lifespan(mount))
|
||||
host.mount("/mcp", verified_by_the_limiter)
|
||||
|
||||
with TestClient(host, base_url="https://backend.agpt.co") as client:
|
||||
response = client.post("/mcp/", headers=HEADERS, json=INITIALIZE)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
validate.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_token_shaped_bearer_is_never_tried_as_an_api_key_over_mcp(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""FastMCP's own check follows the middleware's rule when it runs alone."""
|
||||
validate_key = mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.validate_api_key", new_callable=AsyncMock
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.validate_access_token",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=InvalidTokenError("not found"),
|
||||
)
|
||||
|
||||
assert await ExternalAPITokenVerifier().verify_token("agpt_xt_forged") is None
|
||||
validate_key.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token",
|
||||
["agpt_key", f"{ACCESS_TOKEN_PREFIX}token"],
|
||||
ids=["api-key", "access-token"],
|
||||
)
|
||||
@pytest.mark.parametrize("state, admitted", [("banned", False), ("unbanned", True)])
|
||||
async def test_a_banned_owners_credential_gets_no_mcp_session(
|
||||
mocker: pytest_mock.MockFixture, auth_row, token: str, state: str, admitted: bool
|
||||
) -> None:
|
||||
owner = APIAuthorizationInfo(
|
||||
user_id=auth_row(state), scopes=[], type="api_key", created_at=datetime.now(UTC)
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.validate_api_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=owner,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.validate_access_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(owner, None),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.resolve_credential_tenancy",
|
||||
new_callable=AsyncMock,
|
||||
return_value=("org-1", None),
|
||||
)
|
||||
|
||||
verified = await ExternalAPITokenVerifier().verify_token(token)
|
||||
|
||||
assert (verified is not None) is admitted
|
||||
@@ -0,0 +1,519 @@
|
||||
"""
|
||||
V2 External API - MCP Server Endpoint
|
||||
|
||||
Exposes the platform's Copilot tools as an MCP (Model Context Protocol) server,
|
||||
allowing external MCP clients (Claude Desktop, Cursor, etc.) to interact with
|
||||
agents, runs, library, and other platform features programmatically.
|
||||
|
||||
Uses Streamable HTTP transport with stateless sessions, authenticated via the
|
||||
same API key / OAuth bearer token mechanism as the rest of the external API.
|
||||
|
||||
Tool visibility is scope-aware: ``tools/list`` only returns tools whose
|
||||
required permissions are satisfied by the caller's API key / OAuth token.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, AsyncIterator, Optional, Sequence
|
||||
|
||||
import pydantic
|
||||
from mcp.server.auth.middleware.auth_context import get_access_token
|
||||
from mcp.server.auth.provider import AccessToken, TokenVerifier
|
||||
from mcp.server.auth.settings import AuthSettings
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
from mcp.server.fastmcp.server import Context
|
||||
from mcp.server.fastmcp.tools.base import Tool as MCPTool
|
||||
from mcp.server.fastmcp.utilities.func_metadata import ArgModelBase, FuncMetadata
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from mcp.shared.auth import ProtectedResourceMetadata
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import AnyHttpUrl
|
||||
from starlette import status
|
||||
from starlette.applications import Starlette
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from backend.api.external.middleware import (
|
||||
VerifiedCredential,
|
||||
owner_is_banned,
|
||||
verified_credential,
|
||||
)
|
||||
from backend.api.external.v2.errors import error_response
|
||||
from backend.api.external.v2.mcp_calls import (
|
||||
check_arguments,
|
||||
check_plan,
|
||||
check_spend_allowance,
|
||||
check_tool_limits,
|
||||
input_schema,
|
||||
missing_scopes,
|
||||
)
|
||||
from backend.api.external.v2.mcp_tenancy import check_ids_in_tenant
|
||||
from backend.api.external.v2.tenancy import resolve_credential_tenancy
|
||||
from backend.copilot.context import set_turn_unattended
|
||||
from backend.copilot.model import ChatSession
|
||||
from backend.copilot.sdk.tool_adapter import _execute_tool_sync
|
||||
from backend.copilot.tools import TOOL_REGISTRY
|
||||
from backend.copilot.tools.base import BaseTool
|
||||
from backend.data.auth.api_key import validate_api_key
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import (
|
||||
InvalidClientError,
|
||||
InvalidTokenError,
|
||||
is_access_token,
|
||||
validate_access_token,
|
||||
)
|
||||
from backend.util.exceptions import NotAuthorizedError
|
||||
from backend.util.settings import Settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
META_KEY_REQUIRED_SCOPES = "required_scopes"
|
||||
|
||||
# Registry tools deliberately NOT exposed over MCP, with the reason. Every tool
|
||||
# must either opt in via `BaseTool.allow_external_use` or be listed here;
|
||||
# mcp_server_test.py fails on anything unclassified so new tools can't slip
|
||||
# out of the external surface unnoticed.
|
||||
EXTERNAL_USE_EXCLUSIONS: dict[str, str] = {
|
||||
# Copilot-UI or chat-session-only behaviour
|
||||
"TodoWrite": "drives the copilot task checklist UI",
|
||||
"add_understanding": "copilot's own memory of the user",
|
||||
"ask_question": "needs an interactive chat turn to answer",
|
||||
"decompose_goal": "renders a plan card in the copilot UI",
|
||||
"enter_agent_building_mode": "switches the chat session's mode",
|
||||
"schedule_followup": "schedules a copilot chat follow-up, not an agent",
|
||||
"run_sub_session": "copilot sub-agents live inside a chat session",
|
||||
"get_sub_session_result": "copilot sub-agents live inside a chat session",
|
||||
"browser_navigate": "browser state is bound to the chat session",
|
||||
"browser_act": "browser state is bound to the chat session",
|
||||
"browser_screenshot": "browser state is bound to the chat session",
|
||||
"expert_onboarding": "renders the expert intake card in a chat session",
|
||||
"consult_teammate": "teammate checks run inside a chat session",
|
||||
"find_session": "lists the caller's live chat sessions",
|
||||
"message_session": "messages another live chat session",
|
||||
# Expert team flows: preview in one turn, confirm in the next; MCP calls
|
||||
# get a fresh session each time, and no expert API-key permission exists
|
||||
"hire_expert": "two-step confirm needs chat-session state",
|
||||
"raise_expert": "two-step confirm needs chat-session state",
|
||||
"update_expert": "two-step confirm needs chat-session state",
|
||||
"confirm_expert_change": "two-step confirm needs chat-session state",
|
||||
"update_expert_soul": "two-step confirm needs chat-session state",
|
||||
"confirm_expert_soul_update": "two-step confirm needs chat-session state",
|
||||
"handoff_to_expert": "hands off between experts within a chat session",
|
||||
"delegate_to_expert": "no API-key permission for expert teams yet",
|
||||
"list_team": "no API-key permission for expert teams yet",
|
||||
"list_expert_chats": "no API-key permission for expert teams yet",
|
||||
"read_expert_chat": "no API-key permission for expert teams yet",
|
||||
"install_expert_workflow": "no API-key permission for expert teams yet",
|
||||
"remove_expert_workflow": "no API-key permission for expert teams yet",
|
||||
"list_expert_workflows": "no API-key permission for expert teams yet",
|
||||
"list_expert_credentials": "no API-key permission for expert teams yet",
|
||||
"grant_expert_credential": "no API-key permission for expert teams yet",
|
||||
"revoke_expert_credential": "no API-key permission for expert teams yet",
|
||||
"request_credential_grant": "no API-key permission for expert teams yet",
|
||||
"list_routines": "no API-key permission for expert teams yet",
|
||||
"schedule_routine": "no API-key permission for expert teams yet",
|
||||
# Copilot-internal knowledge stores (feature-flagged / prompt-facing)
|
||||
"memory_store": "copilot Graphiti memory",
|
||||
"memory_search": "copilot Graphiti memory",
|
||||
"memory_forget_search": "copilot Graphiti memory",
|
||||
"memory_forget_confirm": "copilot Graphiti memory",
|
||||
"store_skill": "copilot skill registry feeds copilot prompts",
|
||||
"read_skill": "copilot skill registry feeds copilot prompts",
|
||||
"list_skills": "copilot skill registry feeds copilot prompts",
|
||||
"delete_skill": "copilot skill registry feeds copilot prompts",
|
||||
# Capabilities the v2 REST API also leaves out, by decision
|
||||
"find_capability": "discovery for run_capability, which is not exposed",
|
||||
"describe_capability": "discovery for run_capability, which is not exposed",
|
||||
"run_capability": "runs blocks, MCP servers and any deferred tool by id, "
|
||||
"past each tool's own permission",
|
||||
"resume_capability": "resumes a run_capability review",
|
||||
"connect_integration": "OAuth flows are web-UI only in v2",
|
||||
"post_to_chat_platform": "posts as the platform bot; no permission model yet",
|
||||
"edit_chat_platform_message": "posts as the platform bot; no permission model yet",
|
||||
"list_chat_platform_channels": "posts as the platform bot; no permission model yet",
|
||||
"create_feature_request": "writes to the platform's own Linear workspace, on "
|
||||
"the platform's key",
|
||||
# Security
|
||||
"bash_exec": "sandboxed shell on platform infrastructure",
|
||||
"start_desktop": "sandboxed desktop on platform infrastructure",
|
||||
# Not yet classified for external use
|
||||
"pause_schedule": "not yet reviewed for external use",
|
||||
"resume_schedule": "not yet reviewed for external use",
|
||||
}
|
||||
|
||||
# Tools exposed without any permission, with the reason. A tool that spends
|
||||
# platform money or acts through a platform-owned account must name a permission
|
||||
# in its `allow_external_use` instead; mcp_server_test.py fails on any other tool
|
||||
# that opts in with an empty permission list.
|
||||
UNSCOPED_EXTERNAL_TOOLS: dict[str, str] = {
|
||||
"get_doc_page": "serves a published documentation page",
|
||||
"search_docs": "searches published documentation",
|
||||
"get_agent_building_guide": "returns static guidance text",
|
||||
"validate_agent_graph": "pure check over JSON the caller supplied",
|
||||
"fix_agent_graph": "pure rewrite of JSON the caller supplied",
|
||||
"find_agent": "public marketplace listings, as v2's own /marketplace serves",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Server factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# The path `/mcp` ends up on, counting from the origin: v2 is mounted at
|
||||
# /api/v2, the MCP app inside it.
|
||||
MCP_RESOURCE_PATH = "/api/v2/mcp"
|
||||
|
||||
# RFC 9728 §3.1: the well-known segment is inserted between host and resource
|
||||
# path, so this is the URL a client derives — and it is on the origin root.
|
||||
WELL_KNOWN_PROTECTED_RESOURCE_PATH = (
|
||||
f"/.well-known/oauth-protected-resource{MCP_RESOURCE_PATH}"
|
||||
)
|
||||
|
||||
|
||||
def create_mcp_server() -> FastMCP:
|
||||
"""Create the MCP server with all eligible Copilot tools registered."""
|
||||
base_url = _platform_base_url()
|
||||
|
||||
tools = []
|
||||
for tool in TOOL_REGISTRY.values():
|
||||
allowed, required_perms = tool.allow_external_use
|
||||
if not allowed or required_perms is None:
|
||||
reason = EXTERNAL_USE_EXCLUSIONS.get(tool.name, "unclassified")
|
||||
logger.debug(f"Skipping MCP tool {tool.name}: {reason}")
|
||||
continue
|
||||
tools.append(_mcp_tool(tool, required_perms))
|
||||
|
||||
logger.info(
|
||||
f"MCP server created with {len(tools)} tools: {[t.name for t in tools]}"
|
||||
)
|
||||
return _ScopeAwareMCP(
|
||||
name="autogpt-platform",
|
||||
instructions=(
|
||||
"AutoGPT Platform MCP Server. "
|
||||
"Use these tools to find, create, run, and manage AI agents."
|
||||
),
|
||||
token_verifier=ExternalAPITokenVerifier(),
|
||||
auth=AuthSettings(
|
||||
issuer_url=AnyHttpUrl(base_url),
|
||||
resource_server_url=AnyHttpUrl(f"{base_url}{MCP_RESOURCE_PATH}"),
|
||||
),
|
||||
tools=tools,
|
||||
stateless_http=True,
|
||||
streamable_http_path="/",
|
||||
# FastMCP's default only admits localhost Host headers: DNS-rebinding
|
||||
# protection for an unauthenticated server on a developer's machine.
|
||||
# This one answers on the platform's public hosts and needs a bearer
|
||||
# credential on every request, which a rebinding page cannot supply.
|
||||
transport_security=TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=False
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def create_mcp_app() -> Starlette:
|
||||
"""Create the Starlette ASGI app for the MCP server."""
|
||||
server = create_mcp_server()
|
||||
return server.streamable_http_app()
|
||||
|
||||
|
||||
def protected_resource_metadata() -> ProtectedResourceMetadata:
|
||||
"""RFC 9728 metadata for the MCP server, to be served at the ORIGIN root.
|
||||
|
||||
FastMCP registers this document inside the `/mcp` sub-app, so it lands at
|
||||
`/api/v2/mcp/.well-known/oauth-protected-resource/api/v2/mcp`
|
||||
while a client — following the `resource_metadata` FastMCP itself puts in
|
||||
`WWW-Authenticate` — asks the origin root for it and gets a 404. The root
|
||||
app serves this instead; see `WELL_KNOWN_PROTECTED_RESOURCE_PATH`.
|
||||
"""
|
||||
base_url = _platform_base_url()
|
||||
return ProtectedResourceMetadata(
|
||||
resource=AnyHttpUrl(f"{base_url}{MCP_RESOURCE_PATH}"),
|
||||
authorization_servers=[AnyHttpUrl(base_url)],
|
||||
resource_name="AutoGPT Platform MCP Server",
|
||||
)
|
||||
|
||||
|
||||
def _platform_base_url() -> str:
|
||||
return Settings().config.platform_base_url or "https://platform.agpt.co"
|
||||
|
||||
|
||||
class MCPMount:
|
||||
"""The `/mcp` mount: serves the MCP app the host app's lifespan runs.
|
||||
|
||||
Starlette never runs a mounted app's lifespan, and FastMCP's is what starts
|
||||
the session manager every request goes through; mounted directly, every
|
||||
request was a 500 ("Task group is not initialized"). The host app enters
|
||||
`running()` from its own lifespan instead. Each entry builds a new server,
|
||||
because a session manager can only be run once.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._app: Starlette | None = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def running(self) -> AsyncIterator[None]:
|
||||
"""Serve the MCP app while the host runs; a failed start serves 503.
|
||||
|
||||
The REST API must not go down with its MCP endpoint, so a server that
|
||||
fails to build or start is logged and the mount keeps answering 503.
|
||||
"""
|
||||
async with AsyncExitStack() as stack:
|
||||
try:
|
||||
app = create_mcp_app()
|
||||
await stack.enter_async_context(app.router.lifespan_context(app))
|
||||
except Exception:
|
||||
logger.exception("The MCP server failed to start; /mcp answers 503")
|
||||
yield
|
||||
return
|
||||
self._app = app
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._app = None
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if self._app is None:
|
||||
response = error_response(
|
||||
status.HTTP_503_SERVICE_UNAVAILABLE, "The MCP server is not running"
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
if _ambiguous_credentials(scope):
|
||||
response = error_response(
|
||||
status.HTTP_400_BAD_REQUEST,
|
||||
"Send one credential, as a single `Authorization: Bearer` header",
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
# The v2 rate limiter in front of this mount has verified the
|
||||
# credential already; FastMCP verifies it again by token alone.
|
||||
verified = _request_credential.set(verified_credential(scope))
|
||||
try:
|
||||
await self._app(scope, receive, send)
|
||||
finally:
|
||||
_request_credential.reset(verified)
|
||||
|
||||
|
||||
def _ambiguous_credentials(scope: Scope) -> bool:
|
||||
"""Whether the request names more than one credential.
|
||||
|
||||
FastMCP authenticates the first Authorization header, while the v2 rate
|
||||
limiter reads X-API-Key first; with both, or two Authorization headers,
|
||||
the limit would be charged to one credential and the call made as another.
|
||||
"""
|
||||
names = [name.lower() for name, _ in scope.get("headers", [])]
|
||||
authorizations = names.count(b"authorization")
|
||||
return authorizations > 1 or (authorizations == 1 and b"x-api-key" in names)
|
||||
|
||||
|
||||
mcp_mount = MCPMount()
|
||||
|
||||
# What the v2 middleware verified this request's credential as, so the MCP
|
||||
# token verifier doesn't verify (and for an API key, hash) the same token again.
|
||||
_request_credential: ContextVar[Optional[VerifiedCredential]] = ContextVar(
|
||||
"mcp_request_credential", default=None
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scope-aware FastMCP subclass
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _ScopeAwareMCP(FastMCP):
|
||||
"""FastMCP subclass that filters ``tools/list`` by the caller's scopes."""
|
||||
|
||||
async def list_tools(self) -> list:
|
||||
all_tools = await super().list_tools()
|
||||
access_token = get_access_token()
|
||||
if not access_token:
|
||||
return []
|
||||
caller_scopes = set(access_token.scopes)
|
||||
return [
|
||||
t
|
||||
for t in all_tools
|
||||
if caller_scopes.issuperset(
|
||||
(t.meta or {}).get(META_KEY_REQUIRED_SCOPES, [])
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token verification — reuses existing external API auth infrastructure
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TenantedAccessToken(AccessToken):
|
||||
"""An MCP access token carrying the tenant its credentials act in.
|
||||
|
||||
MCP has no per-request dependency chain, so the tenant is resolved here,
|
||||
once per token verification, instead of by `v2.tenancy`.
|
||||
"""
|
||||
|
||||
organization_id: str
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
class ExternalAPITokenVerifier(TokenVerifier):
|
||||
"""Validates API keys and OAuth tokens via external API auth."""
|
||||
|
||||
async def verify_token(self, token: str) -> AccessToken | None:
|
||||
verified = _request_credential.get()
|
||||
if verified is not None and verified.credential == token:
|
||||
return await self._tenanted(token, verified.auth) if verified.auth else None
|
||||
|
||||
# Try API key first, unless it is in the OAuth token format
|
||||
if not is_access_token(token) and (
|
||||
api_key_info := await validate_api_key(token)
|
||||
):
|
||||
return await self._tenanted(token, api_key_info)
|
||||
|
||||
# Try OAuth bearer token
|
||||
try:
|
||||
token_info, _ = await validate_access_token(token)
|
||||
return await self._tenanted(token, token_info)
|
||||
except (InvalidClientError, InvalidTokenError):
|
||||
return None
|
||||
|
||||
async def _tenanted(
|
||||
self, token: str, auth: APIAuthorizationInfo
|
||||
) -> TenantedAccessToken | None:
|
||||
if await owner_is_banned(auth):
|
||||
logger.warning(f"Rejecting MCP token for banned user {auth.user_id}")
|
||||
return None
|
||||
try:
|
||||
organization_id, team_id = await resolve_credential_tenancy(auth)
|
||||
except NotAuthorizedError as e:
|
||||
logger.warning(f"Rejecting MCP token for {auth.user_id}: {e}")
|
||||
return None
|
||||
return TenantedAccessToken(
|
||||
token=token,
|
||||
client_id=auth.user_id,
|
||||
scopes=[s.value for s in auth.scopes],
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _create_tool_handler(
|
||||
tool: BaseTool,
|
||||
required_scopes: Sequence[str],
|
||||
):
|
||||
"""Create an async MCP tool handler that wraps a BaseTool subclass.
|
||||
|
||||
Before the tool runs, the call is held to the caller's credential: its
|
||||
arguments to the advertised schema, its scopes to ``required_scopes`` plus
|
||||
whatever the arguments add, the ids it names to the credential's
|
||||
organization, the caller to an active plan, and a tool that spends platform
|
||||
money to the caller's allowance.
|
||||
"""
|
||||
|
||||
async def handler(ctx: Context, **kwargs: Any) -> str:
|
||||
# Raised, not returned: a rejection returned as content is reported to
|
||||
# the client as a successful call whose text happens to say "denied".
|
||||
access_token = get_access_token()
|
||||
if not isinstance(access_token, TenantedAccessToken):
|
||||
raise ToolError("Authentication required")
|
||||
|
||||
check_arguments(tool, kwargs)
|
||||
if missing := missing_scopes(
|
||||
tool, required_scopes, kwargs, access_token.scopes
|
||||
):
|
||||
raise ToolError(f"Missing required permission(s): {', '.join(missing)}")
|
||||
|
||||
user_id = access_token.client_id
|
||||
await check_ids_in_tenant(
|
||||
tool.name, kwargs, user_id, access_token.organization_id
|
||||
)
|
||||
await check_tool_limits(tool, kwargs, user_id, access_token.organization_id)
|
||||
await check_plan(user_id)
|
||||
if tool.spends_platform_money:
|
||||
await check_spend_allowance(user_id)
|
||||
|
||||
# An automation session: nobody is watching to answer an approval
|
||||
# card, so the credential's scopes are the authorization, as in REST.
|
||||
session = ChatSession.new(
|
||||
user_id,
|
||||
dry_run=False,
|
||||
origin="automation",
|
||||
organization_id=access_token.organization_id,
|
||||
team_id=access_token.team_id,
|
||||
)
|
||||
session.external_caller = True
|
||||
session.external_permissions = frozenset(
|
||||
APIKeyPermission(scope) for scope in access_token.scopes
|
||||
)
|
||||
# Tools read the turn, not the origin: without this a missing credential
|
||||
# comes back as a setup card, and an ambiguous one as a question.
|
||||
set_turn_unattended(session, scheduled=False)
|
||||
|
||||
try:
|
||||
result = await _execute_tool_sync(tool, user_id, session, kwargs)
|
||||
except Exception:
|
||||
# As in chat: the caller gets no internal exception text.
|
||||
logger.exception(f"MCP tool {tool.name} failed")
|
||||
raise ToolError(f"{tool.name} failed unexpectedly; retry shortly")
|
||||
|
||||
parts = []
|
||||
for block in result.get("content", []):
|
||||
if block.get("type") == "text":
|
||||
parts.append(block["text"])
|
||||
text = "\n".join(parts)
|
||||
# A refusal or failure must reach the client as one, not as a result.
|
||||
if result.get("isError"):
|
||||
raise ToolError(text)
|
||||
return text
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def _mcp_tool(tool: BaseTool, required_perms: Sequence[APIKeyPermission]) -> MCPTool:
|
||||
"""Build an MCP tool from a Copilot tool.
|
||||
|
||||
Constructed directly rather than through `FastMCP.add_tool`, which derives
|
||||
the schema from the handler's signature; ours is `**kwargs` and the real
|
||||
schema comes from the Copilot tool.
|
||||
"""
|
||||
required_scopes = [p.value for p in required_perms]
|
||||
return MCPTool(
|
||||
fn=_create_tool_handler(tool, required_scopes),
|
||||
name=tool.name,
|
||||
title=None,
|
||||
description=tool.description,
|
||||
parameters=input_schema(tool),
|
||||
fn_metadata=_PASSTHROUGH_META,
|
||||
is_async=True,
|
||||
context_kwarg="ctx",
|
||||
annotations=None,
|
||||
meta={META_KEY_REQUIRED_SCOPES: required_scopes},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Passthrough arg model — lets us specify JSON Schema directly instead of
|
||||
# having FastMCP introspect the handler function's signature.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _PassthroughArgs(ArgModelBase):
|
||||
"""Accepts any fields and passes them through as kwargs."""
|
||||
|
||||
model_config = pydantic.ConfigDict(extra="allow")
|
||||
|
||||
def model_dump_one_level(self, **_kwargs: Any) -> dict[str, Any]:
|
||||
return dict(self.__pydantic_extra__ or {})
|
||||
|
||||
|
||||
_PASSTHROUGH_META = FuncMetadata(
|
||||
arg_model=_PassthroughArgs,
|
||||
output_schema=None,
|
||||
output_model=None,
|
||||
wrap_output=False,
|
||||
)
|
||||
@@ -0,0 +1,712 @@
|
||||
"""Keep the MCP tool surface deliberate: every Copilot tool is classified."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException
|
||||
from mcp.server.auth.provider import AccessToken
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
from mcp.types import CallToolRequest, CallToolRequestParams, ListToolsRequest
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette.routing import Match
|
||||
|
||||
from backend.api.external.v2.mcp_calls import TOOL_LIMITERS, _starts_a_paid_run
|
||||
from backend.api.external.v2.mcp_server import (
|
||||
EXTERNAL_USE_EXCLUSIONS,
|
||||
META_KEY_REQUIRED_SCOPES,
|
||||
UNSCOPED_EXTERNAL_TOOLS,
|
||||
WELL_KNOWN_PROTECTED_RESOURCE_PATH,
|
||||
TenantedAccessToken,
|
||||
_ambiguous_credentials,
|
||||
_create_tool_handler,
|
||||
create_mcp_server,
|
||||
protected_resource_metadata,
|
||||
)
|
||||
from backend.api.external.v2.rate_limit import file_upload_limiter
|
||||
from backend.copilot import usage_activation_fixtures
|
||||
from backend.copilot.gate.classifier import Judgement
|
||||
from backend.copilot.model import ChatSession
|
||||
from backend.copilot.rate_limit import RateLimitUnavailable
|
||||
from backend.copilot.tools import TOOL_REGISTRY
|
||||
from backend.copilot.tools.base import CARD_HINTS
|
||||
from backend.copilot.tools.models import ErrorResponse
|
||||
from backend.util.exceptions import GraphValidationError
|
||||
|
||||
|
||||
def _token(*scopes: APIKeyPermission) -> TenantedAccessToken:
|
||||
return TenantedAccessToken(
|
||||
token="agpt_test",
|
||||
client_id="user-1",
|
||||
scopes=[s.value for s in scopes],
|
||||
organization_id="org-1",
|
||||
team_id="team-1",
|
||||
)
|
||||
|
||||
|
||||
def _tool_scopes(tool_name: str) -> list[APIKeyPermission]:
|
||||
return list(TOOL_REGISTRY[tool_name].allow_external_use[1] or [])
|
||||
|
||||
|
||||
async def _call(tool_name: str, arguments: dict[str, Any]):
|
||||
server = create_mcp_server()
|
||||
result = await server._mcp_server.request_handlers[CallToolRequest](
|
||||
CallToolRequest(
|
||||
method="tools/call",
|
||||
params=CallToolRequestParams(name=tool_name, arguments=arguments),
|
||||
)
|
||||
)
|
||||
return result.root
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_tenancy_lookups(mocker: pytest_mock.MockerFixture) -> mock.AsyncMock:
|
||||
"""Calls whose ids are all in the caller's organization."""
|
||||
return mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.check_ids_in_tenant",
|
||||
new_callable=mock.AsyncMock,
|
||||
)
|
||||
|
||||
|
||||
usage_snapshot = usage_activation_fixtures.usage_snapshot
|
||||
|
||||
|
||||
async def test_the_server_carries_every_opted_in_tool_with_its_scopes():
|
||||
"""Registration goes through FastMCP's own constructor, not its internals.
|
||||
|
||||
Writing into `_tool_manager._tools` worked by accident of mcp 1.26.0's
|
||||
layout; this fails the moment the supported path stops carrying the tools.
|
||||
"""
|
||||
server = create_mcp_server()
|
||||
|
||||
exposed = {name for name, t in TOOL_REGISTRY.items() if t.allow_external_use[0]}
|
||||
# The base listing: the subclass's own filters by the caller's scopes.
|
||||
registered = {t.name: t for t in await FastMCP.list_tools(server)}
|
||||
assert set(registered) == exposed
|
||||
|
||||
for name, tool in registered.items():
|
||||
expected = [p.value for p in (TOOL_REGISTRY[name].allow_external_use[1] or [])]
|
||||
assert (tool.meta or {}).get(META_KEY_REQUIRED_SCOPES) == expected
|
||||
|
||||
|
||||
async def test_the_advertised_schema_names_required_arguments_and_nothing_extra():
|
||||
registered = {t.name: t for t in await FastMCP.list_tools(create_mcp_server())}
|
||||
|
||||
move = registered["move_agents_to_folder"].inputSchema
|
||||
assert move["additionalProperties"] is False
|
||||
assert move["required"] == ["agent_ids"]
|
||||
|
||||
# Presets have no v2 permission, and the sandbox paths no sandbox.
|
||||
assert "preset_id" not in registered["run_agent"].inputSchema["properties"]
|
||||
# A blocking wait holds a request and an event-bus connection; poll instead.
|
||||
assert "wait_for_result" not in registered["run_agent"].inputSchema["properties"]
|
||||
# A dry run runs models on the platform's key the allowance doesn't meter.
|
||||
assert "dry_run" not in registered["run_agent"].inputSchema["properties"]
|
||||
assert "wait_if_running" not in (
|
||||
registered["view_agent_output"].inputSchema["properties"]
|
||||
)
|
||||
assert "save_to_path" not in (
|
||||
registered["read_workspace_file"].inputSchema["properties"]
|
||||
)
|
||||
assert "source_path" not in (
|
||||
registered["write_workspace_file"].inputSchema["properties"]
|
||||
)
|
||||
|
||||
|
||||
async def test_the_advertised_schema_carries_none_of_the_approval_cards_hints():
|
||||
"""The card's label and entity hints are for the web app, and `"entity":
|
||||
null` is not JSON Schema a client must accept."""
|
||||
for tool in await FastMCP.list_tools(create_mcp_server()):
|
||||
for name, spec in tool.inputSchema["properties"].items():
|
||||
assert not CARD_HINTS & set(spec), f"{tool.name}.{name}"
|
||||
|
||||
|
||||
def test_create_feature_request_stays_off_the_external_surface():
|
||||
"""It writes to the platform's own Linear workspace on the platform's key."""
|
||||
assert not TOOL_REGISTRY["create_feature_request"].allow_external_use[0]
|
||||
assert "create_feature_request" in EXTERNAL_USE_EXCLUSIONS
|
||||
|
||||
|
||||
def test_the_protected_resource_document_answers_where_clients_look():
|
||||
"""RFC 9728 discovery, at the URL the `WWW-Authenticate` header names.
|
||||
|
||||
FastMCP registers its own copy inside the `/mcp` mount, so the derived URL
|
||||
— well-known segment first, resource path after — 404s unless the root app
|
||||
serves it. This asserts against the real app, not a stand-in.
|
||||
"""
|
||||
from mcp.server.auth.routes import build_resource_metadata_url
|
||||
|
||||
from backend.api.rest_api import app
|
||||
|
||||
metadata = protected_resource_metadata()
|
||||
derived = urlparse(str(build_resource_metadata_url(metadata.resource)))
|
||||
assert derived.path == WELL_KNOWN_PROTECTED_RESOURCE_PATH
|
||||
|
||||
scope = {"type": "http", "method": "GET", "path": derived.path, "headers": []}
|
||||
assert any(
|
||||
route.matches(scope)[0] == Match.FULL for route in app.routes
|
||||
), f"nothing on the root app answers {derived.path}"
|
||||
|
||||
|
||||
async def test_a_rejected_tool_call_is_an_error_not_a_successful_denial(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Returned text is reported to the client as a call that succeeded."""
|
||||
handler = _create_tool_handler(TOOL_REGISTRY["list_folders"], ["READ_GRAPH"])
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token", return_value=None
|
||||
)
|
||||
with pytest.raises(ToolError, match="Authentication required"):
|
||||
await handler(ctx=mock.Mock())
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(APIKeyPermission.READ_RUN),
|
||||
)
|
||||
with pytest.raises(ToolError, match="READ_GRAPH"):
|
||||
await handler(ctx=mock.Mock())
|
||||
|
||||
|
||||
def test_every_tool_is_either_exposed_or_explicitly_excluded():
|
||||
exposed = {name for name, t in TOOL_REGISTRY.items() if t.allow_external_use[0]}
|
||||
unclassified = set(TOOL_REGISTRY) - exposed - set(EXTERNAL_USE_EXCLUSIONS)
|
||||
assert not unclassified, (
|
||||
"Tools neither opted in via allow_external_use nor listed in "
|
||||
f"EXTERNAL_USE_EXCLUSIONS: {sorted(unclassified)}"
|
||||
)
|
||||
|
||||
|
||||
def test_exclusion_list_has_no_stale_or_contradictory_entries():
|
||||
unknown = set(EXTERNAL_USE_EXCLUSIONS) - set(TOOL_REGISTRY)
|
||||
assert not unknown, f"Excluded tools that no longer exist: {sorted(unknown)}"
|
||||
|
||||
contradictory = [
|
||||
name
|
||||
for name in EXTERNAL_USE_EXCLUSIONS
|
||||
if TOOL_REGISTRY[name].allow_external_use[0]
|
||||
]
|
||||
assert (
|
||||
not contradictory
|
||||
), f"Tools both opted in and excluded (drop one): {sorted(contradictory)}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tools_list_leaves_out_tools_the_caller_has_no_scope_for(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
):
|
||||
granted = {APIKeyPermission.READ_LIBRARY}
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=SimpleNamespace(scopes=[p.value for p in granted]),
|
||||
)
|
||||
exposed = {
|
||||
name: set(perms)
|
||||
for name, t in TOOL_REGISTRY.items()
|
||||
for allowed, perms in [t.allow_external_use]
|
||||
if allowed
|
||||
}
|
||||
server = create_mcp_server()
|
||||
|
||||
result = await server._mcp_server.request_handlers[ListToolsRequest](
|
||||
ListToolsRequest(method="tools/list")
|
||||
)
|
||||
|
||||
listed = {t.name for t in result.root.tools}
|
||||
assert listed == {name for name, perms in exposed.items() if perms <= granted}
|
||||
assert listed < set(exposed)
|
||||
|
||||
|
||||
def test_exposed_tools_declare_permissions_as_a_sequence():
|
||||
for name, tool in TOOL_REGISTRY.items():
|
||||
allowed, perms = tool.allow_external_use
|
||||
if allowed:
|
||||
assert perms is not None, f"{name} opted in without a permission list"
|
||||
|
||||
|
||||
def test_no_tool_is_exposed_unscoped_without_a_stated_reason():
|
||||
"""An empty permission list means any key can drive the tool.
|
||||
|
||||
That is right for published docs and public listings and wrong for anything
|
||||
that spends platform money or acts through a platform-owned account, so the
|
||||
open set is enumerated rather than inferred.
|
||||
"""
|
||||
unscoped = {
|
||||
name
|
||||
for name, tool in TOOL_REGISTRY.items()
|
||||
if tool.allow_external_use[0] and not tool.allow_external_use[1]
|
||||
}
|
||||
assert unscoped == set(UNSCOPED_EXTERNAL_TOOLS), (
|
||||
"unlisted tools exposed with no permission: "
|
||||
f"{sorted(unscoped - set(UNSCOPED_EXTERNAL_TOOLS))}; "
|
||||
"listed but no longer unscoped: "
|
||||
f"{sorted(set(UNSCOPED_EXTERNAL_TOOLS) - unscoped)}"
|
||||
)
|
||||
|
||||
|
||||
def test_no_tool_that_spends_platform_money_is_unscoped():
|
||||
for name, tool in TOOL_REGISTRY.items():
|
||||
allowed, perms = tool.allow_external_use
|
||||
if allowed and tool.spends_platform_money:
|
||||
assert perms, f"{name} spends platform money but needs no permission"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token, refusal",
|
||||
[
|
||||
(None, "Authentication required"),
|
||||
# Without a resolved organization the call couldn't be confined to one.
|
||||
(
|
||||
AccessToken(
|
||||
token="agpt_test", client_id="user-1", scopes=["WRITE_LIBRARY"]
|
||||
),
|
||||
"Authentication required",
|
||||
),
|
||||
(_token(), "Missing required permission"),
|
||||
],
|
||||
)
|
||||
async def test_a_call_without_the_right_credentials_reaches_the_client_as_an_error(
|
||||
mocker: pytest_mock.MockerFixture, token, refusal: str
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token", return_value=token
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["delete_folder"]), "_execute")
|
||||
|
||||
result = await _call("delete_folder", {"folder_id": "folder-1"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert refusal in result.content[0].text
|
||||
|
||||
|
||||
async def test_an_account_without_a_plan_cannot_call_a_tool(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
paywalled_user: str,
|
||||
):
|
||||
# Every call, not only one that spends platform money: the tools are
|
||||
# Otto's, and every chat turn needs a plan.
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("delete_folder")),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["delete_folder"]), "_execute")
|
||||
|
||||
result = await _call("delete_folder", {"folder_id": "folder-1"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert "needs an active subscription" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments, unknown",
|
||||
[
|
||||
# Presets have no v2 permission.
|
||||
({"library_agent_id": "agent-1", "preset_id": "preset-1"}, "preset_id"),
|
||||
# A dry run runs models on the platform's key the allowance doesn't meter.
|
||||
({"library_agent_id": "agent-1", "dry_run": True}, "dry_run"),
|
||||
# Not in the schema at all; the tool's own input model has the field.
|
||||
({"library_agent_id": "agent-1", "save_as_preset": True}, "save_as_preset"),
|
||||
# The marker the approval gate passes to a tool it already approved.
|
||||
({"library_agent_id": "agent-1", "_gate_approved": True}, "_gate_approved"),
|
||||
],
|
||||
)
|
||||
async def test_an_argument_the_tool_does_not_advertise_is_refused(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
arguments: dict,
|
||||
unknown: str,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("run_agent")),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["run_agent"]), "_execute")
|
||||
|
||||
result = await _call("run_agent", arguments)
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert f"Unknown argument(s): {unknown}" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments, missing",
|
||||
[
|
||||
({"library_agent_id": "agent-1", "cron": "0 9 * * *"}, "WRITE_SCHEDULE"),
|
||||
({"library_agent_id": "agent-1", "schedule_name": "daily"}, "WRITE_SCHEDULE"),
|
||||
({"username_agent_slug": "someone/agent"}, "WRITE_LIBRARY"),
|
||||
],
|
||||
)
|
||||
async def test_a_branch_needs_the_permission_for_what_it_does(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
arguments: dict,
|
||||
missing: str,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(APIKeyPermission.RUN_AGENT),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["run_agent"]), "_execute")
|
||||
|
||||
result = await _call("run_agent", arguments)
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert f"Missing required permission(s): {missing}" in result.content[0].text
|
||||
|
||||
|
||||
async def test_an_id_in_another_organization_stops_the_call_before_the_tool(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(APIKeyPermission.WRITE_LIBRARY),
|
||||
)
|
||||
library = mocker.patch("backend.api.external.v2.mcp_tenancy.library_db")
|
||||
library.return_value.get_folder = mock.AsyncMock(
|
||||
return_value=SimpleNamespace(organization_id="org-2")
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["delete_folder"]), "_execute")
|
||||
|
||||
result = await _call("delete_folder", {"folder_id": "folder-1"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert "Folder 'folder-1' not found" in result.content[0].text
|
||||
library.return_value.get_folder.assert_awaited_once_with("folder-1", "user-1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"allowance_failure, refusal",
|
||||
[
|
||||
({"paywalled": True}, "needs an active subscription"),
|
||||
(
|
||||
{"rate_limit": RateLimitUnavailable("redis down")},
|
||||
"Usage limits are unavailable",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_a_tool_that_spends_platform_money_checks_the_allowance_first(
|
||||
mocker: pytest_mock.MockerFixture, allowance_failure: dict, refusal: str
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(APIKeyPermission.USE_TOOLS),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.is_user_paywalled",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=allowance_failure.get("paywalled", False),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.get_global_rate_limits",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=(1, 1, None),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.check_rate_limit",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=allowance_failure.get("rate_limit"),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["web_search"]), "_execute")
|
||||
|
||||
result = await _call("web_search", {"query": "weather"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert refusal in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args, paid_run",
|
||||
[
|
||||
({}, True),
|
||||
({"cron": "0 9 * * *"}, False),
|
||||
({"schedule_name": "Daily"}, False),
|
||||
# Read as the tool's own model reads them: a blank cron is no schedule
|
||||
# to the tool, which starts a real run.
|
||||
({"cron": " "}, True),
|
||||
({"schedule_name": " "}, True),
|
||||
({"dry_run": True}, False),
|
||||
# Refused by the tool, so nothing starts.
|
||||
({"cron": ["0 9 * * *"]}, True),
|
||||
],
|
||||
)
|
||||
def test_a_run_agent_call_is_a_paid_run_as_the_tool_reads_it(
|
||||
args: dict[str, Any], paid_run: bool
|
||||
):
|
||||
assert _starts_a_paid_run(args) is paid_run
|
||||
|
||||
|
||||
async def test_a_tool_called_over_mcp_runs_without_waiting_for_an_approval(
|
||||
mocker: pytest_mock.MockerFixture, no_tenancy_lookups: mock.AsyncMock
|
||||
):
|
||||
"""Nobody is watching an MCP call to answer an approval card.
|
||||
|
||||
The session is an automation's, as a scheduled turn's is, so the gate stays
|
||||
out of it and the credential's scopes are the authorization, as in REST.
|
||||
"""
|
||||
tool = TOOL_REGISTRY["delete_folder"]
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("delete_folder")),
|
||||
)
|
||||
mocker.patch("backend.copilot.gate.is_feature_enabled", return_value=True)
|
||||
mocker.patch(
|
||||
"backend.copilot.gate.supervise",
|
||||
return_value=Judgement(allowed=False, reason="Nobody asked for it."),
|
||||
)
|
||||
open_review = mocker.patch(
|
||||
"backend.copilot.gate.review_store.open_review", return_value=True
|
||||
)
|
||||
run = mocker.patch.object(
|
||||
type(tool),
|
||||
"_execute",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=ErrorResponse(message="stub", session_id=None),
|
||||
)
|
||||
|
||||
await _call("delete_folder", {"folder_id": "folder-1"})
|
||||
|
||||
open_review.assert_not_called()
|
||||
run.assert_awaited_once()
|
||||
session = run.await_args.args[1]
|
||||
assert isinstance(session, ChatSession)
|
||||
assert session.metadata.origin == "automation"
|
||||
assert session.external_caller
|
||||
assert session.external_permissions == frozenset(_tool_scopes("delete_folder"))
|
||||
assert (session.organization_id, session.team_id) == ("org-1", "team-1")
|
||||
no_tenancy_lookups.assert_awaited_once_with(
|
||||
"delete_folder", {"folder_id": "folder-1"}, "user-1", "org-1"
|
||||
)
|
||||
|
||||
|
||||
async def test_a_missing_credential_comes_back_as_an_error_not_a_setup_card(
|
||||
mocker: pytest_mock.MockerFixture, no_tenancy_lookups: mock.AsyncMock
|
||||
):
|
||||
"""Nobody can answer a setup card an MCP call puts to the user, so a run
|
||||
whose credentials have gone must fail naming what to connect."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("run_agent")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
TOOL_LIMITERS["run_agent"], "check", new_callable=mock.AsyncMock
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.get_credit_model",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(get_credits=mock.AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.copilot.tools.run_agent.build_missing_credentials_from_graph",
|
||||
return_value={"credentials": {"provider": "firecrawl"}},
|
||||
)
|
||||
graph = mock.MagicMock(id="graph-1", version=1, input_schema={})
|
||||
graph.name = "Daily Scraper"
|
||||
|
||||
async def credentials_gone(self, user_id, session, **kwargs):
|
||||
return await self._handle_graph_validation_race(
|
||||
error=GraphValidationError(
|
||||
message="Graph is invalid",
|
||||
node_errors={"node": {"credentials": "These credentials are required"}},
|
||||
),
|
||||
graph=graph,
|
||||
user_id=user_id,
|
||||
session_id=session.session_id,
|
||||
action_verb="running",
|
||||
)
|
||||
|
||||
mocker.patch.object(type(TOOL_REGISTRY["run_agent"]), "_execute", credentials_gone)
|
||||
|
||||
result = await _call("run_agent", {"library_agent_id": "agent-1"})
|
||||
|
||||
reply = json.loads(result.content[0].text)
|
||||
assert reply["type"] == "error"
|
||||
assert "firecrawl" in reply["message"]
|
||||
|
||||
|
||||
async def test_a_tool_is_held_to_the_cap_of_the_rest_route_it_mirrors(
|
||||
mocker: pytest_mock.MockerFixture, no_tenancy_lookups: mock.AsyncMock
|
||||
):
|
||||
"""REST caps runs at 60 a minute; MCP's global 200 must not be a way round."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("run_agent")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
TOOL_LIMITERS["run_agent"],
|
||||
"check",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=HTTPException(status_code=429, detail="Rate limit exceeded"),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["run_agent"]), "_execute")
|
||||
|
||||
result = await _call("run_agent", {"library_agent_id": "agent-1"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert "Rate limit exceeded" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_name, arguments",
|
||||
[
|
||||
("fix_agent_graph", {"agent_json": {"nodes": []}, "write_to": "fixed.json"}),
|
||||
("find_library_agent", {"query": "mail", "write_graph_to": "graph.json"}),
|
||||
],
|
||||
)
|
||||
async def test_an_argument_that_writes_a_file_is_held_to_the_upload_cap(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
):
|
||||
"""REST caps uploads at 20 per 5 minutes; writing the file through a tool
|
||||
argument must not be a way round it."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes(tool_name), APIKeyPermission.WRITE_FILES),
|
||||
)
|
||||
for limiter in TOOL_LIMITERS.values():
|
||||
mocker.patch.object(limiter, "check", new_callable=mock.AsyncMock)
|
||||
mocker.patch.object(
|
||||
file_upload_limiter,
|
||||
"check",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=HTTPException(status_code=429, detail="Rate limit exceeded"),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY[tool_name]), "_execute")
|
||||
|
||||
result = await _call(tool_name, arguments)
|
||||
|
||||
run.assert_not_called()
|
||||
assert result.isError
|
||||
assert "Rate limit exceeded" in result.content[0].text
|
||||
|
||||
|
||||
async def test_a_call_that_writes_no_file_is_not_charged_an_upload(
|
||||
mocker: pytest_mock.MockerFixture, no_tenancy_lookups: mock.AsyncMock
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("find_library_agent")),
|
||||
)
|
||||
search = mocker.patch.object(
|
||||
TOOL_LIMITERS["find_library_agent"], "check", new_callable=mock.AsyncMock
|
||||
)
|
||||
upload = mocker.patch.object(
|
||||
file_upload_limiter, "check", new_callable=mock.AsyncMock
|
||||
)
|
||||
mocker.patch.object(
|
||||
type(TOOL_REGISTRY["find_library_agent"]),
|
||||
"_execute",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=ErrorResponse(message="No agents found"),
|
||||
)
|
||||
|
||||
await _call("find_library_agent", {"query": "mail"})
|
||||
|
||||
search.assert_awaited_once()
|
||||
upload.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_run_on_a_zero_balance_is_refused_as_rest_refuses_it(
|
||||
mocker: pytest_mock.MockerFixture, no_tenancy_lookups: mock.AsyncMock
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("run_agent")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
TOOL_LIMITERS["run_agent"], "check", new_callable=mock.AsyncMock
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.get_credit_model",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(get_credits=mock.AsyncMock(return_value=0)),
|
||||
)
|
||||
run = mocker.patch.object(type(TOOL_REGISTRY["run_agent"]), "_execute")
|
||||
|
||||
result = await _call("run_agent", {"library_agent_id": "agent-1"})
|
||||
|
||||
run.assert_not_called()
|
||||
assert "Insufficient balance" in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_name", ["validate_agent_graph", "fix_agent_graph"])
|
||||
async def test_a_malformed_graph_is_refused_as_such_not_as_an_internal_error(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
tool_name: str,
|
||||
):
|
||||
"""The schema lets `agent_json` be a string, parsed only inside the tool;
|
||||
a node that isn't an object reached the validator, whose crash text came
|
||||
back to the client."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes(tool_name)),
|
||||
)
|
||||
|
||||
result = await _call(tool_name, {"agent_json": '{"nodes": [1], "links": []}'})
|
||||
|
||||
text = result.content[0].text
|
||||
assert "must be a list of objects" in text
|
||||
assert "attribute" not in text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
[(b"authorization", b"Bearer agpt_a"), (b"x-api-key", b"agpt_b")],
|
||||
[(b"authorization", b"Bearer agpt_a"), (b"authorization", b"Bearer agpt_b")],
|
||||
],
|
||||
)
|
||||
def test_a_request_naming_two_credentials_is_refused(headers: list) -> None:
|
||||
"""The limiter and FastMCP would each pick a different one."""
|
||||
assert _ambiguous_credentials({"type": "http", "headers": headers})
|
||||
|
||||
|
||||
def test_a_single_bearer_credential_is_accepted() -> None:
|
||||
assert not _ambiguous_credentials(
|
||||
{"type": "http", "headers": [(b"authorization", b"Bearer agpt_a")]}
|
||||
)
|
||||
|
||||
|
||||
async def test_a_tool_that_fails_unexpectedly_returns_no_internal_text(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
no_tenancy_lookups: mock.AsyncMock,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""The failure is logged here; the client gets a generic error."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server.get_access_token",
|
||||
return_value=_token(*_tool_scopes("get_agent_building_guide")),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_calls.is_user_paywalled",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=False,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.mcp_server._execute_tool_sync",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=RuntimeError("relation /srv/internal.sql does not exist"),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
result = await _call("get_agent_building_guide", {})
|
||||
|
||||
assert result.isError
|
||||
assert "internal.sql" not in result.content[0].text
|
||||
assert "internal.sql" in caplog.text
|
||||
@@ -0,0 +1,148 @@
|
||||
"""Keep the ids an MCP call names inside the organization it acts in.
|
||||
|
||||
The Copilot tools look rows up by user. A v2 credential is bound to one
|
||||
organization, and the REST API answers 404 for a row tagged with another
|
||||
(`tenancy.in_tenant`); this gives an MCP call the same answer, before the tool
|
||||
runs. Tools that list or search on their own filter what they return through
|
||||
`copilot.tools.external_scope`.
|
||||
|
||||
Every argument of an exposed tool that looks like an id is classified below;
|
||||
mcp_tenancy_test.py fails on one that isn't.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
|
||||
from backend.copilot.tools.external_scope import in_tenant
|
||||
from backend.data.db_accessors import execution_db, graph_db, library_db
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IdKind(Enum):
|
||||
LIBRARY_AGENT = "Library agent"
|
||||
FOLDER = "Folder"
|
||||
RUN = "Run"
|
||||
PRESET = "Preset"
|
||||
|
||||
|
||||
AGENT = IdKind.LIBRARY_AGENT
|
||||
FOLDER = IdKind.FOLDER
|
||||
RUN = IdKind.RUN
|
||||
PRESET = IdKind.PRESET
|
||||
|
||||
# Per tool, the arguments that name an organization's row. An agent id may be
|
||||
# a library agent id or a graph id, as the tools accept either.
|
||||
TENANTED_ARGS: dict[str, dict[str, IdKind]] = {
|
||||
"create_agent": {"library_agent_ids": AGENT, "folder_id": FOLDER},
|
||||
"customize_agent": {"library_agent_ids": AGENT, "folder_id": FOLDER},
|
||||
"edit_agent": {"agent_id": AGENT, "library_agent_ids": AGENT},
|
||||
"find_library_agent": {"agent_id": AGENT},
|
||||
"run_agent": {"library_agent_id": AGENT},
|
||||
"view_agent_output": {"library_agent_id": AGENT, "execution_id": RUN},
|
||||
"list_schedules": {"library_agent_id": AGENT},
|
||||
"create_folder": {"parent_id": FOLDER},
|
||||
"list_folders": {"parent_id": FOLDER},
|
||||
"update_folder": {"folder_id": FOLDER},
|
||||
"delete_folder": {"folder_id": FOLDER},
|
||||
"move_folder": {"folder_id": FOLDER, "target_parent_id": FOLDER},
|
||||
"move_agents_to_folder": {"agent_ids": AGENT, "folder_id": FOLDER},
|
||||
"list_presets": {"library_agent_id": AGENT, "graph_id": AGENT},
|
||||
"update_preset": {"preset_id": PRESET},
|
||||
"delete_preset": {"preset_id": PRESET},
|
||||
"list_agent_triggers": {"library_agent_id": AGENT},
|
||||
"setup_agent_webhook_trigger": {"library_agent_id": AGENT, "graph_id": AGENT},
|
||||
}
|
||||
|
||||
# Arguments that look like ids but name no organization's row, with the reason.
|
||||
UNTENANTED_ARGS: dict[str, dict[str, str]] = {
|
||||
"delete_schedule": {
|
||||
"schedule_id": "the schedule tools only find the organization's schedules",
|
||||
},
|
||||
"list_schedules": {
|
||||
"graph_id": "the schedule tools only list the organization's schedules",
|
||||
},
|
||||
"read_workspace_file": {"file_id": "workspace files belong to the user"},
|
||||
"delete_workspace_file": {"file_id": "workspace files belong to the user"},
|
||||
"list_workspace_files": {"folder_id": "a workspace folder, the user's own"},
|
||||
"run_agent": {"username_agent_slug": "names a public marketplace listing"},
|
||||
"view_agent_output": {"store_slug": "names a public marketplace listing"},
|
||||
}
|
||||
|
||||
|
||||
async def check_ids_in_tenant(
|
||||
tool_name: str, args: dict[str, Any], user_id: str, organization_id: str
|
||||
) -> None:
|
||||
"""Refuse a call naming a row tagged with another organization.
|
||||
|
||||
An id that names nothing passes; the tool says it wasn't found.
|
||||
"""
|
||||
for arg, kind in TENANTED_ARGS.get(tool_name, {}).items():
|
||||
for value in _ids(args.get(arg)):
|
||||
try:
|
||||
allowed = await _in_tenant(kind, value, user_id, organization_id)
|
||||
except Exception as exc:
|
||||
logger.exception(f"MCP tenancy check failed on {tool_name}.{arg}")
|
||||
raise ToolError("Couldn't check access; retry shortly") from exc
|
||||
if not allowed:
|
||||
raise ToolError(f"{kind.value} '{value}' not found")
|
||||
|
||||
|
||||
def _ids(value: Any) -> list[str]:
|
||||
"""The ids an argument holds, normalised the way the tools read them."""
|
||||
values = value if isinstance(value, list) else [value]
|
||||
return [v.strip() for v in values if isinstance(v, str) and v.strip()]
|
||||
|
||||
|
||||
async def _in_tenant(
|
||||
kind: IdKind, value: str, user_id: str, organization_id: str
|
||||
) -> bool:
|
||||
match kind:
|
||||
case IdKind.LIBRARY_AGENT:
|
||||
organizations = await _agent_organizations(value, user_id)
|
||||
case IdKind.FOLDER:
|
||||
organizations = await _folder_organizations(value, user_id)
|
||||
case IdKind.RUN:
|
||||
run = await execution_db().get_graph_execution_meta(
|
||||
user_id=user_id, execution_id=value
|
||||
)
|
||||
organizations = [run.organization_id] if run else []
|
||||
case IdKind.PRESET:
|
||||
preset = await library_db().get_preset(user_id=user_id, preset_id=value)
|
||||
organizations = [preset.organization_id] if preset else []
|
||||
return all(in_tenant(org, organization_id) for org in organizations)
|
||||
|
||||
|
||||
async def _agent_organizations(value: str, user_id: str) -> list[str | None]:
|
||||
"""The organization of every row the id could name, as the tools resolve it.
|
||||
|
||||
A library agent by its own id or by its graph's; failing both, a graph of
|
||||
the user's own with no library entry, which ``edit_agent`` opens directly.
|
||||
"""
|
||||
lib = library_db()
|
||||
try:
|
||||
by_id = await lib.get_library_agent(value, user_id)
|
||||
except NotFoundError:
|
||||
by_id = None
|
||||
by_graph = await lib.get_library_agent_by_graph_id(
|
||||
user_id, value, include_archived=True
|
||||
)
|
||||
organizations = [a.organization_id for a in (by_id, by_graph) if a is not None]
|
||||
if organizations:
|
||||
return organizations
|
||||
graph = await graph_db().get_graph(value, None, user_id=user_id)
|
||||
if graph is not None and graph.user_id == user_id:
|
||||
return [graph.organization_id]
|
||||
return []
|
||||
|
||||
|
||||
async def _folder_organizations(value: str, user_id: str) -> list[str | None]:
|
||||
try:
|
||||
folder = await library_db().get_folder(value, user_id)
|
||||
except NotFoundError:
|
||||
return []
|
||||
return [folder.organization_id]
|
||||
@@ -0,0 +1,219 @@
|
||||
"""An MCP call can't name a row in another organization than its credential's."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from mcp.server.fastmcp.exceptions import ToolError
|
||||
|
||||
from backend.api.external.v2.mcp_tenancy import (
|
||||
TENANTED_ARGS,
|
||||
UNTENANTED_ARGS,
|
||||
check_ids_in_tenant,
|
||||
)
|
||||
from backend.copilot.tools import TOOL_REGISTRY
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
USER = "user-1"
|
||||
ORG = "org-1"
|
||||
|
||||
|
||||
def _looks_like_an_id(argument: str) -> bool:
|
||||
return argument.endswith(("_id", "_ids", "_slug"))
|
||||
|
||||
|
||||
def _exposed_tools() -> dict[str, set[str]]:
|
||||
return {
|
||||
name: set(tool.external_parameters.get("properties", {}))
|
||||
for name, tool in TOOL_REGISTRY.items()
|
||||
if tool.allow_external_use[0]
|
||||
}
|
||||
|
||||
|
||||
def test_every_id_an_exposed_tool_takes_is_classified():
|
||||
unclassified = sorted(
|
||||
f"{name}.{arg}"
|
||||
for name, arguments in _exposed_tools().items()
|
||||
for arg in arguments
|
||||
if _looks_like_an_id(arg)
|
||||
and arg not in TENANTED_ARGS.get(name, {})
|
||||
and arg not in UNTENANTED_ARGS.get(name, {})
|
||||
)
|
||||
assert not unclassified, (
|
||||
"Id arguments neither checked against the organization (TENANTED_ARGS) "
|
||||
f"nor listed with a reason (UNTENANTED_ARGS): {unclassified}"
|
||||
)
|
||||
|
||||
|
||||
def test_the_classification_names_no_stale_tool_or_argument():
|
||||
exposed = _exposed_tools()
|
||||
stale = sorted(
|
||||
f"{name}.{arg}"
|
||||
for table in (TENANTED_ARGS, UNTENANTED_ARGS)
|
||||
for name, arguments in table.items()
|
||||
for arg in arguments
|
||||
if arg not in exposed.get(name, set())
|
||||
)
|
||||
assert not stale, f"Classified arguments no exposed tool takes: {stale}"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def library(mocker: pytest_mock.MockerFixture) -> MagicMock:
|
||||
lib = MagicMock()
|
||||
lib.get_library_agent = AsyncMock(side_effect=NotFoundError("no"))
|
||||
lib.get_library_agent_by_graph_id = AsyncMock(return_value=None)
|
||||
lib.get_folder = AsyncMock(side_effect=NotFoundError("no"))
|
||||
mocker.patch("backend.api.external.v2.mcp_tenancy.library_db", return_value=lib)
|
||||
return lib
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def graphs(mocker: pytest_mock.MockerFixture) -> MagicMock:
|
||||
db = MagicMock()
|
||||
db.get_graph = AsyncMock(return_value=None)
|
||||
mocker.patch("backend.api.external.v2.mcp_tenancy.graph_db", return_value=db)
|
||||
return db
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runs(mocker: pytest_mock.MockerFixture) -> MagicMock:
|
||||
db = MagicMock()
|
||||
db.get_graph_execution_meta = AsyncMock(return_value=None)
|
||||
mocker.patch("backend.api.external.v2.mcp_tenancy.execution_db", return_value=db)
|
||||
return db
|
||||
|
||||
|
||||
def _row(organization_id: str | None, **fields) -> SimpleNamespace:
|
||||
return SimpleNamespace(organization_id=organization_id, **fields)
|
||||
|
||||
|
||||
async def test_a_folder_in_another_organization_is_not_found(library: MagicMock):
|
||||
library.get_folder = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError, match="Folder 'folder-1' not found"):
|
||||
await check_ids_in_tenant("delete_folder", {"folder_id": "folder-1"}, USER, ORG)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("organization_id", [ORG, None])
|
||||
async def test_a_folder_in_the_organization_or_untagged_passes(
|
||||
library: MagicMock, organization_id: str | None
|
||||
):
|
||||
library.get_folder = AsyncMock(return_value=_row(organization_id))
|
||||
|
||||
await check_ids_in_tenant("delete_folder", {"folder_id": "folder-1"}, USER, ORG)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool", ["update_preset", "delete_preset"])
|
||||
async def test_a_preset_in_another_organization_is_not_found(
|
||||
library: MagicMock, tool: str
|
||||
):
|
||||
library.get_preset = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError, match="Preset 'preset-1' not found"):
|
||||
await check_ids_in_tenant(tool, {"preset_id": "preset-1"}, USER, ORG)
|
||||
library.get_preset.assert_awaited_once_with(user_id=USER, preset_id="preset-1")
|
||||
|
||||
|
||||
async def test_an_id_naming_nothing_is_left_to_the_tool(
|
||||
library: MagicMock, graphs: MagicMock, runs: MagicMock
|
||||
):
|
||||
await check_ids_in_tenant(
|
||||
"view_agent_output",
|
||||
{"library_agent_id": "agent-1", "execution_id": "run-1"},
|
||||
USER,
|
||||
ORG,
|
||||
)
|
||||
|
||||
|
||||
async def test_a_library_agent_in_another_organization_is_not_found(
|
||||
library: MagicMock, graphs: MagicMock
|
||||
):
|
||||
library.get_library_agent = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError, match="Library agent 'agent-1' not found"):
|
||||
await check_ids_in_tenant(
|
||||
"run_agent", {"library_agent_id": "agent-1"}, USER, ORG
|
||||
)
|
||||
graphs.get_graph.assert_not_called()
|
||||
|
||||
|
||||
async def test_an_agent_named_by_its_graph_id_is_checked_too(
|
||||
library: MagicMock, graphs: MagicMock
|
||||
):
|
||||
library.get_library_agent_by_graph_id = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError, match="Library agent 'graph-1' not found"):
|
||||
await check_ids_in_tenant("edit_agent", {"agent_id": "graph-1"}, USER, ORG)
|
||||
library.get_library_agent_by_graph_id.assert_awaited_once_with(
|
||||
USER, "graph-1", include_archived=True
|
||||
)
|
||||
|
||||
|
||||
async def test_an_own_graph_without_a_library_entry_is_checked_by_its_own_tag(
|
||||
library: MagicMock, graphs: MagicMock
|
||||
):
|
||||
graphs.get_graph = AsyncMock(return_value=_row("org-2", user_id=USER))
|
||||
|
||||
with pytest.raises(ToolError, match="Library agent 'graph-1' not found"):
|
||||
await check_ids_in_tenant("edit_agent", {"agent_id": "graph-1"}, USER, ORG)
|
||||
|
||||
|
||||
async def test_someone_elses_graph_is_left_to_the_tool(
|
||||
library: MagicMock, graphs: MagicMock
|
||||
):
|
||||
"""A marketplace graph carries its publisher's organization, not the caller's."""
|
||||
graphs.get_graph = AsyncMock(return_value=_row("org-2", user_id="publisher"))
|
||||
|
||||
await check_ids_in_tenant("edit_agent", {"agent_id": "graph-1"}, USER, ORG)
|
||||
|
||||
|
||||
async def test_every_id_in_a_list_is_checked(library: MagicMock, graphs: MagicMock):
|
||||
library.get_library_agent = AsyncMock(
|
||||
side_effect=[_row(ORG), _row("org-2")],
|
||||
)
|
||||
|
||||
with pytest.raises(ToolError, match="Library agent 'agent-2' not found"):
|
||||
await check_ids_in_tenant(
|
||||
"move_agents_to_folder",
|
||||
{"agent_ids": ["agent-1", "agent-2"]},
|
||||
USER,
|
||||
ORG,
|
||||
)
|
||||
|
||||
|
||||
async def test_an_id_is_checked_the_way_the_tool_reads_it(library: MagicMock):
|
||||
"""The tools strip the ids they're given, so the check does too."""
|
||||
library.get_folder = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError):
|
||||
await check_ids_in_tenant(
|
||||
"delete_folder", {"folder_id": " folder-1 "}, USER, ORG
|
||||
)
|
||||
library.get_folder.assert_awaited_once_with("folder-1", USER)
|
||||
|
||||
|
||||
async def test_a_run_in_another_organization_is_not_found(runs: MagicMock):
|
||||
runs.get_graph_execution_meta = AsyncMock(return_value=_row("org-2"))
|
||||
|
||||
with pytest.raises(ToolError, match="Run 'run-1' not found"):
|
||||
await check_ids_in_tenant(
|
||||
"view_agent_output", {"execution_id": "run-1"}, USER, ORG
|
||||
)
|
||||
runs.get_graph_execution_meta.assert_awaited_once_with(
|
||||
user_id=USER, execution_id="run-1"
|
||||
)
|
||||
|
||||
|
||||
async def test_a_failed_lookup_refuses_without_leaking_the_error(library: MagicMock):
|
||||
library.get_folder = AsyncMock(side_effect=RuntimeError("query engine down"))
|
||||
|
||||
with pytest.raises(ToolError) as refusal:
|
||||
await check_ids_in_tenant("delete_folder", {"folder_id": "folder-1"}, USER, ORG)
|
||||
assert "query engine" not in str(refusal.value)
|
||||
|
||||
|
||||
async def test_a_tool_without_ids_needs_no_lookup(library: MagicMock):
|
||||
await check_ids_in_tenant("search_docs", {"query": "folder_id"}, USER, ORG)
|
||||
|
||||
library.get_folder.assert_not_called()
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
A member sees a colleague's org-home runs and graphs through v2, but the writes
|
||||
beneath these routes are the owner's, so each refuses a colleague before acting.
|
||||
|
||||
These drive the real data layer: the defect sits between an org-visible read and
|
||||
an owner-scoped write, and a mocked database absorbs it.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
from uuid import uuid4
|
||||
|
||||
import fastapi
|
||||
import httpx
|
||||
import prisma.models
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data.user import get_or_create_user
|
||||
from backend.util.test import SpinTestServer
|
||||
|
||||
from . import graphs
|
||||
from .errors import add_v2_exception_handlers
|
||||
from .models import GraphCreateRequest
|
||||
from .runs import runs_router
|
||||
from .tenancy import TenantContext, require_auth
|
||||
|
||||
|
||||
class Colleagues(BaseModel):
|
||||
owner: TenantContext
|
||||
mate: TenantContext
|
||||
graph_id: str
|
||||
run_id: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def colleagues(server: SpinTestServer) -> Colleagues:
|
||||
org_id = str(uuid4())
|
||||
owner, mate = [await _member(org_id) for _ in range(2)]
|
||||
graph = await graphs.create_graph(
|
||||
GraphCreateRequest(name="Shared", nodes=[], links=[]), auth=owner
|
||||
)
|
||||
run = await execution_db.create_graph_execution(
|
||||
graph_id=graph.id,
|
||||
graph_version=graph.version,
|
||||
starting_nodes_input=[],
|
||||
inputs={},
|
||||
user_id=owner.user_id,
|
||||
organization_id=org_id,
|
||||
)
|
||||
return Colleagues(owner=owner, mate=mate, graph_id=graph.id, run_id=run.id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_colleagues_run_is_not_deleted(colleagues: Colleagues) -> None:
|
||||
refused = await _call("DELETE", f"/runs/{colleagues.run_id}", colleagues.mate)
|
||||
|
||||
assert refused.status_code == 403, refused.text
|
||||
assert not (await _run_row(colleagues.run_id)).isDeleted
|
||||
|
||||
deleted = await _call("DELETE", f"/runs/{colleagues.run_id}", colleagues.owner)
|
||||
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
assert (await _run_row(colleagues.run_id)).isDeleted
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_colleagues_run_is_not_cancelled(
|
||||
colleagues: Colleagues, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
queue = Mock(publish_message=AsyncMock())
|
||||
mocker.patch(
|
||||
"backend.executor.utils.get_async_execution_queue",
|
||||
new_callable=AsyncMock,
|
||||
return_value=queue,
|
||||
)
|
||||
|
||||
refused = await _call("POST", f"/runs/{colleagues.run_id}/stop", colleagues.mate)
|
||||
|
||||
assert refused.status_code == 403, refused.text
|
||||
queue.publish_message.assert_not_awaited()
|
||||
|
||||
stopped = await _call("POST", f"/runs/{colleagues.run_id}/stop", colleagues.owner)
|
||||
|
||||
assert stopped.status_code == 202, stopped.text
|
||||
queue.publish_message.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_colleague_cannot_append_a_version(colleagues: Colleagues) -> None:
|
||||
body = {"name": "Shared", "nodes": [], "links": []}
|
||||
|
||||
refused = await _call(
|
||||
"PUT", f"/graphs/{colleagues.graph_id}", colleagues.mate, json=body
|
||||
)
|
||||
|
||||
assert refused.status_code == 403, refused.text
|
||||
assert await _version_owners(colleagues.graph_id) == {1: colleagues.owner.user_id}
|
||||
|
||||
updated = await _call(
|
||||
"PUT", f"/graphs/{colleagues.graph_id}", colleagues.owner, json=body
|
||||
)
|
||||
|
||||
assert updated.status_code == 200, updated.text
|
||||
assert await _version_owners(colleagues.graph_id) == {
|
||||
1: colleagues.owner.user_id,
|
||||
2: colleagues.owner.user_id,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_colleague_cannot_activate_a_version(
|
||||
colleagues: Colleagues, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
activate = mocker.spy(graphs, "before_graph_activate")
|
||||
path = f"/graphs/{colleagues.graph_id}/versions/active"
|
||||
body = {"active_graph_version": 1}
|
||||
|
||||
refused = await _call("PUT", path, colleagues.mate, json=body)
|
||||
|
||||
assert refused.status_code == 403, refused.text
|
||||
activate.assert_not_called()
|
||||
|
||||
activated = await _call("PUT", path, colleagues.owner, json=body)
|
||||
|
||||
assert activated.status_code == 204, activated.text
|
||||
activate.assert_called_once()
|
||||
|
||||
|
||||
async def _member(org_id: str) -> TenantContext:
|
||||
user_id = str(uuid4())
|
||||
await get_or_create_user({"sub": user_id, "email": f"{user_id}@example.com"})
|
||||
return TenantContext(
|
||||
user_id=user_id,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=org_id,
|
||||
)
|
||||
|
||||
|
||||
async def _call(
|
||||
method: str, path: str, auth: TenantContext, json: dict | None = None
|
||||
) -> httpx.Response:
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(runs_router, prefix="/runs")
|
||||
app.include_router(graphs.graphs_router, prefix="/graphs")
|
||||
add_v2_exception_handlers(app)
|
||||
app.dependency_overrides[require_auth] = lambda: auth
|
||||
# Without raise_app_exceptions=False a 500 surfaces as the exception, not a response.
|
||||
transport = httpx.ASGITransport(app=app, raise_app_exceptions=False)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://v2") as client:
|
||||
return await client.request(method, path, json=json)
|
||||
|
||||
|
||||
async def _run_row(run_id: str) -> prisma.models.AgentGraphExecution:
|
||||
return await prisma.models.AgentGraphExecution.prisma().find_unique_or_raise(
|
||||
where={"id": run_id}
|
||||
)
|
||||
|
||||
|
||||
async def _version_owners(graph_id: str) -> dict[int, str]:
|
||||
versions = await prisma.models.AgentGraph.prisma().find_many(where={"id": graph_id})
|
||||
return {v.version: v.userId for v in versions}
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
V2 External API - Pagination
|
||||
|
||||
One envelope and one set of query parameters for every list endpoint:
|
||||
`?limit=&cursor=` in, `{"items", "next_cursor", "total_count"}` out.
|
||||
|
||||
The cursor is opaque. It carries a version and a kind, so a cursor from another
|
||||
endpoint, a stale format, or a hand-made one is rejected rather than silently
|
||||
read as page 1 — which is the failure a later move to keyset pagination would
|
||||
otherwise hit.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
from typing import Annotated, Any, Generic, Optional, Sequence, TypeVar
|
||||
|
||||
from fastapi import Depends, HTTPException, Query
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from starlette import status
|
||||
|
||||
from .common import DEFAULT_PAGE_SIZE, MAX_PAGE_SIZE
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
CURSOR_VERSION = 1
|
||||
|
||||
# An offset this deep is a forged cursor, not a caller walking a list. Without a
|
||||
# bound, Prisma raises a DataError on the skip and the caller gets a 500.
|
||||
MAX_PAGE = 1_000_000
|
||||
|
||||
|
||||
class Page(BaseModel, Generic[T]):
|
||||
"""The response envelope of every v2 list endpoint."""
|
||||
|
||||
items: list[T]
|
||||
next_cursor: Optional[str] = Field(
|
||||
description="Pass as `cursor` to fetch the next page. `null` on the last page."
|
||||
)
|
||||
total_count: Optional[int] = Field(
|
||||
description=(
|
||||
"Items matching the request across all pages, or `null` where the "
|
||||
"underlying source cannot report one. Always present."
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class PageRequest(BaseModel):
|
||||
"""The `limit`/`cursor` query parameters of every v2 list endpoint."""
|
||||
|
||||
limit: int
|
||||
cursor: Optional[str] = None
|
||||
position: dict[str, Any] = Field(default_factory=dict, exclude=True)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _decode_cursor(self) -> "PageRequest":
|
||||
# Decoded here rather than where it is read, so a malformed cursor is a
|
||||
# 400 even on an endpoint that ignores the value.
|
||||
self.position = _decode(self.cursor)
|
||||
return self
|
||||
|
||||
@property
|
||||
def page(self) -> int:
|
||||
"""1-indexed page, for the offset-paginated sources."""
|
||||
if not self.position:
|
||||
return 1
|
||||
self.require_cursor_kind("p", "a page")
|
||||
page = self.position.get("p")
|
||||
if not isinstance(page, int) or not 1 <= page <= MAX_PAGE:
|
||||
raise _malformed_cursor()
|
||||
return page
|
||||
|
||||
@property
|
||||
def token(self) -> Optional[str]:
|
||||
"""Opaque position, for the keyset-paginated sources."""
|
||||
if not self.position:
|
||||
return None
|
||||
self.require_cursor_kind("t", "a keyset token")
|
||||
token = self.position.get("t")
|
||||
if not isinstance(token, str):
|
||||
raise _malformed_cursor()
|
||||
return token
|
||||
|
||||
def paged(self, items: Sequence[T], total_count: int) -> Page[T]:
|
||||
"""For an offset-paginated source that reports a total."""
|
||||
return Page[T](
|
||||
items=list(items),
|
||||
next_cursor=(
|
||||
encode_page_cursor(self.page + 1)
|
||||
if self.page * self.limit < total_count
|
||||
else None
|
||||
),
|
||||
total_count=total_count,
|
||||
)
|
||||
|
||||
def slice(self, items: Sequence[T]) -> Page[T]:
|
||||
"""For a source that returns everything it has; paginate in memory."""
|
||||
start = (self.page - 1) * self.limit
|
||||
return self.paged(items[start : start + self.limit], len(items))
|
||||
|
||||
def keyset(
|
||||
self,
|
||||
items: Sequence[T],
|
||||
next_token: Optional[str],
|
||||
total_count: Optional[int] = None,
|
||||
) -> Page[T]:
|
||||
"""For a keyset-paginated source, which may not be able to report a total."""
|
||||
# Checked here too: a caller that never reads `token` would otherwise
|
||||
# accept a page cursor and answer with the first page.
|
||||
self.require_cursor_kind("t", "a keyset token")
|
||||
return Page[T](
|
||||
items=list(items),
|
||||
next_cursor=encode_token_cursor(next_token) if next_token else None,
|
||||
total_count=total_count,
|
||||
)
|
||||
|
||||
def uncounted(self, items: Sequence[T]) -> Page[T]:
|
||||
"""For a source that answers in one shot: no second page, no total."""
|
||||
_reject_any_cursor(self)
|
||||
return Page[T](items=list(items), next_cursor=None, total_count=None)
|
||||
|
||||
def require_cursor_kind(self, kind: str, expected: str) -> None:
|
||||
"""No-op without a cursor; a cursor of another kind is a 400."""
|
||||
if self.position and self.position.get("k") != kind:
|
||||
raise _wrong_cursor(expected)
|
||||
|
||||
|
||||
def page_request(
|
||||
limit: Annotated[
|
||||
int,
|
||||
Query(
|
||||
ge=1,
|
||||
le=MAX_PAGE_SIZE,
|
||||
description=f"Items per page (max {MAX_PAGE_SIZE})",
|
||||
),
|
||||
] = DEFAULT_PAGE_SIZE,
|
||||
cursor: Annotated[
|
||||
Optional[str],
|
||||
Query(description="`next_cursor` from the previous response"),
|
||||
] = None,
|
||||
) -> PageRequest:
|
||||
return PageRequest(limit=limit, cursor=cursor)
|
||||
|
||||
|
||||
def single_page_request(
|
||||
request: Annotated[PageRequest, Depends(page_request)],
|
||||
) -> PageRequest:
|
||||
"""For endpoints whose source answers in one shot.
|
||||
|
||||
Rejects a cursor here rather than in the route body, so a request carrying
|
||||
one never reaches the upstream call it would then discard.
|
||||
"""
|
||||
_reject_any_cursor(request)
|
||||
return request
|
||||
|
||||
|
||||
def encode_page_cursor(page: int) -> str:
|
||||
return _encode({"k": "p", "p": page})
|
||||
|
||||
|
||||
def encode_token_cursor(token: str) -> str:
|
||||
return _encode({"k": "t", "t": token})
|
||||
|
||||
|
||||
def _encode(payload: dict[str, Any]) -> str:
|
||||
raw = json.dumps({"v": CURSOR_VERSION, **payload}, separators=(",", ":")).encode()
|
||||
return base64.urlsafe_b64encode(raw).decode().rstrip("=")
|
||||
|
||||
|
||||
def _decode(cursor: Optional[str]) -> dict[str, Any]:
|
||||
if not cursor:
|
||||
return {}
|
||||
try:
|
||||
payload = json.loads(
|
||||
base64.urlsafe_b64decode(cursor + "=" * (-len(cursor) % 4))
|
||||
)
|
||||
except (binascii.Error, UnicodeDecodeError, ValueError, TypeError):
|
||||
raise _malformed_cursor()
|
||||
if not isinstance(payload, dict) or payload.get("v") != CURSOR_VERSION:
|
||||
raise _malformed_cursor()
|
||||
return payload
|
||||
|
||||
|
||||
def _reject_any_cursor(request: PageRequest) -> None:
|
||||
if request.position:
|
||||
raise _wrong_cursor("no cursor — this endpoint returns a single page")
|
||||
|
||||
|
||||
def _malformed_cursor() -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Malformed cursor. Pass back the `next_cursor` of a previous response.",
|
||||
)
|
||||
|
||||
|
||||
def _wrong_cursor(expected: str) -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=(
|
||||
f"This endpoint expects {expected}. "
|
||||
"Cursors are not interchangeable between endpoints."
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Every v2 route either needs an active plan or is listed here as open.
|
||||
|
||||
The host app's ``backend/api/paywall_routes_test.py`` sees only its own routes,
|
||||
and this app is mounted inside it, so it keeps its own list. The split follows
|
||||
the internal API's: anything that runs compute, calls an outside service or
|
||||
stores new files or rows needs a plan.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pytest
|
||||
from fastapi.dependencies.models import Dependant
|
||||
from fastapi.routing import APIRoute
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.api.external.v2.app import v2_app
|
||||
from backend.api.external.v2.errors import add_v2_exception_handlers
|
||||
from backend.api.external.v2.routes import v2_router
|
||||
from backend.api.external.v2.tenancy import (
|
||||
TenantContext,
|
||||
require_auth,
|
||||
require_paid_permission,
|
||||
)
|
||||
|
||||
# Every route's check is a closure over this one code object.
|
||||
PLAN_CHECK = require_paid_permission().__code__
|
||||
|
||||
OPEN_WHILE_PAYWALLED = {
|
||||
# Reading the account's own data, or the catalog.
|
||||
"GET /blocks",
|
||||
"GET /files",
|
||||
"GET /files/{file_id}",
|
||||
"GET /files/{file_id}/download",
|
||||
"GET /graphs",
|
||||
"GET /graphs/{graph_id}",
|
||||
"GET /graphs/{graph_id}/blocks",
|
||||
"GET /graphs/{graph_id}/credentials",
|
||||
"GET /graphs/{graph_id}/library-agent",
|
||||
"GET /graphs/{graph_id}/marketplace-listing",
|
||||
"GET /graphs/{graph_id}/versions",
|
||||
"GET /integrations/credentials",
|
||||
"GET /library/agents",
|
||||
"GET /library/agents/{agent_id}",
|
||||
"GET /library/agents/{agent_id}/credentials",
|
||||
"GET /library/folders",
|
||||
"GET /library/folders/tree",
|
||||
"GET /library/folders/{folder_id}",
|
||||
"GET /library/presets",
|
||||
"GET /library/presets/{preset_id}",
|
||||
"GET /marketplace/agents",
|
||||
"GET /marketplace/agents/by-version/{version_id}",
|
||||
"GET /marketplace/agents/{username}/{agent_name}",
|
||||
"GET /marketplace/creators",
|
||||
"GET /marketplace/creators/{username}",
|
||||
"GET /marketplace/profile",
|
||||
"GET /marketplace/submissions",
|
||||
"GET /me",
|
||||
"GET /runs",
|
||||
"GET /runs/reviews",
|
||||
"GET /runs/{run_id}",
|
||||
"GET /schedules",
|
||||
# Seeing what the account is billed.
|
||||
"GET /credits",
|
||||
"GET /credits/cost-summary",
|
||||
"GET /credits/invoices",
|
||||
"GET /credits/subscription",
|
||||
"GET /credits/transactions",
|
||||
# Stopping, deleting or revoking.
|
||||
"DELETE /files/{file_id}",
|
||||
"DELETE /integrations/credentials/{credential_id}",
|
||||
"DELETE /library/agents/{agent_id}",
|
||||
"DELETE /library/folders/{folder_id}",
|
||||
"DELETE /library/presets/{preset_id}",
|
||||
"DELETE /marketplace/submissions/{version_id}",
|
||||
"DELETE /runs/{run_id}",
|
||||
"DELETE /runs/{run_id}/share",
|
||||
"POST /runs/{run_id}/stop",
|
||||
"DELETE /schedules/{schedule_id}",
|
||||
# Organising what the account already has: names, places, settings, sharing.
|
||||
"PATCH /graphs/{graph_id}/settings",
|
||||
"PATCH /library/agents/{agent_id}",
|
||||
"POST /library/folders",
|
||||
"PATCH /library/folders/{folder_id}",
|
||||
"POST /library/folders/{folder_id}/move",
|
||||
"PATCH /marketplace/profile",
|
||||
"POST /runs/{run_id}/share",
|
||||
}
|
||||
|
||||
# One closed route per router, as a client calls it.
|
||||
CLOSED_SAMPLES = [
|
||||
"POST /files/upload",
|
||||
"POST /graphs",
|
||||
"POST /integrations/credentials",
|
||||
"POST /library/agents/a-1/runs",
|
||||
"POST /library/presets/p-1/execute",
|
||||
"POST /marketplace/submissions",
|
||||
"POST /runs/r-1/reviews",
|
||||
"POST /schedules",
|
||||
"GET /search?query=weather",
|
||||
]
|
||||
|
||||
|
||||
def test_every_route_needs_a_plan_or_is_listed_open():
|
||||
unlisted, checked_but_listed = [], []
|
||||
for key, route in _routes():
|
||||
checked = _depends_on_plan_check(route.dependant)
|
||||
listed = key in OPEN_WHILE_PAYWALLED
|
||||
if checked and listed:
|
||||
checked_but_listed.append(key)
|
||||
elif not checked and not listed:
|
||||
unlisted.append(key)
|
||||
|
||||
assert not unlisted, (
|
||||
"Authenticate these with require_paid_permission, or list them in "
|
||||
f"OPEN_WHILE_PAYWALLED with the reason they stay open: {sorted(unlisted)}"
|
||||
)
|
||||
assert not checked_but_listed, (
|
||||
"These routes need a plan, so drop them from the open list: "
|
||||
f"{sorted(checked_but_listed)}"
|
||||
)
|
||||
|
||||
|
||||
def test_the_open_list_names_only_routes_that_exist():
|
||||
assert not OPEN_WHILE_PAYWALLED - {key for key, _ in _routes()}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call", CLOSED_SAMPLES)
|
||||
def test_an_account_without_a_plan_gets_402_from_a_closed_route(
|
||||
call: str, paywalled_user: str
|
||||
):
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(v2_router)
|
||||
add_v2_exception_handlers(app)
|
||||
app.dependency_overrides[require_auth] = lambda: TenantContext(
|
||||
user_id=paywalled_user,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
method, path = call.split(" ", 1)
|
||||
|
||||
response = fastapi.testclient.TestClient(app).request(method, path)
|
||||
|
||||
assert response.status_code == 402, response.text
|
||||
assert response.json()["error"]["code"] == "payment_required"
|
||||
|
||||
|
||||
def _routes():
|
||||
for route in v2_app.routes:
|
||||
if isinstance(route, APIRoute):
|
||||
for method in route.methods:
|
||||
yield f"{method} {route.path}", route
|
||||
|
||||
|
||||
def _depends_on_plan_check(dependant: Dependant) -> bool:
|
||||
return any(
|
||||
(inspect.isfunction(dep.call) and dep.call.__code__ is PLAN_CHECK)
|
||||
or _depends_on_plan_check(dep)
|
||||
for dep in dependant.dependencies
|
||||
)
|
||||
@@ -0,0 +1,466 @@
|
||||
"""
|
||||
Tests for what a v2 credential is allowed to reach.
|
||||
|
||||
Least privilege is the point: a key granted one scope must not be a way to read
|
||||
data, spend platform money, or destroy platform state that its scopes never
|
||||
mentioned. Each test here pins one such hole shut.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
from unittest import mock
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission, ContentType
|
||||
from pydantic import SecretStr
|
||||
|
||||
from backend.api.external.v2.errors import add_v2_exception_handlers
|
||||
from backend.api.external.v2.models import SearchContentType
|
||||
from backend.api.external.v2.pagination import PageRequest
|
||||
from backend.api.external.v2.tenancy import TenantContext, require_auth
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.auth.oauth import ACCESS_TOKEN_PREFIX
|
||||
from backend.data.model import APIKeyCredentials, is_sdk_default
|
||||
from backend.integrations.credentials_store import SYSTEM_CREDENTIAL_IDS
|
||||
from backend.integrations.providers import ProviderName
|
||||
from backend.util.exceptions import NeedConfirmation
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG_ID = "org-1"
|
||||
_REQUEST = mock.Mock(spec=fastapi.Request)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Search: a private content type costs the scope that guards it elsewhere
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type,scope",
|
||||
[
|
||||
(SearchContentType.LIBRARY_AGENT, APIKeyPermission.READ_LIBRARY),
|
||||
(SearchContentType.WORKSPACE_FILE, APIKeyPermission.READ_FILES),
|
||||
],
|
||||
)
|
||||
async def test_searching_private_content_without_its_scope_is_refused(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
content_type: SearchContentType,
|
||||
scope: APIKeyPermission,
|
||||
) -> None:
|
||||
"""A zero-scope key asked for these and got the caller's own rows back."""
|
||||
hybrid = _mock_hybrid_search(mocker)
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
await search_with(mocker, [content_type], scopes=[])
|
||||
|
||||
assert raised.value.status_code == 403
|
||||
assert scope.value in str(raised.value.detail)
|
||||
hybrid.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type,scope",
|
||||
[
|
||||
(SearchContentType.LIBRARY_AGENT, APIKeyPermission.READ_LIBRARY),
|
||||
(SearchContentType.WORKSPACE_FILE, APIKeyPermission.READ_FILES),
|
||||
],
|
||||
)
|
||||
async def test_searching_private_content_with_its_scope_is_allowed(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
content_type: SearchContentType,
|
||||
scope: APIKeyPermission,
|
||||
) -> None:
|
||||
hybrid = _mock_hybrid_search(mocker)
|
||||
|
||||
await search_with(mocker, [content_type], scopes=[scope])
|
||||
|
||||
assert hybrid.await_args.kwargs["content_types"] == [
|
||||
ContentType(content_type.value)
|
||||
]
|
||||
|
||||
|
||||
async def test_search_without_content_types_asks_only_for_public_ones(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""v2 names the default rather than inheriting whatever the engine picks."""
|
||||
hybrid = _mock_hybrid_search(mocker)
|
||||
|
||||
await search_with(mocker, None, scopes=[])
|
||||
|
||||
assert hybrid.await_args.kwargs["content_types"] == [
|
||||
ContentType.STORE_AGENT,
|
||||
ContentType.BLOCK,
|
||||
ContentType.INTEGRATION,
|
||||
ContentType.DOCUMENTATION,
|
||||
]
|
||||
|
||||
|
||||
def test_chat_sessions_are_not_a_searchable_content_type() -> None:
|
||||
"""No v2 scope covers chat, so the parameter cannot name it at all."""
|
||||
assert "CHAT_SESSION" in {t.value for t in ContentType}
|
||||
assert "CHAT_SESSION" not in {t.value for t in SearchContentType}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("content_type", list(SearchContentType))
|
||||
async def test_search_results_carry_their_content_type(
|
||||
mocker: pytest_mock.MockFixture, content_type: SearchContentType
|
||||
) -> None:
|
||||
hybrid = _mock_hybrid_search(mocker)
|
||||
hybrid.return_value = ([_search_row(content_type.value)], 1)
|
||||
|
||||
page = await search_with(mocker, [content_type], scopes=list(APIKeyPermission))
|
||||
|
||||
assert [r.content_type for r in page.items] == [content_type]
|
||||
assert page.items[0].content_id == "id-1"
|
||||
|
||||
|
||||
async def test_a_result_outside_the_v2_types_fails_rather_than_leaking(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The engine is only asked for v2 types; a chat session it returns anyway
|
||||
is a server fault, never a row handed to the caller."""
|
||||
hybrid = _mock_hybrid_search(mocker)
|
||||
hybrid.return_value = ([_search_row("CHAT_SESSION")], 1)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await search_with(mocker, None, scopes=list(APIKeyPermission))
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Sharing: both directions cost the same scopes
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_share_and_unshare_require_the_same_permissions() -> None:
|
||||
"""Unsharing needed only SHARE_RUN while sharing needed READ_RUN as well."""
|
||||
assert _published_scopes("/runs/{run_id}/share", "post") == _published_scopes(
|
||||
"/runs/{run_id}/share", "delete"
|
||||
)
|
||||
assert _published_scopes("/runs/{run_id}/share", "delete") == {
|
||||
"READ_RUN",
|
||||
"SHARE_RUN",
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Credentials: the platform's own credentials are not the caller's to delete
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def test_deleting_a_system_credential_is_refused(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""`delete_creds_by_id` raises ValueError on these, which v2 reported as a
|
||||
400 blaming the caller. The internal route answers 403."""
|
||||
from .integrations.credentials import delete_credential
|
||||
|
||||
system_id = sorted(SYSTEM_CREDENTIAL_IDS)[0]
|
||||
deleter = _mock_creds_store(mocker, _credential(system_id))
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
await delete_credential(_REQUEST, credential_id=system_id, auth=_tenant())
|
||||
|
||||
assert raised.value.status_code == 403
|
||||
deleter.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_deleting_a_managed_credential_is_refused(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .integrations.credentials import delete_credential
|
||||
|
||||
deleter = _mock_creds_store(mocker, _credential("managed-1", is_managed=True))
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
await delete_credential(_REQUEST, credential_id="managed-1", auth=_tenant())
|
||||
|
||||
assert raised.value.status_code == 403
|
||||
deleter.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_deleting_an_sdk_default_credential_is_not_found(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .integrations.credentials import delete_credential
|
||||
|
||||
cred_id = "openai-default"
|
||||
assert is_sdk_default(cred_id)
|
||||
deleter = _mock_creds_store(mocker, _credential(cred_id))
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
await delete_credential(_REQUEST, credential_id=cred_id, auth=_tenant())
|
||||
|
||||
assert raised.value.status_code == 404
|
||||
deleter.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_deleting_an_ordinary_credential_still_works(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The guards above must not have closed the endpoint's actual job."""
|
||||
from .integrations.credentials import delete_credential
|
||||
|
||||
credential = _credential("mine-1")
|
||||
deleter = _mock_creds_store(mocker, credential)
|
||||
|
||||
await delete_credential(_REQUEST, credential_id="mine-1", auth=_tenant())
|
||||
|
||||
# The internal route's deletion: webhooks torn down first, OAuth revoked.
|
||||
deleter.assert_awaited_once_with(
|
||||
_REQUEST, USER_ID, credential, ProviderName(credential.provider), False
|
||||
)
|
||||
|
||||
|
||||
async def test_a_credential_a_trigger_still_uses_needs_force(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .integrations.credentials import delete_credential
|
||||
|
||||
deleter = _mock_creds_store(mocker, _credential("mine-1"))
|
||||
deleter.side_effect = NeedConfirmation("A webhook is still in use")
|
||||
|
||||
with pytest.raises(fastapi.HTTPException) as raised:
|
||||
await delete_credential(_REQUEST, credential_id="mine-1", auth=_tenant())
|
||||
|
||||
assert raised.value.status_code == 409
|
||||
assert "force=true" in raised.value.detail
|
||||
|
||||
|
||||
def test_credential_listing_says_which_credentials_are_the_platform_s() -> None:
|
||||
"""Without this, managed and user credentials are indistinguishable."""
|
||||
from .models import CredentialInfo
|
||||
|
||||
assert CredentialInfo.from_internal(_credential("a", is_managed=True)).is_managed
|
||||
assert not CredentialInfo.from_internal(_credential("b")).is_managed
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Auth: verified once per request, and the limiter never blocks on its own bugs
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def test_a_request_verifies_its_credential_once() -> None:
|
||||
"""The limiter and the route both resolve the caller; an API key costs a
|
||||
Scrypt hash to verify, and it was paid twice per request."""
|
||||
calls = _resolve_calls()
|
||||
|
||||
with mock.patch("backend.api.external.middleware.validate_api_key", new=calls):
|
||||
response = _client().get("/probe", headers={"X-API-Key": "k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert calls.count == 1
|
||||
|
||||
|
||||
async def test_a_transient_auth_failure_in_the_limiter_does_not_fail_the_request() -> (
|
||||
None
|
||||
):
|
||||
"""The limiter caught only HTTPException around auth resolution, so any
|
||||
other error propagated out of it and the route never ran — a blip in the
|
||||
credential store 500'd a request the route itself could have served."""
|
||||
calls = _resolve_calls(fail_first=True)
|
||||
|
||||
with mock.patch("backend.api.external.middleware.validate_api_key", new=calls):
|
||||
response = _client().get("/probe", headers={"X-API-Key": "k"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert calls.count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
{"X-API-Key": "k"},
|
||||
{"Authorization": "Bearer k"},
|
||||
{"Authorization": f"Bearer {ACCESS_TOKEN_PREFIX}t"},
|
||||
],
|
||||
ids=["api-key", "api-key-as-bearer", "access-token"],
|
||||
)
|
||||
@pytest.mark.parametrize("state, status", [("banned", 403), ("unbanned", 200)])
|
||||
def test_a_banned_owners_credential_reaches_no_v2_route(
|
||||
auth_row, headers: dict[str, str], state: str, status: int
|
||||
) -> None:
|
||||
owner = _owner(auth_row(state))
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
new=mock.AsyncMock(return_value=owner),
|
||||
),
|
||||
mock.patch(
|
||||
"backend.api.external.middleware.validate_access_token",
|
||||
new=mock.AsyncMock(return_value=(owner, None)),
|
||||
),
|
||||
):
|
||||
response = _client().get("/probe", headers=headers)
|
||||
|
||||
assert response.status_code == status, response.text
|
||||
if status == 403:
|
||||
assert response.json()["error"]["details"]["code"] == "account_banned"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _owner(user_id: str) -> APIAuthorizationInfo:
|
||||
return APIAuthorizationInfo(
|
||||
user_id=user_id,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
created_at=datetime.now(tz=timezone.utc),
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
|
||||
def _tenant(scopes: Optional[list[APIKeyPermission]] = None) -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission) if scopes is None else scopes,
|
||||
type="api_key",
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
|
||||
def _page() -> PageRequest:
|
||||
return PageRequest(limit=20, page=1, cursor=None)
|
||||
|
||||
|
||||
def _mock_hybrid_search(mocker: pytest_mock.MockFixture):
|
||||
return mocker.patch(
|
||||
"backend.api.external.v2.search.unified_hybrid_search",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=([], 0),
|
||||
)
|
||||
|
||||
|
||||
def _search_row(content_type: str) -> dict:
|
||||
return {
|
||||
"content_type": content_type,
|
||||
"content_id": "id-1",
|
||||
"searchable_text": "text",
|
||||
"metadata": {},
|
||||
"updated_at": datetime.now(tz=timezone.utc),
|
||||
"combined_score": 0.5,
|
||||
}
|
||||
|
||||
|
||||
async def search_with(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
content_types: Optional[list[SearchContentType]],
|
||||
scopes: list[APIKeyPermission],
|
||||
):
|
||||
from .search import search
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.search.search_limiter.check",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
# Every library agent a row names is the caller's, in this organization.
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.search.library_db.get_library_agent_organizations",
|
||||
new_callable=mock.AsyncMock,
|
||||
side_effect=lambda _user, ids: {i: ORG_ID for i in ids},
|
||||
)
|
||||
return await search(
|
||||
response=fastapi.Response(),
|
||||
query="q",
|
||||
content_types=content_types,
|
||||
category=None,
|
||||
page=_page(),
|
||||
auth=_tenant(scopes=scopes),
|
||||
)
|
||||
|
||||
|
||||
def _published_scopes(path: str, method: str) -> set[str]:
|
||||
"""The scopes the OpenAPI document tells a caller the endpoint needs."""
|
||||
from .app import v2_app
|
||||
|
||||
operation = v2_app.openapi()["paths"][path][method]
|
||||
return {
|
||||
scope
|
||||
for requirement in operation["security"]
|
||||
for scopes in requirement.values()
|
||||
for scope in scopes
|
||||
}
|
||||
|
||||
|
||||
def _credential(cred_id: str, is_managed: bool = False) -> APIKeyCredentials:
|
||||
return APIKeyCredentials(
|
||||
id=cred_id,
|
||||
provider="openai",
|
||||
title="t",
|
||||
api_key=SecretStr("secret"),
|
||||
is_managed=is_managed,
|
||||
)
|
||||
|
||||
|
||||
def _mock_creds_store(mocker: pytest_mock.MockFixture, credential: APIKeyCredentials):
|
||||
"""Returns the shared deletion the route hands an allowed delete to."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.integrations.credentials.creds_manager.store"
|
||||
".get_creds_by_id",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=credential,
|
||||
)
|
||||
return mocker.patch(
|
||||
"backend.api.external.v2.integrations.credentials"
|
||||
".delete_credentials_and_dependents",
|
||||
new_callable=mock.AsyncMock,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_calls(fail_first: bool = False):
|
||||
"""An API-key validator that counts how often the request verified the key."""
|
||||
|
||||
async def validate(key: str) -> APIAuthorizationInfo:
|
||||
validate.count += 1
|
||||
if fail_first and validate.count == 1:
|
||||
raise RuntimeError("credential store is briefly unreachable")
|
||||
return APIAuthorizationInfo(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
created_at=datetime.now(tz=timezone.utc),
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
validate.count = 0
|
||||
return validate
|
||||
|
||||
|
||||
def _client() -> fastapi.testclient.TestClient:
|
||||
"""A one-route v2-shaped app: the real middleware over the real dependency."""
|
||||
from .global_rate_limit import GlobalRateLimitMiddleware
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
|
||||
@app.get("/probe")
|
||||
async def probe(auth: TenantContext = fastapi.Security(require_auth)) -> str:
|
||||
return auth.user_id
|
||||
|
||||
app.add_middleware(GlobalRateLimitMiddleware)
|
||||
add_v2_exception_handlers(app)
|
||||
return fastapi.testclient.TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _offline(mocker: pytest_mock.MockFixture) -> None:
|
||||
"""No test here is about Redis or the org tables."""
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.RateLimiter.check",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.RateLimiter.exhausted",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=False,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.tenancy.resolve_credential_tenancy",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=(ORG_ID, None),
|
||||
)
|
||||
@@ -0,0 +1,24 @@
|
||||
"""V2 External API - Per-Endpoint Rate Limiters"""
|
||||
|
||||
from fastapi import Response
|
||||
|
||||
from backend.api.utils.rate_limit import RateLimiter
|
||||
|
||||
media_upload_limiter = RateLimiter(
|
||||
"v2:media_upload", max_requests=10, window_seconds=300
|
||||
)
|
||||
search_limiter = RateLimiter("v2:search", max_requests=30, window_seconds=60)
|
||||
graph_exec_limiter = RateLimiter("v2:graph_exec", max_requests=60, window_seconds=60)
|
||||
file_upload_limiter = RateLimiter("v2:file_upload", max_requests=20, window_seconds=300)
|
||||
# Every call fans out to uncached Stripe reads (proration, period end, pending
|
||||
# change), so the cap is the internal endpoint's: 60/min/user.
|
||||
subscription_limiter = RateLimiter(
|
||||
"v2:subscription", max_requests=60, window_seconds=60
|
||||
)
|
||||
|
||||
|
||||
async def enforce(limiter: RateLimiter, user_id: str, response: Response) -> None:
|
||||
"""Apply `limiter` and report its window instead of the global one: it is the
|
||||
narrower cap, so it is the one the caller has to back off on."""
|
||||
if state := await limiter.check(user_id):
|
||||
response.headers.update(state.headers())
|
||||
@@ -0,0 +1,487 @@
|
||||
"""Rate limiting as a contract: what a client is told, and what it is keyed on.
|
||||
|
||||
A cap the caller cannot see forces it to retry blind, and an anonymous bucket
|
||||
keyed on a header the caller writes is not a cap at all.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException, Response
|
||||
|
||||
from backend.api.external.middleware import resolve_request_auth
|
||||
from backend.api.external.v2 import credits, global_rate_limit
|
||||
from backend.api.external.v2.global_rate_limit import (
|
||||
GlobalRateLimitMiddleware,
|
||||
client_ip,
|
||||
)
|
||||
from backend.api.utils.rate_limit import RateLimiter
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.util.settings import Config
|
||||
|
||||
PEER = "10.0.0.9"
|
||||
MID_MINUTE = datetime(2026, 1, 1, 12, 0, 30, tzinfo=UTC)
|
||||
|
||||
|
||||
async def test_a_request_under_the_cap_is_told_where_it_stands(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
redis.incr.return_value = 3
|
||||
|
||||
state = await RateLimiter("t", max_requests=10, window_seconds=60).check("u1")
|
||||
|
||||
assert state is not None
|
||||
assert state.headers()["X-RateLimit-Limit"] == "10"
|
||||
assert state.headers()["X-RateLimit-Remaining"] == "7"
|
||||
assert 1 <= int(state.headers()["X-RateLimit-Reset"]) <= 60
|
||||
|
||||
|
||||
async def test_a_blocked_request_carries_retry_after(redis: mock.AsyncMock) -> None:
|
||||
"""Without it a client retries blind, adding load to what the cap protects."""
|
||||
redis.incr.return_value = 11
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await RateLimiter("t", max_requests=10, window_seconds=60).check("u1")
|
||||
|
||||
headers = raised.value.headers or {}
|
||||
assert raised.value.status_code == 429
|
||||
assert 1 <= int(headers["Retry-After"]) <= 60
|
||||
assert headers["X-RateLimit-Remaining"] == "0"
|
||||
|
||||
|
||||
async def test_an_unmeasurable_window_publishes_no_numbers(
|
||||
redis: mock.AsyncMock,
|
||||
) -> None:
|
||||
"""Redis is down: fail open, but do not report a count nobody measured."""
|
||||
redis.incr.side_effect = ConnectionError("redis is gone")
|
||||
|
||||
assert (
|
||||
await RateLimiter("t", max_requests=10, window_seconds=60).check("u1") is None
|
||||
)
|
||||
|
||||
|
||||
async def test_the_response_carries_the_callers_window_position(
|
||||
redis: mock.AsyncMock, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""Behind the hosted platform's three hops: the client's bucket, and told so."""
|
||||
mocker.patch.object(
|
||||
global_rate_limit.settings.config,
|
||||
"trusted_proxy_count",
|
||||
Config.model_fields["trusted_proxy_count"].default,
|
||||
)
|
||||
redis.incr.return_value = 1
|
||||
sent: list[dict] = []
|
||||
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
|
||||
async def send(message):
|
||||
sent.append(message)
|
||||
|
||||
forwarded = b"203.0.113.7, 172.70.1.1, 34.1.1.1"
|
||||
scope = _scope() | {"headers": [(b"x-forwarded-for", forwarded)]}
|
||||
await GlobalRateLimitMiddleware(app)(scope, _receive, send)
|
||||
|
||||
assert redis.incr.await_args.args[0].startswith("rl:v2:global:anon:203.0.113.7:")
|
||||
headers = dict(sent[0]["headers"])
|
||||
assert headers[b"x-ratelimit-limit"] == b"5"
|
||||
assert headers[b"x-ratelimit-remaining"] == b"4"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"hops, forwarded, expected",
|
||||
[
|
||||
(1, "203.0.113.7", "203.0.113.7"),
|
||||
# The spoofed entry sits left of the one our proxy appended.
|
||||
(1, "1.2.3.4, 203.0.113.7", "203.0.113.7"),
|
||||
(2, "203.0.113.7, 198.51.100.1", "203.0.113.7"),
|
||||
(2, "9.9.9.9, 203.0.113.7, 198.51.100.1", "203.0.113.7"),
|
||||
# Fewer hops than configured means the header did not come through
|
||||
# our own proxies; the socket peer is the only trustworthy value.
|
||||
(2, "203.0.113.7", PEER),
|
||||
(1, "", PEER),
|
||||
(0, "203.0.113.7", PEER),
|
||||
],
|
||||
)
|
||||
def test_the_anonymous_bucket_ignores_caller_written_hops(
|
||||
mocker: pytest_mock.MockFixture, hops: int, forwarded: str, expected: str
|
||||
) -> None:
|
||||
"""A caller that can set the key picks its own bucket, so there is no cap."""
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.settings.config.trusted_proxy_count",
|
||||
hops,
|
||||
)
|
||||
|
||||
headers = [(b"x-forwarded-for", forwarded.encode())] if forwarded else []
|
||||
assert client_ip(_scope() | {"headers": headers}) == expected
|
||||
|
||||
|
||||
async def test_the_subscription_read_is_capped_before_it_reaches_stripe(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Uncached Stripe reads on every call, so the cap has to gate the fan-out."""
|
||||
user = mocker.patch.object(credits, "get_user_by_id", new=mock.AsyncMock())
|
||||
mocker.patch.object(
|
||||
credits.subscription_limiter,
|
||||
"check",
|
||||
new=mock.AsyncMock(side_effect=HTTPException(status_code=429, detail="nope")),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await credits.get_subscription_status(
|
||||
response=Response(), auth=mock.Mock(user_id="u1")
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 429
|
||||
user.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis(mocker: pytest_mock.MockFixture) -> mock.AsyncMock:
|
||||
client = mock.AsyncMock()
|
||||
client.set.return_value = True
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.get_redis_async",
|
||||
new=mock.AsyncMock(return_value=client),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(return_value=None),
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
def _scope() -> dict:
|
||||
return {"type": "http", "client": (PEER, 4242), "headers": []}
|
||||
|
||||
|
||||
async def _receive() -> dict:
|
||||
return {"type": "http.request"}
|
||||
|
||||
|
||||
class _Counters:
|
||||
"""A fake Redis window store: counts per key, readable as the limiter reads them."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.counts: dict[str, int] = {}
|
||||
|
||||
async def set(self, name: str, value: int, ex: int = 0, nx: bool = False) -> bool:
|
||||
if nx and name in self.counts:
|
||||
return False
|
||||
self.counts[name] = int(value)
|
||||
return True
|
||||
|
||||
async def incr(self, name: str) -> int:
|
||||
self.counts[name] = self.counts.get(name, 0) + 1
|
||||
return self.counts[name]
|
||||
|
||||
async def get(self, name: str) -> bytes | None:
|
||||
return str(self.counts[name]).encode() if name in self.counts else None
|
||||
|
||||
def total(self, fragment: str) -> int:
|
||||
return sum(n for key, n in self.counts.items() if fragment in key)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def counters(mocker: pytest_mock.MockFixture) -> _Counters:
|
||||
"""Its limiters' clock stays mid-minute, so no test's calls split across
|
||||
two windows."""
|
||||
store = _Counters()
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.get_redis_async",
|
||||
new=mock.AsyncMock(return_value=store),
|
||||
)
|
||||
mocker.patch("backend.api.utils.rate_limit.datetime", _MidMinute)
|
||||
return store
|
||||
|
||||
|
||||
class _MidMinute(datetime):
|
||||
@classmethod
|
||||
def now(cls, tz=None):
|
||||
return MID_MINUTE.astimezone(tz)
|
||||
|
||||
|
||||
async def _send_with_key(key: bytes, verify: mock.AsyncMock) -> int:
|
||||
"""One request carrying `key`; the status the middleware answered with."""
|
||||
return await _send([(b"x-api-key", key)])
|
||||
|
||||
|
||||
async def _send_with_bearer(token: bytes) -> int:
|
||||
"""One request carrying `token` as a bearer; the status it was answered with."""
|
||||
return await _send([(b"authorization", b"Bearer " + token)])
|
||||
|
||||
|
||||
async def _send(headers: list[tuple[bytes, bytes]]) -> int:
|
||||
statuses: list[int] = []
|
||||
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
|
||||
async def send(message):
|
||||
if message["type"] == "http.response.start":
|
||||
statuses.append(message["status"])
|
||||
|
||||
scope = _scope()
|
||||
scope["headers"] = headers
|
||||
await GlobalRateLimitMiddleware(app)(scope, _receive, send)
|
||||
return statuses[0]
|
||||
|
||||
|
||||
async def test_a_valid_key_flooding_the_api_is_refused_before_the_hash(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
"""The per-user cap answers only after the hash it exists to bound."""
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(user_id="user-1")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
statuses = [await _send_with_key(b"agpt_validkey", verify) for _ in range(310)]
|
||||
|
||||
assert statuses.count(200) == 300
|
||||
assert statuses[-1] == 429
|
||||
assert verify.await_count == 300
|
||||
|
||||
|
||||
async def test_a_failing_key_is_refused_unhashed_without_locking_out_other_keys(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
"""A bad key behind a shared address throttles its own head, not the address."""
|
||||
rejection = HTTPException(status_code=401, detail="Invalid API key")
|
||||
|
||||
async def resolve(scope, api_key, bearer):
|
||||
if api_key.startswith("agpt_bad"):
|
||||
raise rejection
|
||||
return mock.Mock(user_id="user-1")
|
||||
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(side_effect=resolve),
|
||||
)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
for _ in range(30):
|
||||
await _send_with_key(b"agpt_badkey", verify)
|
||||
hashed = verify.await_count
|
||||
|
||||
assert await _send_with_key(b"agpt_badkey", verify) == 429
|
||||
assert verify.await_count == hashed
|
||||
assert await _send_with_key(b"agpt_goodkey", verify) == 200
|
||||
|
||||
|
||||
async def test_an_address_failing_across_many_heads_is_refused_unhashed(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(side_effect=HTTPException(status_code=401, detail="no")),
|
||||
)
|
||||
counters.counts[global_rate_limit._failed_auth_limiter._key(PEER, MID_MINUTE)] = 300
|
||||
|
||||
assert await _send_with_key(b"agpt_fresh", verify) == 429
|
||||
verify.assert_not_awaited()
|
||||
|
||||
|
||||
async def test_a_token_shaped_api_key_header_is_counted_as_a_key(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
"""Sent as X-API-Key it is tried as a key, so it can cost a hash."""
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(user_id="user-1")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
await _send_with_key(b"agpt_xt_value", verify)
|
||||
|
||||
assert counters.total("credential-presented") == 1
|
||||
|
||||
|
||||
async def test_a_valid_oauth_token_flooding_the_api_is_refused_before_the_lookup(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
"""A token costs a database lookup, which the per-user cap counts only after."""
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(return_value=mock.Mock(user_id="user-1")),
|
||||
)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
statuses = [await _send_with_bearer(b"agpt_xt_token") for _ in range(310)]
|
||||
|
||||
assert statuses.count(200) == 300
|
||||
assert statuses[-1] == 429
|
||||
assert verify.await_count == 300
|
||||
assert "agpt_xt_token" not in "".join(counters.counts)
|
||||
|
||||
|
||||
async def test_forged_oauth_tokens_stop_costing_lookups_once_the_address_keeps_failing(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
"""Every forged token is new to its own counter; the address's is what fills."""
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(side_effect=HTTPException(status_code=401, detail="no")),
|
||||
)
|
||||
|
||||
for n in range(300):
|
||||
await _send_with_bearer(b"agpt_xt_forged-%d" % n)
|
||||
|
||||
assert await _send_with_bearer(b"agpt_xt_forged-next") == 429
|
||||
assert verify.await_count == 300
|
||||
|
||||
|
||||
async def test_one_failing_oauth_token_does_not_lock_out_another(
|
||||
mocker: pytest_mock.MockFixture, counters: _Counters
|
||||
) -> None:
|
||||
rejection = HTTPException(status_code=401, detail="Invalid token")
|
||||
|
||||
async def resolve(scope, api_key, bearer):
|
||||
if bearer.credentials == "agpt_xt_expired":
|
||||
raise rejection
|
||||
return mock.Mock(user_id="user-1")
|
||||
|
||||
verify = mocker.patch(
|
||||
"backend.api.external.v2.global_rate_limit.resolve_request_auth",
|
||||
new=mock.AsyncMock(side_effect=resolve),
|
||||
)
|
||||
mocker.patch.object(
|
||||
global_rate_limit._authenticated_limiter, "check", return_value=None
|
||||
)
|
||||
|
||||
for _ in range(30):
|
||||
await _send_with_bearer(b"agpt_xt_expired")
|
||||
looked_up = verify.await_count
|
||||
|
||||
assert await _send_with_bearer(b"agpt_xt_expired") == 429
|
||||
assert verify.await_count == looked_up
|
||||
assert await _send_with_bearer(b"agpt_xt_current") == 200
|
||||
|
||||
|
||||
async def test_a_request_verifies_its_credential_once_even_when_rejected(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The rate limiter and the route both ask; an invalid key is hashed once."""
|
||||
validate = mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
new=mock.AsyncMock(return_value=None),
|
||||
)
|
||||
scope: dict = {"type": "http"}
|
||||
|
||||
for _ in range(2):
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await resolve_request_auth(scope, api_key="agpt_wrongkey", bearer=None)
|
||||
assert raised.value.status_code == 401
|
||||
|
||||
validate.assert_awaited_once()
|
||||
|
||||
|
||||
async def test_a_cached_verification_is_reused_only_for_the_same_credential(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
validate = mocker.patch(
|
||||
"backend.api.external.middleware.validate_api_key",
|
||||
new=mock.AsyncMock(side_effect=[_principal("a"), _principal("b")]),
|
||||
)
|
||||
scope: dict = {"type": "http"}
|
||||
|
||||
first = await resolve_request_auth(scope, api_key="agpt_a", bearer=None)
|
||||
again = await resolve_request_auth(scope, api_key="agpt_a", bearer=None)
|
||||
other = await resolve_request_auth(scope, api_key="agpt_b", bearer=None)
|
||||
|
||||
assert first is again
|
||||
assert other.user_id == "b"
|
||||
assert validate.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path, bucket",
|
||||
[
|
||||
("/api/v2/openapi.json", "anon-docs:"),
|
||||
("/api/v2/docs", "anon-docs:"),
|
||||
("/api/v2/library/agents", "anon:"),
|
||||
],
|
||||
)
|
||||
async def test_reading_the_docs_without_a_key_has_its_own_bucket(
|
||||
redis: mock.AsyncMock, path: str, bucket: str
|
||||
) -> None:
|
||||
"""The docs send agents to the spec before they have a key; reading it must
|
||||
not use up the five calls a minute everything else anonymous gets."""
|
||||
redis.incr.return_value = 1
|
||||
|
||||
async def app(scope, receive, send):
|
||||
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||
|
||||
async def send(message):
|
||||
pass
|
||||
|
||||
scope = _scope()
|
||||
scope.update(path=path, root_path="/api/v2")
|
||||
await GlobalRateLimitMiddleware(app)(scope, _receive, send)
|
||||
|
||||
assert redis.incr.await_args.args[0].startswith(f"rl:v2:global:{bucket}")
|
||||
|
||||
|
||||
def _principal(user_id: str) -> APIAuthorizationInfo:
|
||||
return APIAuthorizationInfo(
|
||||
user_id=user_id, scopes=[], type="api_key", created_at=datetime.now(UTC)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call", ["check", "exhausted"])
|
||||
async def test_a_limiter_that_fails_open_logs_no_subject(
|
||||
mocker: pytest_mock.MockFixture, caplog: pytest.LogCaptureFixture, call: str
|
||||
):
|
||||
"""A pre-verification bucket names an API key by its head; a Redis outage
|
||||
must not put that in the log."""
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.get_redis_async",
|
||||
new=mock.AsyncMock(side_effect=ConnectionError("down")),
|
||||
)
|
||||
limiter = RateLimiter("probe", max_requests=1, window_seconds=60)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="backend.api.utils.rate_limit"):
|
||||
await getattr(limiter, call)("203.0.113.7:agpt_abc")
|
||||
|
||||
assert "failed open" in caplog.text
|
||||
assert "agpt_abc" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call, open_answer", [("check", None), ("exhausted", False)])
|
||||
async def test_a_hanging_redis_fails_the_limiter_open_fast(
|
||||
mocker: pytest_mock.MockFixture, call: str, open_answer: object
|
||||
):
|
||||
"""Every v2 request waits on these checks."""
|
||||
mocker.patch(
|
||||
"backend.api.utils.rate_limit.get_redis_async",
|
||||
new=mock.AsyncMock(return_value=_HangingRedis()),
|
||||
)
|
||||
limiter = RateLimiter("probe", max_requests=1, window_seconds=60)
|
||||
|
||||
answer = await asyncio.wait_for(getattr(limiter, call)("user-1"), timeout=5)
|
||||
|
||||
assert answer is open_answer
|
||||
|
||||
|
||||
class _HangingRedis:
|
||||
async def set(self, *args, **kwargs):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def incr(self, *args, **kwargs):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def get(self, *args, **kwargs):
|
||||
await asyncio.Event().wait()
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
V2 External API Routes
|
||||
|
||||
This module defines the main v2 router that aggregates all v2 API endpoints.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter
|
||||
from starlette import status
|
||||
|
||||
from .blocks import blocks_router
|
||||
from .credits import credits_router
|
||||
from .errors import ErrorResponse
|
||||
from .files import file_workspace_router
|
||||
from .graphs import graphs_router
|
||||
from .identity import identity_router
|
||||
from .integrations import integrations_router
|
||||
from .library import library_router
|
||||
from .marketplace import marketplace_router
|
||||
from .runs import runs_router
|
||||
from .schedules import schedules_router
|
||||
from .search import search_router
|
||||
|
||||
# Declared once, on the parent router, so every v2 operation advertises the one
|
||||
# error body shape instead of FastAPI's default `{"detail": ...}`.
|
||||
ERROR_RESPONSES: dict[int | str, dict] = {
|
||||
code: {"model": ErrorResponse, "description": description}
|
||||
for code, description in {
|
||||
status.HTTP_400_BAD_REQUEST: "Malformed request",
|
||||
status.HTTP_401_UNAUTHORIZED: "Missing or invalid credentials",
|
||||
status.HTTP_403_FORBIDDEN: "Credentials lack a required scope",
|
||||
status.HTTP_404_NOT_FOUND: "No such resource",
|
||||
status.HTTP_413_CONTENT_TOO_LARGE: "Request body too large",
|
||||
status.HTTP_422_UNPROCESSABLE_CONTENT: "Request failed validation",
|
||||
status.HTTP_429_TOO_MANY_REQUESTS: "Rate limit exceeded",
|
||||
status.HTTP_500_INTERNAL_SERVER_ERROR: "Unhandled server error",
|
||||
}.items()
|
||||
}
|
||||
|
||||
v2_router = APIRouter(responses=ERROR_RESPONSES)
|
||||
|
||||
# Include all sub-routers
|
||||
v2_router.include_router(blocks_router, prefix="/blocks")
|
||||
v2_router.include_router(credits_router, prefix="/credits")
|
||||
v2_router.include_router(file_workspace_router, prefix="/files")
|
||||
v2_router.include_router(graphs_router, prefix="/graphs")
|
||||
v2_router.include_router(identity_router)
|
||||
v2_router.include_router(integrations_router, prefix="/integrations")
|
||||
v2_router.include_router(library_router, prefix="/library")
|
||||
v2_router.include_router(marketplace_router, prefix="/marketplace")
|
||||
v2_router.include_router(runs_router, prefix="/runs")
|
||||
v2_router.include_router(schedules_router, prefix="/schedules")
|
||||
v2_router.include_router(search_router, prefix="/search")
|
||||
@@ -0,0 +1,119 @@
|
||||
"""A run's `credentials_inputs` reach the executor as credential models.
|
||||
|
||||
`add_graph_execution` maps each credential onto the graph's nodes with
|
||||
`.model_dump()`. Typed as a plain dict, every v2 run of an agent that needs a
|
||||
credential reached it as a dict and failed with a 500.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import Response
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import ValidationError
|
||||
|
||||
from backend.data.execution import ExecutionStatus, GraphExecutionMeta
|
||||
from backend.data.model import CredentialsMetaInput
|
||||
|
||||
from .models import AgentRunRequest, AgentRunScheduleCreateRequest
|
||||
from .tenancy import TenantContext
|
||||
|
||||
FIELD = "openai_api_key_credentials"
|
||||
CREDENTIAL = {"id": "cred-1", "title": "Mine", "provider": "openai", "type": "api_key"}
|
||||
|
||||
|
||||
def test_a_run_request_parses_credentials_into_models() -> None:
|
||||
request = AgentRunRequest.model_validate(
|
||||
{"credentials_inputs": {FIELD: CREDENTIAL}}
|
||||
)
|
||||
|
||||
parsed = request.credentials_inputs[FIELD]
|
||||
assert isinstance(parsed, CredentialsMetaInput)
|
||||
assert parsed.model_dump(exclude_none=True) == CREDENTIAL
|
||||
|
||||
|
||||
def test_a_schedule_request_parses_credentials_into_models() -> None:
|
||||
request = AgentRunScheduleCreateRequest.model_validate(
|
||||
{
|
||||
"graph_id": "graph-1",
|
||||
"name": "Daily",
|
||||
"cron": "0 9 * * *",
|
||||
"credentials_inputs": {FIELD: CREDENTIAL},
|
||||
}
|
||||
)
|
||||
|
||||
assert isinstance(request.credentials_inputs[FIELD], CredentialsMetaInput)
|
||||
|
||||
|
||||
def test_a_credential_without_an_id_is_a_validation_error() -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
AgentRunRequest.model_validate(
|
||||
{"credentials_inputs": {FIELD: {"provider": "openai", "type": "api_key"}}}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_run_hands_the_executor_credential_models(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .library.agents import execute_agent
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.agents.graph_exec_limiter.check",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(
|
||||
organization_id="org-1",
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
input_schema={},
|
||||
),
|
||||
)
|
||||
add_execution = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=GraphExecutionMeta.model_construct(
|
||||
id="run-1",
|
||||
user_id="user-1",
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
preset_id=None,
|
||||
status=ExecutionStatus.QUEUED,
|
||||
started_at=None,
|
||||
ended_at=None,
|
||||
inputs={},
|
||||
credential_inputs=None,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
stats=None,
|
||||
organization_id="org-1",
|
||||
team_id=None,
|
||||
),
|
||||
)
|
||||
|
||||
await execute_agent(
|
||||
response=Response(),
|
||||
request=AgentRunRequest.model_validate(
|
||||
{"credentials_inputs": {FIELD: CREDENTIAL}}
|
||||
),
|
||||
agent_id="agent-1",
|
||||
auth=TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=[APIKeyPermission.RUN_AGENT],
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
),
|
||||
)
|
||||
|
||||
passed = add_execution.await_args.kwargs["graph_credentials_inputs"]
|
||||
assert passed[FIELD].model_dump(exclude_none=True) == CREDENTIAL
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Run start and schedule create refuse inputs the graph's `input_schema` rejects.
|
||||
|
||||
The executor runs a graph with a misspelled or missing input anyway, and the run
|
||||
ends COMPLETED with no outputs and no error, so the request is the only place a
|
||||
client can learn its inputs were wrong.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
import fastapi
|
||||
import fastapi.testclient
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.blocks.io import AgentInputBlock
|
||||
from backend.data.execution import ExecutionStatus, GraphExecutionMeta
|
||||
from backend.data.graph import BaseGraph
|
||||
|
||||
from .errors import add_v2_exception_handlers
|
||||
from .idempotency import IDEMPOTENCY_HEADER
|
||||
from .library import library_router
|
||||
from .schedules import schedules_router
|
||||
from .tenancy import TenantContext, require_auth
|
||||
|
||||
ORG_ID = "org-1"
|
||||
INPUT_SCHEMA = BaseGraph._generate_schema(
|
||||
(AgentInputBlock.Input, {"name": "topic"}),
|
||||
(AgentInputBlock.Input, {"name": "tone", "value": "neutral"}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/library/agents/agent-1/runs", "/schedules"])
|
||||
@pytest.mark.parametrize(
|
||||
"inputs,bad_input,error_type",
|
||||
[
|
||||
({"topic": "AI", "topik": "AI"}, "topik", "extra_forbidden"),
|
||||
({"tone": "dry"}, "topic", "missing"),
|
||||
({"topic": None}, "topic", "missing"),
|
||||
],
|
||||
)
|
||||
def test_inputs_the_graph_rejects_are_a_422_naming_them(
|
||||
client: fastapi.testclient.TestClient,
|
||||
started: mock.AsyncMock,
|
||||
path: str,
|
||||
inputs: dict[str, Any],
|
||||
bad_input: str,
|
||||
error_type: str,
|
||||
) -> None:
|
||||
response = client.post(path, json=_body(path, inputs))
|
||||
|
||||
assert response.status_code == 422, response.text
|
||||
error = response.json()["error"]
|
||||
assert error["code"] == "validation_error"
|
||||
assert [(e["loc"], e["type"]) for e in error["details"]["errors"]] == [
|
||||
(["body", "inputs", bad_input], error_type)
|
||||
]
|
||||
started.assert_not_awaited()
|
||||
|
||||
|
||||
def test_inputs_the_graph_accepts_start_the_run(
|
||||
client: fastapi.testclient.TestClient, started: mock.AsyncMock
|
||||
) -> None:
|
||||
response = client.post(
|
||||
"/library/agents/agent-1/runs", json={"inputs": {"topic": "AI"}}
|
||||
)
|
||||
|
||||
assert response.status_code == 202, response.text
|
||||
assert started.await_args.kwargs["inputs"] == {"topic": "AI"}
|
||||
|
||||
|
||||
def test_a_422_does_not_hold_the_idempotency_key(
|
||||
client: fastapi.testclient.TestClient, started: mock.AsyncMock
|
||||
) -> None:
|
||||
"""Held, the key would answer the corrected retry with a 409 for 24 hours."""
|
||||
headers = {IDEMPOTENCY_HEADER: "retry-me"}
|
||||
path = "/library/agents/agent-1/runs"
|
||||
|
||||
refused = client.post(path, json={"inputs": {}}, headers=headers)
|
||||
retried = client.post(path, json={"inputs": {"topic": "AI"}}, headers=headers)
|
||||
|
||||
assert refused.status_code == 422, refused.text
|
||||
assert retried.status_code == 202, retried.text
|
||||
started.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path,inputs,replayed",
|
||||
[
|
||||
("/library/agents/agent-1/runs", {"tone": "dry", "topic": "AI"}, True),
|
||||
("/library/agents/agent-2/runs", {"topic": "AI", "tone": "dry"}, False),
|
||||
("/library/agents/agent-1/runs", {"topic": "ML", "tone": "dry"}, False),
|
||||
],
|
||||
ids=["same-request", "other-agent", "other-inputs"],
|
||||
)
|
||||
def test_a_key_reused_on_a_different_request_is_a_422_not_the_earlier_run(
|
||||
client: fastapi.testclient.TestClient,
|
||||
started: mock.AsyncMock,
|
||||
mocker: pytest_mock.MockFixture,
|
||||
path: str,
|
||||
inputs: dict[str, Any],
|
||||
replayed: bool,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.data.execution.get_graph_execution",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
headers = {IDEMPOTENCY_HEADER: "reused"}
|
||||
first = client.post(
|
||||
"/library/agents/agent-1/runs",
|
||||
json={"inputs": {"topic": "AI", "tone": "dry"}},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
second = client.post(path, json={"inputs": inputs}, headers=headers)
|
||||
|
||||
assert first.status_code == 202, first.text
|
||||
assert second.status_code == (202 if replayed else 422), second.text
|
||||
started.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(mocker: pytest_mock.MockFixture) -> fastapi.testclient.TestClient:
|
||||
store: dict[str, str] = {}
|
||||
|
||||
async def set_(name: str, value: str, nx: bool = False, ex: int = 0) -> bool:
|
||||
if nx and name in store:
|
||||
return False
|
||||
store[name] = value
|
||||
return True
|
||||
|
||||
async def delete(name: str) -> None:
|
||||
store.pop(name, None)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.idempotency.get_redis_async",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(
|
||||
set=mock.AsyncMock(side_effect=set_),
|
||||
get=mock.AsyncMock(side_effect=store.get),
|
||||
delete=mock.AsyncMock(side_effect=delete),
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.agents.graph_exec_limiter.check",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(get_credits=mock.AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(
|
||||
organization_id=ORG_ID,
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
input_schema=INPUT_SCHEMA,
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=mock.Mock(version=1, input_schema=INPUT_SCHEMA),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_active_graph_version",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=1,
|
||||
)
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(library_router, prefix="/library")
|
||||
app.include_router(schedules_router, prefix="/schedules")
|
||||
app.dependency_overrides[require_auth] = lambda: TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
add_v2_exception_handlers(app)
|
||||
return fastapi.testclient.TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def started(mocker: pytest_mock.MockFixture) -> mock.AsyncMock:
|
||||
"""Every way a request can start work: a run now, or a schedule of them."""
|
||||
run = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=mock.AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
scheduler = mocker.patch(
|
||||
"backend.api.external.v2.schedules.get_scheduler_client"
|
||||
).return_value
|
||||
scheduler.add_execution_schedule = run
|
||||
return run
|
||||
|
||||
|
||||
def _body(path: str, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
if path == "/schedules":
|
||||
return {
|
||||
"graph_id": "graph-1",
|
||||
"name": "Daily",
|
||||
"cron": "0 9 * * *",
|
||||
"timezone": "UTC",
|
||||
"inputs": inputs,
|
||||
}
|
||||
return {"inputs": inputs}
|
||||
|
||||
|
||||
def _run() -> GraphExecutionMeta:
|
||||
return GraphExecutionMeta.model_construct(
|
||||
id="run-1",
|
||||
user_id="user-1",
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
preset_id=None,
|
||||
status=ExecutionStatus.QUEUED,
|
||||
started_at=None,
|
||||
ended_at=None,
|
||||
inputs={},
|
||||
credential_inputs=None,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
stats=None,
|
||||
organization_id=ORG_ID,
|
||||
team_id=None,
|
||||
)
|
||||
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
V2 External API - Runs Endpoints
|
||||
|
||||
Provides access to agent runs and human-in-the-loop reviews.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Query, Security
|
||||
from prisma.enums import APIKeyPermission, ReviewStatus
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.graph_executions import sharing
|
||||
from backend.api.features.graph_executions.review.model import ReviewItem
|
||||
from backend.api.features.graph_executions.review.service import process_reviews
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data import human_review as review_db
|
||||
from backend.data.execution import ExecutionStatus, GraphExecution
|
||||
from backend.executor import utils as execution_utils
|
||||
|
||||
from .models import (
|
||||
AgentGraphRun,
|
||||
AgentGraphRunDetails,
|
||||
AgentRunReview,
|
||||
AgentRunReviewsSubmitRequest,
|
||||
AgentRunReviewsSubmitResponse,
|
||||
AgentRunReviewStatus,
|
||||
AgentRunShareResponse,
|
||||
RunStatus,
|
||||
)
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .tenancy import (
|
||||
TenantContext,
|
||||
in_tenant,
|
||||
owned_by_caller,
|
||||
require_paid_permission,
|
||||
require_permission,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
runs_router = APIRouter(tags=["runs"])
|
||||
|
||||
|
||||
# Registered before the `/{run_id}` routes below: Starlette matches in
|
||||
# registration order, so `/{run_id}` would otherwise swallow `/reviews`.
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints - Reviews (Human-in-the-loop)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@runs_router.get(
|
||||
path="/reviews",
|
||||
summary="List agent run human-in-the-loop reviews",
|
||||
operation_id="listAgentRunReviews",
|
||||
)
|
||||
async def list_reviews(
|
||||
run_id: Optional[str] = Query(
|
||||
default=None, description="Filter by graph execution ID"
|
||||
),
|
||||
review_status: Optional[AgentRunReviewStatus] = Query(
|
||||
default=None,
|
||||
alias="status",
|
||||
description="Filter by review status",
|
||||
),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_RUN_REVIEW)
|
||||
),
|
||||
) -> Page[AgentRunReview]:
|
||||
"""
|
||||
List human-in-the-loop reviews for agent runs.
|
||||
|
||||
Returns reviews of all statuses if no status filter is given.
|
||||
"""
|
||||
reviews, pagination = await review_db.get_reviews(
|
||||
user_id=auth.user_id,
|
||||
graph_exec_id=run_id,
|
||||
status=ReviewStatus(review_status) if review_status else None,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
organization_id=auth.organization_id,
|
||||
graph_runs_only=True,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[AgentRunReview.from_internal(r) for r in reviews],
|
||||
total_count=pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@runs_router.post(
|
||||
path="/{run_id}/reviews",
|
||||
summary="Submit agent run human-in-the-loop reviews",
|
||||
operation_id="submitAgentRunReviews",
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def submit_reviews(
|
||||
request: AgentRunReviewsSubmitRequest,
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(APIKeyPermission.WRITE_RUN_REVIEW)
|
||||
),
|
||||
) -> AgentRunReviewsSubmitResponse:
|
||||
"""
|
||||
Submit responses to all pending human-in-the-loop reviews for a run.
|
||||
|
||||
All pending reviews for the run must be included in the request.
|
||||
Approving a review continues execution; rejecting terminates that branch.
|
||||
"""
|
||||
# Reviews carry no organization of their own; the run they belong to does.
|
||||
await _assert_own_run_in_tenant(run_id, auth)
|
||||
|
||||
outcome = await process_reviews(
|
||||
auth.user_id,
|
||||
[
|
||||
ReviewItem(
|
||||
node_exec_id=decision.node_exec_id,
|
||||
approved=decision.approved,
|
||||
reviewed_data=decision.edited_payload,
|
||||
message=decision.message,
|
||||
auto_approve_future=decision.auto_approve_future,
|
||||
)
|
||||
for decision in request.reviews
|
||||
],
|
||||
graph_exec_id=run_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
)
|
||||
|
||||
return AgentRunReviewsSubmitResponse(
|
||||
run_id=run_id,
|
||||
approved_count=outcome.approved_count,
|
||||
rejected_count=outcome.rejected_count,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints - Runs
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@runs_router.get(
|
||||
path="",
|
||||
summary="List agent runs",
|
||||
operation_id="listAgentRuns",
|
||||
)
|
||||
async def list_runs(
|
||||
graph_id: Optional[str] = Query(default=None, description="Filter by graph ID"),
|
||||
# `Annotated`, not `= Query(...)`: the latter leaves the sentinel as the
|
||||
# Python default, and this handler is also called directly in tests.
|
||||
statuses: Annotated[
|
||||
Optional[list[RunStatus]],
|
||||
Query(description="Filter by run status; repeat to match several"),
|
||||
] = None,
|
||||
started_after: Annotated[
|
||||
Optional[datetime], Query(description="Only runs created at or after this time")
|
||||
] = None,
|
||||
started_before: Annotated[
|
||||
Optional[datetime],
|
||||
Query(description="Only runs created at or before this time"),
|
||||
] = None,
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_RUN)),
|
||||
) -> Page[AgentGraphRun]:
|
||||
"""List agent runs, optionally filtered by graph, status and creation time."""
|
||||
result = await execution_db.get_graph_executions_paginated(
|
||||
user_id=auth.user_id,
|
||||
graph_id=graph_id,
|
||||
statuses=[ExecutionStatus(s) for s in statuses] if statuses else None,
|
||||
created_time_gte=started_after,
|
||||
created_time_lte=started_before,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
return page.paged(
|
||||
[AgentGraphRun.from_internal(e) for e in result.executions],
|
||||
total_count=result.pagination.total_items,
|
||||
)
|
||||
|
||||
|
||||
@runs_router.get(
|
||||
path="/{run_id}",
|
||||
summary="Get agent run details",
|
||||
operation_id="getAgentRunDetails",
|
||||
)
|
||||
async def get_run(
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_RUN)),
|
||||
) -> AgentGraphRunDetails:
|
||||
"""Get detailed information about a specific run."""
|
||||
result = await execution_db.get_graph_execution(
|
||||
user_id=auth.user_id,
|
||||
execution_id=run_id,
|
||||
include_node_executions=True,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
|
||||
if not result:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Run #{run_id} not found",
|
||||
)
|
||||
|
||||
return AgentGraphRunDetails.from_internal(result)
|
||||
|
||||
|
||||
@runs_router.post(
|
||||
path="/{run_id}/stop",
|
||||
summary="Stop agent run",
|
||||
operation_id="stopAgentRun",
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def stop_run(
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_RUN)),
|
||||
) -> AgentGraphRun:
|
||||
"""
|
||||
Stop a run that hasn't finished: one that is incomplete, queued, running,
|
||||
or waiting for a review.
|
||||
|
||||
Waits up to 15 seconds for the run to stop, then returns it with its status
|
||||
at that moment. A run that has already finished answers `409`.
|
||||
"""
|
||||
# The cancel is published before stop_graph_execution checks the owner.
|
||||
run = await _assert_own_run_in_tenant(run_id, auth)
|
||||
if run.status in _FINISHED:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"Run #{run_id} has already finished ({run.status.value})",
|
||||
)
|
||||
|
||||
try:
|
||||
await execution_utils.stop_graph_execution(
|
||||
graph_exec_id=run_id,
|
||||
user_id=auth.user_id,
|
||||
wait_timeout=_STOP_WAIT_SECONDS,
|
||||
)
|
||||
except TimeoutError:
|
||||
# The cancel is published; the executor is still winding the run
|
||||
# down, which the returned status shows.
|
||||
logger.warning(f"Run #{run_id} did not stop within {_STOP_WAIT_SECONDS}s")
|
||||
|
||||
return AgentGraphRun.from_internal(await _assert_own_run_in_tenant(run_id, auth))
|
||||
|
||||
|
||||
@runs_router.delete(
|
||||
path="/{run_id}",
|
||||
summary="Delete agent run",
|
||||
operation_id="deleteAgentRun",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_run(
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_RUN)),
|
||||
) -> None:
|
||||
"""Delete an agent run. A shared run stops being downloadable too."""
|
||||
await _assert_own_run_in_tenant(run_id, auth)
|
||||
|
||||
await sharing.delete_execution(auth.user_id, run_id)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints - Sharing
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@runs_router.post(
|
||||
path="/{run_id}/share",
|
||||
summary="Enable sharing for an agent run",
|
||||
operation_id="enableAgentRunShare",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def enable_sharing(
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_RUN, APIKeyPermission.SHARE_RUN)
|
||||
),
|
||||
) -> AgentRunShareResponse:
|
||||
"""Enable public sharing for a run.
|
||||
|
||||
Sharing again issues a new token, and links from the earlier share stop
|
||||
working.
|
||||
"""
|
||||
await _assert_own_run_in_tenant(run_id, auth)
|
||||
|
||||
share_token = await sharing.share_execution(auth.user_id, run_id)
|
||||
|
||||
return AgentRunShareResponse(
|
||||
share_url=sharing.share_url(share_token), share_token=share_token
|
||||
)
|
||||
|
||||
|
||||
@runs_router.delete(
|
||||
path="/{run_id}/share",
|
||||
summary="Disable sharing for an agent run",
|
||||
operation_id="disableAgentRunShare",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def disable_sharing(
|
||||
run_id: str = Path(description="Graph Execution ID"),
|
||||
auth: TenantContext = Security(
|
||||
require_permission(APIKeyPermission.READ_RUN, APIKeyPermission.SHARE_RUN)
|
||||
),
|
||||
) -> None:
|
||||
"""Disable public sharing for a run, and the file downloads it allowed."""
|
||||
await _assert_own_run_in_tenant(run_id, auth)
|
||||
|
||||
await sharing.unshare_execution(auth.user_id, run_id)
|
||||
|
||||
|
||||
_FINISHED = frozenset(
|
||||
{ExecutionStatus.COMPLETED, ExecutionStatus.FAILED, ExecutionStatus.TERMINATED}
|
||||
)
|
||||
_STOP_WAIT_SECONDS = 15.0
|
||||
|
||||
|
||||
async def _assert_own_run_in_tenant(run_id: str, auth: TenantContext) -> GraphExecution:
|
||||
"""Also 403 for an org-mate's run, which the owner-scoped writes cannot reach."""
|
||||
return owned_by_caller(
|
||||
await _assert_run_in_tenant(run_id, auth), auth, f"Run #{run_id}"
|
||||
)
|
||||
|
||||
|
||||
async def _assert_run_in_tenant(run_id: str, auth: TenantContext) -> GraphExecution:
|
||||
"""404 before acting on a run the credentials cannot reach."""
|
||||
return in_tenant(
|
||||
await execution_db.get_graph_execution(
|
||||
user_id=auth.user_id,
|
||||
execution_id=run_id,
|
||||
organization_id=auth.organization_id,
|
||||
),
|
||||
auth,
|
||||
f"Run #{run_id}",
|
||||
)
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Writes to a run act only on the caller's own run, and sharing goes through
|
||||
the same service as the web app, so a revoked share can't keep its files
|
||||
downloadable."""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data.execution import ExecutionStatus, GraphExecution
|
||||
from backend.util.exceptions import NotAuthorizedError
|
||||
|
||||
from .models import AgentRunReviewDecision, AgentRunReviewsSubmitRequest
|
||||
from .runs import delete_run, disable_sharing, enable_sharing, stop_run, submit_reviews
|
||||
from .tenancy import TenantContext
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG_ID = "org-1"
|
||||
|
||||
|
||||
def _auth() -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
|
||||
def _run(status: ExecutionStatus = ExecutionStatus.RUNNING) -> GraphExecution:
|
||||
now = datetime.now(timezone.utc)
|
||||
return GraphExecution(
|
||||
id="run-1",
|
||||
user_id=USER_ID,
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
inputs={},
|
||||
credential_inputs={},
|
||||
nodes_input_masks={},
|
||||
preset_id=None,
|
||||
status=status,
|
||||
started_at=now,
|
||||
ended_at=None,
|
||||
stats=None,
|
||||
outputs={},
|
||||
organization_id=ORG_ID,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def run_db(mocker: pytest_mock.MockFixture) -> SimpleNamespace:
|
||||
"""The run lookups and writes; `calls` records the order of the writes."""
|
||||
calls: list[str] = []
|
||||
|
||||
def recorder(name: str) -> AsyncMock:
|
||||
return AsyncMock(side_effect=lambda **_: calls.append(name))
|
||||
|
||||
db = SimpleNamespace(
|
||||
calls=calls,
|
||||
get_graph_execution=AsyncMock(return_value=_run()),
|
||||
delete_shared_execution_files=recorder("delete_files"),
|
||||
update_graph_execution_share_status=recorder("update_share"),
|
||||
create_shared_execution_files=recorder("create_files"),
|
||||
delete_graph_execution=recorder("delete_run"),
|
||||
)
|
||||
for name, value in vars(db).items():
|
||||
if name != "calls":
|
||||
mocker.patch.object(execution_db, name, new=value)
|
||||
return db
|
||||
|
||||
|
||||
async def test_sharing_rebuilds_the_file_allowlist(run_db) -> None:
|
||||
response = await enable_sharing(run_id="run-1", auth=_auth())
|
||||
|
||||
assert run_db.calls == ["delete_files", "update_share", "create_files"]
|
||||
assert response.share_url.endswith(response.share_token)
|
||||
|
||||
|
||||
async def test_unsharing_revokes_the_file_downloads(run_db) -> None:
|
||||
await disable_sharing(run_id="run-1", auth=_auth())
|
||||
|
||||
assert run_db.calls == ["delete_files", "update_share"]
|
||||
assert (
|
||||
run_db.update_graph_execution_share_status.await_args.kwargs["is_shared"]
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
async def test_deleting_a_run_revokes_its_file_downloads(run_db) -> None:
|
||||
await delete_run(run_id="run-1", auth=_auth())
|
||||
|
||||
assert run_db.calls == ["delete_run", "delete_files"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"write",
|
||||
[
|
||||
lambda: enable_sharing(run_id="run-1", auth=_auth()),
|
||||
lambda: disable_sharing(run_id="run-1", auth=_auth()),
|
||||
lambda: delete_run(run_id="run-1", auth=_auth()),
|
||||
lambda: stop_run(run_id="run-1", auth=_auth()),
|
||||
lambda: submit_reviews(
|
||||
AgentRunReviewsSubmitRequest(
|
||||
reviews=[AgentRunReviewDecision(node_exec_id="node-1", approved=True)]
|
||||
),
|
||||
run_id="run-1",
|
||||
auth=_auth(),
|
||||
),
|
||||
],
|
||||
ids=["share", "unshare", "delete", "stop", "review"],
|
||||
)
|
||||
async def test_writes_refuse_a_teammates_run(
|
||||
run_db, mocker: pytest_mock.MockFixture, write
|
||||
) -> None:
|
||||
"""A teammate's run in the same organization is visible to reads, but the
|
||||
writes beneath these routes are its owner's: none happens, and the caller
|
||||
gets 403."""
|
||||
run_db.get_graph_execution.return_value = _run().model_copy(
|
||||
update={"user_id": "teammate"}
|
||||
)
|
||||
stop = mocker.patch(
|
||||
"backend.executor.utils.stop_graph_execution", new_callable=AsyncMock
|
||||
)
|
||||
reviews = mocker.patch(
|
||||
"backend.api.external.v2.runs.process_reviews", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await write()
|
||||
|
||||
assert run_db.calls == []
|
||||
stop.assert_not_awaited()
|
||||
reviews.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"finished",
|
||||
[ExecutionStatus.COMPLETED, ExecutionStatus.FAILED, ExecutionStatus.TERMINATED],
|
||||
)
|
||||
async def test_stopping_a_finished_run_is_a_conflict(
|
||||
run_db, mocker: pytest_mock.MockFixture, finished: ExecutionStatus
|
||||
) -> None:
|
||||
run_db.get_graph_execution.return_value = _run(finished)
|
||||
stop = mocker.patch(
|
||||
"backend.executor.utils.stop_graph_execution", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as caught:
|
||||
await stop_run(run_id="run-1", auth=_auth())
|
||||
|
||||
assert caught.value.status_code == 409
|
||||
stop.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stoppable",
|
||||
[
|
||||
ExecutionStatus.INCOMPLETE,
|
||||
ExecutionStatus.QUEUED,
|
||||
ExecutionStatus.RUNNING,
|
||||
ExecutionStatus.REVIEW,
|
||||
],
|
||||
)
|
||||
async def test_stopping_an_unfinished_run_returns_its_status(
|
||||
run_db, mocker: pytest_mock.MockFixture, stoppable: ExecutionStatus
|
||||
) -> None:
|
||||
run_db.get_graph_execution.side_effect = [
|
||||
_run(stoppable),
|
||||
_run(ExecutionStatus.TERMINATED),
|
||||
]
|
||||
stop = mocker.patch(
|
||||
"backend.executor.utils.stop_graph_execution", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
run = await stop_run(run_id="run-1", auth=_auth())
|
||||
|
||||
assert run.status == "TERMINATED"
|
||||
assert stop.await_args.kwargs["graph_exec_id"] == "run-1"
|
||||
|
||||
|
||||
async def test_a_slow_stop_returns_the_run_instead_of_failing(
|
||||
run_db, mocker: pytest_mock.MockFixture
|
||||
) -> None:
|
||||
"""The cancel is already published when the wait times out; the caller
|
||||
gets the run as it stands, not a 500."""
|
||||
run_db.get_graph_execution.side_effect = [_run(), _run()]
|
||||
mocker.patch(
|
||||
"backend.executor.utils.stop_graph_execution",
|
||||
new=AsyncMock(side_effect=TimeoutError()),
|
||||
)
|
||||
|
||||
run = await stop_run(run_id="run-1", auth=_auth())
|
||||
|
||||
assert run.status == "RUNNING"
|
||||
@@ -0,0 +1,177 @@
|
||||
"""
|
||||
V2 External API - Schedules Endpoints
|
||||
|
||||
Provides endpoints for managing execution schedules.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Security
|
||||
from prisma.enums import APIKeyPermission
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.experts import experts_db
|
||||
from backend.api.features.schedule_visibility import visible_graph_schedules
|
||||
from backend.data import graph as graph_db
|
||||
from backend.data.tenancy import get_user_team_ids
|
||||
from backend.data.user import get_user_by_id
|
||||
from backend.executor.scheduler import GraphExecutionJobInfo
|
||||
from backend.util.clients import get_scheduler_client
|
||||
from backend.util.exceptions import NotFoundError
|
||||
from backend.util.timezone_utils import get_user_timezone_or_utc
|
||||
|
||||
from .library.helpers import assert_inputs_match
|
||||
from .models import AgentRunSchedule, AgentRunScheduleCreateRequest
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .tenancy import TenantContext, require_paid_permission, require_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
schedules_router = APIRouter(tags=["schedules"])
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@schedules_router.get(
|
||||
path="",
|
||||
summary="List run schedules",
|
||||
operation_id="listGraphRunSchedules",
|
||||
)
|
||||
async def list_all_schedules(
|
||||
graph_id: Optional[str] = Query(default=None, description="Filter by graph ID"),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.READ_SCHEDULE)),
|
||||
) -> Page[AgentRunSchedule]:
|
||||
"""List schedules for the authenticated user.
|
||||
|
||||
A paused schedule is listed with no `next_run_time`.
|
||||
"""
|
||||
schedules = await _visible_schedules(auth, graph_id=graph_id)
|
||||
# The scheduler keeps a schedule the caller owns even when it belongs to
|
||||
# another org, so the org filter is applied here rather than trusted there.
|
||||
return page.slice(
|
||||
[AgentRunSchedule.from_internal(s) for s in schedules if _in_tenant(s, auth)]
|
||||
)
|
||||
|
||||
|
||||
@schedules_router.delete(
|
||||
path="/{schedule_id}",
|
||||
summary="Delete run schedule",
|
||||
operation_id="deleteGraphRunSchedule",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_schedule(
|
||||
schedule_id: str,
|
||||
auth: TenantContext = Security(require_permission(APIKeyPermission.WRITE_SCHEDULE)),
|
||||
) -> None:
|
||||
"""Delete an execution schedule."""
|
||||
# The scheduler deletes by user id alone, so the org is checked here: a
|
||||
# schedule the key's org cannot see must not be deletable through it.
|
||||
await _assert_schedule_in_tenant(schedule_id, auth)
|
||||
|
||||
try:
|
||||
await get_scheduler_client().delete_schedule(
|
||||
schedule_id=schedule_id,
|
||||
user_id=auth.user_id,
|
||||
)
|
||||
except NotFoundError:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Schedule #{schedule_id} not found",
|
||||
)
|
||||
|
||||
|
||||
@schedules_router.post(
|
||||
path="",
|
||||
summary="Create run schedule",
|
||||
operation_id="createGraphRunSchedule",
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
async def create_schedule(
|
||||
request: AgentRunScheduleCreateRequest,
|
||||
auth: TenantContext = Security(
|
||||
require_paid_permission(
|
||||
APIKeyPermission.WRITE_SCHEDULE, APIKeyPermission.RUN_AGENT
|
||||
)
|
||||
),
|
||||
) -> AgentRunSchedule:
|
||||
"""Create a new execution schedule for a graph."""
|
||||
graph_id = request.graph_id
|
||||
# Without a version get_graph returns the newest, which can be a draft.
|
||||
version = request.graph_version
|
||||
if version is None:
|
||||
version = await graph_db.get_active_graph_version(
|
||||
graph_id, auth.user_id, organization_id=auth.organization_id
|
||||
)
|
||||
graph = await graph_db.get_graph(
|
||||
graph_id=graph_id,
|
||||
version=version,
|
||||
user_id=auth.user_id,
|
||||
organization_id=auth.organization_id,
|
||||
)
|
||||
if not graph:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph #{graph_id} v{request.graph_version} not found.",
|
||||
)
|
||||
assert_inputs_match(graph.input_schema, request.inputs)
|
||||
|
||||
# Determine timezone
|
||||
if request.timezone:
|
||||
user_timezone = request.timezone
|
||||
else:
|
||||
user = await get_user_by_id(auth.user_id)
|
||||
user_timezone = get_user_timezone_or_utc(user.timezone if user else None)
|
||||
|
||||
result = await get_scheduler_client().add_execution_schedule(
|
||||
user_id=auth.user_id,
|
||||
graph_id=graph_id,
|
||||
graph_version=graph.version,
|
||||
name=request.name,
|
||||
cron=request.cron,
|
||||
input_data=request.inputs,
|
||||
input_credentials=request.credentials_inputs,
|
||||
user_timezone=user_timezone,
|
||||
organization_id=auth.organization_id,
|
||||
team_id=auth.team_id,
|
||||
# As on the internal route: a graph one hired expert installed runs as
|
||||
# that expert's, under its budget and pause.
|
||||
expert_id=await experts_db.resolve_expert_for_graph(auth.user_id, graph_id),
|
||||
)
|
||||
|
||||
return AgentRunSchedule.from_internal(result)
|
||||
|
||||
|
||||
async def _visible_schedules(
|
||||
auth: TenantContext, graph_id: Optional[str] = None
|
||||
) -> list[GraphExecutionJobInfo]:
|
||||
"""The schedules the internal routes show: paused ones included, except
|
||||
those an archived expert's removal paused, which only a re-hire may touch.
|
||||
"""
|
||||
schedules = await get_scheduler_client().get_graph_execution_schedules(
|
||||
user_id=auth.user_id,
|
||||
graph_id=graph_id,
|
||||
organization_id=auth.organization_id,
|
||||
team_ids=await get_user_team_ids(auth.user_id, auth.organization_id),
|
||||
include_paused=True,
|
||||
)
|
||||
return await visible_graph_schedules(schedules, auth.user_id)
|
||||
|
||||
|
||||
async def _assert_schedule_in_tenant(schedule_id: str, auth: TenantContext) -> None:
|
||||
"""404 for a schedule outside the organization the credentials act in."""
|
||||
visible = await _visible_schedules(auth)
|
||||
if not any(
|
||||
schedule.id == schedule_id and _in_tenant(schedule, auth)
|
||||
for schedule in visible
|
||||
):
|
||||
raise NotFoundError(f"Schedule #{schedule_id} not found")
|
||||
|
||||
|
||||
def _in_tenant(schedule: GraphExecutionJobInfo, auth: TenantContext) -> bool:
|
||||
"""Schedules store an untagged org as "", so normalise before comparing."""
|
||||
return (schedule.organization_id or None) in (None, auth.organization_id)
|
||||
@@ -0,0 +1,160 @@
|
||||
"""v2 schedules show what the internal routes show, paused ones included."""
|
||||
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import ValidationError
|
||||
|
||||
from backend.util.exceptions import NotFoundError
|
||||
|
||||
from .models import AgentRunScheduleCreateRequest
|
||||
from .pagination import PageRequest
|
||||
from .schedules import create_schedule, delete_schedule, list_all_schedules
|
||||
from .tenancy import TenantContext
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-1",
|
||||
)
|
||||
|
||||
|
||||
def _schedule(schedule_id: str, *, paused: bool, expert_id: str | None = None):
|
||||
return Mock(
|
||||
id=schedule_id,
|
||||
organization_id="org-1",
|
||||
expert_id=expert_id,
|
||||
next_run_time="" if paused else "2026-10-10T09:00:00+00:00",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scheduler(mocker: pytest_mock.MockFixture) -> Mock:
|
||||
client = Mock(
|
||||
get_graph_execution_schedules=AsyncMock(
|
||||
return_value=[
|
||||
_schedule("running", paused=False),
|
||||
_schedule("paused", paused=True),
|
||||
# Paused when its expert was archived; only a re-hire restores it.
|
||||
_schedule("archived", paused=True, expert_id="gone"),
|
||||
]
|
||||
),
|
||||
delete_schedule=AsyncMock(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.schedules.get_scheduler_client", return_value=client
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.schedules.get_user_team_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.schedule_visibility.experts_db",
|
||||
return_value=Mock(active_expert_ids=AsyncMock(return_value=set())),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.models.AgentRunSchedule.from_internal",
|
||||
side_effect=lambda s: s,
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
async def test_paused_schedules_are_listed(scheduler: Mock) -> None:
|
||||
listed = await list_all_schedules(
|
||||
graph_id=None, page=PageRequest(limit=25), auth=_AUTH
|
||||
)
|
||||
|
||||
assert [s.id for s in listed.items] == ["running", "paused"]
|
||||
assert scheduler.get_graph_execution_schedules.await_args.kwargs["include_paused"]
|
||||
|
||||
|
||||
async def test_a_paused_schedule_can_be_deleted(scheduler: Mock) -> None:
|
||||
await delete_schedule(schedule_id="paused", auth=_AUTH)
|
||||
|
||||
scheduler.delete_schedule.assert_awaited_once_with(
|
||||
schedule_id="paused", user_id="user-1"
|
||||
)
|
||||
|
||||
|
||||
async def test_an_archived_experts_paused_schedule_stays_out_of_reach(
|
||||
scheduler: Mock,
|
||||
) -> None:
|
||||
with pytest.raises(NotFoundError):
|
||||
await delete_schedule(schedule_id="archived", auth=_AUTH)
|
||||
|
||||
scheduler.delete_schedule.assert_not_awaited()
|
||||
|
||||
|
||||
def test_an_unknown_timezone_is_refused_before_it_reaches_the_scheduler() -> None:
|
||||
with pytest.raises(ValidationError, match="Unknown timezone"):
|
||||
AgentRunScheduleCreateRequest(
|
||||
graph_id="graph-1", name="daily", cron="0 9 * * *", timezone="Mars/Olympus"
|
||||
)
|
||||
|
||||
request = AgentRunScheduleCreateRequest(
|
||||
graph_id="graph-1", name="daily", cron="0 9 * * *", timezone="Europe/Paris"
|
||||
)
|
||||
assert request.timezone == "Europe/Paris"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("requested, scheduled", [(None, 3), (5, 5)])
|
||||
async def test_a_schedule_without_a_version_runs_the_active_one(
|
||||
mocker: pytest_mock.MockFixture, requested: int | None, scheduled: int
|
||||
) -> None:
|
||||
"""The newest version can be a draft; the request says the active one."""
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_active_graph_version",
|
||||
new_callable=AsyncMock,
|
||||
return_value=3,
|
||||
)
|
||||
get_graph = mocker.patch(
|
||||
"backend.data.graph.get_graph", new_callable=AsyncMock, return_value=None
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await create_schedule(
|
||||
request=AgentRunScheduleCreateRequest(
|
||||
graph_id="graph-1",
|
||||
name="daily",
|
||||
cron="0 9 * * *",
|
||||
graph_version=requested,
|
||||
),
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
assert get_graph.await_args.kwargs["version"] == scheduled
|
||||
|
||||
|
||||
async def test_a_schedule_for_a_graph_a_hired_expert_installed_is_that_expert_s(
|
||||
mocker: pytest_mock.MockFixture, scheduler: Mock
|
||||
) -> None:
|
||||
"""As on the internal route, so its runs meet the expert's budget and pause."""
|
||||
mocker.patch(
|
||||
"backend.api.features.experts.experts_db.resolve_expert_for_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value="expert-1",
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(input_schema={"properties": {}}, version=1),
|
||||
)
|
||||
scheduler.add_execution_schedule = AsyncMock(return_value=Mock())
|
||||
|
||||
await create_schedule(
|
||||
request=AgentRunScheduleCreateRequest(
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
name="daily",
|
||||
cron="0 9 * * *",
|
||||
timezone="Europe/Paris",
|
||||
),
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
assert scheduler.add_execution_schedule.await_args.kwargs["expert_id"] == "expert-1"
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
V2 External API - Search Endpoints
|
||||
|
||||
Cross-domain hybrid search across agents, blocks, and documentation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security
|
||||
from prisma.enums import APIKeyPermission, ContentType
|
||||
from starlette import status
|
||||
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.features.search.hybrid_search import (
|
||||
HybridSearchRow,
|
||||
unified_hybrid_search,
|
||||
)
|
||||
|
||||
from .models import MarketplaceSearchResult, SearchContentType
|
||||
from .pagination import Page, PageRequest, page_request
|
||||
from .rate_limit import enforce, search_limiter
|
||||
from .tenancy import TenantContext, require_paid_permission
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
search_router = APIRouter(tags=["search"])
|
||||
|
||||
|
||||
@search_router.get(
|
||||
path="",
|
||||
summary="Search content and capabilities of the platform",
|
||||
operation_id="search",
|
||||
)
|
||||
async def search(
|
||||
response: Response,
|
||||
query: str = Query(description="Search query"),
|
||||
content_types: Optional[list[SearchContentType]] = Query(
|
||||
default=None, description="Content types to filter by"
|
||||
),
|
||||
category: Optional[str] = Query(default=None, description="Filter by category"),
|
||||
page: PageRequest = Depends(page_request),
|
||||
auth: TenantContext = Security(require_paid_permission()),
|
||||
) -> Page[MarketplaceSearchResult]:
|
||||
"""
|
||||
Search the platform's content and capabilities (hybrid search: literal + semantic).
|
||||
|
||||
Searches public agents, blocks and documentation by default. The caller's own
|
||||
library agents and workspace files are searchable too, each requiring the same
|
||||
permission that its own endpoints require.
|
||||
|
||||
**Rate limit:** 30 requests per minute per user.
|
||||
"""
|
||||
await enforce(search_limiter, auth.user_id, response)
|
||||
|
||||
requested = content_types or PUBLIC_CONTENT_TYPES
|
||||
for content_type in requested:
|
||||
if (scope := PRIVATE_CONTENT_TYPE_SCOPES.get(content_type)) and (
|
||||
scope not in auth.scopes
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"Searching {content_type.value} requires the "
|
||||
f"{scope.value} permission",
|
||||
)
|
||||
|
||||
results, total_count = await unified_hybrid_search(
|
||||
query=query,
|
||||
content_types=[ContentType(t.value) for t in requested],
|
||||
category=category,
|
||||
page=page.page,
|
||||
page_size=page.limit,
|
||||
user_id=auth.user_id,
|
||||
)
|
||||
results = await _in_tenant_results(results, auth)
|
||||
|
||||
return page.paged(
|
||||
[
|
||||
MarketplaceSearchResult(
|
||||
content_type=SearchContentType(r["content_type"]),
|
||||
content_id=r["content_id"],
|
||||
searchable_text=r["searchable_text"],
|
||||
metadata=r.get("metadata"),
|
||||
updated_at=r.get("updated_at"),
|
||||
combined_score=r.get("combined_score"),
|
||||
)
|
||||
for r in results
|
||||
],
|
||||
total_count=total_count,
|
||||
)
|
||||
|
||||
|
||||
async def _in_tenant_results(
|
||||
results: list[HybridSearchRow], auth: TenantContext
|
||||
) -> list[HybridSearchRow]:
|
||||
"""Drop library agents the credential's organization can't see.
|
||||
|
||||
The search index records whose a row is but not its organization, so a
|
||||
match from the user's other organization (or a deleted entry) is removed
|
||||
here, as a direct read of it would answer 404.
|
||||
"""
|
||||
agent_ids = [
|
||||
r["content_id"]
|
||||
for r in results
|
||||
if r["content_type"] == SearchContentType.LIBRARY_AGENT.value
|
||||
]
|
||||
if not agent_ids:
|
||||
return results
|
||||
organizations = await library_db.get_library_agent_organizations(
|
||||
auth.user_id, agent_ids
|
||||
)
|
||||
return [
|
||||
r
|
||||
for r in results
|
||||
if r["content_type"] != SearchContentType.LIBRARY_AGENT.value
|
||||
or (
|
||||
r["content_id"] in organizations
|
||||
and organizations[r["content_id"]] in (None, auth.organization_id)
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
# Searching these reaches the caller's own rows, so each costs the permission
|
||||
# that guards the same data elsewhere in v2. Everything else is public content.
|
||||
PRIVATE_CONTENT_TYPE_SCOPES: dict[SearchContentType, APIKeyPermission] = {
|
||||
SearchContentType.LIBRARY_AGENT: APIKeyPermission.READ_LIBRARY,
|
||||
SearchContentType.WORKSPACE_FILE: APIKeyPermission.READ_FILES,
|
||||
}
|
||||
|
||||
# Named here rather than left to the internal default, so v2's contract does not
|
||||
# change when that default does.
|
||||
PUBLIC_CONTENT_TYPES = [
|
||||
t for t in SearchContentType if t not in PRIVATE_CONTENT_TYPE_SCOPES
|
||||
]
|
||||
@@ -0,0 +1,67 @@
|
||||
"""A search answers with what a direct read of each result would."""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest_mock
|
||||
from fastapi import Response
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from .models import SearchContentType
|
||||
from .pagination import PageRequest
|
||||
from .search import search
|
||||
from .tenancy import TenantContext
|
||||
|
||||
_AUTH = TenantContext(
|
||||
user_id="user-1",
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id="org-a",
|
||||
)
|
||||
|
||||
|
||||
def _row(content_type: str, content_id: str) -> dict:
|
||||
return {
|
||||
"content_type": content_type,
|
||||
"content_id": content_id,
|
||||
"searchable_text": "text",
|
||||
"metadata": None,
|
||||
"updated_at": None,
|
||||
"combined_score": 1.0,
|
||||
}
|
||||
|
||||
|
||||
async def test_library_agents_outside_the_organization_are_left_out(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The index records the user, not the organization, of a library agent."""
|
||||
mocker.patch("backend.api.external.v2.search.enforce", new_callable=AsyncMock)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.search.unified_hybrid_search",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
[
|
||||
_row("LIBRARY_AGENT", "mine-in-a"),
|
||||
_row("LIBRARY_AGENT", "mine-in-b"),
|
||||
_row("LIBRARY_AGENT", "untagged"),
|
||||
_row("LIBRARY_AGENT", "deleted"),
|
||||
_row("BLOCK", "block-1"),
|
||||
],
|
||||
5,
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.search.library_db.get_library_agent_organizations",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"mine-in-a": "org-a", "mine-in-b": "org-b", "untagged": None},
|
||||
)
|
||||
|
||||
found = await search(
|
||||
response=Response(),
|
||||
query="report",
|
||||
content_types=[SearchContentType.LIBRARY_AGENT, SearchContentType.BLOCK],
|
||||
category=None,
|
||||
page=PageRequest(limit=25),
|
||||
auth=_AUTH,
|
||||
)
|
||||
|
||||
assert [r.content_id for r in found.items] == ["mine-in-a", "untagged", "block-1"]
|
||||
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
V2 External API - Tenancy resolution
|
||||
|
||||
Every v2 request acts inside exactly one organization, and optionally inside one
|
||||
team within it. The tenant is bound to the credential: an API key is minted in
|
||||
an org (`organizationId`) and may be pinned to a team (`teamIdRestriction`),
|
||||
which is the same pair v1 external already honours. A key with no org — every
|
||||
key minted before orgs existed — resolves to the caller's personal org, so
|
||||
personal keys keep behaving exactly as they did.
|
||||
|
||||
A key that is not pinned to a team may select one per request with `X-Team-Id`,
|
||||
matching the header the internal API uses. Without it the request runs at
|
||||
org-home, which is what the web app does when no team is selected.
|
||||
|
||||
Handlers depend on `require_permission`/`require_auth` from this module rather
|
||||
than from `..middleware`, so no v2 endpoint can reach the database without a
|
||||
membership-verified tenant.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Literal, Optional, Protocol, TypeVar
|
||||
|
||||
from autogpt_libs.auth.dependencies import TEAM_HEADER_NAME
|
||||
from fastapi import Header, HTTPException, Request, Security
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
from prisma.enums import APIKeyPermission
|
||||
from pydantic import BaseModel
|
||||
from starlette import status
|
||||
|
||||
from backend.api.external import middleware
|
||||
from backend.copilot.rate_limit import enforce_payment_paywall
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.db import prisma
|
||||
from backend.util.exceptions import NotAuthorizedError, NotFoundError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T", bound="TenantedRow")
|
||||
OwnedT = TypeVar("OwnedT", bound="OwnedRow")
|
||||
|
||||
|
||||
class TenantContext(BaseModel):
|
||||
"""An authenticated caller with its organization and team resolved.
|
||||
|
||||
Deliberately not an `APIAuthorizationInfo` subclass: `organization_id` is
|
||||
optional there and always resolved here, and a handler has no business
|
||||
reading a credential's hash metadata.
|
||||
"""
|
||||
|
||||
user_id: str
|
||||
scopes: list[APIKeyPermission]
|
||||
type: Literal["oauth", "api_key"]
|
||||
organization_id: str
|
||||
team_id: Optional[str] = None
|
||||
|
||||
|
||||
async def require_auth(
|
||||
request: Request,
|
||||
api_key: Optional[str] = Security(middleware.api_key_header),
|
||||
bearer: Optional[HTTPAuthorizationCredentials] = Security(middleware.bearer_auth),
|
||||
requested_team: Optional[str] = Header(
|
||||
default=None,
|
||||
alias=TEAM_HEADER_NAME,
|
||||
description=(
|
||||
"Team to act in, within the organization the credentials belong to. "
|
||||
"Rejected for credentials already pinned to a team."
|
||||
),
|
||||
),
|
||||
) -> TenantContext:
|
||||
"""Require valid credentials, then resolve the tenant they act in.
|
||||
|
||||
Every v2 route reaches this one dependency, so tests override it once and
|
||||
`tenancy_test` can assert no handler bypasses it.
|
||||
"""
|
||||
auth = await middleware.resolve_request_auth(request.scope, api_key, bearer)
|
||||
if auth is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing authentication. Provide API key or access token.",
|
||||
)
|
||||
return await resolve_tenant(auth, requested_team)
|
||||
|
||||
|
||||
def require_paid_permission(*permissions: APIKeyPermission):
|
||||
"""`require_permission`, and a 402 for an account without an active plan.
|
||||
|
||||
For routes that start work, store files or call out, as on the internal API.
|
||||
"""
|
||||
check_permissions = require_permission(*permissions)
|
||||
|
||||
async def check_plan(
|
||||
tenant: TenantContext = Security(check_permissions),
|
||||
) -> TenantContext:
|
||||
await enforce_payment_paywall(tenant.user_id)
|
||||
return tenant
|
||||
|
||||
return check_plan
|
||||
|
||||
|
||||
def require_permission(*permissions: APIKeyPermission):
|
||||
"""Require the given scopes, and resolve the credential's tenant."""
|
||||
|
||||
async def check_permissions(
|
||||
tenant: TenantContext = Security(
|
||||
require_auth, scopes=[p.value for p in permissions]
|
||||
),
|
||||
) -> TenantContext:
|
||||
if missing := [p for p in permissions if p not in tenant.scopes]:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Missing required permission(s): "
|
||||
f"{', '.join(p.value for p in missing)}",
|
||||
)
|
||||
return tenant
|
||||
|
||||
return check_permissions
|
||||
|
||||
|
||||
def require_billing_permission(*permissions: APIKeyPermission):
|
||||
"""`require_permission`, for the organization's own wallet.
|
||||
|
||||
An organization's credits are pooled, so the web app shows their balance,
|
||||
ledger and invoices only to its owner and billing managers; a credential
|
||||
acting in the organization is held to the same roles. In a personal
|
||||
organization the caller is always its owner.
|
||||
"""
|
||||
|
||||
async def check_billing_role(
|
||||
tenant: TenantContext = Security(require_permission(*permissions)),
|
||||
) -> TenantContext:
|
||||
member = await prisma.orgmember.find_unique(
|
||||
where={
|
||||
"orgId_userId": {
|
||||
"orgId": tenant.organization_id,
|
||||
"userId": tenant.user_id,
|
||||
}
|
||||
}
|
||||
)
|
||||
if member is None or not (member.isOwner or member.isBillingManager):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=(
|
||||
"Only the organization's owner or a billing manager can read "
|
||||
"its credits"
|
||||
),
|
||||
)
|
||||
return tenant
|
||||
|
||||
return check_billing_role
|
||||
|
||||
|
||||
async def resolve_tenant(
|
||||
auth: APIAuthorizationInfo, requested_team: Optional[str] = None
|
||||
) -> TenantContext:
|
||||
"""Resolve the org and team this request acts in, verifying membership."""
|
||||
organization_id, team_id = await resolve_credential_tenancy(auth)
|
||||
|
||||
if requested_team := (requested_team or "").strip():
|
||||
if auth.team_id_restriction and requested_team != auth.team_id_restriction:
|
||||
raise NotAuthorizedError(
|
||||
f"These credentials are restricted to team {auth.team_id_restriction}"
|
||||
)
|
||||
await _assert_active_team_member(auth.user_id, organization_id, requested_team)
|
||||
team_id = requested_team
|
||||
|
||||
return TenantContext(
|
||||
user_id=auth.user_id,
|
||||
scopes=auth.scopes,
|
||||
type=auth.type,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
class TenantedRow(Protocol):
|
||||
"""Any resource row that records which organization it belongs to."""
|
||||
|
||||
organization_id: Optional[str]
|
||||
|
||||
|
||||
def in_tenant(resource: Optional[T], auth: TenantContext, what: str) -> T:
|
||||
"""Return the resource, or 404 if it is missing or in another organization.
|
||||
|
||||
Resolving the tenant is not enough on its own: the by-id lookups underneath
|
||||
v2 match on user id, so a key minted in one org would otherwise reach the
|
||||
same user's rows in another. 404 rather than 403 — whether a resource
|
||||
exists elsewhere is not this credential's business. Untagged rows (created
|
||||
before org tagging) stay visible to their owner, as everywhere else.
|
||||
"""
|
||||
if resource is None or (
|
||||
resource.organization_id is not None
|
||||
and resource.organization_id != auth.organization_id
|
||||
):
|
||||
raise NotFoundError(f"{what} not found")
|
||||
return resource
|
||||
|
||||
|
||||
class OwnedRow(Protocol):
|
||||
"""Any resource row that records the user who owns it."""
|
||||
|
||||
user_id: str
|
||||
|
||||
|
||||
def owned_by_caller(resource: OwnedT, auth: TenantContext, what: str) -> OwnedT:
|
||||
"""Return the resource, or 403 if another member of the org owns it.
|
||||
|
||||
Org-visible reads show a colleague's rows; the writes beneath v2 match on
|
||||
the owner's id, so refuse before any of them runs.
|
||||
"""
|
||||
if resource.user_id != auth.user_id:
|
||||
raise NotAuthorizedError(
|
||||
f"{what} belongs to another member of this organization"
|
||||
)
|
||||
return resource
|
||||
|
||||
|
||||
async def resolve_credential_tenancy(
|
||||
auth: APIAuthorizationInfo,
|
||||
) -> tuple[str, Optional[str]]:
|
||||
"""The org and team a credential is bound to, independent of any request.
|
||||
|
||||
Also used by the MCP server, which authenticates outside the dependency
|
||||
chain and so cannot go through `resolve_tenant`.
|
||||
"""
|
||||
if auth.organization_id:
|
||||
await _assert_active_org_member(auth.user_id, auth.organization_id)
|
||||
if auth.team_id_restriction:
|
||||
await _assert_active_team_member(
|
||||
auth.user_id, auth.organization_id, auth.team_id_restriction
|
||||
)
|
||||
return auth.organization_id, auth.team_id_restriction
|
||||
return await _personal_tenancy(auth.user_id)
|
||||
|
||||
|
||||
async def _assert_active_org_member(user_id: str, organization_id: str) -> None:
|
||||
member = await prisma.orgmember.find_unique(
|
||||
where={"orgId_userId": {"orgId": organization_id, "userId": user_id}},
|
||||
include={"Org": True},
|
||||
)
|
||||
if member is None or member.status != "ACTIVE":
|
||||
raise NotAuthorizedError(
|
||||
"These credentials belong to an organization you are not an active "
|
||||
"member of"
|
||||
)
|
||||
if member.Org is not None and member.Org.deletedAt is not None:
|
||||
raise NotAuthorizedError("This organization has been deleted")
|
||||
|
||||
|
||||
async def _assert_active_team_member(
|
||||
user_id: str, organization_id: str, team_id: str
|
||||
) -> None:
|
||||
member = await prisma.teammember.find_unique(
|
||||
where={"teamId_userId": {"teamId": team_id, "userId": user_id}},
|
||||
include={"Team": True},
|
||||
)
|
||||
if member is None or member.status != "ACTIVE":
|
||||
raise NotAuthorizedError(f"You are not an active member of team {team_id}")
|
||||
if member.Team is None or member.Team.orgId != organization_id:
|
||||
raise NotAuthorizedError(
|
||||
f"Team {team_id} does not belong to the organization these credentials "
|
||||
"act in"
|
||||
)
|
||||
|
||||
|
||||
async def _personal_tenancy(user_id: str) -> tuple[str, Optional[str]]:
|
||||
"""Org for a credential minted before orgs, or outside one.
|
||||
|
||||
Org-home, not the personal org's default team: `get_graph` and friends
|
||||
treat a team as an exact-match filter, so pinning a pre-org key to a team
|
||||
would hide every graph it could read yesterday. Bootstraps a missing
|
||||
personal org rather than leaving the account unable to call the API.
|
||||
"""
|
||||
from backend.api.features.orgs.db import get_user_default_team
|
||||
|
||||
organization_id, _ = await get_user_default_team(user_id)
|
||||
if organization_id is None:
|
||||
logger.warning(
|
||||
f"User {user_id} has no personal org and bootstrap failed — "
|
||||
"account in inconsistent state"
|
||||
)
|
||||
raise NotAuthorizedError(
|
||||
"No organization context available. Your account may be in an "
|
||||
"inconsistent state — please contact support."
|
||||
)
|
||||
return organization_id, None
|
||||
@@ -0,0 +1,827 @@
|
||||
"""
|
||||
Tests for v2 tenancy: which organization and team a request acts in.
|
||||
|
||||
Two things must hold. Structurally, every v2 route resolves its tenant through
|
||||
`tenancy.require_auth` — a handler that skipped it would silently fall back to
|
||||
"the user, across every org they belong to". Behaviourally, the tenant reaching
|
||||
the database is the one bound to the credential: a key minted in org A never
|
||||
passes org B, and a key with no org passes the caller's personal org.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import re
|
||||
import textwrap
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException, Response
|
||||
from fastapi.routing import APIRoute
|
||||
from prisma.enums import APIKeyPermission
|
||||
|
||||
from backend.data.auth.base import APIAuthorizationInfo
|
||||
from backend.data.execution import ExecutionStatus, GraphExecutionMeta
|
||||
from backend.util.exceptions import NotAuthorizedError, NotFoundError
|
||||
|
||||
from . import runs, tenancy
|
||||
from .pagination import PageRequest
|
||||
from .routes import v2_router
|
||||
from .tenancy import TenantContext, resolve_tenant
|
||||
|
||||
USER_ID = "user-1"
|
||||
ORG_A = "org-a"
|
||||
ORG_B = "org-b"
|
||||
TEAM_A = "team-a"
|
||||
PERSONAL_ORG = "personal-org"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Structural: no handler can bypass tenant resolution
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_every_v2_route_resolves_the_tenant() -> None:
|
||||
"""Every route's dependency tree contains the one tenancy dependency."""
|
||||
bypassing = [
|
||||
f"{sorted(route.methods)[0]} {route.path}"
|
||||
for route in v2_router.routes
|
||||
if isinstance(route, APIRoute)
|
||||
and not _depends_on_tenancy(route.dependant.dependencies)
|
||||
]
|
||||
assert not bypassing, f"routes that never resolve a tenant: {bypassing}"
|
||||
|
||||
|
||||
def test_v2_router_has_routes_to_check() -> None:
|
||||
"""Guards the test above against passing on an empty router."""
|
||||
assert len([r for r in v2_router.routes if isinstance(r, APIRoute)]) > 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Behavioural: the tenant reaching the database is the credential's
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_library_listing_is_scoped_to_the_key_org(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""A key minted in org A lists org A's library, never org B's."""
|
||||
from .library.agents import list_library_agents
|
||||
|
||||
list_db = mocker.patch(
|
||||
"backend.api.features.library.db.list_library_agents",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(agents=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
|
||||
await list_library_agents(
|
||||
published=None, favorite=None, page=_page(), auth=_key_for(ORG_A)
|
||||
)
|
||||
|
||||
assert list_db.await_args.kwargs["organization_id"] == ORG_A
|
||||
assert list_db.await_args.kwargs["organization_id"] != ORG_B
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_listing_is_scoped_to_the_key_org(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .runs import list_runs
|
||||
|
||||
list_db = mocker.patch(
|
||||
"backend.data.execution.get_graph_executions_paginated",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(executions=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
|
||||
await list_runs(graph_id=None, page=_page(), auth=_key_for(ORG_B))
|
||||
|
||||
assert list_db.await_args.kwargs["organization_id"] == ORG_B
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_personal_key_passes_the_personal_org_not_none(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""A key minted before orgs existed acts in the caller's personal org.
|
||||
|
||||
Passing None instead would list the user's rows across every org they
|
||||
belong to — the pre-tenancy behaviour this replaces.
|
||||
"""
|
||||
from .library.agents import list_library_agents
|
||||
|
||||
list_db = mocker.patch(
|
||||
"backend.api.features.library.db.list_library_agents",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(agents=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
|
||||
await list_library_agents(
|
||||
published=None, favorite=None, page=_page(), auth=_key_for(PERSONAL_ORG)
|
||||
)
|
||||
|
||||
assert list_db.await_args.kwargs["organization_id"] == PERSONAL_ORG
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_creation_carries_org_and_team(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""A team-pinned key tenants the runs it starts to that team."""
|
||||
from .library.agents import execute_agent
|
||||
from .models import AgentRunRequest
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(
|
||||
organization_id=ORG_A, graph_id="graph-1", graph_version=1, input_schema={}
|
||||
),
|
||||
)
|
||||
add_execution = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
|
||||
await execute_agent(
|
||||
response=Response(),
|
||||
request=AgentRunRequest(inputs={}),
|
||||
agent_id="agent-1",
|
||||
auth=_key_for(ORG_A, team_id=TEAM_A),
|
||||
)
|
||||
|
||||
assert add_execution.await_args.kwargs["organization_id"] == ORG_A
|
||||
assert add_execution.await_args.kwargs["team_id"] == TEAM_A
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_billing_uses_the_org_credit_model(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""An org key bills the org ledger, not the owner's personal wallet."""
|
||||
from .credits import get_balance
|
||||
|
||||
credit_model = mocker.patch(
|
||||
"backend.api.external.v2.credits.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=42)),
|
||||
)
|
||||
|
||||
balance = await get_balance(auth=_key_for(ORG_A))
|
||||
|
||||
assert credit_model.await_args.args == (USER_ID, ORG_A)
|
||||
assert balance.balance_cents == 42
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_marketplace_submissions_are_scoped_to_the_key_org(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .marketplace import list_submissions
|
||||
|
||||
list_db = mocker.patch(
|
||||
"backend.api.features.store.db.get_store_submissions",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(submissions=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
|
||||
await list_submissions(page=_page(), auth=_key_for(ORG_A))
|
||||
|
||||
assert list_db.await_args.kwargs["organization_id"] == ORG_A
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_read_from_an_org_home_key_is_not_team_filtered(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""`get_graph` treats a team as an exact-match filter, so v2 read paths
|
||||
pass none at all: with one, a key stops finding org-home graphs."""
|
||||
from .graphs import get_graph
|
||||
|
||||
graph_db = mocker.patch(
|
||||
"backend.data.graph.get_graph",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await get_graph(graph_id="graph-1", version=None, auth=_key_for(ORG_A))
|
||||
|
||||
assert graph_db.await_args.kwargs["organization_id"] == ORG_A
|
||||
assert "team_id" not in graph_db.await_args.kwargs
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Resolution: what the credential and the X-Team-Id header are allowed to say
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_key_resolves_to_its_own_org(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
_mock_membership(mocker, org_id=ORG_A)
|
||||
|
||||
tenant = await resolve_tenant(_credential(organization_id=ORG_A))
|
||||
|
||||
assert (tenant.organization_id, tenant.team_id) == (ORG_A, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_pinned_key_resolves_to_its_team(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
_mock_membership(mocker, org_id=ORG_A)
|
||||
|
||||
tenant = await resolve_tenant(
|
||||
_credential(organization_id=ORG_A, team_id_restriction=TEAM_A)
|
||||
)
|
||||
|
||||
assert (tenant.organization_id, tenant.team_id) == (ORG_A, TEAM_A)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"team_member",
|
||||
[
|
||||
None,
|
||||
Mock(status="INACTIVE", Team=Mock(orgId=ORG_A)),
|
||||
Mock(status="ACTIVE", Team=Mock(orgId=ORG_B)),
|
||||
],
|
||||
ids=["removed", "inactive", "team-moved-orgs"],
|
||||
)
|
||||
async def test_team_pinned_key_is_rejected_once_its_team_membership_ends(
|
||||
mocker: pytest_mock.MockFixture, team_member: Mock | None
|
||||
) -> None:
|
||||
"""Leaving a team must disable the keys pinned to it."""
|
||||
prisma = _mock_membership(mocker, org_id=ORG_A)
|
||||
prisma.teammember.find_unique = AsyncMock(return_value=team_member)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(
|
||||
_credential(organization_id=ORG_A, team_id_restriction=TEAM_A)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_without_org_falls_back_to_the_personal_org_at_org_home(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Org-home, never the default team: a team is an exact-match read filter,
|
||||
so a team here would hide graphs a pre-org key could read yesterday."""
|
||||
mocker.patch(
|
||||
"backend.api.features.orgs.db.get_user_default_team",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(PERSONAL_ORG, "personal-team"),
|
||||
)
|
||||
|
||||
tenant = await resolve_tenant(_credential())
|
||||
|
||||
assert (tenant.organization_id, tenant.team_id) == (PERSONAL_ORG, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_is_rejected_once_membership_ends(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Leaving an org must disable the keys minted in it."""
|
||||
_mock_membership(mocker, org_id=ORG_A, status="INACTIVE")
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(_credential(organization_id=ORG_A))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_is_rejected_when_the_org_is_deleted(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
_mock_membership(mocker, org_id=ORG_A, deleted_at=datetime.now(tz=timezone.utc))
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(_credential(organization_id=ORG_A))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_header_selects_a_team_the_caller_belongs_to(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
prisma = _mock_membership(mocker, org_id=ORG_A)
|
||||
prisma.teammember.find_unique = AsyncMock(
|
||||
return_value=Mock(status="ACTIVE", Team=Mock(orgId=ORG_A))
|
||||
)
|
||||
|
||||
tenant = await resolve_tenant(
|
||||
_credential(organization_id=ORG_A), requested_team=TEAM_A
|
||||
)
|
||||
|
||||
assert tenant.team_id == TEAM_A
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_header_cannot_override_a_pinned_key(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
_mock_membership(mocker, org_id=ORG_A)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(
|
||||
_credential(organization_id=ORG_A, team_id_restriction=TEAM_A),
|
||||
requested_team="other-team",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_header_cannot_reach_another_orgs_team(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
prisma = _mock_membership(mocker, org_id=ORG_A)
|
||||
prisma.teammember.find_unique = AsyncMock(
|
||||
return_value=Mock(status="ACTIVE", Team=Mock(orgId=ORG_B))
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(
|
||||
_credential(organization_id=ORG_A), requested_team="team-b"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_account_without_a_personal_org_is_rejected(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
mocker.patch(
|
||||
"backend.api.features.orgs.db.get_user_default_team",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(None, None),
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await resolve_tenant(_credential())
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Cross-organization access: an org-A key must not reach org-B resources
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_read_an_org_b_library_agent(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""404, not 403: whether it exists elsewhere is not this key's business."""
|
||||
from .library.agents import get_library_agent
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B),
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await get_library_agent(agent_id="la-1", auth=_key_for(ORG_A))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_run_an_org_b_library_agent(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The run would be tenanted, billed and visible in org A."""
|
||||
from .library.agents import execute_agent
|
||||
from .models import AgentRunRequest
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B, graph_id="g1", graph_version=1),
|
||||
)
|
||||
add_execution = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await execute_agent(
|
||||
response=Response(),
|
||||
request=AgentRunRequest(inputs={}),
|
||||
agent_id="la-1",
|
||||
auth=_key_for(ORG_A),
|
||||
)
|
||||
add_execution.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_delete_an_org_b_library_agent(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .library.agents import delete_library_agent
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B),
|
||||
)
|
||||
delete = mocker.patch(
|
||||
"backend.api.features.library.db.delete_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await delete_library_agent(agent_id="la-1", auth=_key_for(ORG_A))
|
||||
delete.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_execute_an_org_b_preset(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .library.presets import execute_preset
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.library.helpers.get_credit_model",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(get_credits=AsyncMock(return_value=100)),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B),
|
||||
)
|
||||
add_execution = mocker.patch(
|
||||
"backend.executor.utils.add_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(),
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await execute_preset(response=Response(), preset_id="p-1", auth=_key_for(ORG_A))
|
||||
add_execution.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call, write",
|
||||
[
|
||||
("get", None),
|
||||
("update", "update_triggered_preset"),
|
||||
("delete", "delete_preset_with_webhook_cleanup"),
|
||||
],
|
||||
)
|
||||
async def test_org_a_key_cannot_reach_an_org_b_preset(
|
||||
mocker: pytest_mock.MockFixture, call: str, write: str | None
|
||||
) -> None:
|
||||
from .library import presets
|
||||
from .models import AgentPresetUpdateRequest
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_preset",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B),
|
||||
)
|
||||
writes = {
|
||||
name: mocker.patch(
|
||||
f"backend.api.features.library.triggers.{name}", new_callable=AsyncMock
|
||||
)
|
||||
for name in ("update_triggered_preset", "delete_preset_with_webhook_cleanup")
|
||||
}
|
||||
handler = {
|
||||
"get": lambda: presets.get_preset(preset_id="p-1", auth=_key_for(ORG_A)),
|
||||
"update": lambda: presets.update_preset(
|
||||
request=AgentPresetUpdateRequest(is_active=True),
|
||||
preset_id="p-1",
|
||||
auth=_key_for(ORG_A),
|
||||
),
|
||||
"delete": lambda: presets.delete_preset(preset_id="p-1", auth=_key_for(ORG_A)),
|
||||
}[call]
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await handler()
|
||||
for mock in writes.values():
|
||||
mock.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call", ["create", "setup_trigger"])
|
||||
async def test_org_a_key_cannot_put_a_preset_on_an_org_b_graph(
|
||||
mocker: pytest_mock.MockFixture, call: str
|
||||
) -> None:
|
||||
"""A preset takes its graph's organization, so the graph must be in the key's."""
|
||||
from .library import presets
|
||||
from .models import AgentPresetCreateRequest, AgentTriggerSetupRequest
|
||||
|
||||
get_graph = mocker.patch(
|
||||
"backend.data.graph.get_graph", new_callable=AsyncMock, return_value=None
|
||||
)
|
||||
create = mocker.patch(
|
||||
"backend.api.features.library.db.create_preset", new_callable=AsyncMock
|
||||
)
|
||||
setup = mocker.patch(
|
||||
"backend.api.features.library.triggers.setup_triggered_preset",
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
fields = {"name": "p", "graph_id": "graph-b", "graph_version": 1}
|
||||
handler = {
|
||||
"create": lambda: presets.create_preset(
|
||||
request=AgentPresetCreateRequest(**fields), auth=_key_for(ORG_A)
|
||||
),
|
||||
"setup_trigger": lambda: presets.setup_trigger(
|
||||
request=AgentTriggerSetupRequest(**fields, trigger_config={}),
|
||||
auth=_key_for(ORG_A),
|
||||
),
|
||||
}[call]
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await handler()
|
||||
assert get_graph.await_args.kwargs["organization_id"] == ORG_A
|
||||
create.assert_not_awaited()
|
||||
setup.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_delete_an_org_b_folder(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .library.folders import delete_folder
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_folder",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(organization_id=ORG_B),
|
||||
)
|
||||
delete = mocker.patch(
|
||||
"backend.api.features.library.db.delete_folder", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await delete_folder(folder_id="f-1", auth=_key_for(ORG_A))
|
||||
delete.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_delete_an_org_b_run(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
from .runs import delete_run
|
||||
|
||||
mocker.patch(
|
||||
"backend.data.execution.get_graph_execution",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_run(organization_id=ORG_B),
|
||||
)
|
||||
delete = mocker.patch(
|
||||
"backend.data.execution.delete_graph_execution", new_callable=AsyncMock
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await delete_run(run_id="run-1", auth=_key_for(ORG_A))
|
||||
delete.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_org_a_key_cannot_see_or_delete_an_org_b_schedule(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The scheduler keeps a schedule its owner made in any org, so v2 filters
|
||||
on the schedule's own org rather than trusting that listing."""
|
||||
from .schedules import delete_schedule, list_all_schedules
|
||||
|
||||
scheduler = Mock(
|
||||
get_graph_execution_schedules=AsyncMock(
|
||||
return_value=[_schedule("s-a", ORG_A), _schedule("s-b", ORG_B)]
|
||||
),
|
||||
delete_schedule=AsyncMock(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.schedules.get_scheduler_client",
|
||||
return_value=scheduler,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.schedules.get_user_team_ids",
|
||||
new_callable=AsyncMock,
|
||||
return_value=[],
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.schedules.visible_graph_schedules",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=lambda schedules, _user_id: schedules,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.external.v2.models.AgentRunSchedule.from_internal",
|
||||
side_effect=lambda s: s,
|
||||
)
|
||||
|
||||
listed = await list_all_schedules(graph_id=None, page=_page(), auth=_key_for(ORG_A))
|
||||
|
||||
assert [s.id for s in listed.items] == ["s-a"]
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await delete_schedule(schedule_id="s-b", auth=_key_for(ORG_A))
|
||||
scheduler.delete_schedule.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_untagged_rows_stay_reachable(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""Rows created before org tagging belong to their owner, not to nobody."""
|
||||
from .library.agents import get_library_agent
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_library_agent(organization_id=None),
|
||||
)
|
||||
from_internal = mocker.patch(
|
||||
"backend.api.external.v2.models.LibraryAgent.from_internal"
|
||||
)
|
||||
|
||||
await get_library_agent(agent_id="la-1", auth=_key_for(ORG_A))
|
||||
|
||||
assert from_internal.call_args.args[0].organization_id is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preset_and_review_listing_are_scoped_to_the_key_org(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""These list by user id underneath, so the org has to reach the query."""
|
||||
from .library.presets import list_presets
|
||||
from .runs import list_reviews
|
||||
|
||||
presets = mocker.patch(
|
||||
"backend.api.features.library.db.list_presets",
|
||||
new_callable=AsyncMock,
|
||||
return_value=Mock(presets=[], pagination=Mock(total_items=0)),
|
||||
)
|
||||
reviews = mocker.patch(
|
||||
"backend.data.human_review.get_reviews",
|
||||
new_callable=AsyncMock,
|
||||
return_value=([], Mock(total_items=0)),
|
||||
)
|
||||
|
||||
await list_presets(graph_id=None, page=_page(), auth=_key_for(ORG_A))
|
||||
await list_reviews(
|
||||
run_id=None, review_status=None, page=_page(), auth=_key_for(ORG_A)
|
||||
)
|
||||
|
||||
assert presets.await_args.kwargs["organization_id"] == ORG_A
|
||||
assert reviews.await_args.kwargs["organization_id"] == ORG_A
|
||||
|
||||
|
||||
def test_every_item_handler_checks_the_tenant() -> None:
|
||||
"""A handler that fetches by id must enforce the tenant on what it fetched.
|
||||
|
||||
Resolving the tenant is only half of it: several lookups underneath match
|
||||
on user id alone, so a handler that skips the check reads across
|
||||
organizations. Either form counts — `in_tenant` on the row, or passing the
|
||||
org into a query that filters on it.
|
||||
"""
|
||||
unchecked = [
|
||||
f"{sorted(route.methods)[0]} {route.path}"
|
||||
for route in v2_router.routes
|
||||
if isinstance(route, APIRoute)
|
||||
and "{" in route.path
|
||||
and route.endpoint.__name__ not in _ITEM_HANDLERS_WITHOUT_TENANTED_ROWS
|
||||
and not _checks_tenant(inspect.getsource(route.endpoint))
|
||||
]
|
||||
assert not unchecked, f"item handlers that never check the tenant: {unchecked}"
|
||||
|
||||
|
||||
def _checks_tenant(source: str) -> bool:
|
||||
return (
|
||||
"in_tenant" in source
|
||||
or "organization_id=auth.organization_id" in source
|
||||
or any(f"{helper.__name__}(" in source for helper in _TENANT_CHECKING_HELPERS)
|
||||
)
|
||||
|
||||
|
||||
# Helpers that fetch a row by id and check its tenant before returning it.
|
||||
_TENANT_CHECKING_HELPERS = [runs._assert_run_in_tenant, runs._assert_own_run_in_tenant]
|
||||
|
||||
|
||||
def test_the_tenant_checking_helpers_check_the_tenant() -> None:
|
||||
"""Each calls `in_tenant`, or another of them that does."""
|
||||
for helper in _TENANT_CHECKING_HELPERS:
|
||||
checks = ["in_tenant"] + [
|
||||
other.__name__ for other in _TENANT_CHECKING_HELPERS if other is not helper
|
||||
]
|
||||
body = _body(helper)
|
||||
assert any(re.search(rf"\b{name}\(", body) for name in checks), helper.__name__
|
||||
|
||||
|
||||
def _body(function) -> str:
|
||||
"""The function's statements without its signature, whose name says "tenant"."""
|
||||
definition = ast.parse(textwrap.dedent(inspect.getsource(function))).body[0]
|
||||
assert isinstance(definition, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
return "\n".join(ast.unparse(statement) for statement in definition.body)
|
||||
|
||||
|
||||
# Item handlers that touch nothing an organization owns: workspace files and
|
||||
# integration credentials are stored per user, and marketplace listings and
|
||||
# creator profiles are public — a published listing is readable by anyone.
|
||||
_ITEM_HANDLERS_WITHOUT_TENANTED_ROWS = {
|
||||
"get_file",
|
||||
"delete_file",
|
||||
"download_file",
|
||||
"delete_credential",
|
||||
"get_agent_by_version",
|
||||
"get_agent_details",
|
||||
"get_creator_details",
|
||||
"get_marketplace_listing_for_graph",
|
||||
}
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _depends_on_tenancy(dependencies) -> bool:
|
||||
return any(
|
||||
dep.call is tenancy.require_auth or _depends_on_tenancy(dep.dependencies)
|
||||
for dep in dependencies
|
||||
)
|
||||
|
||||
|
||||
def _key_for(organization_id: str, team_id: str | None = None) -> TenantContext:
|
||||
return TenantContext(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
def _credential(
|
||||
organization_id: str | None = None, team_id_restriction: str | None = None
|
||||
) -> APIAuthorizationInfo:
|
||||
return APIAuthorizationInfo(
|
||||
user_id=USER_ID,
|
||||
scopes=list(APIKeyPermission),
|
||||
type="api_key",
|
||||
created_at=datetime.now(tz=timezone.utc),
|
||||
organization_id=organization_id,
|
||||
team_id_restriction=team_id_restriction,
|
||||
)
|
||||
|
||||
|
||||
def _mock_membership(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
org_id: str,
|
||||
status: str = "ACTIVE",
|
||||
deleted_at: datetime | None = None,
|
||||
) -> Mock:
|
||||
prisma = mocker.patch("backend.api.external.v2.tenancy.prisma")
|
||||
prisma.orgmember.find_unique = AsyncMock(
|
||||
return_value=Mock(status=status, Org=Mock(id=org_id, deletedAt=deleted_at))
|
||||
)
|
||||
prisma.teammember.find_unique = AsyncMock(
|
||||
return_value=Mock(status="ACTIVE", Team=Mock(orgId=org_id))
|
||||
)
|
||||
return prisma
|
||||
|
||||
|
||||
def _page() -> PageRequest:
|
||||
return PageRequest(limit=20)
|
||||
|
||||
|
||||
def _schedule(schedule_id: str, organization_id: str) -> Mock:
|
||||
return Mock(id=schedule_id, organization_id=organization_id)
|
||||
|
||||
|
||||
def _library_agent(organization_id: str | None) -> Mock:
|
||||
return Mock(id="la-1", organization_id=organization_id, graph_id="g1")
|
||||
|
||||
|
||||
def _run(organization_id: str = ORG_A) -> GraphExecutionMeta:
|
||||
return GraphExecutionMeta.model_construct(
|
||||
id="run-1",
|
||||
user_id=USER_ID,
|
||||
graph_id="graph-1",
|
||||
graph_version=1,
|
||||
preset_id=None,
|
||||
status=ExecutionStatus.QUEUED,
|
||||
started_at=None,
|
||||
ended_at=None,
|
||||
inputs={},
|
||||
credential_inputs=None,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
stats=None,
|
||||
organization_id=organization_id,
|
||||
team_id=TEAM_A,
|
||||
)
|
||||
@@ -13,7 +13,7 @@ from backend.data.credit import (
|
||||
admin_get_user_history,
|
||||
get_user_credit_model,
|
||||
)
|
||||
from backend.data.model import UserTransaction
|
||||
from backend.data.model import UserCreditTransactionAdminView
|
||||
from backend.data.platform_cost import (
|
||||
COPILOT_USAGE_EXPORT_MAX_DAYS,
|
||||
CopilotWeeklyUsageRow,
|
||||
@@ -90,7 +90,7 @@ async def admin_get_all_user_history(
|
||||
|
||||
|
||||
class CreditTransactionsExportResponse(BaseModel):
|
||||
transactions: list[UserTransaction]
|
||||
transactions: list[UserCreditTransactionAdminView]
|
||||
total_rows: int
|
||||
window_days: int
|
||||
max_window_days: int
|
||||
|
||||
@@ -9,7 +9,7 @@ import pytest_mock
|
||||
from autogpt_libs.auth.jwt_utils import get_jwt_payload
|
||||
from pytest_snapshot.plugin import Snapshot
|
||||
|
||||
from backend.data.model import UserTransaction
|
||||
from backend.data.model import UserCreditTransactionAdminView
|
||||
from backend.util.json import SafeJson
|
||||
from backend.util.models import Pagination
|
||||
|
||||
@@ -127,14 +127,14 @@ def test_get_user_history_success(
|
||||
# Mock the admin_get_user_history function
|
||||
mock_history_response = UserHistoryResponse(
|
||||
history=[
|
||||
UserTransaction(
|
||||
UserCreditTransactionAdminView(
|
||||
user_id="user-1",
|
||||
user_email="[email protected]",
|
||||
amount=1000,
|
||||
reason="Initial grant",
|
||||
transaction_type=prisma.enums.CreditTransactionType.GRANT,
|
||||
),
|
||||
UserTransaction(
|
||||
UserCreditTransactionAdminView(
|
||||
user_id="user-2",
|
||||
user_email="[email protected]",
|
||||
amount=-50,
|
||||
@@ -178,7 +178,7 @@ def test_get_user_history_with_filters(
|
||||
# Mock the admin_get_user_history function
|
||||
mock_history_response = UserHistoryResponse(
|
||||
history=[
|
||||
UserTransaction(
|
||||
UserCreditTransactionAdminView(
|
||||
user_id="user-3",
|
||||
user_email="[email protected]",
|
||||
amount=500,
|
||||
@@ -333,7 +333,7 @@ def test_export_credit_transactions_success(
|
||||
) -> None:
|
||||
mock_export = AsyncMock(
|
||||
return_value=[
|
||||
UserTransaction(
|
||||
UserCreditTransactionAdminView(
|
||||
user_id="user-1",
|
||||
user_email="[email protected]",
|
||||
amount=1000,
|
||||
@@ -342,7 +342,7 @@ def test_export_credit_transactions_success(
|
||||
admin_email="[email protected]",
|
||||
transaction_type=prisma.enums.CreditTransactionType.GRANT,
|
||||
),
|
||||
UserTransaction(
|
||||
UserCreditTransactionAdminView(
|
||||
user_id="user-2",
|
||||
user_email="[email protected]",
|
||||
amount=-50,
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
from pydantic import BaseModel
|
||||
|
||||
from backend.data.model import UserTransaction
|
||||
from backend.data.model import UserCreditTransactionAdminView
|
||||
from backend.util.models import Pagination
|
||||
|
||||
|
||||
class UserHistoryResponse(BaseModel):
|
||||
"""Response model for listings with version history"""
|
||||
|
||||
history: list[UserTransaction]
|
||||
history: list[UserCreditTransactionAdminView]
|
||||
pagination: Pagination
|
||||
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ class CreateAPIKeyRequest(pydantic.BaseModel):
|
||||
name: str
|
||||
permissions: list[APIKeyPermission]
|
||||
description: Optional[str] = None
|
||||
team_id: Optional[str] = None
|
||||
|
||||
|
||||
class CreateAPIKeyResponse(pydantic.BaseModel):
|
||||
|
||||
@@ -22,12 +22,21 @@ async def create_api_key(
|
||||
ctx: Annotated[RequestContext, Security(get_request_context)],
|
||||
) -> CreateAPIKeyResponse:
|
||||
"""Create a new API key"""
|
||||
if request.team_id and request.team_id not in await get_user_team_ids(
|
||||
user_id, ctx.org_id
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="You are not an active member of that team",
|
||||
)
|
||||
|
||||
api_key_info, plain_text_key = await api_key_db.create_api_key(
|
||||
name=request.name,
|
||||
user_id=user_id,
|
||||
permissions=request.permissions,
|
||||
description=request.description,
|
||||
organization_id=ctx.org_id,
|
||||
team_id_restriction=request.team_id,
|
||||
)
|
||||
return CreateAPIKeyResponse(api_key=api_key_info, plain_text_key=plain_text_key)
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ from pytest_snapshot.plugin import Snapshot
|
||||
|
||||
from backend.api.features.billing.credits.routes import get_credit_history, router
|
||||
from backend.api.rest_api import app as real_app
|
||||
from backend.api.rest_api import handle_internal_http_error
|
||||
from backend.api.utils.exceptions import add_exception_handlers
|
||||
from backend.data.credit import AutoTopUpConfig, UserCreditBase
|
||||
from backend.data.model import TransactionHistory
|
||||
|
||||
@@ -561,7 +561,7 @@ HISTORY_CONTEXT = RequestContext(
|
||||
def history_client():
|
||||
history_app = fastapi.FastAPI()
|
||||
history_app.include_router(router)
|
||||
history_app.add_exception_handler(ValueError, handle_internal_http_error(400))
|
||||
add_exception_handlers(history_app)
|
||||
history_app.dependency_overrides[requires_user] = lambda: None
|
||||
history_app.dependency_overrides[get_user_id] = lambda: "user"
|
||||
history_app.dependency_overrides[get_request_context] = lambda: HISTORY_CONTEXT
|
||||
|
||||
@@ -9,7 +9,7 @@ import pytest
|
||||
import pytest_mock
|
||||
|
||||
from backend.api.features.chat import speech as speech_routes
|
||||
from backend.api.rest_api import handle_internal_http_error
|
||||
from backend.api.utils.exceptions import add_exception_handlers
|
||||
from backend.copilot import speech as speech_module
|
||||
from backend.copilot.rate_limit import RateLimitExceeded, RateLimitUnavailable
|
||||
from backend.copilot.usage_activation_fixtures import usage_snapshot
|
||||
@@ -21,7 +21,7 @@ app = fastapi.FastAPI()
|
||||
app.include_router(speech_routes.router)
|
||||
# The route lets ValueError through to the app's global handler; without this
|
||||
# the test app would 500 where production 400s.
|
||||
app.add_exception_handler(ValueError, handle_internal_http_error(400))
|
||||
add_exception_handlers(app)
|
||||
client = fastapi.testclient.TestClient(app)
|
||||
|
||||
AUDIO = b"ID3-fake-mp3-bytes"
|
||||
@@ -82,9 +82,9 @@ def test_speech_refuses_a_paywalled_caller_before_spending(
|
||||
new=AsyncMock(side_effect=UserPaywalledError("no active subscription")),
|
||||
)
|
||||
|
||||
with pytest.raises(UserPaywalledError):
|
||||
client.post("/speech", json={"text": "hello"})
|
||||
response = client.post("/speech", json={"text": "hello"})
|
||||
|
||||
assert response.status_code == 402
|
||||
speech_module._speech_client.assert_not_called()
|
||||
record_usage.assert_not_awaited()
|
||||
|
||||
|
||||
+94
-67
@@ -7,6 +7,7 @@ import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import pytest_mock
|
||||
from fastapi import HTTPException
|
||||
from prisma.enums import ReviewStatus
|
||||
from pytest_snapshot.plugin import Snapshot
|
||||
|
||||
@@ -40,7 +41,8 @@ from backend.data.graph import (
|
||||
)
|
||||
from backend.util.exceptions import BlockExecutionError
|
||||
|
||||
from .model import PendingHumanReviewModel
|
||||
from .model import PendingHumanReviewModel, ReviewItem
|
||||
from .service import process_reviews
|
||||
|
||||
# Using a fixed timestamp for reproducible tests
|
||||
FIXED_NOW = datetime.datetime(2023, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc)
|
||||
@@ -179,6 +181,7 @@ async def test_get_pending_reviews_for_execution_not_available(
|
||||
|
||||
|
||||
_ROUTES = "backend.api.features.graph_executions.review.routes"
|
||||
_SERVICE = "backend.api.features.graph_executions.review.service"
|
||||
|
||||
|
||||
def _chat_review(user_id: str, status: ReviewStatus) -> PendingHumanReviewModel:
|
||||
@@ -242,18 +245,18 @@ async def test_approving_a_chat_review_resumes_no_graph_and_scopes_auto_approval
|
||||
) -> None:
|
||||
waiting = _chat_review(test_user_id, ReviewStatus.WAITING)
|
||||
mocker.patch(
|
||||
f"{_ROUTES}.get_reviews_by_node_exec_ids",
|
||||
f"{_SERVICE}.get_reviews_by_node_exec_ids",
|
||||
return_value={waiting.node_exec_id: waiting},
|
||||
)
|
||||
mocker.patch(
|
||||
f"{_ROUTES}.process_all_reviews_for_execution",
|
||||
f"{_SERVICE}.process_all_reviews_for_execution",
|
||||
return_value={
|
||||
waiting.node_exec_id: _chat_review(test_user_id, ReviewStatus.APPROVED)
|
||||
},
|
||||
)
|
||||
auto_approve = mocker.patch(f"{_ROUTES}.create_auto_approval_record")
|
||||
graph_exec_meta = mocker.patch(f"{_ROUTES}.get_graph_execution_meta")
|
||||
resume = mocker.patch(f"{_ROUTES}.add_graph_execution")
|
||||
auto_approve = mocker.patch(f"{_SERVICE}.create_auto_approval_record")
|
||||
graph_exec_meta = mocker.patch(f"{_SERVICE}.get_graph_execution_meta")
|
||||
resume = mocker.patch(f"{_SERVICE}.add_graph_execution")
|
||||
|
||||
response = await client.post(
|
||||
"/api/review/action",
|
||||
@@ -296,10 +299,10 @@ async def test_one_request_cannot_act_on_reviews_from_two_scopes(
|
||||
chat = _chat_review(test_user_id, ReviewStatus.WAITING)
|
||||
second = chat.model_copy(update={"node_exec_id": "second", **other})
|
||||
mocker.patch(
|
||||
f"{_ROUTES}.get_reviews_by_node_exec_ids",
|
||||
f"{_SERVICE}.get_reviews_by_node_exec_ids",
|
||||
return_value={chat.node_exec_id: chat, second.node_exec_id: second},
|
||||
)
|
||||
process = mocker.patch(f"{_ROUTES}.process_all_reviews_for_execution")
|
||||
process = mocker.patch(f"{_SERVICE}.process_all_reviews_for_execution")
|
||||
|
||||
response = await client.post(
|
||||
"/api/review/action",
|
||||
@@ -315,6 +318,30 @@ async def test_one_request_cannot_act_on_reviews_from_two_scopes(
|
||||
process.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_a_run_pinned_request_cannot_act_on_a_chat_review(
|
||||
mocker: pytest_mock.MockerFixture,
|
||||
test_user_id: str,
|
||||
) -> None:
|
||||
chat = _chat_review(test_user_id, ReviewStatus.WAITING)
|
||||
mocker.patch(
|
||||
f"{_SERVICE}.get_reviews_by_node_exec_ids",
|
||||
return_value={chat.node_exec_id: chat},
|
||||
)
|
||||
mocker.patch(f"{_SERVICE}._assert_awaiting_review")
|
||||
process = mocker.patch(f"{_SERVICE}.process_all_reviews_for_execution")
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await process_reviews(
|
||||
test_user_id,
|
||||
[ReviewItem(node_exec_id=chat.node_exec_id, approved=True)],
|
||||
graph_exec_id="ge-1",
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 404
|
||||
process.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="session")
|
||||
async def test_process_review_action_approve_success(
|
||||
client: httpx.AsyncClient,
|
||||
@@ -327,7 +354,7 @@ async def test_process_review_action_approve_success(
|
||||
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {"test_node_123": sample_pending_review}
|
||||
|
||||
@@ -337,7 +364,7 @@ async def test_process_review_action_approve_success(
|
||||
mock_get_reviews_for_execution.return_value = [sample_pending_review]
|
||||
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
# Create approved review for return
|
||||
approved_review = PendingHumanReviewModel(
|
||||
@@ -361,19 +388,19 @@ async def test_process_review_action_approve_success(
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
mock_get_graph_exec.return_value = mock_graph_exec_meta
|
||||
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.add_graph_execution"
|
||||
"backend.api.features.graph_executions.review.service.add_graph_execution"
|
||||
)
|
||||
|
||||
request_data = {
|
||||
@@ -409,13 +436,13 @@ async def test_process_review_action_reject_success(
|
||||
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {"test_node_123": sample_pending_review}
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -427,7 +454,7 @@ async def test_process_review_action_reject_success(
|
||||
mock_get_reviews_for_execution.return_value = [sample_pending_review]
|
||||
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
rejected_review = PendingHumanReviewModel(
|
||||
node_exec_id="test_node_123",
|
||||
@@ -449,7 +476,7 @@ async def test_process_review_action_reject_success(
|
||||
mock_process_all_reviews.return_value = {"test_node_123": rejected_review}
|
||||
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
@@ -504,7 +531,7 @@ async def test_process_review_action_mixed_success(
|
||||
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {
|
||||
"test_node_123": sample_pending_review,
|
||||
@@ -517,7 +544,7 @@ async def test_process_review_action_mixed_success(
|
||||
mock_get_reviews_for_execution.return_value = [sample_pending_review, second_review]
|
||||
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
# Create approved version of first review
|
||||
approved_review = PendingHumanReviewModel(
|
||||
@@ -562,14 +589,14 @@ async def test_process_review_action_mixed_success(
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
mock_get_graph_exec.return_value = mock_graph_exec_meta
|
||||
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
@@ -628,14 +655,14 @@ async def test_process_review_action_review_not_found(
|
||||
"""Test error when review is not found"""
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
# Return empty dict to simulate review not found
|
||||
mock_get_reviews_for_user.return_value = {}
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -649,7 +676,7 @@ async def test_process_review_action_review_not_found(
|
||||
|
||||
# Mock process_all_reviews to simulate not finding reviews
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
# This should raise a ValueError with "Reviews not found" message based on the data/human_review.py logic
|
||||
mock_process_all_reviews.side_effect = ValueError(
|
||||
@@ -682,13 +709,13 @@ async def test_process_review_action_partial_failure(
|
||||
"""Test handling of partial failures in review processing"""
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {"test_node_123": sample_pending_review}
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -702,7 +729,7 @@ async def test_process_review_action_partial_failure(
|
||||
|
||||
# Mock partial failure in processing
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
mock_process_all_reviews.side_effect = ValueError("Some reviews failed validation")
|
||||
|
||||
@@ -732,14 +759,14 @@ async def test_process_review_action_invalid_node_exec_id(
|
||||
"""Test failure when trying to process review with invalid node execution ID"""
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
# Return empty dict to simulate review not found
|
||||
mock_get_reviews_for_user.return_value = {}
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -772,13 +799,13 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
"""Test that auto_approve_future_actions flag creates auto-approval records"""
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {"test_node_123": sample_pending_review}
|
||||
|
||||
# Mock process_all_reviews
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
approved_review = PendingHumanReviewModel(
|
||||
node_exec_id="test_node_123",
|
||||
@@ -801,7 +828,7 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
|
||||
# Mock get_node_executions to return node_id mapping
|
||||
mock_get_node_executions = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_node_executions"
|
||||
"backend.api.features.graph_executions.review.service.get_node_executions"
|
||||
)
|
||||
mock_node_exec = mocker.Mock(spec=NodeExecutionResult)
|
||||
mock_node_exec.node_exec_id = "test_node_123"
|
||||
@@ -810,12 +837,12 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
|
||||
# Mock create_auto_approval_record
|
||||
mock_create_auto_approval = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.create_auto_approval_record"
|
||||
"backend.api.features.graph_executions.review.service.create_auto_approval_record"
|
||||
)
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -823,13 +850,13 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
|
||||
# Mock has_pending_reviews_for_graph_exec
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
# Mock get_graph_settings to return custom settings
|
||||
mock_get_settings = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_settings"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_settings"
|
||||
)
|
||||
mock_get_settings.return_value = GraphSettings(
|
||||
human_in_the_loop_safe_mode=True,
|
||||
@@ -838,7 +865,7 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
|
||||
# Mock get_user_by_id to prevent database access
|
||||
mock_get_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_user_by_id"
|
||||
"backend.api.features.graph_executions.review.service.get_user_by_id"
|
||||
)
|
||||
mock_user = mocker.Mock()
|
||||
mock_user.timezone = "UTC"
|
||||
@@ -846,7 +873,7 @@ async def test_process_review_action_auto_approve_creates_auto_approval_records(
|
||||
|
||||
# Mock add_graph_execution
|
||||
mock_add_execution = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.add_graph_execution"
|
||||
"backend.api.features.graph_executions.review.service.add_graph_execution"
|
||||
)
|
||||
|
||||
request_data = {
|
||||
@@ -903,13 +930,13 @@ async def test_process_review_action_without_auto_approve_still_loads_settings(
|
||||
"""Test that execution context is created with settings even without auto-approve"""
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
mock_get_reviews_for_user.return_value = {"test_node_123": sample_pending_review}
|
||||
|
||||
# Mock process_all_reviews
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
approved_review = PendingHumanReviewModel(
|
||||
node_exec_id="test_node_123",
|
||||
@@ -932,12 +959,12 @@ async def test_process_review_action_without_auto_approve_still_loads_settings(
|
||||
|
||||
# Mock create_auto_approval_record - should NOT be called when auto_approve is False
|
||||
mock_create_auto_approval = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.create_auto_approval_record"
|
||||
"backend.api.features.graph_executions.review.service.create_auto_approval_record"
|
||||
)
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -945,13 +972,13 @@ async def test_process_review_action_without_auto_approve_still_loads_settings(
|
||||
|
||||
# Mock has_pending_reviews_for_graph_exec
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
# Mock get_graph_settings with sensitive_action_safe_mode enabled
|
||||
mock_get_settings = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_settings"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_settings"
|
||||
)
|
||||
mock_get_settings.return_value = GraphSettings(
|
||||
human_in_the_loop_safe_mode=False,
|
||||
@@ -960,7 +987,7 @@ async def test_process_review_action_without_auto_approve_still_loads_settings(
|
||||
|
||||
# Mock get_user_by_id to prevent database access
|
||||
mock_get_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_user_by_id"
|
||||
"backend.api.features.graph_executions.review.service.get_user_by_id"
|
||||
)
|
||||
mock_user = mocker.Mock()
|
||||
mock_user.timezone = "UTC"
|
||||
@@ -968,7 +995,7 @@ async def test_process_review_action_without_auto_approve_still_loads_settings(
|
||||
|
||||
# Mock add_graph_execution
|
||||
mock_add_execution = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.add_graph_execution"
|
||||
"backend.api.features.graph_executions.review.service.add_graph_execution"
|
||||
)
|
||||
|
||||
# Request WITHOUT auto_approve_future (defaults to False)
|
||||
@@ -1051,7 +1078,7 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock get_reviews_by_node_exec_ids (called to find the graph_exec_id)
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
# Need to return both reviews in WAITING state (before processing)
|
||||
approved_review_waiting = PendingHumanReviewModel(
|
||||
@@ -1091,7 +1118,7 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock process_all_reviews
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
mock_process_all_reviews.return_value = {
|
||||
"node_exec_approved": approved_review,
|
||||
@@ -1100,7 +1127,7 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock get_node_executions to return node_id mapping
|
||||
mock_get_node_executions = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_node_executions"
|
||||
"backend.api.features.graph_executions.review.service.get_node_executions"
|
||||
)
|
||||
mock_node_exec = mocker.Mock(spec=NodeExecutionResult)
|
||||
mock_node_exec.node_exec_id = "node_exec_approved"
|
||||
@@ -1109,12 +1136,12 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock create_auto_approval_record
|
||||
mock_create_auto_approval = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.create_auto_approval_record"
|
||||
"backend.api.features.graph_executions.review.service.create_auto_approval_record"
|
||||
)
|
||||
|
||||
# Mock get_graph_execution_meta to return execution in REVIEW status
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -1122,19 +1149,19 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock has_pending_reviews_for_graph_exec
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
# Mock get_graph_settings
|
||||
mock_get_settings = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_settings"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_settings"
|
||||
)
|
||||
mock_get_settings.return_value = GraphSettings()
|
||||
|
||||
# Mock get_user_by_id to prevent database access
|
||||
mock_get_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_user_by_id"
|
||||
"backend.api.features.graph_executions.review.service.get_user_by_id"
|
||||
)
|
||||
mock_user = mocker.Mock()
|
||||
mock_user.timezone = "UTC"
|
||||
@@ -1142,7 +1169,7 @@ async def test_process_review_action_auto_approve_only_applies_to_approved_revie
|
||||
|
||||
# Mock add_graph_execution
|
||||
mock_add_execution = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.add_graph_execution"
|
||||
"backend.api.features.graph_executions.review.service.add_graph_execution"
|
||||
)
|
||||
|
||||
request_data = {
|
||||
@@ -1198,7 +1225,7 @@ async def test_process_review_action_per_review_auto_approve_granularity(
|
||||
"""Test that auto-approval can be set per-review (granular control)"""
|
||||
# Mock get_reviews_by_node_exec_ids - return different reviews based on node_exec_id
|
||||
mock_get_reviews_for_user = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_reviews_by_node_exec_ids"
|
||||
"backend.api.features.graph_executions.review.service.get_reviews_by_node_exec_ids"
|
||||
)
|
||||
|
||||
# Create a mapping of node_exec_id to review
|
||||
@@ -1255,7 +1282,7 @@ async def test_process_review_action_per_review_auto_approve_granularity(
|
||||
|
||||
# Mock process_all_reviews - return 3 approved reviews
|
||||
mock_process_all_reviews = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.process_all_reviews_for_execution"
|
||||
"backend.api.features.graph_executions.review.service.process_all_reviews_for_execution"
|
||||
)
|
||||
mock_process_all_reviews.return_value = {
|
||||
"node_1_auto": PendingHumanReviewModel(
|
||||
@@ -1313,7 +1340,7 @@ async def test_process_review_action_per_review_auto_approve_granularity(
|
||||
|
||||
# Mock get_node_executions to return batch node data
|
||||
mock_get_node_executions = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_node_executions"
|
||||
"backend.api.features.graph_executions.review.service.get_node_executions"
|
||||
)
|
||||
# Create mock node executions for each review
|
||||
mock_node_execs = []
|
||||
@@ -1326,12 +1353,12 @@ async def test_process_review_action_per_review_auto_approve_granularity(
|
||||
|
||||
# Mock create_auto_approval_record
|
||||
mock_create_auto_approval = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.create_auto_approval_record"
|
||||
"backend.api.features.graph_executions.review.service.create_auto_approval_record"
|
||||
)
|
||||
|
||||
# Mock get_graph_execution_meta
|
||||
mock_get_graph_exec = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_execution_meta"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_execution_meta"
|
||||
)
|
||||
mock_graph_exec_meta = mocker.Mock()
|
||||
mock_graph_exec_meta.status = ExecutionStatus.REVIEW
|
||||
@@ -1339,22 +1366,22 @@ async def test_process_review_action_per_review_auto_approve_granularity(
|
||||
|
||||
# Mock has_pending_reviews_for_graph_exec
|
||||
mock_has_pending = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.has_pending_reviews_for_graph_exec"
|
||||
"backend.api.features.graph_executions.review.service.has_pending_reviews_for_graph_exec"
|
||||
)
|
||||
mock_has_pending.return_value = False
|
||||
|
||||
# Mock settings and execution
|
||||
mock_get_settings = mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_settings"
|
||||
"backend.api.features.graph_executions.review.service.get_graph_settings"
|
||||
)
|
||||
mock_get_settings.return_value = GraphSettings(
|
||||
human_in_the_loop_safe_mode=False, sensitive_action_safe_mode=False
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.add_graph_execution"
|
||||
"backend.api.features.graph_executions.review.service.add_graph_execution"
|
||||
)
|
||||
mocker.patch("backend.api.features.graph_executions.review.routes.get_user_by_id")
|
||||
mocker.patch("backend.api.features.graph_executions.review.service.get_user_by_id")
|
||||
|
||||
# Request with granular auto-approval:
|
||||
# - node_1_auto: auto_approve_future=True
|
||||
@@ -1422,7 +1449,7 @@ async def test_an_answer_on_a_chat_card_wakes_that_chat(
|
||||
review = sample_pending_review.model_copy(
|
||||
update={"graph_exec_id": graph_exec_id, "session_id": session_id}
|
||||
)
|
||||
routes = "backend.api.features.graph_executions.review.routes"
|
||||
routes = _SERVICE
|
||||
mocker.patch(
|
||||
f"{routes}.get_reviews_by_node_exec_ids",
|
||||
return_value={"test_node_123": review},
|
||||
@@ -1509,7 +1536,7 @@ async def test_an_approved_chat_card_sets_the_rule_it_asked_for(
|
||||
mocker.patch(
|
||||
"backend.copilot.gate.held._held", side_effect=lambda _: dict(held_calls)
|
||||
)
|
||||
routes = "backend.api.features.graph_executions.review.routes"
|
||||
routes = _SERVICE
|
||||
mocker.patch(
|
||||
f"{routes}.get_reviews_by_node_exec_ids",
|
||||
return_value={"test_node_123": review},
|
||||
@@ -1599,7 +1626,7 @@ async def test_a_rejected_irreversible_action_never_runs_when_the_run_resumes(
|
||||
await create_graph(graph, user_id=test_user_id)
|
||||
# The user's library setting that asks for irreversible actions to pause.
|
||||
mocker.patch(
|
||||
"backend.api.features.graph_executions.review.routes.get_graph_settings",
|
||||
"backend.api.features.graph_executions.review.service.get_graph_settings",
|
||||
new=mocker.AsyncMock(
|
||||
return_value=GraphSettings(sensitive_action_safe_mode=True)
|
||||
),
|
||||
|
||||
@@ -1,46 +1,22 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any, List
|
||||
from typing import List
|
||||
|
||||
import autogpt_libs.auth as autogpt_auth_lib
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Security, status
|
||||
from prisma.enums import ReviewStatus
|
||||
|
||||
from backend.copilot.constants import legacy_chat_session_id, parse_node_id_from_exec_id
|
||||
from backend.copilot.gate.chat_rules import set_answer_rules as set_chat_rules
|
||||
from backend.copilot.constants import legacy_chat_session_id
|
||||
from backend.copilot.gate.held import answered_cards as held_answered_cards
|
||||
from backend.copilot.gate.held import subject_keys as held_subject_keys
|
||||
from backend.copilot.gate.held import wake as wake_for_held_calls
|
||||
from backend.copilot.gate.workflow import (
|
||||
close_step_cards,
|
||||
finish_step_cards,
|
||||
is_step_card,
|
||||
step_of,
|
||||
step_rule_keys,
|
||||
)
|
||||
from backend.copilot.rate_limit import enforce_payment_paywall
|
||||
from backend.data.execution import (
|
||||
ExecutionContext,
|
||||
ExecutionStatus,
|
||||
get_graph_execution_meta,
|
||||
get_node_executions,
|
||||
)
|
||||
from backend.data.graph import get_graph_settings
|
||||
from backend.data.execution import get_graph_execution_meta
|
||||
from backend.data.human_review import (
|
||||
create_auto_approval_record,
|
||||
get_pending_reviews_for_chat_session,
|
||||
get_pending_reviews_for_execution,
|
||||
get_pending_reviews_for_user,
|
||||
get_reviews_by_node_exec_ids,
|
||||
has_pending_reviews_for_graph_exec,
|
||||
process_all_reviews_for_execution,
|
||||
)
|
||||
from backend.data.model import USER_TIMEZONE_NOT_SET
|
||||
from backend.data.user import get_user_by_id
|
||||
from backend.data.workspace import get_or_create_workspace
|
||||
from backend.executor.utils import add_graph_execution
|
||||
|
||||
from .model import PendingHumanReviewModel, ReviewRequest, ReviewResponse
|
||||
from .service import process_reviews
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -51,37 +27,6 @@ router = APIRouter(
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_node_ids(
|
||||
node_exec_ids: list[str],
|
||||
graph_exec_id: str | None,
|
||||
) -> dict[str, str]:
|
||||
"""Resolve node_exec_id -> node_id for auto-approval records.
|
||||
|
||||
A chat review's id encodes its node id as "{node_id}:{random}"; graph
|
||||
executions look up node_id from NodeExecution records.
|
||||
"""
|
||||
if not node_exec_ids:
|
||||
return {}
|
||||
|
||||
if graph_exec_id is None:
|
||||
return {neid: parse_node_id_from_exec_id(neid) for neid in node_exec_ids}
|
||||
|
||||
node_execs = await get_node_executions(
|
||||
graph_exec_id=graph_exec_id, include_exec_data=False
|
||||
)
|
||||
node_exec_map = {ne.node_exec_id: ne.node_id for ne in node_execs}
|
||||
|
||||
result = {}
|
||||
for neid in node_exec_ids:
|
||||
if neid in node_exec_map:
|
||||
result[neid] = node_exec_map[neid]
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed to resolve node_id for {neid}: Node execution not found."
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@router.get(
|
||||
"/pending",
|
||||
summary="Get Pending Reviews",
|
||||
@@ -217,252 +162,24 @@ async def process_review_action(
|
||||
user_id: str = Security(autogpt_auth_lib.get_user_id),
|
||||
) -> ReviewResponse:
|
||||
"""Process reviews with approve or reject actions."""
|
||||
|
||||
# Collect all node exec IDs from the request
|
||||
all_request_node_ids = {review.node_exec_id for review in request.reviews}
|
||||
|
||||
if not all_request_node_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="At least one review must be provided",
|
||||
)
|
||||
|
||||
# Batch fetch all requested reviews (regardless of status for idempotent handling)
|
||||
reviews_map = await get_reviews_by_node_exec_ids(
|
||||
list(all_request_node_ids), user_id
|
||||
outcome = await process_reviews(
|
||||
user_id, request.reviews, wait_for_other_cards=request.wait_for_other_cards
|
||||
)
|
||||
|
||||
# Validate all reviews were found (must exist, any status is OK for now)
|
||||
missing_ids = all_request_node_ids - set(reviews_map.keys())
|
||||
if missing_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Review(s) not found: {', '.join(missing_ids)}",
|
||||
)
|
||||
|
||||
# Validate all reviews belong to the same execution or chat
|
||||
scopes = {
|
||||
(review.graph_exec_id, review.session_id) for review in reviews_map.values()
|
||||
}
|
||||
if len(scopes) > 1:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="All reviews in a single request must belong to the same execution.",
|
||||
)
|
||||
|
||||
graph_exec_id, chat_session_id = next(iter(scopes))
|
||||
|
||||
# A run still running would finish its pass with nothing waiting, the step
|
||||
# unrun; a queued or ended run is safe to answer.
|
||||
for review in reviews_map.values():
|
||||
step = step_of(review)
|
||||
if step is not None and await _running(user_id, step.graph_exec_id):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="This workflow is still running other steps. Answer once it pauses.",
|
||||
)
|
||||
|
||||
# Validate execution status for graph executions; a chat has none
|
||||
if graph_exec_id is not None:
|
||||
graph_exec_meta = await get_graph_execution_meta(
|
||||
user_id=user_id, execution_id=graph_exec_id
|
||||
)
|
||||
if not graph_exec_meta:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph execution #{graph_exec_id} not found",
|
||||
)
|
||||
if graph_exec_meta.status not in (
|
||||
ExecutionStatus.REVIEW,
|
||||
ExecutionStatus.INCOMPLETE,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"Cannot process reviews while execution status is {graph_exec_meta.status}",
|
||||
)
|
||||
|
||||
# Build review decisions map and track which reviews requested auto-approval
|
||||
# Auto-approved reviews use original data (no modifications allowed)
|
||||
review_decisions = {}
|
||||
auto_approve_requests = {} # Map node_exec_id -> auto_approve_future flag
|
||||
|
||||
for review in request.reviews:
|
||||
review_status = (
|
||||
ReviewStatus.APPROVED if review.approved else ReviewStatus.REJECTED
|
||||
)
|
||||
# If this review requested auto-approval, don't allow data modifications
|
||||
reviewed_data = None if review.auto_approve_future else review.reviewed_data
|
||||
review_decisions[review.node_exec_id] = (
|
||||
review_status,
|
||||
reviewed_data,
|
||||
review.message,
|
||||
)
|
||||
auto_approve_requests[review.node_exec_id] = review.auto_approve_future
|
||||
|
||||
# Read before the answer lands: a turn it wakes claims the held calls.
|
||||
ruled = [r.node_exec_id for r in request.reviews if r.chat_rule]
|
||||
chat_rule_keys = (
|
||||
await held_subject_keys(chat_session_id, ruled)
|
||||
| step_rule_keys(reviews_map[node_exec_id] for node_exec_id in ruled)
|
||||
if chat_session_id is not None
|
||||
else {}
|
||||
)
|
||||
|
||||
# Process all reviews
|
||||
updated_reviews = await process_all_reviews_for_execution(
|
||||
user_id=user_id,
|
||||
review_decisions=review_decisions,
|
||||
)
|
||||
|
||||
# Create auto-approval records for approved reviews that requested it
|
||||
# Deduplicate by node_id to avoid race conditions when multiple reviews
|
||||
# for the same node are processed in parallel
|
||||
async def create_auto_approval_for_node(
|
||||
node_id: str, review_result
|
||||
) -> tuple[str, bool]:
|
||||
"""
|
||||
Create auto-approval record for a node.
|
||||
Returns (node_id, success) tuple for tracking failures.
|
||||
"""
|
||||
try:
|
||||
await create_auto_approval_record(
|
||||
user_id=user_id,
|
||||
node_id=node_id,
|
||||
payload=review_result.payload,
|
||||
graph_exec_id=review_result.graph_exec_id,
|
||||
graph_id=review_result.graph_id,
|
||||
graph_version=review_result.graph_version,
|
||||
chat_session_id=review_result.session_id,
|
||||
)
|
||||
return (node_id, True)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to create auto-approval record for node {node_id}",
|
||||
exc_info=e,
|
||||
)
|
||||
return (node_id, False)
|
||||
|
||||
# Collect node_exec_ids that need auto-approval and resolve their node_ids
|
||||
node_exec_ids_needing_auto_approval = [
|
||||
node_exec_id
|
||||
for node_exec_id, review_result in updated_reviews.items()
|
||||
if review_result.status == ReviewStatus.APPROVED
|
||||
and auto_approve_requests.get(node_exec_id, False)
|
||||
]
|
||||
|
||||
node_id_map = await _resolve_node_ids(
|
||||
node_exec_ids_needing_auto_approval, graph_exec_id
|
||||
)
|
||||
|
||||
# Deduplicate by node_id — one auto-approval per node
|
||||
nodes_needing_auto_approval: dict[str, Any] = {}
|
||||
for node_exec_id in node_exec_ids_needing_auto_approval:
|
||||
node_id = node_id_map.get(node_exec_id)
|
||||
if node_id and node_id not in nodes_needing_auto_approval:
|
||||
nodes_needing_auto_approval[node_id] = updated_reviews[node_exec_id]
|
||||
|
||||
# Execute all auto-approval creations in parallel (deduplicated by node_id)
|
||||
auto_approval_results = await asyncio.gather(
|
||||
*[
|
||||
create_auto_approval_for_node(node_id, review_result)
|
||||
for node_id, review_result in nodes_needing_auto_approval.items()
|
||||
],
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
# Count auto-approval failures
|
||||
auto_approval_failed_count = 0
|
||||
for result in auto_approval_results:
|
||||
if isinstance(result, Exception):
|
||||
auto_approval_failed_count += 1
|
||||
logger.error(
|
||||
f"Unexpected exception during auto-approval creation: {result}"
|
||||
)
|
||||
elif isinstance(result, tuple) and len(result) == 2 and not result[1]:
|
||||
auto_approval_failed_count += 1
|
||||
|
||||
# Count results
|
||||
approved_count = sum(
|
||||
1
|
||||
for review in updated_reviews.values()
|
||||
if review.status == ReviewStatus.APPROVED
|
||||
)
|
||||
rejected_count = sum(
|
||||
1
|
||||
for review in updated_reviews.values()
|
||||
if review.status == ReviewStatus.REJECTED
|
||||
)
|
||||
|
||||
if chat_session_id is not None and chat_rule_keys:
|
||||
await set_chat_rules(
|
||||
chat_session_id,
|
||||
user_id,
|
||||
updated_reviews,
|
||||
{review.node_exec_id: review.chat_rule for review in request.reviews},
|
||||
chat_rule_keys,
|
||||
{review.node_exec_id: review.chat_rule_scope for review in request.reviews},
|
||||
)
|
||||
|
||||
# A held call finishes on its own: an answer starts the chat's next turn, or
|
||||
# from the chat's own queue, its last answer does.
|
||||
awaiting_other_cards = False
|
||||
if chat_session_id is not None and updated_reviews:
|
||||
# A workflow step's card answers the step's own row; its run resumes
|
||||
# below, so the card never wakes the chat.
|
||||
finished = await finish_step_cards(user_id, updated_reviews.values())
|
||||
for run in {step.graph_exec_id: step for step in finished.answered}.values():
|
||||
meta = await get_graph_execution_meta(
|
||||
user_id=user_id, execution_id=run.graph_exec_id
|
||||
)
|
||||
if (
|
||||
meta is not None
|
||||
and meta.status in (ExecutionStatus.REVIEW, ExecutionStatus.INCOMPLETE)
|
||||
and not await has_pending_reviews_for_graph_exec(run.graph_exec_id)
|
||||
):
|
||||
await resume_reviewed_execution(
|
||||
user_id, run.graph_exec_id, run.graph_id, run.graph_version
|
||||
)
|
||||
awaiting_other_cards = not await wake_for_held_calls(
|
||||
user_id,
|
||||
chat_session_id,
|
||||
[r for r in updated_reviews.values() if not is_step_card(r)],
|
||||
until_last=request.wait_for_other_cards,
|
||||
)
|
||||
if finished.failed:
|
||||
# Each card is kept, so answering it again finishes the job.
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="The answer could not reach its workflow step. Try again.",
|
||||
)
|
||||
|
||||
# Resume graph execution only for real graph executions; a chat is resumed
|
||||
# by the LLM calling resume_capability with the review_id
|
||||
if graph_exec_id is not None and updated_reviews:
|
||||
await close_step_cards(user_id, updated_reviews)
|
||||
if not await has_pending_reviews_for_graph_exec(graph_exec_id):
|
||||
first_review = next(iter(updated_reviews.values()))
|
||||
assert first_review.graph_id and first_review.graph_version is not None
|
||||
await resume_reviewed_execution(
|
||||
user_id,
|
||||
graph_exec_id,
|
||||
first_review.graph_id,
|
||||
first_review.graph_version,
|
||||
)
|
||||
|
||||
# Build error message if auto-approvals failed
|
||||
error_message = None
|
||||
if auto_approval_failed_count > 0:
|
||||
if outcome.auto_approval_failed_count:
|
||||
error_message = (
|
||||
f"{auto_approval_failed_count} auto-approval setting(s) could not be saved. "
|
||||
f"You may need to manually approve these reviews in future executions."
|
||||
f"{outcome.auto_approval_failed_count} auto-approval setting(s) could "
|
||||
f"not be saved. You may need to manually approve these reviews in "
|
||||
f"future executions."
|
||||
)
|
||||
|
||||
return ReviewResponse(
|
||||
approved_count=approved_count,
|
||||
rejected_count=rejected_count,
|
||||
failed_count=auto_approval_failed_count,
|
||||
approved_count=outcome.approved_count,
|
||||
rejected_count=outcome.rejected_count,
|
||||
failed_count=outcome.auto_approval_failed_count,
|
||||
error=error_message,
|
||||
awaiting_other_cards=awaiting_other_cards,
|
||||
awaiting_other_cards=outcome.awaiting_other_cards,
|
||||
)
|
||||
|
||||
|
||||
@@ -479,39 +196,3 @@ async def continue_with_answered_reviews(
|
||||
"""Start the chat's turn with the cards answered so far; the rest stay open
|
||||
and the turn is told how many still wait."""
|
||||
await wake_for_held_calls(user_id, chat_session_id)
|
||||
|
||||
|
||||
async def _running(user_id: str, graph_exec_id: str) -> bool:
|
||||
meta = await get_graph_execution_meta(user_id=user_id, execution_id=graph_exec_id)
|
||||
return meta is not None and meta.status == ExecutionStatus.RUNNING
|
||||
|
||||
|
||||
async def resume_reviewed_execution(
|
||||
user_id: str, graph_exec_id: str, graph_id: str, graph_version: int
|
||||
) -> None:
|
||||
"""Re-queue a run whose reviews are all answered; its paused steps re-run
|
||||
and find their answers."""
|
||||
try:
|
||||
user = await get_user_by_id(user_id)
|
||||
settings = await get_graph_settings(
|
||||
user_id=user_id, graph_id=graph_id, graph_version=graph_version
|
||||
)
|
||||
user_timezone = (
|
||||
user.timezone if user.timezone != USER_TIMEZONE_NOT_SET else "UTC"
|
||||
)
|
||||
workspace = await get_or_create_workspace(user_id)
|
||||
execution_context = ExecutionContext(
|
||||
human_in_the_loop_safe_mode=settings.human_in_the_loop_safe_mode,
|
||||
sensitive_action_safe_mode=settings.sensitive_action_safe_mode,
|
||||
user_timezone=user_timezone,
|
||||
workspace_id=workspace.id,
|
||||
)
|
||||
await add_graph_execution(
|
||||
graph_id=graph_id,
|
||||
user_id=user_id,
|
||||
graph_exec_id=graph_exec_id,
|
||||
execution_context=execution_context,
|
||||
)
|
||||
logger.info(f"Resumed execution {graph_exec_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to resume execution {graph_exec_id}: {str(e)}")
|
||||
@@ -0,0 +1,387 @@
|
||||
"""Human-in-the-loop review processing, shared by the internal and v2 APIs.
|
||||
|
||||
One implementation so the two surfaces cannot disagree about which reviews a
|
||||
request may touch, what a run's status has to be, or when execution resumes.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from prisma.enums import ReviewStatus
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from backend.copilot.constants import parse_node_id_from_exec_id
|
||||
from backend.copilot.gate.chat_rules import set_answer_rules as set_chat_rules
|
||||
from backend.copilot.gate.held import subject_keys as held_subject_keys
|
||||
from backend.copilot.gate.held import wake as wake_for_held_calls
|
||||
from backend.copilot.gate.workflow import (
|
||||
close_step_cards,
|
||||
finish_step_cards,
|
||||
is_step_card,
|
||||
step_of,
|
||||
step_rule_keys,
|
||||
)
|
||||
from backend.data.execution import (
|
||||
ExecutionContext,
|
||||
ExecutionStatus,
|
||||
get_graph_execution_meta,
|
||||
get_node_executions,
|
||||
)
|
||||
from backend.data.graph import get_graph_settings
|
||||
from backend.data.human_review import (
|
||||
create_auto_approval_record,
|
||||
get_reviews_by_node_exec_ids,
|
||||
has_pending_reviews_for_graph_exec,
|
||||
process_all_reviews_for_execution,
|
||||
)
|
||||
from backend.data.model import USER_TIMEZONE_NOT_SET
|
||||
from backend.data.user import get_user_by_id
|
||||
from backend.data.workspace import get_or_create_workspace
|
||||
from backend.executor.utils import add_graph_execution
|
||||
|
||||
from .model import ReviewItem
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ReviewOutcome(BaseModel):
|
||||
"""What processing a batch of review decisions did."""
|
||||
|
||||
graph_exec_id: Optional[str] = Field(
|
||||
description="The run the reviews belonged to; None for an AutoPilot chat's"
|
||||
)
|
||||
approved_count: int
|
||||
rejected_count: int
|
||||
auto_approval_failed_count: int = Field(
|
||||
description="Approved reviews whose auto-approve-future record failed to save"
|
||||
)
|
||||
awaiting_other_cards: bool = Field(
|
||||
default=False,
|
||||
description="The answers started no chat turn: other cards in the chat wait",
|
||||
)
|
||||
|
||||
|
||||
async def process_reviews(
|
||||
user_id: str,
|
||||
reviews: Sequence[ReviewItem],
|
||||
*,
|
||||
graph_exec_id: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
team_id: Optional[str] = None,
|
||||
wait_for_other_cards: bool = False,
|
||||
) -> ReviewOutcome:
|
||||
"""Approve or reject reviews, then resume the run once none are pending.
|
||||
|
||||
`graph_exec_id` pins the run the decisions must belong to, for a caller
|
||||
that already knows it from the URL; without it the run is taken from the
|
||||
reviews themselves, which must then all belong to one. `wait_for_other_cards`
|
||||
is a chat queue's answer, which starts no turn while another card waits.
|
||||
"""
|
||||
if not reviews:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="At least one review must be provided",
|
||||
)
|
||||
|
||||
reviews_map = await get_reviews_by_node_exec_ids(
|
||||
[review.node_exec_id for review in reviews], user_id
|
||||
)
|
||||
missing_ids = {r.node_exec_id for r in reviews} - set(reviews_map)
|
||||
if missing_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Review(s) not found: {', '.join(sorted(missing_ids))}",
|
||||
)
|
||||
|
||||
graph_exec_id, chat_session_id = _one_scope(reviews_map.values(), graph_exec_id)
|
||||
# A run still running would finish its pass with nothing waiting, the step
|
||||
# unrun; a queued or ended run is safe to answer.
|
||||
for review in reviews_map.values():
|
||||
step = step_of(review)
|
||||
if step is not None and await _running(user_id, step.graph_exec_id):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="This workflow is still running other steps. Answer once it pauses.",
|
||||
)
|
||||
if graph_exec_id is not None:
|
||||
await _assert_awaiting_review(user_id, graph_exec_id)
|
||||
|
||||
# Read before the answer lands: a turn it wakes claims the held calls.
|
||||
ruled = [r.node_exec_id for r in reviews if r.chat_rule]
|
||||
chat_rule_keys = (
|
||||
await held_subject_keys(chat_session_id, ruled)
|
||||
| step_rule_keys(reviews_map[node_exec_id] for node_exec_id in ruled)
|
||||
if chat_session_id is not None
|
||||
else {}
|
||||
)
|
||||
|
||||
# An auto-approved review takes the original data: approving future runs of
|
||||
# a block is not the place to also edit this one's payload.
|
||||
updated_reviews = await process_all_reviews_for_execution(
|
||||
user_id=user_id,
|
||||
review_decisions={
|
||||
review.node_exec_id: (
|
||||
ReviewStatus.APPROVED if review.approved else ReviewStatus.REJECTED,
|
||||
None if review.auto_approve_future else review.reviewed_data,
|
||||
review.message,
|
||||
)
|
||||
for review in reviews
|
||||
},
|
||||
)
|
||||
|
||||
failed = await _record_auto_approvals(
|
||||
user_id,
|
||||
{r.node_exec_id: r.auto_approve_future for r in reviews},
|
||||
updated_reviews,
|
||||
graph_exec_id,
|
||||
)
|
||||
|
||||
if chat_session_id is not None and chat_rule_keys:
|
||||
await set_chat_rules(
|
||||
chat_session_id,
|
||||
user_id,
|
||||
updated_reviews,
|
||||
{review.node_exec_id: review.chat_rule for review in reviews},
|
||||
chat_rule_keys,
|
||||
{review.node_exec_id: review.chat_rule_scope for review in reviews},
|
||||
)
|
||||
|
||||
# A held call finishes on its own: an answer starts the chat's next turn, or
|
||||
# from the chat's own queue, its last answer does.
|
||||
awaiting_other_cards = False
|
||||
if chat_session_id is not None and updated_reviews:
|
||||
# A workflow step's card answers the step's own row; its run resumes
|
||||
# below, so the card never wakes the chat.
|
||||
finished = await finish_step_cards(user_id, updated_reviews.values())
|
||||
for run in {step.graph_exec_id: step for step in finished.answered}.values():
|
||||
meta = await get_graph_execution_meta(
|
||||
user_id=user_id, execution_id=run.graph_exec_id
|
||||
)
|
||||
if (
|
||||
meta is not None
|
||||
and meta.status in (ExecutionStatus.REVIEW, ExecutionStatus.INCOMPLETE)
|
||||
and not await has_pending_reviews_for_graph_exec(run.graph_exec_id)
|
||||
):
|
||||
await _resume(
|
||||
user_id,
|
||||
run.graph_exec_id,
|
||||
run.graph_id,
|
||||
run.graph_version,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
awaiting_other_cards = not await wake_for_held_calls(
|
||||
user_id,
|
||||
chat_session_id,
|
||||
[r for r in updated_reviews.values() if not is_step_card(r)],
|
||||
until_last=wait_for_other_cards,
|
||||
)
|
||||
if finished.failed:
|
||||
# Each card is kept, so answering it again finishes the job.
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="The answer could not reach its workflow step. Try again.",
|
||||
)
|
||||
|
||||
if graph_exec_id is not None and updated_reviews:
|
||||
await close_step_cards(user_id, updated_reviews)
|
||||
await _resume_if_nothing_pending(
|
||||
user_id,
|
||||
graph_exec_id,
|
||||
updated_reviews,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
return ReviewOutcome(
|
||||
graph_exec_id=graph_exec_id,
|
||||
approved_count=sum(
|
||||
1 for r in updated_reviews.values() if r.status == ReviewStatus.APPROVED
|
||||
),
|
||||
rejected_count=sum(
|
||||
1 for r in updated_reviews.values() if r.status == ReviewStatus.REJECTED
|
||||
),
|
||||
auto_approval_failed_count=failed,
|
||||
awaiting_other_cards=awaiting_other_cards,
|
||||
)
|
||||
|
||||
|
||||
def _one_scope(
|
||||
reviews, graph_exec_id: Optional[str]
|
||||
) -> tuple[Optional[str], Optional[str]]:
|
||||
"""The one run, or the one chat, every decision belongs to."""
|
||||
scopes = {(review.graph_exec_id, review.session_id) for review in reviews}
|
||||
if graph_exec_id is not None:
|
||||
if {g for g, _ in scopes} != {graph_exec_id}:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Review(s) not found for run {graph_exec_id}",
|
||||
)
|
||||
return graph_exec_id, None
|
||||
|
||||
if len(scopes) > 1:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="All reviews in a single request must belong to the same execution.",
|
||||
)
|
||||
return next(iter(scopes))
|
||||
|
||||
|
||||
async def _assert_awaiting_review(user_id: str, graph_exec_id: str) -> None:
|
||||
graph_exec_meta = await get_graph_execution_meta(
|
||||
user_id=user_id, execution_id=graph_exec_id
|
||||
)
|
||||
if not graph_exec_meta:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=f"Graph execution #{graph_exec_id} not found",
|
||||
)
|
||||
if graph_exec_meta.status not in (
|
||||
ExecutionStatus.REVIEW,
|
||||
ExecutionStatus.INCOMPLETE,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=(
|
||||
"Cannot process reviews while execution status is "
|
||||
f"{graph_exec_meta.status}"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _record_auto_approvals(
|
||||
user_id: str,
|
||||
requested: dict[str, bool],
|
||||
updated_reviews: dict[str, Any],
|
||||
graph_exec_id: Optional[str],
|
||||
) -> int:
|
||||
"""Save one auto-approval per node, returning how many could not be saved."""
|
||||
node_exec_ids = [
|
||||
node_exec_id
|
||||
for node_exec_id, review in updated_reviews.items()
|
||||
if review.status == ReviewStatus.APPROVED and requested.get(node_exec_id, False)
|
||||
]
|
||||
node_id_map = await _resolve_node_ids(node_exec_ids, graph_exec_id)
|
||||
|
||||
# Deduplicate by node_id — concurrent reviews of one node would otherwise race.
|
||||
by_node: dict[str, Any] = {}
|
||||
for node_exec_id in node_exec_ids:
|
||||
node_id = node_id_map.get(node_exec_id)
|
||||
if node_id and node_id not in by_node:
|
||||
by_node[node_id] = updated_reviews[node_exec_id]
|
||||
|
||||
results = await asyncio.gather(
|
||||
*[
|
||||
_create_auto_approval(user_id, node_id, review)
|
||||
for node_id, review in by_node.items()
|
||||
],
|
||||
return_exceptions=True,
|
||||
)
|
||||
return sum(1 for saved in results if isinstance(saved, BaseException) or not saved)
|
||||
|
||||
|
||||
async def _create_auto_approval(user_id: str, node_id: str, review) -> bool:
|
||||
try:
|
||||
await create_auto_approval_record(
|
||||
user_id=user_id,
|
||||
graph_exec_id=review.graph_exec_id,
|
||||
graph_id=review.graph_id,
|
||||
graph_version=review.graph_version,
|
||||
node_id=node_id,
|
||||
payload=review.payload,
|
||||
chat_session_id=review.session_id,
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to create auto-approval record for node {node_id}", exc_info=e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _resolve_node_ids(
|
||||
node_exec_ids: list[str], graph_exec_id: Optional[str]
|
||||
) -> dict[str, str]:
|
||||
"""Resolve node_exec_id -> node_id for auto-approval records.
|
||||
|
||||
A chat review's id encodes its node id as "{node_id}:{random}".
|
||||
"""
|
||||
if not node_exec_ids:
|
||||
return {}
|
||||
if graph_exec_id is None:
|
||||
return {neid: parse_node_id_from_exec_id(neid) for neid in node_exec_ids}
|
||||
|
||||
node_execs = await get_node_executions(
|
||||
graph_exec_id=graph_exec_id, include_exec_data=False
|
||||
)
|
||||
by_exec_id = {ne.node_exec_id: ne.node_id for ne in node_execs}
|
||||
for missing in [neid for neid in node_exec_ids if neid not in by_exec_id]:
|
||||
logger.error(f"Failed to resolve node_id for {missing}: execution not found.")
|
||||
return {neid: by_exec_id[neid] for neid in node_exec_ids if neid in by_exec_id}
|
||||
|
||||
|
||||
async def _resume_if_nothing_pending(
|
||||
user_id: str,
|
||||
graph_exec_id: str,
|
||||
updated_reviews: dict[str, Any],
|
||||
*,
|
||||
organization_id: Optional[str],
|
||||
team_id: Optional[str],
|
||||
) -> None:
|
||||
if await has_pending_reviews_for_graph_exec(graph_exec_id):
|
||||
return
|
||||
|
||||
first_review = next(iter(updated_reviews.values()))
|
||||
await _resume(
|
||||
user_id,
|
||||
graph_exec_id,
|
||||
first_review.graph_id,
|
||||
first_review.graph_version,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
async def _resume(
|
||||
user_id: str,
|
||||
graph_exec_id: str,
|
||||
graph_id: Optional[str],
|
||||
graph_version: Optional[int],
|
||||
*,
|
||||
organization_id: Optional[str],
|
||||
team_id: Optional[str],
|
||||
) -> None:
|
||||
"""Re-queue a run whose reviews are all answered; its paused steps re-run
|
||||
and find their answers."""
|
||||
try:
|
||||
assert graph_id and graph_version is not None
|
||||
user = await get_user_by_id(user_id)
|
||||
settings = await get_graph_settings(
|
||||
user_id=user_id, graph_id=graph_id, graph_version=graph_version
|
||||
)
|
||||
workspace = await get_or_create_workspace(user_id)
|
||||
await add_graph_execution(
|
||||
graph_id=graph_id,
|
||||
user_id=user_id,
|
||||
graph_exec_id=graph_exec_id,
|
||||
execution_context=ExecutionContext(
|
||||
human_in_the_loop_safe_mode=settings.human_in_the_loop_safe_mode,
|
||||
sensitive_action_safe_mode=settings.sensitive_action_safe_mode,
|
||||
user_timezone=(
|
||||
user.timezone if user.timezone != USER_TIMEZONE_NOT_SET else "UTC"
|
||||
),
|
||||
workspace_id=workspace.id,
|
||||
),
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
logger.info(f"Resumed execution {graph_exec_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to resume execution {graph_exec_id}: {e}")
|
||||
|
||||
|
||||
async def _running(user_id: str, graph_exec_id: str) -> bool:
|
||||
meta = await get_graph_execution_meta(user_id=user_id, execution_id=graph_exec_id)
|
||||
return meta is not None and meta.status == ExecutionStatus.RUNNING
|
||||
+2
-2
@@ -45,7 +45,7 @@ from backend.data.redis_client import get_redis_async
|
||||
from backend.util.exceptions import BlockExecutionError
|
||||
from backend.util.feature_flag import Flag
|
||||
|
||||
_ROUTES = "backend.api.features.graph_executions.review.routes"
|
||||
_SERVICE = "backend.api.features.graph_executions.review.service"
|
||||
_BLOCK = SendDiscordMessageBlock()
|
||||
_INPUT = {"channel_name": "general", "message_content": "The digest is out"}
|
||||
_KEY = f"block:{_BLOCK.id}"
|
||||
@@ -78,7 +78,7 @@ def flags():
|
||||
|
||||
@pytest.fixture
|
||||
def resumed():
|
||||
with patch(f"{_ROUTES}.add_graph_execution", AsyncMock()) as add:
|
||||
with patch(f"{_SERVICE}.add_graph_execution", AsyncMock()) as add:
|
||||
yield add
|
||||
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from autogpt_libs.auth import get_request_context, get_user_id, requires_user
|
||||
@@ -7,6 +7,7 @@ from autogpt_libs.auth.models import RequestContext
|
||||
from fastapi import APIRouter, Body, HTTPException, Path, Query, Response, Security
|
||||
from starlette.status import HTTP_204_NO_CONTENT, HTTP_404_NOT_FOUND
|
||||
|
||||
from backend.api.features.graph_executions import sharing
|
||||
from backend.api.features.graph_executions.model import (
|
||||
ExecutionShareRequest,
|
||||
ExecutionShareResponse,
|
||||
@@ -23,13 +24,10 @@ from backend.data.onboarding import (
|
||||
complete_onboarding_step,
|
||||
get_user_onboarding,
|
||||
)
|
||||
from backend.data.sharing.tokens import SHARE_TOKEN_PATTERN, generate_share_token
|
||||
from backend.data.sharing.tokens import SHARE_TOKEN_PATTERN
|
||||
from backend.data.workspace import get_workspace_file_by_id
|
||||
from backend.executor import utils as execution_utils
|
||||
from backend.util.exceptions import NotFoundError
|
||||
from backend.util.settings import Settings
|
||||
|
||||
settings = Settings()
|
||||
|
||||
# No router-level auth dependency: the two /public/shared routes are
|
||||
# deliberately unauthenticated, so each route keeps its own.
|
||||
@@ -227,9 +225,7 @@ async def delete_graph_execution(
|
||||
user_id: Annotated[str, Security(get_user_id)],
|
||||
ctx: Annotated[RequestContext, Security(get_request_context)],
|
||||
) -> None:
|
||||
await execution_db.delete_graph_execution(
|
||||
graph_exec_id=graph_exec_id, user_id=user_id
|
||||
)
|
||||
await sharing.delete_execution(user_id, graph_exec_id)
|
||||
|
||||
|
||||
@router.post(
|
||||
@@ -244,49 +240,15 @@ async def enable_execution_sharing(
|
||||
_body: ExecutionShareRequest = Body(default=ExecutionShareRequest()),
|
||||
) -> ExecutionShareResponse:
|
||||
"""Enable sharing for a graph execution."""
|
||||
# Verify the execution belongs to the user
|
||||
execution = await execution_db.get_graph_execution(
|
||||
user_id=user_id, execution_id=graph_exec_id
|
||||
)
|
||||
if not execution:
|
||||
raise HTTPException(status_code=404, detail="Execution not found")
|
||||
|
||||
# Generate a unique share token
|
||||
share_token = generate_share_token()
|
||||
|
||||
# Remove stale allowlist records before updating the token — prevents a
|
||||
# window where old records + new token could coexist.
|
||||
await execution_db.delete_shared_execution_files(execution_id=graph_exec_id)
|
||||
|
||||
# Update the execution with share info — the underlying update_many
|
||||
# also enforces (id, user_id) at the DB layer, so a TOCTOU delete
|
||||
# between the pre-check above and this write surfaces as 404 rather
|
||||
# than a silent no-op.
|
||||
try:
|
||||
await execution_db.update_graph_execution_share_status(
|
||||
execution_id=graph_exec_id,
|
||||
user_id=user_id,
|
||||
is_shared=True,
|
||||
share_token=share_token,
|
||||
shared_at=datetime.now(timezone.utc),
|
||||
)
|
||||
share_token = await sharing.share_execution(user_id, graph_exec_id)
|
||||
except NotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc))
|
||||
|
||||
# Create allowlist of workspace files referenced in outputs
|
||||
await execution_db.create_shared_execution_files(
|
||||
execution_id=graph_exec_id,
|
||||
share_token=share_token,
|
||||
user_id=user_id,
|
||||
outputs=execution.outputs,
|
||||
return ExecutionShareResponse(
|
||||
share_url=sharing.share_url(share_token), share_token=share_token
|
||||
)
|
||||
|
||||
# Return the share URL
|
||||
frontend_url = settings.config.frontend_base_url or "http://localhost:3000"
|
||||
share_url = f"{frontend_url}/share/{share_token}"
|
||||
|
||||
return ExecutionShareResponse(share_url=share_url, share_token=share_token)
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/graphs/{graph_id}/executions/{graph_exec_id}/share",
|
||||
@@ -300,26 +262,8 @@ async def disable_execution_sharing(
|
||||
ctx: Annotated[RequestContext, Security(get_request_context)],
|
||||
) -> None:
|
||||
"""Disable sharing for a graph execution."""
|
||||
# Verify the execution belongs to the user
|
||||
execution = await execution_db.get_graph_execution(
|
||||
user_id=user_id, execution_id=graph_exec_id
|
||||
)
|
||||
if not execution:
|
||||
raise HTTPException(status_code=404, detail="Execution not found")
|
||||
|
||||
# Remove shared file allowlist records
|
||||
await execution_db.delete_shared_execution_files(execution_id=graph_exec_id)
|
||||
|
||||
# Remove share info — owner-gated at the DB layer; TOCTOU delete
|
||||
# after the pre-check surfaces as 404.
|
||||
try:
|
||||
await execution_db.update_graph_execution_share_status(
|
||||
execution_id=graph_exec_id,
|
||||
user_id=user_id,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
shared_at=None,
|
||||
)
|
||||
await sharing.unshare_execution(user_id, graph_exec_id)
|
||||
except NotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc))
|
||||
|
||||
|
||||
@@ -374,6 +374,11 @@ def test_delete_execution_returns_204(mocker: pytest_mock.MockFixture) -> None:
|
||||
"backend.api.features.graph_executions.routes.execution_db.delete_graph_execution",
|
||||
AsyncMock(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.graph_executions.routes.execution_db"
|
||||
".delete_shared_execution_files",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
response = client.delete("/executions/exec-1")
|
||||
|
||||
@@ -381,6 +386,26 @@ def test_delete_execution_returns_204(mocker: pytest_mock.MockFixture) -> None:
|
||||
assert deleted.await_args.kwargs["graph_exec_id"] == "exec-1"
|
||||
|
||||
|
||||
def test_deleting_an_execution_revokes_its_shared_file_downloads(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
"""The public file download checks only the allowlist, so it has to go
|
||||
with the run; the shared page already hides deleted runs."""
|
||||
mocker.patch(
|
||||
"backend.api.features.graph_executions.routes.execution_db.delete_graph_execution",
|
||||
AsyncMock(),
|
||||
)
|
||||
delete_files = mocker.patch(
|
||||
"backend.api.features.graph_executions.routes.execution_db"
|
||||
".delete_shared_execution_files",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
client.delete("/executions/exec-1")
|
||||
|
||||
delete_files.assert_awaited_once_with(execution_id="exec-1")
|
||||
|
||||
|
||||
def test_get_shared_execution_returns_404_for_an_unknown_token(
|
||||
mocker: pytest_mock.MockFixture,
|
||||
) -> None:
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Turning a run's public share link on and off, and deleting a run.
|
||||
|
||||
A run's share state, its share token and the allowlist of output files the
|
||||
public link may download have to change together: the public file download
|
||||
checks only the allowlist. The web app's routes and External API v2 both come
|
||||
through here, so neither can update one and leave the others behind.
|
||||
|
||||
Every function loads the caller's own run first and raises `NotFoundError`
|
||||
otherwise, so nothing is written for a run someone else owns.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from backend.data import execution as execution_db
|
||||
from backend.data.sharing.tokens import generate_share_token
|
||||
from backend.util.exceptions import NotFoundError
|
||||
from backend.util.settings import Settings
|
||||
|
||||
settings = Settings()
|
||||
|
||||
|
||||
async def share_execution(user_id: str, execution_id: str) -> str:
|
||||
"""Share the caller's run publicly and return its new share token.
|
||||
|
||||
A previous share's file allowlist is dropped before the new token is
|
||||
written, so links from an earlier share stop working.
|
||||
"""
|
||||
execution = await _own_execution(user_id, execution_id)
|
||||
share_token = generate_share_token()
|
||||
|
||||
await execution_db.delete_shared_execution_files(execution_id=execution_id)
|
||||
# Owner-gated at the DB layer too: a delete racing the check above
|
||||
# surfaces as NotFoundError rather than a silent no-op.
|
||||
await execution_db.update_graph_execution_share_status(
|
||||
execution_id=execution_id,
|
||||
user_id=user_id,
|
||||
is_shared=True,
|
||||
share_token=share_token,
|
||||
shared_at=datetime.now(timezone.utc),
|
||||
)
|
||||
await execution_db.create_shared_execution_files(
|
||||
execution_id=execution_id,
|
||||
share_token=share_token,
|
||||
user_id=user_id,
|
||||
outputs=execution.outputs,
|
||||
)
|
||||
return share_token
|
||||
|
||||
|
||||
async def unshare_execution(user_id: str, execution_id: str) -> None:
|
||||
"""Revoke the run's public link and every file download it allowed."""
|
||||
await _own_execution(user_id, execution_id)
|
||||
|
||||
await execution_db.delete_shared_execution_files(execution_id=execution_id)
|
||||
await execution_db.update_graph_execution_share_status(
|
||||
execution_id=execution_id,
|
||||
user_id=user_id,
|
||||
is_shared=False,
|
||||
share_token=None,
|
||||
shared_at=None,
|
||||
)
|
||||
|
||||
|
||||
async def delete_execution(user_id: str, execution_id: str) -> None:
|
||||
"""Delete the caller's run, revoking the file downloads a share allowed.
|
||||
|
||||
The shared page already hides deleted runs, but the file download checks
|
||||
only the allowlist, which would otherwise outlive the run.
|
||||
"""
|
||||
await execution_db.delete_graph_execution(
|
||||
graph_exec_id=execution_id, user_id=user_id
|
||||
)
|
||||
await execution_db.delete_shared_execution_files(execution_id=execution_id)
|
||||
|
||||
|
||||
def share_url(share_token: str) -> str:
|
||||
frontend_url = settings.config.frontend_base_url or "http://localhost:3000"
|
||||
return f"{frontend_url}/share/{share_token}"
|
||||
|
||||
|
||||
async def _own_execution(
|
||||
user_id: str, execution_id: str
|
||||
) -> execution_db.GraphExecution:
|
||||
execution = await execution_db.get_graph_execution(
|
||||
user_id=user_id, execution_id=execution_id
|
||||
)
|
||||
if execution is None:
|
||||
raise NotFoundError(f"Run #{execution_id} not found")
|
||||
return execution
|
||||
@@ -42,14 +42,14 @@ async def list_graphs(
|
||||
user_id: Annotated[str, Security(get_user_id)],
|
||||
ctx: Annotated[RequestContext, Security(get_request_context)],
|
||||
) -> Sequence[graph_db.GraphMeta]:
|
||||
paginated_result = await graph_db.list_graphs_paginated(
|
||||
graphs, _ = await graph_db.list_graphs_paginated(
|
||||
user_id=user_id,
|
||||
page=1,
|
||||
page_size=250,
|
||||
filter_by="active",
|
||||
organization_id=ctx.org_id,
|
||||
)
|
||||
return paginated_result.graphs
|
||||
return graphs
|
||||
|
||||
|
||||
@router.get(
|
||||
@@ -209,8 +209,8 @@ async def update_graph(
|
||||
|
||||
skipped_webhook_presets: list[library_model.SkippedWebhookPreset] = []
|
||||
if new_graph_version.is_active:
|
||||
await library_db.update_library_agent_version_and_settings(
|
||||
user_id, new_graph_version
|
||||
await library_db.update_agent_version_in_library(
|
||||
user_id, new_graph_version.id, new_graph_version.version
|
||||
)
|
||||
await graph_db.set_graph_active_version(
|
||||
graph_id=graph_id, version=new_graph_version.version, user_id=user_id
|
||||
@@ -279,8 +279,8 @@ async def set_graph_active_version(
|
||||
)
|
||||
|
||||
# Keep the library agent up to date with the new active version
|
||||
await library_db.update_library_agent_version_and_settings(
|
||||
user_id, new_active_graph
|
||||
await library_db.update_agent_version_in_library(
|
||||
user_id, new_active_graph.id, new_active_graph.version
|
||||
)
|
||||
|
||||
if current_active_graph and current_active_graph.version != new_active_version:
|
||||
|
||||
@@ -13,7 +13,7 @@ from pytest_snapshot.plugin import Snapshot
|
||||
|
||||
from backend.api.features.graphs.routes import router
|
||||
from backend.api.rest_api import app as real_app
|
||||
from backend.api.rest_api import handle_internal_http_error
|
||||
from backend.api.utils.exceptions import add_exception_handlers
|
||||
from backend.blocks.basic import StoreValueBlock
|
||||
from backend.blocks.generic_webhook.triggers import GenericWebhookTriggerBlock
|
||||
from backend.blocks.google.sheets import GoogleSheetsReadBlock
|
||||
@@ -23,14 +23,13 @@ from backend.data import execution as execution_db
|
||||
from backend.data.execution import ExecutionStatus
|
||||
from backend.data.graph import GraphModel, NodeModel
|
||||
from backend.data.onboarding import OnboardingStep
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import GraphActivationError
|
||||
from backend.util.exceptions import GraphValidationError
|
||||
|
||||
app = fastapi.FastAPI()
|
||||
app.include_router(router)
|
||||
# Mirror rest_api.py's GraphActivationError -> 400 mapping so the atomicity
|
||||
# tests verify the same behaviour the real app exposes.
|
||||
app.add_exception_handler(GraphActivationError, handle_internal_http_error(400))
|
||||
# Same handler set as the real app (GraphActivationError -> 400 included), so
|
||||
# the atomicity tests verify the behaviour it actually exposes.
|
||||
add_exception_handlers(app)
|
||||
client = fastapi.testclient.TestClient(app)
|
||||
|
||||
|
||||
@@ -229,7 +228,7 @@ def test_get_graphs(
|
||||
|
||||
mocker.patch(
|
||||
"backend.data.graph.list_graphs_paginated",
|
||||
return_value=Mock(graphs=[mock_graph]),
|
||||
return_value=([mock_graph], Mock()),
|
||||
)
|
||||
|
||||
response = client.get("/graphs")
|
||||
@@ -571,7 +570,7 @@ def test_update_graph_returns_400_and_persists_nothing_on_activation_error(
|
||||
"backend.api.features.graphs.routes.graph_db.create_graph", new=AsyncMock()
|
||||
)
|
||||
update_lib_agent_mock = mocker.patch(
|
||||
"backend.api.features.graphs.routes.library_db.update_library_agent_version_and_settings",
|
||||
"backend.api.features.graphs.routes.library_db.update_agent_version_in_library",
|
||||
new=AsyncMock(),
|
||||
)
|
||||
|
||||
|
||||
@@ -1179,15 +1179,43 @@ async def delete_credentials(
|
||||
)
|
||||
|
||||
try:
|
||||
await remove_all_webhooks_for_credentials(user_id, creds, force)
|
||||
tokens_revoked = await delete_credentials_and_dependents(
|
||||
request, user_id, creds, provider, force
|
||||
)
|
||||
except NeedConfirmation as e:
|
||||
return CredentialsDeletionNeedsConfirmationResponse(message=str(e))
|
||||
|
||||
return CredentialsDeletionResponse(revoked=tokens_revoked)
|
||||
|
||||
|
||||
async def delete_credentials_and_dependents(
|
||||
request: Request,
|
||||
user_id: str,
|
||||
creds: Credentials,
|
||||
provider: ProviderName,
|
||||
force: bool,
|
||||
) -> bool | None:
|
||||
"""Delete a user's credentials with what hangs off them.
|
||||
|
||||
Removes the webhooks registered with them first, then the credentials,
|
||||
then revokes OAuth tokens at the provider. The External API's delete
|
||||
calls this too, so neither surface can leave a live webhook or token
|
||||
behind. Returns whether the provider revoked the tokens (None when there
|
||||
were none to revoke).
|
||||
|
||||
Raises:
|
||||
NeedConfirmation: a webhook is still in use and ``force`` is False;
|
||||
nothing has been deleted.
|
||||
"""
|
||||
await remove_all_webhooks_for_credentials(user_id, creds, force)
|
||||
|
||||
tokens_revoked = None
|
||||
if provider == ProviderName.CODEX:
|
||||
tokens_revoked = await revoke_codex_credentials(creds_manager, user_id, cred_id)
|
||||
tokens_revoked = await revoke_codex_credentials(
|
||||
creds_manager, user_id, creds.id
|
||||
)
|
||||
else:
|
||||
await creds_manager.delete(user_id, cred_id)
|
||||
await creds_manager.delete(user_id, creds.id)
|
||||
|
||||
if isinstance(creds, OAuth2Credentials) and provider != ProviderName.CODEX:
|
||||
if provider_matches(provider.value, ProviderName.MCP.value):
|
||||
@@ -1206,7 +1234,7 @@ async def delete_credentials(
|
||||
handler = _get_provider_oauth_handler(request, provider)
|
||||
tokens_revoked = await handler.revoke_tokens(creds)
|
||||
|
||||
return CredentialsDeletionResponse(revoked=tokens_revoked)
|
||||
return tokens_revoked
|
||||
|
||||
|
||||
# ------------------------- WEBHOOK STUFF -------------------------- #
|
||||
|
||||
@@ -22,6 +22,7 @@ from backend.util.exceptions import NotFoundError
|
||||
from backend.util.json import SafeJson
|
||||
|
||||
from ._schedule_info import _fetch_schedule_info
|
||||
from .exceptions import LibraryAgentInAnotherOrganizationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -144,8 +145,13 @@ def _marketplace_metadata(
|
||||
async def restore_existing_library_agent(
|
||||
store_listing_version: prisma.models.StoreListingVersion,
|
||||
user_id: str,
|
||||
organization_id: str | None = None,
|
||||
) -> library_model.LibraryAgent | None:
|
||||
"""Restore and return an existing library entry without loading its graph."""
|
||||
"""Restore and return an existing library entry without loading its graph.
|
||||
|
||||
With ``organization_id``, an entry tagged with another organization is
|
||||
left alone and ``LibraryAgentInAnotherOrganizationError`` raised.
|
||||
"""
|
||||
ag = _require_graph(store_listing_version)
|
||||
|
||||
client = prisma.models.LibraryAgent.prisma()
|
||||
@@ -159,6 +165,13 @@ async def restore_existing_library_agent(
|
||||
existing = await client.find_unique(where=where)
|
||||
if existing is None:
|
||||
return None
|
||||
if organization_id is not None and existing.organizationId not in (
|
||||
None,
|
||||
organization_id,
|
||||
):
|
||||
raise LibraryAgentInAnotherOrganizationError(
|
||||
"This agent is already in your library in another organization"
|
||||
)
|
||||
|
||||
restored = await client.update(
|
||||
where=where,
|
||||
@@ -187,6 +200,8 @@ async def add_graph_to_library(
|
||||
store_listing_version: prisma.models.StoreListingVersion,
|
||||
*,
|
||||
tx: prisma.Prisma | None = None,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
"""Check existing / restore soft-deleted / create new LibraryAgent.
|
||||
|
||||
@@ -195,13 +210,13 @@ async def add_graph_to_library(
|
||||
atomic with the caller's other writes; catching a uniqueness error inside
|
||||
PostgreSQL's transaction would leave that transaction aborted.
|
||||
"""
|
||||
settings_json = SafeJson(GraphSettings.from_graph(graph_model).model_dump())
|
||||
settings_json = SafeJson(GraphSettings().model_dump())
|
||||
_include = library_agent_include(
|
||||
user_id, include_nodes=False, include_executions=False
|
||||
)
|
||||
marketplace = _marketplace_metadata(store_listing_version)
|
||||
create_data, update_data = await _library_agent_payloads(
|
||||
graph_model, user_id, settings_json, marketplace
|
||||
graph_model, user_id, settings_json, marketplace, organization_id, team_id
|
||||
)
|
||||
|
||||
if tx is not None:
|
||||
@@ -230,8 +245,11 @@ async def _library_agent_payloads(
|
||||
user_id: str,
|
||||
settings_json: SafeJson,
|
||||
marketplace: dict[str, str | None],
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> tuple[dict, dict]:
|
||||
organization_id, team_id = await resolve_default_tenancy(user_id)
|
||||
if organization_id is None:
|
||||
organization_id, team_id = await resolve_default_tenancy(user_id)
|
||||
create_data = {
|
||||
"User": {"connect": {"id": user_id}},
|
||||
"AgentGraph": {
|
||||
|
||||
@@ -9,7 +9,7 @@ import backend.data.graph as graph_db
|
||||
from backend.data.includes import library_agent_include
|
||||
|
||||
from ._schedule_info import _fetch_schedule_info
|
||||
from .db import _fetch_execution_counts, _fetch_marketplace_details
|
||||
from .db import _fetch_execution_counts, _fetch_store_listing
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -60,7 +60,6 @@ async def get_library_agents_by_ids(
|
||||
detail.row,
|
||||
sub_graphs=detail.sub_graphs,
|
||||
store_listing=detail.store_listing,
|
||||
profile=detail.profile,
|
||||
schedule_info=schedule_info,
|
||||
execution_count_override=execution_counts.get(detail.row.agentGraphId),
|
||||
)
|
||||
@@ -72,7 +71,6 @@ async def get_library_agents_by_ids(
|
||||
class _AgentDetails(BaseModel):
|
||||
row: prisma.models.LibraryAgent
|
||||
store_listing: prisma.models.StoreListing | None
|
||||
profile: prisma.models.Profile | None
|
||||
sub_graphs: list[prisma.models.AgentGraph]
|
||||
|
||||
|
||||
@@ -83,13 +81,11 @@ async def _load_details(
|
||||
return None
|
||||
try:
|
||||
async with limit:
|
||||
(store_listing, profile), sub_graphs = await asyncio.gather(
|
||||
_fetch_marketplace_details(row.AgentGraph.id),
|
||||
store_listing, sub_graphs = await asyncio.gather(
|
||||
_fetch_store_listing(row.AgentGraph.id),
|
||||
graph_db.get_sub_graphs(row.AgentGraph),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading LibraryAgent #{row.id}: {e}")
|
||||
return None
|
||||
return _AgentDetails(
|
||||
row=row, store_listing=store_listing, profile=profile, sub_graphs=sub_graphs
|
||||
)
|
||||
return _AgentDetails(row=row, store_listing=store_listing, sub_graphs=sub_graphs)
|
||||
@@ -4,6 +4,7 @@ import logging
|
||||
from typing import Literal, LiteralString, Optional, cast
|
||||
|
||||
import fastapi
|
||||
import prisma.enums
|
||||
import prisma.errors
|
||||
import prisma.models
|
||||
import prisma.types
|
||||
@@ -16,8 +17,12 @@ import backend.data.integrations as integrations_db
|
||||
from backend.api.features.library.exceptions import (
|
||||
FolderAlreadyExistsError,
|
||||
FolderValidationError,
|
||||
LibraryAgentInAnotherOrganizationError,
|
||||
)
|
||||
from backend.api.features.orgs.db import resolve_member_tenancy
|
||||
from backend.api.features.store.store_listing_versions import (
|
||||
installable_store_version_where,
|
||||
)
|
||||
from backend.data.db import get_database_schema, transaction
|
||||
from backend.data.execution import get_graph_execution
|
||||
from backend.data.expert_attribution import resolve_attributable_expert
|
||||
@@ -28,6 +33,7 @@ from backend.data.includes import (
|
||||
library_agent_include,
|
||||
)
|
||||
from backend.data.model import CredentialsMetaInput, GraphInput
|
||||
from backend.data.tenancy import organization_filter, writable_organization_ids
|
||||
from backend.integrations.creds_manager import IntegrationCredentialsManager
|
||||
from backend.integrations.webhooks.graph_lifecycle_hooks import (
|
||||
before_graph_activate,
|
||||
@@ -35,7 +41,12 @@ from backend.integrations.webhooks.graph_lifecycle_hooks import (
|
||||
on_graph_deactivate,
|
||||
)
|
||||
from backend.util.clients import get_scheduler_client
|
||||
from backend.util.exceptions import InvalidInputError, MissingConfigError, NotFoundError
|
||||
from backend.util.exceptions import (
|
||||
InvalidInputError,
|
||||
MissingConfigError,
|
||||
NotAuthorizedError,
|
||||
NotFoundError,
|
||||
)
|
||||
from backend.util.json import SafeJson
|
||||
from backend.util.models import Pagination
|
||||
from backend.util.product_analytics import track_listing_added_to_library
|
||||
@@ -75,28 +86,11 @@ async def _fetch_execution_counts(user_id: str, graph_ids: list[str]) -> dict[st
|
||||
}
|
||||
|
||||
|
||||
async def _fetch_marketplace_details(
|
||||
graph_id: str,
|
||||
) -> tuple[prisma.models.StoreListing | None, prisma.models.Profile | None]:
|
||||
store_listing = await prisma.models.StoreListing.prisma().find_first(
|
||||
where={
|
||||
"agentGraphId": graph_id,
|
||||
"isDeleted": False,
|
||||
"hasApprovedVersion": True,
|
||||
},
|
||||
include={"ActiveVersion": True},
|
||||
)
|
||||
profile = None
|
||||
if store_listing and store_listing.ActiveVersion and store_listing.owningUserId:
|
||||
profile = await prisma.models.Profile.prisma().find_first(
|
||||
where={"userId": store_listing.owningUserId}
|
||||
)
|
||||
return store_listing, profile
|
||||
|
||||
|
||||
async def list_library_agents(
|
||||
user_id: str,
|
||||
search_term: Optional[str] = None,
|
||||
published: Optional[bool] = None,
|
||||
favorite: Optional[bool] = None,
|
||||
sort_by: library_model.LibraryAgentSort = library_model.LibraryAgentSort.UPDATED_AT,
|
||||
page: int = 1,
|
||||
page_size: int = 50,
|
||||
@@ -113,6 +107,8 @@ async def list_library_agents(
|
||||
Args:
|
||||
user_id: The ID of the user whose LibraryAgents we want to retrieve.
|
||||
search_term: Optional string to filter agents by name/description.
|
||||
published: Allows filtering by marketplace publish status;
|
||||
`True` -> only published agents, `False` -> only unpublished agents.
|
||||
sort_by: Sorting field (createdAt, updatedAt, isFavorite, isCreatedByUser).
|
||||
page: Current page (1-indexed).
|
||||
page_size: Number of items per page.
|
||||
@@ -199,6 +195,18 @@ async def list_library_agents(
|
||||
},
|
||||
]
|
||||
|
||||
# Filter by favorite status
|
||||
if favorite is not None:
|
||||
where_clause["isFavorite"] = favorite
|
||||
|
||||
# By graph, as a single read reports it: the entry follows every save of a
|
||||
# listed graph, so its own version may never have been submitted.
|
||||
if published is not None:
|
||||
listed = await _listed_graph_ids(where_clause)
|
||||
where_clause["agentGraphId"] = (
|
||||
{"in": listed} if published else {"not_in": listed}
|
||||
)
|
||||
|
||||
order_by: (
|
||||
prisma.types.LibraryAgentOrderByInput
|
||||
| list[prisma.types.LibraryAgentOrderByInput]
|
||||
@@ -206,7 +214,8 @@ async def list_library_agents(
|
||||
) = None
|
||||
|
||||
if sort_by == library_model.LibraryAgentSort.CREATED_AT:
|
||||
order_by = {"createdAt": "asc"}
|
||||
# The id breaks ties, so paging by offset neither skips nor repeats.
|
||||
order_by = [{"createdAt": "asc"}, {"id": "asc"}]
|
||||
elif sort_by == library_model.LibraryAgentSort.UPDATED_AT:
|
||||
order_by = {"updatedAt": "desc"}
|
||||
elif sort_by == library_model.LibraryAgentSort.LAST_RUN:
|
||||
@@ -270,6 +279,21 @@ async def list_library_agents(
|
||||
)
|
||||
|
||||
|
||||
async def _listed_graph_ids(where: prisma.types.LibraryAgentWhereInput) -> list[str]:
|
||||
"""The graphs among these library entries with a live, approved listing."""
|
||||
entries = await prisma.models.LibraryAgent.prisma().find_many(
|
||||
where=where, distinct=["agentGraphId"]
|
||||
)
|
||||
listings = await prisma.models.StoreListing.prisma().find_many(
|
||||
where={
|
||||
"agentGraphId": {"in": [entry.agentGraphId for entry in entries]},
|
||||
"isDeleted": False,
|
||||
"hasApprovedVersion": True,
|
||||
}
|
||||
)
|
||||
return [listing.agentGraphId for listing in listings]
|
||||
|
||||
|
||||
async def list_favorite_library_agents(
|
||||
user_id: str,
|
||||
page: int = 1,
|
||||
@@ -387,26 +411,33 @@ async def get_library_agent(id: str, user_id: str) -> library_model.LibraryAgent
|
||||
if not library_agent.AgentGraph:
|
||||
raise NotFoundError(f"Agent graph for library agent #{id} not found")
|
||||
|
||||
(
|
||||
marketplace_details,
|
||||
schedule_info,
|
||||
sub_graphs,
|
||||
) = await asyncio.gather(
|
||||
_fetch_marketplace_details(library_agent.AgentGraph.id),
|
||||
store_listing, schedule_info, sub_graphs = await asyncio.gather(
|
||||
_fetch_store_listing(library_agent.AgentGraph.id),
|
||||
_fetch_schedule_info(user_id, graph_id=library_agent.AgentGraph.id),
|
||||
graph_db.get_sub_graphs(library_agent.AgentGraph),
|
||||
)
|
||||
store_listing, profile = marketplace_details
|
||||
|
||||
return library_model.LibraryAgent.from_db(
|
||||
library_agent,
|
||||
sub_graphs=sub_graphs,
|
||||
store_listing=store_listing,
|
||||
profile=profile,
|
||||
schedule_info=schedule_info,
|
||||
)
|
||||
|
||||
|
||||
async def _fetch_store_listing(graph_id: str) -> prisma.models.StoreListing | None:
|
||||
# Keyed on the graph, not its version: the owner's library agent follows
|
||||
# every save, and a later version is still the listed agent.
|
||||
return await prisma.models.StoreListing.prisma().find_first(
|
||||
where={
|
||||
"agentGraphId": graph_id,
|
||||
"isDeleted": False,
|
||||
"hasApprovedVersion": True,
|
||||
},
|
||||
include={"ActiveVersion": True, "CreatorProfile": True},
|
||||
)
|
||||
|
||||
|
||||
async def get_library_agent_by_store_version_id(
|
||||
store_listing_version_id: str,
|
||||
user_id: str,
|
||||
@@ -445,6 +476,21 @@ async def get_library_agent_by_store_version_id(
|
||||
return library_model.LibraryAgent.from_db(agent, schedule_info=schedule_info)
|
||||
|
||||
|
||||
async def get_library_agent_organizations(
|
||||
user_id: str, library_agent_ids: list[str]
|
||||
) -> dict[str, str | None]:
|
||||
"""The organization of each of the user's live library agents among these ids.
|
||||
|
||||
Deleted entries and ids that aren't the user's are left out.
|
||||
"""
|
||||
if not library_agent_ids:
|
||||
return {}
|
||||
rows = await prisma.models.LibraryAgent.prisma().find_many(
|
||||
where={"id": {"in": library_agent_ids}, "userId": user_id, "isDeleted": False}
|
||||
)
|
||||
return {row.id: row.organizationId for row in rows}
|
||||
|
||||
|
||||
async def get_library_agent_id_by_graph_id(user_id: str, graph_id: str) -> str | None:
|
||||
"""Id-only lookup for building deep links (e.g. the expert run post) —
|
||||
no relation includes, no schedule info, unlike the full getter below."""
|
||||
@@ -653,9 +699,8 @@ async def create_library_agent(
|
||||
}
|
||||
},
|
||||
settings=SafeJson(
|
||||
GraphSettings.from_graph(
|
||||
graph_entry,
|
||||
hitl_safe_mode=hitl_safe_mode,
|
||||
GraphSettings(
|
||||
human_in_the_loop_safe_mode=hitl_safe_mode,
|
||||
sensitive_action_safe_mode=sensitive_action_safe_mode,
|
||||
).model_dump()
|
||||
),
|
||||
@@ -687,9 +732,8 @@ async def create_library_agent(
|
||||
else {}
|
||||
),
|
||||
"settings": SafeJson(
|
||||
GraphSettings.from_graph(
|
||||
graph_entry,
|
||||
hitl_safe_mode=hitl_safe_mode,
|
||||
GraphSettings(
|
||||
human_in_the_loop_safe_mode=hitl_safe_mode,
|
||||
sensitive_action_safe_mode=sensitive_action_safe_mode,
|
||||
).model_dump()
|
||||
),
|
||||
@@ -726,6 +770,9 @@ async def update_agent_version_in_library(
|
||||
user_id: str,
|
||||
agent_graph_id: str,
|
||||
agent_graph_version: int,
|
||||
*,
|
||||
library_agent_id: Optional[str] = None,
|
||||
organization_ids: Optional[frozenset[str | None]] = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
"""
|
||||
Updates the agent version in the library for any agent owned by the user.
|
||||
@@ -734,25 +781,52 @@ async def update_agent_version_in_library(
|
||||
user_id: Owner of the LibraryAgent.
|
||||
agent_graph_id: The agent graph's ID to update.
|
||||
agent_graph_version: The new version of the agent graph.
|
||||
library_agent_id: The entry to move; without it, the user's first
|
||||
entry for the graph.
|
||||
organization_ids: Confine the change to entries tagged with these
|
||||
organizations (None for untagged ones), so that a call acting in
|
||||
one organization moves only its own entry and removes none of
|
||||
another's, and onto none of the user's own versions outside them.
|
||||
|
||||
Raises:
|
||||
DatabaseError: If there's an error with the update.
|
||||
NotFoundError: If no library agent is found for this user and agent.
|
||||
NotFoundError: If no library agent is found for this user and agent,
|
||||
or the target version doesn't exist or is another user's that
|
||||
isn't approved in the marketplace.
|
||||
LibraryAgentInAnotherOrganizationError: If the target version is in
|
||||
the user's library in an organization outside ``organization_ids``.
|
||||
NotAuthorizedError: If the target version is one of the user's own,
|
||||
outside ``organization_ids``.
|
||||
"""
|
||||
logger.debug(
|
||||
f"Updating agent version in library for user #{user_id}, "
|
||||
f"agent #{agent_graph_id} v{agent_graph_version}"
|
||||
)
|
||||
entry: prisma.types.LibraryAgentWhereInput = {
|
||||
"userId": user_id,
|
||||
"agentGraphId": agent_graph_id,
|
||||
}
|
||||
if library_agent_id is not None:
|
||||
entry["id"] = library_agent_id
|
||||
if organization_ids is not None:
|
||||
entry["AND"] = [
|
||||
cast(
|
||||
prisma.types.LibraryAgentWhereInput,
|
||||
organization_filter(organization_ids),
|
||||
)
|
||||
]
|
||||
async with transaction() as tx:
|
||||
await _assert_target_version(
|
||||
tx, user_id, agent_graph_id, agent_graph_version, organization_ids
|
||||
)
|
||||
library_agent = await prisma.models.LibraryAgent.prisma(tx).find_first_or_raise(
|
||||
where={
|
||||
"userId": user_id,
|
||||
"agentGraphId": agent_graph_id,
|
||||
},
|
||||
where=entry
|
||||
)
|
||||
|
||||
# Delete any conflicting LibraryAgent for the target version
|
||||
await prisma.models.LibraryAgent.prisma(tx).delete_many(
|
||||
# An entry is unique per user, graph and version, so one already on
|
||||
# the target version gives way; another organization's isn't this
|
||||
# call's to remove.
|
||||
conflicts = await prisma.models.LibraryAgent.prisma(tx).find_many(
|
||||
where={
|
||||
"userId": user_id,
|
||||
"agentGraphId": agent_graph_id,
|
||||
@@ -760,6 +834,17 @@ async def update_agent_version_in_library(
|
||||
"id": {"not": library_agent.id},
|
||||
}
|
||||
)
|
||||
if organization_ids is not None and any(
|
||||
conflict.organizationId not in organization_ids for conflict in conflicts
|
||||
):
|
||||
raise LibraryAgentInAnotherOrganizationError(
|
||||
f"Version {agent_graph_version} of this agent is already in your "
|
||||
"library in another organization"
|
||||
)
|
||||
if conflicts:
|
||||
await prisma.models.LibraryAgent.prisma(tx).delete_many(
|
||||
where={"id": {"in": [conflict.id for conflict in conflicts]}}
|
||||
)
|
||||
|
||||
lib = await prisma.models.LibraryAgent.prisma(tx).update(
|
||||
where={"id": library_agent.id},
|
||||
@@ -773,7 +858,7 @@ async def update_agent_version_in_library(
|
||||
},
|
||||
},
|
||||
},
|
||||
include={"AgentGraph": True},
|
||||
include=library_agent_include(user_id),
|
||||
)
|
||||
|
||||
if lib is None:
|
||||
@@ -781,8 +866,20 @@ async def update_agent_version_in_library(
|
||||
f"Failed to update library agent for {agent_graph_id} v{agent_graph_version}"
|
||||
)
|
||||
|
||||
sub_graphs = (
|
||||
await graph_db.get_sub_graphs(lib.AgentGraph) if lib.AgentGraph else None
|
||||
)
|
||||
schedule_info = await _fetch_schedule_info(user_id, graph_id=agent_graph_id)
|
||||
return library_model.LibraryAgent.from_db(lib, schedule_info=schedule_info)
|
||||
library_agent = library_model.LibraryAgent.from_db(
|
||||
lib, sub_graphs=sub_graphs, schedule_info=schedule_info
|
||||
)
|
||||
|
||||
# Re-embed so name/description/instructions changes are reflected in
|
||||
# similarity search results before the next create-time gate runs.
|
||||
assert lib.AgentGraph is not None # for type checker
|
||||
schedule_library_agent_embedding(library_agent.id, user_id, lib.AgentGraph)
|
||||
|
||||
return library_agent
|
||||
|
||||
|
||||
async def create_graph_in_library(
|
||||
@@ -790,8 +887,14 @@ async def create_graph_in_library(
|
||||
user_id: str,
|
||||
folder_id: str | None = None,
|
||||
is_hidden: bool = False,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> tuple[graph_db.GraphModel, library_model.LibraryAgent]:
|
||||
"""Create a new graph and add it to the user's library."""
|
||||
"""Create a new graph and add it to the user's library.
|
||||
|
||||
Without ``organization_id`` the graph is untagged and the library entry
|
||||
goes to the user's default team.
|
||||
"""
|
||||
graph.version = 1
|
||||
graph_model = graph_db.make_graph_model(graph, user_id)
|
||||
graph_model.reassign_ids(user_id=user_id, reassign_graph_id=True)
|
||||
@@ -805,7 +908,9 @@ async def create_graph_in_library(
|
||||
else:
|
||||
await clear_unowned_auto_credentials(graph_model, user_id)
|
||||
|
||||
created_graph = await graph_db.create_graph(graph_model, user_id)
|
||||
created_graph = await graph_db.create_graph(
|
||||
graph_model, user_id, organization_id=organization_id, team_id=team_id
|
||||
)
|
||||
|
||||
library_agents = await create_library_agent(
|
||||
graph=created_graph,
|
||||
@@ -814,6 +919,8 @@ async def create_graph_in_library(
|
||||
create_library_agents_for_sub_graphs=False,
|
||||
folder_id=folder_id,
|
||||
is_hidden=is_hidden,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
return created_graph, library_agents[0]
|
||||
@@ -822,8 +929,28 @@ async def create_graph_in_library(
|
||||
async def update_graph_in_library(
|
||||
graph: graph_db.Graph,
|
||||
user_id: str,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
move_webhook_presets: bool = True,
|
||||
) -> tuple[graph_db.GraphModel, library_model.LibraryAgent]:
|
||||
"""Create a new version of an existing graph and update the library entry."""
|
||||
"""Create a new version of an existing graph and update the library entry.
|
||||
|
||||
``organization_id``/``team_id`` tag the new version; None leaves it untagged.
|
||||
A call acting in ``organization_id`` may only version a graph whose
|
||||
versions are all the user's, within the organizations it may change
|
||||
(``writable_organization_ids``), and moves only their library entries and
|
||||
webhook presets.
|
||||
``move_webhook_presets`` False leaves the user's webhook presets on the
|
||||
versions they run now.
|
||||
"""
|
||||
organization_ids = (
|
||||
await writable_organization_ids(user_id, organization_id)
|
||||
if organization_id is not None
|
||||
else None
|
||||
)
|
||||
if organization_ids is not None:
|
||||
await graph_db.assert_graph_within(graph.id, user_id, organization_ids)
|
||||
|
||||
existing_versions = await graph_db.get_graph_all_versions(graph.id, user_id)
|
||||
current_active_version = (
|
||||
next((v for v in existing_versions if v.is_active), None)
|
||||
@@ -844,7 +971,9 @@ async def update_graph_in_library(
|
||||
else:
|
||||
await clear_unowned_auto_credentials(graph_model, user_id)
|
||||
|
||||
created_graph = await graph_db.create_graph(graph_model, user_id)
|
||||
created_graph = await graph_db.create_graph(
|
||||
graph_model, user_id, organization_id=organization_id, team_id=team_id
|
||||
)
|
||||
|
||||
library_agent = await get_library_agent_by_graph_id(
|
||||
user_id, created_graph.id, include_archived=True
|
||||
@@ -852,8 +981,11 @@ async def update_graph_in_library(
|
||||
if not library_agent:
|
||||
raise NotFoundError(f"Library agent not found for graph {created_graph.id}")
|
||||
|
||||
library_agent = await update_library_agent_version_and_settings(
|
||||
user_id, created_graph
|
||||
library_agent = await update_agent_version_in_library(
|
||||
user_id,
|
||||
created_graph.id,
|
||||
created_graph.version,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
|
||||
if created_graph.is_active:
|
||||
@@ -872,38 +1004,55 @@ async def update_graph_in_library(
|
||||
# so the migration result is intentionally discarded here. Skipped
|
||||
# presets are surfaced on the interactive graph-activation endpoints
|
||||
# (update_graph / set_graph_active_version) instead.
|
||||
if created_graph.webhook_input_node:
|
||||
if move_webhook_presets and created_graph.webhook_input_node:
|
||||
await migrate_webhook_presets_to_new_version(
|
||||
user_id=user_id,
|
||||
new_graph=created_graph,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
|
||||
return created_graph, library_agent
|
||||
|
||||
|
||||
async def update_library_agent_version_and_settings(
|
||||
user_id: str, agent_graph: graph_db.GraphModel
|
||||
) -> library_model.LibraryAgent:
|
||||
"""Update library agent to point to new graph version and sync settings."""
|
||||
library = await update_agent_version_in_library(
|
||||
user_id, agent_graph.id, agent_graph.version
|
||||
async def _assert_target_version(
|
||||
tx: prisma.Prisma,
|
||||
user_id: str,
|
||||
graph_id: str,
|
||||
version: int,
|
||||
organization_ids: Optional[frozenset[str | None]],
|
||||
) -> None:
|
||||
"""Refuse a graph version the user may not point a library entry at.
|
||||
|
||||
Any of their own, within *organization_ids* when given: running the entry
|
||||
runs that version, billed where the entry is. Of another user's, a
|
||||
marketplace agent's, only one that adding it to the library would install
|
||||
(the same predicate): a version in the library that was ever submitted can
|
||||
be read, run and forked, so pointing at a pending or rejected one would
|
||||
unlock it. A refused version of another user's answers as a missing one
|
||||
does, so neither shows which versions exist.
|
||||
"""
|
||||
not_found = NotFoundError(f"Graph #{graph_id} v{version} not found")
|
||||
target = await prisma.models.AgentGraph.prisma(tx).find_unique(
|
||||
where={"graphVersionId": {"id": graph_id, "version": version}}
|
||||
)
|
||||
updated_settings = GraphSettings.from_graph(
|
||||
graph=agent_graph,
|
||||
hitl_safe_mode=library.settings.human_in_the_loop_safe_mode,
|
||||
sensitive_action_safe_mode=library.settings.sensitive_action_safe_mode,
|
||||
builder_chat_session_id=library.settings.builder_chat_session_id,
|
||||
)
|
||||
if updated_settings != library.settings:
|
||||
library = await update_library_agent(
|
||||
library_agent_id=library.id,
|
||||
user_id=user_id,
|
||||
settings=updated_settings,
|
||||
if target is None:
|
||||
raise not_found
|
||||
if target.userId != user_id:
|
||||
installable = await prisma.models.StoreListingVersion.prisma(tx).find_first(
|
||||
where={
|
||||
"agentGraphId": graph_id,
|
||||
"agentGraphVersion": version,
|
||||
**installable_store_version_where(),
|
||||
}
|
||||
)
|
||||
if installable is None:
|
||||
raise not_found
|
||||
return
|
||||
if organization_ids is not None and target.organizationId not in organization_ids:
|
||||
raise NotAuthorizedError(
|
||||
f"Version {version} of this agent is in another organization; "
|
||||
"choose it with a credential for that organization"
|
||||
)
|
||||
# Re-embed so name/description/instructions changes are reflected in
|
||||
# similarity search results before the next create-time gate runs.
|
||||
schedule_library_agent_embedding(library.id, user_id, agent_graph)
|
||||
return library
|
||||
|
||||
|
||||
async def update_library_agent(
|
||||
@@ -917,6 +1066,7 @@ async def update_library_agent(
|
||||
is_deleted: Optional[Literal[False]] = None,
|
||||
settings: Optional[GraphSettings] = None,
|
||||
folder_id: Optional[str] = None,
|
||||
organization_ids: Optional[frozenset[str | None]] = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
"""
|
||||
Updates the specified LibraryAgent record.
|
||||
@@ -930,6 +1080,8 @@ async def update_library_agent(
|
||||
is_archived: Whether this agent is archived.
|
||||
settings: User-specific settings for this library agent.
|
||||
folder_id: Folder ID to move agent to (None to skip).
|
||||
organization_ids: With ``graph_version``, the organizations whose
|
||||
entries the change may touch; see ``update_agent_version_in_library``.
|
||||
|
||||
Returns:
|
||||
The updated LibraryAgent.
|
||||
@@ -987,6 +1139,8 @@ async def update_library_agent(
|
||||
user_id=user_id,
|
||||
agent_graph_id=agent.graph_id,
|
||||
agent_graph_version=graph_version,
|
||||
library_agent_id=library_agent_id,
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
|
||||
# Otherwise, just update the simple fields
|
||||
@@ -1225,14 +1379,20 @@ async def add_store_agent_to_library(
|
||||
user_id: str,
|
||||
*,
|
||||
track_listing_added: bool = True,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
"""Adds a marketplace agent to the user’s library.
|
||||
"""Adds a marketplace agent to the user's library.
|
||||
|
||||
``track_listing_added`` sends ``listing_added_to_library`` on a first-time
|
||||
add. Pass ``False`` when the system installs the agent on the user's
|
||||
behalf (an expert's preloads or workflows), which is not the user adding
|
||||
a listing.
|
||||
|
||||
With ``organization_id`` a new entry is tagged with it (and ``team_id``),
|
||||
and an existing entry in another organization raises
|
||||
``LibraryAgentInAnotherOrganizationError`` rather than being restored.
|
||||
|
||||
See also: `add_store_agent_to_library_as_admin()` which uses
|
||||
`get_graph_as_admin` to bypass marketplace status checks for admin review.
|
||||
"""
|
||||
@@ -1241,6 +1401,8 @@ async def add_store_agent_to_library(
|
||||
user_id,
|
||||
tx=None,
|
||||
track_listing_added=track_listing_added,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
@@ -1265,6 +1427,8 @@ async def _add_store_agent_to_library(
|
||||
*,
|
||||
tx: prisma.Prisma | None,
|
||||
track_listing_added: bool,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
logger.debug(
|
||||
"Adding agent from store listing version #%s to library for user #%s",
|
||||
@@ -1277,14 +1441,21 @@ async def _add_store_agent_to_library(
|
||||
if tx is None:
|
||||
# The transactional path upserts instead, so an existing entry is
|
||||
# restored atomically with the caller's other writes.
|
||||
existing = await restore_existing_library_agent(store_listing_version, user_id)
|
||||
existing = await restore_existing_library_agent(
|
||||
store_listing_version, user_id, organization_id=organization_id
|
||||
)
|
||||
if existing is not None:
|
||||
return existing
|
||||
graph_model = await resolve_graph_model_for_library(
|
||||
store_listing_version, user_id, admin=False
|
||||
)
|
||||
library_agent = await add_graph_to_library(
|
||||
graph_model, user_id, store_listing_version, tx=tx
|
||||
graph_model,
|
||||
user_id,
|
||||
store_listing_version,
|
||||
tx=tx,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
if track_listing_added:
|
||||
track_listing_added_to_library(
|
||||
@@ -1322,41 +1493,11 @@ async def add_store_agent_to_library_as_admin(
|
||||
##############################################
|
||||
|
||||
|
||||
async def _fetch_user_folders(
|
||||
user_id: str,
|
||||
extra_where: Optional[prisma.types.LibraryFolderWhereInput] = None,
|
||||
include_relations: bool = True,
|
||||
) -> list[prisma.models.LibraryFolder]:
|
||||
"""
|
||||
Shared helper to fetch folders for a user with consistent query params.
|
||||
|
||||
Args:
|
||||
user_id: The ID of the user.
|
||||
extra_where: Additional where-clause filters to merge in.
|
||||
include_relations: Whether to include LibraryAgents and Children relations
|
||||
(used to derive counts via len(); Prisma Python has no _count include).
|
||||
|
||||
Returns:
|
||||
A list of raw Prisma LibraryFolder records.
|
||||
"""
|
||||
where_clause: prisma.types.LibraryFolderWhereInput = {
|
||||
"userId": user_id,
|
||||
"isDeleted": False,
|
||||
}
|
||||
if extra_where:
|
||||
where_clause.update(extra_where)
|
||||
|
||||
return await prisma.models.LibraryFolder.prisma().find_many(
|
||||
where=where_clause,
|
||||
order={"createdAt": "asc"},
|
||||
include=LIBRARY_FOLDER_INCLUDE if include_relations else None,
|
||||
)
|
||||
|
||||
|
||||
async def list_folders(
|
||||
user_id: str,
|
||||
parent_id: Optional[str] = None,
|
||||
include_relations: bool = True,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> list[library_model.LibraryFolder]:
|
||||
"""
|
||||
Lists folders for a user, optionally filtered by parent.
|
||||
@@ -1376,6 +1517,7 @@ async def list_folders(
|
||||
user_id,
|
||||
extra_where={"parentId": parent_id},
|
||||
include_relations=include_relations,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
|
||||
return [
|
||||
@@ -1390,6 +1532,7 @@ async def list_folders(
|
||||
|
||||
async def get_folder_tree(
|
||||
user_id: str,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> list[library_model.LibraryFolderTree]:
|
||||
"""
|
||||
Gets the full folder tree for a user.
|
||||
@@ -1403,7 +1546,7 @@ async def get_folder_tree(
|
||||
logger.debug(f"Getting folder tree for user #{user_id}")
|
||||
|
||||
# Fetch all folders for the user
|
||||
all_folders = await _fetch_user_folders(user_id)
|
||||
all_folders = await _fetch_user_folders(user_id, organization_id=organization_id)
|
||||
|
||||
# Build a map of folder ID to folder data
|
||||
folder_map: dict[str, library_model.LibraryFolderTree] = {
|
||||
@@ -1430,6 +1573,44 @@ async def get_folder_tree(
|
||||
return root_folders
|
||||
|
||||
|
||||
async def _fetch_user_folders(
|
||||
user_id: str,
|
||||
extra_where: Optional[prisma.types.LibraryFolderWhereInput] = None,
|
||||
include_relations: bool = True,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> list[prisma.models.LibraryFolder]:
|
||||
"""
|
||||
Shared helper to fetch folders for a user with consistent query params.
|
||||
|
||||
Args:
|
||||
user_id: The ID of the user.
|
||||
extra_where: Additional where-clause filters to merge in.
|
||||
include_relations: Whether to include LibraryAgents and Children relations
|
||||
(used to derive counts via len(); Prisma Python has no _count include).
|
||||
|
||||
Returns:
|
||||
A list of raw Prisma LibraryFolder records.
|
||||
"""
|
||||
where_clause: prisma.types.LibraryFolderWhereInput = {
|
||||
"userId": user_id,
|
||||
"isDeleted": False,
|
||||
}
|
||||
if extra_where:
|
||||
where_clause.update(extra_where)
|
||||
if organization_id is not None:
|
||||
# Orgless rows show under every org of their owner only (userId is always
|
||||
# in the where); create_folder still writes them, so they are not just legacy.
|
||||
where_clause["AND"] = [
|
||||
{"OR": [{"organizationId": organization_id}, {"organizationId": None}]}
|
||||
]
|
||||
|
||||
return await prisma.models.LibraryFolder.prisma().find_many(
|
||||
where=where_clause,
|
||||
order={"createdAt": "asc"},
|
||||
include=LIBRARY_FOLDER_INCLUDE if include_relations else None,
|
||||
)
|
||||
|
||||
|
||||
async def get_folder(
|
||||
folder_id: str,
|
||||
user_id: str,
|
||||
@@ -1466,49 +1647,14 @@ async def get_folder(
|
||||
)
|
||||
|
||||
|
||||
async def _is_descendant_of(
|
||||
folder_id: str,
|
||||
potential_ancestor_id: str,
|
||||
user_id: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if folder_id is a descendant of (or equal to) potential_ancestor_id.
|
||||
|
||||
Fetches all user folders in a single query and walks the parent chain
|
||||
in memory to avoid N database round-trips.
|
||||
|
||||
Args:
|
||||
folder_id: The ID of the folder to check.
|
||||
potential_ancestor_id: The ID of the potential ancestor.
|
||||
user_id: The ID of the user.
|
||||
|
||||
Returns:
|
||||
True if folder_id is a descendant of (or equal to) potential_ancestor_id.
|
||||
"""
|
||||
all_folders = await prisma.models.LibraryFolder.prisma().find_many(
|
||||
where={"userId": user_id, "isDeleted": False},
|
||||
)
|
||||
parent_map = {f.id: f.parentId for f in all_folders}
|
||||
|
||||
visited: set[str] = set()
|
||||
current_id: str | None = folder_id
|
||||
while current_id:
|
||||
if current_id == potential_ancestor_id:
|
||||
return True
|
||||
if current_id in visited:
|
||||
break # cycle detected
|
||||
visited.add(current_id)
|
||||
current_id = parent_map.get(current_id)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def create_folder(
|
||||
user_id: str,
|
||||
name: str,
|
||||
parent_id: Optional[str] = None,
|
||||
icon: Optional[str] = None,
|
||||
color: Optional[str] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
team_id: Optional[str] = None,
|
||||
) -> library_model.LibraryFolder:
|
||||
"""
|
||||
Creates a new folder for the user.
|
||||
@@ -1519,6 +1665,8 @@ async def create_folder(
|
||||
parent_id: Optional parent folder ID.
|
||||
icon: Optional icon identifier.
|
||||
color: Optional hex color code.
|
||||
organization_id, team_id: Tenancy to tag the folder with; without
|
||||
them it is untagged, which every organization of the user sees.
|
||||
|
||||
Returns:
|
||||
The created LibraryFolder.
|
||||
@@ -1545,6 +1693,10 @@ async def create_folder(
|
||||
create_data["color"] = color
|
||||
if parent_id:
|
||||
create_data["Parent"] = {"connect": {"id": parent_id}}
|
||||
if organization_id is not None:
|
||||
create_data["organizationId"] = organization_id
|
||||
if team_id is not None:
|
||||
create_data["Team"] = {"connect": {"id": team_id}}
|
||||
|
||||
try:
|
||||
folder = await prisma.models.LibraryFolder.prisma().create(data=create_data)
|
||||
@@ -1714,6 +1866,43 @@ async def move_folder(
|
||||
)
|
||||
|
||||
|
||||
async def _is_descendant_of(
|
||||
folder_id: str,
|
||||
potential_ancestor_id: str,
|
||||
user_id: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if folder_id is a descendant of (or equal to) potential_ancestor_id.
|
||||
|
||||
Fetches all user folders in a single query and walks the parent chain
|
||||
in memory to avoid N database round-trips.
|
||||
|
||||
Args:
|
||||
folder_id: The ID of the folder to check.
|
||||
potential_ancestor_id: The ID of the potential ancestor.
|
||||
user_id: The ID of the user.
|
||||
|
||||
Returns:
|
||||
True if folder_id is a descendant of (or equal to) potential_ancestor_id.
|
||||
"""
|
||||
all_folders = await prisma.models.LibraryFolder.prisma().find_many(
|
||||
where={"userId": user_id, "isDeleted": False},
|
||||
)
|
||||
parent_map = {f.id: f.parentId for f in all_folders}
|
||||
|
||||
visited: set[str] = set()
|
||||
current_id: str | None = folder_id
|
||||
while current_id:
|
||||
if current_id == potential_ancestor_id:
|
||||
return True
|
||||
if current_id in visited:
|
||||
break # cycle detected
|
||||
visited.add(current_id)
|
||||
current_id = parent_map.get(current_id)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def delete_folder(
|
||||
folder_id: str,
|
||||
user_id: str,
|
||||
@@ -1926,13 +2115,16 @@ def collect_tree_ids(
|
||||
|
||||
|
||||
async def get_folder_agent_summaries(
|
||||
user_id: str, folder_id: str
|
||||
user_id: str, folder_id: str, organization_id: Optional[str] = None
|
||||
) -> list[dict[str, str | None]]:
|
||||
"""Get a lightweight list of agents in a folder (id, name, description)."""
|
||||
all_agents: list[library_model.LibraryAgent] = []
|
||||
for page in itertools.count(1):
|
||||
resp = await list_library_agents(
|
||||
user_id=user_id, folder_id=folder_id, page=page
|
||||
user_id=user_id,
|
||||
folder_id=folder_id,
|
||||
page=page,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
all_agents.extend(resp.agents)
|
||||
if page >= resp.pagination.total_pages:
|
||||
@@ -1944,12 +2136,16 @@ async def get_folder_agent_summaries(
|
||||
|
||||
async def get_root_agent_summaries(
|
||||
user_id: str,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> list[dict[str, str | None]]:
|
||||
"""Get a lightweight list of root-level agents (folderId IS NULL)."""
|
||||
all_agents: list[library_model.LibraryAgent] = []
|
||||
for page in itertools.count(1):
|
||||
resp = await list_library_agents(
|
||||
user_id=user_id, include_root_only=True, page=page
|
||||
user_id=user_id,
|
||||
include_root_only=True,
|
||||
page=page,
|
||||
organization_id=organization_id,
|
||||
)
|
||||
all_agents.extend(resp.agents)
|
||||
if page >= resp.pagination.total_pages:
|
||||
@@ -1960,11 +2156,14 @@ async def get_root_agent_summaries(
|
||||
|
||||
|
||||
async def get_folder_agents_map(
|
||||
user_id: str, folder_ids: list[str]
|
||||
user_id: str, folder_ids: list[str], organization_id: Optional[str] = None
|
||||
) -> dict[str, list[dict[str, str | None]]]:
|
||||
"""Get agent summaries for multiple folders concurrently."""
|
||||
results = await asyncio.gather(
|
||||
*(get_folder_agent_summaries(user_id, fid) for fid in folder_ids)
|
||||
*(
|
||||
get_folder_agent_summaries(user_id, fid, organization_id)
|
||||
for fid in folder_ids
|
||||
)
|
||||
)
|
||||
return dict(zip(folder_ids, results))
|
||||
|
||||
@@ -1981,6 +2180,7 @@ async def list_presets(
|
||||
graph_id: Optional[str] = None,
|
||||
expert_id: str | None = None,
|
||||
filter_by_expert: bool = False,
|
||||
organization_id: Optional[str] = None,
|
||||
) -> library_model.LibraryAgentPresetResponse:
|
||||
"""
|
||||
Retrieves a paginated list of AgentPresets for the specified user.
|
||||
@@ -2018,6 +2218,11 @@ async def list_presets(
|
||||
query_filter["agentGraphId"] = graph_id
|
||||
if filter_by_expert:
|
||||
query_filter["expertId"] = expert_id
|
||||
if organization_id is not None:
|
||||
# Own rows in this org, plus rows created before org tagging.
|
||||
query_filter["AND"] = [
|
||||
{"OR": [{"organizationId": organization_id}, {"organizationId": None}]}
|
||||
]
|
||||
|
||||
presets_records = await prisma.models.AgentPreset.prisma().find_many(
|
||||
where=query_filter,
|
||||
@@ -2320,8 +2525,9 @@ async def update_preset(
|
||||
update_data["description"] = description
|
||||
if is_active is not None:
|
||||
update_data["isActive"] = is_active
|
||||
if inputs or credentials:
|
||||
if not (inputs and credentials):
|
||||
if inputs is not None or credentials is not None:
|
||||
# Either may be empty: no credentials needed, or every input cleared.
|
||||
if inputs is None or credentials is None:
|
||||
raise ValueError(
|
||||
"Preset inputs and credentials must be provided together"
|
||||
)
|
||||
@@ -2403,6 +2609,8 @@ async def set_preset_webhook(
|
||||
async def migrate_webhook_presets_to_new_version(
|
||||
user_id: str,
|
||||
new_graph: graph_db.GraphModel,
|
||||
*,
|
||||
organization_ids: Optional[frozenset[str | None]] = None,
|
||||
) -> library_model.WebhookPresetMigrationResult:
|
||||
"""
|
||||
Migrates webhook-attached presets for a graph to a newly activated version.
|
||||
@@ -2436,6 +2644,8 @@ async def migrate_webhook_presets_to_new_version(
|
||||
Args:
|
||||
user_id: The owner of the presets.
|
||||
new_graph: The newly activated graph version to migrate presets to.
|
||||
organization_ids: Confine the migration to presets tagged with these
|
||||
organizations (None for untagged ones).
|
||||
|
||||
Returns:
|
||||
The migration outcome: the number of presets migrated and the presets
|
||||
@@ -2446,14 +2656,27 @@ async def migrate_webhook_presets_to_new_version(
|
||||
# New version has no webhook trigger to migrate presets onto.
|
||||
return library_model.WebhookPresetMigrationResult()
|
||||
|
||||
in_organizations = (
|
||||
[
|
||||
cast(
|
||||
prisma.types.AgentPresetWhereInput,
|
||||
organization_filter(organization_ids),
|
||||
)
|
||||
]
|
||||
if organization_ids is not None
|
||||
else None
|
||||
)
|
||||
candidates_where: prisma.types.AgentPresetWhereInput = {
|
||||
"userId": user_id,
|
||||
"agentGraphId": new_graph.id,
|
||||
"agentGraphVersion": {"lt": new_graph.version},
|
||||
"webhookId": {"not": None},
|
||||
"isDeleted": False,
|
||||
}
|
||||
if in_organizations:
|
||||
candidates_where["AND"] = in_organizations
|
||||
candidates = await prisma.models.AgentPreset.prisma().find_many(
|
||||
where={
|
||||
"userId": user_id,
|
||||
"agentGraphId": new_graph.id,
|
||||
"agentGraphVersion": {"lt": new_graph.version},
|
||||
"webhookId": {"not": None},
|
||||
"isDeleted": False,
|
||||
},
|
||||
where=candidates_where
|
||||
)
|
||||
if not candidates:
|
||||
return library_model.WebhookPresetMigrationResult()
|
||||
@@ -2501,18 +2724,21 @@ async def migrate_webhook_presets_to_new_version(
|
||||
skipped_presets=skipped_presets
|
||||
)
|
||||
|
||||
# Preserve candidate order for a deterministic query. Re-assert userId and
|
||||
# the version guard so a concurrent activation that already bumped a preset
|
||||
# past new_graph.version can't be downgraded between the find_many above
|
||||
# and this update.
|
||||
# Preserve candidate order for a deterministic query. Re-assert userId,
|
||||
# the organizations and the version guard so a concurrent activation that
|
||||
# already bumped a preset past new_graph.version can't be downgraded
|
||||
# between the find_many above and this update.
|
||||
ids_to_migrate = [preset.id for preset in candidates if preset.id in compatible_ids]
|
||||
migrate_where: prisma.types.AgentPresetWhereInput = {
|
||||
"id": {"in": ids_to_migrate},
|
||||
"userId": user_id,
|
||||
"agentGraphVersion": {"lt": new_graph.version},
|
||||
"isDeleted": False,
|
||||
}
|
||||
if in_organizations:
|
||||
migrate_where["AND"] = in_organizations
|
||||
count = await prisma.models.AgentPreset.prisma().update_many(
|
||||
where={
|
||||
"id": {"in": ids_to_migrate},
|
||||
"userId": user_id,
|
||||
"agentGraphVersion": {"lt": new_graph.version},
|
||||
"isDeleted": False,
|
||||
},
|
||||
where=migrate_where,
|
||||
data={"agentGraphVersion": new_graph.version},
|
||||
)
|
||||
if count > 0:
|
||||
@@ -2545,7 +2771,10 @@ async def delete_preset(user_id: str, preset_id: str) -> None:
|
||||
|
||||
|
||||
async def fork_library_agent(
|
||||
library_agent_id: str, user_id: str
|
||||
library_agent_id: str,
|
||||
user_id: str,
|
||||
organization_id: str | None = None,
|
||||
team_id: str | None = None,
|
||||
) -> library_model.LibraryAgent:
|
||||
"""
|
||||
Clones a library agent and its underyling graph and nodes (with new ids) for the given user.
|
||||
@@ -2553,6 +2782,8 @@ async def fork_library_agent(
|
||||
Args:
|
||||
library_agent_id: The ID of the library agent to fork.
|
||||
user_id: The ID of the user who owns the library agent.
|
||||
organization_id, team_id: Tenancy for the copy; without them the graph
|
||||
is untagged and the library entry goes to the user's default team.
|
||||
|
||||
Returns:
|
||||
The forked parent (if it has sub-graphs) LibraryAgent.
|
||||
@@ -2579,7 +2810,11 @@ async def fork_library_agent(
|
||||
# GraphActivationError, but the forked graph row exists; callers should
|
||||
# surface that as a 400 to the user.
|
||||
new_graph = await graph_db.fork_graph(
|
||||
original_agent.graph_id, original_agent.graph_version, user_id
|
||||
original_agent.graph_id,
|
||||
original_agent.graph_version,
|
||||
user_id,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
new_graph = await before_graph_activate(new_graph, user_id=user_id)
|
||||
|
||||
@@ -2590,6 +2825,8 @@ async def fork_library_agent(
|
||||
user_id,
|
||||
hitl_safe_mode=original_agent.settings.human_in_the_loop_safe_mode,
|
||||
sensitive_action_safe_mode=original_agent.settings.sensitive_action_safe_mode,
|
||||
organization_id=organization_id,
|
||||
team_id=team_id,
|
||||
)
|
||||
)[0]
|
||||
|
||||
|
||||
@@ -10,12 +10,21 @@ import prisma.models
|
||||
import pytest
|
||||
|
||||
from backend.api.features.experts import experts_db
|
||||
from backend.api.features.store.store_listing_versions import (
|
||||
installable_store_version_where,
|
||||
)
|
||||
from backend.data.db import connect
|
||||
from backend.data.includes import library_agent_include
|
||||
from backend.util.exceptions import MissingConfigError, NotFoundError
|
||||
from backend.util.exceptions import (
|
||||
MissingConfigError,
|
||||
NotAuthorizedError,
|
||||
NotFoundError,
|
||||
)
|
||||
|
||||
from . import db
|
||||
from . import model as library_model
|
||||
from ._add_to_library import restore_existing_library_agent
|
||||
from .exceptions import LibraryAgentInAnotherOrganizationError
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expert_id", [None, "expert-1"])
|
||||
@@ -281,17 +290,20 @@ async def test_add_agent_to_library(mocker):
|
||||
return_value=mock_library_agent_data
|
||||
)
|
||||
|
||||
# Mock graph_db.get_graph function that's called in resolve_graph_for_library
|
||||
# (lives in _add_to_library.py after refactor, not db.py)
|
||||
# Mock _fetch_schedule_info used by add_graph_to_library
|
||||
mocker.patch(
|
||||
"backend.api.features.library._add_to_library._fetch_schedule_info",
|
||||
new=mocker.AsyncMock(return_value={}),
|
||||
)
|
||||
|
||||
# Mock graph_db.get_graph, called by resolve_graph_model_for_library
|
||||
mock_graph_db = mocker.patch(
|
||||
"backend.api.features.library._add_to_library.graph_db"
|
||||
)
|
||||
mock_graph_model = mocker.Mock()
|
||||
mock_graph_model.id = "agent1"
|
||||
mock_graph_model.version = 1
|
||||
mock_graph_model.nodes = (
|
||||
[]
|
||||
) # Empty list so _has_human_in_the_loop_blocks returns False
|
||||
mock_graph_model.nodes = []
|
||||
mock_graph_db.get_graph = mocker.AsyncMock(return_value=mock_graph_model)
|
||||
|
||||
# Mock the model conversion
|
||||
@@ -592,7 +604,7 @@ async def test_update_graph_in_library_allows_archived_library_agent(mocker):
|
||||
new=mocker.AsyncMock(return_value=current_library_agent),
|
||||
)
|
||||
mock_update_library_agent = mocker.patch(
|
||||
"backend.api.features.library.db.update_library_agent_version_and_settings",
|
||||
"backend.api.features.library.db.update_agent_version_in_library",
|
||||
new=mocker.AsyncMock(return_value=updated_library_agent),
|
||||
)
|
||||
mock_clear_unowned = mocker.patch(
|
||||
@@ -619,7 +631,102 @@ async def test_update_graph_in_library_allows_archived_library_agent(mocker):
|
||||
"graph-id",
|
||||
include_archived=True,
|
||||
)
|
||||
mock_update_library_agent.assert_awaited_once_with("test-user", created_graph)
|
||||
mock_update_library_agent.assert_awaited_once_with(
|
||||
"test-user", "graph-id", 2, organization_ids=None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("move", [True, False])
|
||||
async def test_update_graph_in_library_moves_webhook_presets_only_when_asked(
|
||||
mocker, move: bool
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.orgs.db.get_user_default_team",
|
||||
AsyncMock(return_value=("org-fallback", "team-fallback")),
|
||||
)
|
||||
graph_model = mocker.Mock(is_active=True)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.get_graph_all_versions",
|
||||
new=mocker.AsyncMock(return_value=[mocker.Mock(version=1, is_active=True)]),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.make_graph_model",
|
||||
return_value=graph_model,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.before_graph_activate",
|
||||
new=mocker.AsyncMock(return_value=graph_model),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.create_graph",
|
||||
new=mocker.AsyncMock(
|
||||
return_value=mocker.Mock(
|
||||
id="graph-id", version=2, is_active=True, webhook_input_node=object()
|
||||
)
|
||||
),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.get_library_agent_by_graph_id",
|
||||
new=mocker.AsyncMock(return_value=mocker.Mock()),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.update_agent_version_in_library",
|
||||
new=mocker.AsyncMock(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.set_graph_active_version",
|
||||
new=mocker.AsyncMock(),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.on_graph_deactivate", new=mocker.AsyncMock()
|
||||
)
|
||||
migrate = mocker.patch(
|
||||
"backend.api.features.library.db.migrate_webhook_presets_to_new_version",
|
||||
new=mocker.AsyncMock(),
|
||||
)
|
||||
|
||||
await db.update_graph_in_library(
|
||||
mocker.Mock(id="graph-id"), "test-user", move_webhook_presets=move
|
||||
)
|
||||
|
||||
assert migrate.await_count == (1 if move else 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_graph_in_library_tags_the_graph_and_entry_with_the_tenant(
|
||||
mocker,
|
||||
):
|
||||
"""An External API call creates the agent in its credential's organization."""
|
||||
graph_model = mocker.Mock(is_active=False)
|
||||
created_graph = mocker.Mock(id="graph-id", version=1)
|
||||
library_agent = mocker.Mock()
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.make_graph_model",
|
||||
return_value=graph_model,
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.clear_unowned_auto_credentials",
|
||||
new=mocker.AsyncMock(),
|
||||
)
|
||||
create_graph = mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.create_graph",
|
||||
new=mocker.AsyncMock(return_value=created_graph),
|
||||
)
|
||||
create_entry = mocker.patch(
|
||||
"backend.api.features.library.db.create_library_agent",
|
||||
new=mocker.AsyncMock(return_value=[library_agent]),
|
||||
)
|
||||
|
||||
await db.create_graph_in_library(
|
||||
mocker.Mock(), "test-user", organization_id="org-1", team_id="team-1"
|
||||
)
|
||||
|
||||
create_graph.assert_awaited_once_with(
|
||||
graph_model, "test-user", organization_id="org-1", team_id="team-1"
|
||||
)
|
||||
assert create_entry.await_args.kwargs["organization_id"] == "org-1"
|
||||
assert create_entry.await_args.kwargs["team_id"] == "team-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -786,6 +893,37 @@ async def test_get_library_agent_rejects_missing_graph(mocker):
|
||||
mock_from_db.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_library_agent_loads_the_graphs_store_listing(mocker):
|
||||
"""The listing is looked up by graph id with the relations `from_db`
|
||||
reads, so the library page keeps its marketplace info."""
|
||||
library_agent = MagicMock(id="library-id", AgentGraph=MagicMock(id="graph-1"))
|
||||
mock_library_agent = mocker.patch("prisma.models.LibraryAgent.prisma")
|
||||
mock_library_agent.return_value.find_first = mocker.AsyncMock(
|
||||
return_value=library_agent
|
||||
)
|
||||
listing = MagicMock(name="store-listing")
|
||||
mock_store_listing = mocker.patch("prisma.models.StoreListing.prisma")
|
||||
mock_store_listing.return_value.find_first = mocker.AsyncMock(return_value=listing)
|
||||
mocker.patch.object(db, "_fetch_schedule_info", mocker.AsyncMock(return_value={}))
|
||||
mocker.patch.object(
|
||||
db.graph_db, "get_sub_graphs", mocker.AsyncMock(return_value=[])
|
||||
)
|
||||
mock_from_db = mocker.patch.object(library_model.LibraryAgent, "from_db")
|
||||
|
||||
await db.get_library_agent("library-id", "test-user")
|
||||
|
||||
mock_store_listing.return_value.find_first.assert_awaited_once_with(
|
||||
where={
|
||||
"agentGraphId": "graph-1",
|
||||
"isDeleted": False,
|
||||
"hasApprovedVersion": True,
|
||||
},
|
||||
include={"ActiveVersion": True, "CreatorProfile": True},
|
||||
)
|
||||
assert mock_from_db.call_args.kwargs["store_listing"] is listing
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_library_agents_skips_failed_agent(mocker):
|
||||
"""Agents that fail parsing should be skipped — covers the except branch."""
|
||||
@@ -1594,6 +1732,51 @@ async def test_create_expert_preset_rejects_shared_webhook(mocker):
|
||||
mock_preset_client.create.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"inputs, credentials, saved",
|
||||
[({"topic": "x"}, {}, ["topic"]), ({}, {}, []), ({"topic": "x"}, None, None)],
|
||||
)
|
||||
async def test_new_inputs_need_credentials_alongside_which_may_be_empty(
|
||||
mocker, inputs, credentials, saved
|
||||
):
|
||||
"""Inputs and credentials share their rows, so one is never replaced alone;
|
||||
either may be `{}`, as for a preset that needs no credentials."""
|
||||
mocker.patch.object(
|
||||
db, "get_preset", new=AsyncMock(return_value=MagicMock(expert_id=None))
|
||||
)
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_tx():
|
||||
yield None
|
||||
|
||||
preset_client = AsyncMock()
|
||||
mocker.patch.object(prisma.models.AgentPreset, "prisma", return_value=preset_client)
|
||||
mocker.patch.object(
|
||||
prisma.models.AgentNodeExecutionInputOutput, "prisma", return_value=AsyncMock()
|
||||
)
|
||||
mocker.patch.object(db, "transaction", fake_tx)
|
||||
mocker.patch.object(
|
||||
library_model.LibraryAgentPreset, "from_db", return_value=MagicMock()
|
||||
)
|
||||
|
||||
update = db.update_preset(
|
||||
user_id="owner",
|
||||
preset_id="preset-1",
|
||||
inputs=inputs,
|
||||
credentials=credentials,
|
||||
)
|
||||
|
||||
if saved is None:
|
||||
with pytest.raises(ValueError):
|
||||
await update
|
||||
preset_client.update.assert_not_awaited()
|
||||
return
|
||||
await update
|
||||
rows = preset_client.update.await_args.kwargs["data"]["InputPresets"]["create"]
|
||||
assert [row["name"] for row in rows] == saved
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_expert_preset_rehomes_on_metadata_only_update(mocker):
|
||||
current = MagicMock(
|
||||
@@ -2209,3 +2392,236 @@ async def test_cleanup_trigger_agents_processes_each_independently(mocker):
|
||||
recursive_delete.assert_awaited_once_with(
|
||||
library_agent_id="sole", user_id="test-user", soft_delete=True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_re_adding_a_listing_leaves_another_organizations_entry_alone(mocker):
|
||||
"""The entry is unique per user, so a re-add from another org would restore
|
||||
and hand back a row that org can't see."""
|
||||
delegate = MagicMock()
|
||||
delegate.find_unique = AsyncMock(return_value=MagicMock(organizationId="org-b"))
|
||||
delegate.update = AsyncMock()
|
||||
mocker.patch("prisma.models.LibraryAgent.prisma", return_value=delegate)
|
||||
listing_version = MagicMock(AgentGraph=MagicMock(id="graph-1", version=1))
|
||||
|
||||
with pytest.raises(LibraryAgentInAnotherOrganizationError):
|
||||
await restore_existing_library_agent(
|
||||
listing_version, "user-1", organization_id="org-a"
|
||||
)
|
||||
delegate.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def library_rows(mocker):
|
||||
"""The LibraryAgent queries inside update_agent_version_in_library."""
|
||||
rows = MagicMock()
|
||||
rows.find_first_or_raise = AsyncMock(return_value=MagicMock(id="entry-a"))
|
||||
rows.find_many = AsyncMock(return_value=[])
|
||||
rows.delete_many = AsyncMock()
|
||||
rows.update = AsyncMock(return_value=MagicMock())
|
||||
mocker.patch("prisma.models.LibraryAgent.prisma", return_value=rows)
|
||||
# The target version: the user's own, in org-a.
|
||||
rows.versions = MagicMock(
|
||||
find_unique=AsyncMock(
|
||||
return_value=MagicMock(userId="test-user", organizationId="org-a")
|
||||
)
|
||||
)
|
||||
mocker.patch("prisma.models.AgentGraph.prisma", return_value=rows.versions)
|
||||
# The marketplace version another user's target must have: none.
|
||||
rows.listings = MagicMock(find_first=AsyncMock(return_value=None))
|
||||
mocker.patch("prisma.models.StoreListingVersion.prisma", return_value=rows.listings)
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction():
|
||||
yield None
|
||||
|
||||
mocker.patch("backend.api.features.library.db.transaction", transaction)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db._fetch_schedule_info",
|
||||
new=AsyncMock(return_value={}),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.get_sub_graphs",
|
||||
new=AsyncMock(return_value=[]),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.library_model.LibraryAgent.from_db",
|
||||
return_value=MagicMock(id="entry-a"),
|
||||
)
|
||||
mocker.patch("backend.api.features.library.db.schedule_library_agent_embedding")
|
||||
return rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_version_moves_only_the_named_entry_in_the_callers_organizations(
|
||||
library_rows,
|
||||
):
|
||||
"""The user's first entry for the graph could be another organization's."""
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user",
|
||||
"graph-id",
|
||||
2,
|
||||
library_agent_id="entry-a",
|
||||
organization_ids=frozenset({"org-a"}),
|
||||
)
|
||||
|
||||
assert library_rows.update.await_args.kwargs["where"] == {"id": "entry-a"}
|
||||
where = library_rows.find_first_or_raise.await_args.kwargs["where"]
|
||||
assert where["id"] == "entry-a"
|
||||
assert where["AND"] == [{"OR": [{"organizationId": {"in": ["org-a"]}}]}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_another_organizations_entry_on_the_version_is_not_replaced(
|
||||
library_rows,
|
||||
):
|
||||
"""An entry is unique per user, graph and version, and the one already on
|
||||
the version isn't this call's to remove."""
|
||||
library_rows.find_many.return_value = [
|
||||
MagicMock(id="entry-b", organizationId="org-b")
|
||||
]
|
||||
|
||||
with pytest.raises(LibraryAgentInAnotherOrganizationError):
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user", "graph-id", 2, organization_ids=frozenset({"org-a"})
|
||||
)
|
||||
|
||||
library_rows.delete_many.assert_not_awaited()
|
||||
library_rows.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"organization_id, organization_ids",
|
||||
[
|
||||
("org-b", frozenset({"org-a"})),
|
||||
# An untagged version is the personal org's.
|
||||
(None, frozenset({"org-a"})),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_entry_is_not_pointed_at_the_users_version_outside_its_organizations(
|
||||
library_rows, organization_id, organization_ids
|
||||
):
|
||||
"""Running the entry would run that version, billed where the entry is."""
|
||||
library_rows.versions.find_unique.return_value = MagicMock(
|
||||
userId="test-user", organizationId=organization_id
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user",
|
||||
"graph-id",
|
||||
2,
|
||||
library_agent_id="entry-a",
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
|
||||
library_rows.find_first_or_raise.assert_not_awaited()
|
||||
library_rows.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("organization_ids", [None, frozenset({"org-a"})])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_marketplace_agent_follows_only_a_version_it_could_be_installed_at(
|
||||
library_rows, organization_ids
|
||||
):
|
||||
"""A version in the library that was ever submitted can be read, run and
|
||||
forked, so pointing a creator's agent at a pending or rejected one would
|
||||
unlock it. This takes what adding the agent takes, by the same predicate."""
|
||||
library_rows.versions.find_unique.return_value = MagicMock(
|
||||
userId="creator", organizationId="creators-org"
|
||||
)
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user",
|
||||
"graph-id",
|
||||
2,
|
||||
library_agent_id="entry-a",
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
library_rows.update.assert_not_awaited()
|
||||
|
||||
library_rows.listings.find_first.return_value = MagicMock()
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user",
|
||||
"graph-id",
|
||||
2,
|
||||
library_agent_id="entry-a",
|
||||
organization_ids=organization_ids,
|
||||
)
|
||||
library_rows.update.assert_awaited_once()
|
||||
assert library_rows.listings.find_first.await_args.kwargs["where"] == {
|
||||
"agentGraphId": "graph-id",
|
||||
"agentGraphVersion": 2,
|
||||
**installable_store_version_where(),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_organizations_the_users_own_versions_move_as_before(
|
||||
library_rows,
|
||||
):
|
||||
"""The web app's own saves and version picks pass no organization."""
|
||||
library_rows.versions.find_unique.return_value = MagicMock(
|
||||
userId="test-user", organizationId="org-b"
|
||||
)
|
||||
|
||||
await db.update_agent_version_in_library("test-user", "graph-id", 2)
|
||||
|
||||
library_rows.update.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("organization_ids", [None, frozenset({"org-a"})])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_version_that_does_not_exist_is_not_found(
|
||||
library_rows, organization_ids
|
||||
):
|
||||
library_rows.versions.find_unique.return_value = None
|
||||
|
||||
with pytest.raises(NotFoundError):
|
||||
await db.update_agent_version_in_library(
|
||||
"test-user", "graph-id", 9, organization_ids=organization_ids
|
||||
)
|
||||
|
||||
library_rows.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_presets_move_only_within_the_callers_organizations(mocker):
|
||||
presets = MagicMock(find_many=AsyncMock(return_value=[]))
|
||||
mocker.patch("prisma.models.AgentPreset.prisma", return_value=presets)
|
||||
|
||||
await db.migrate_webhook_presets_to_new_version(
|
||||
"test-user",
|
||||
MagicMock(id="graph-id", version=2),
|
||||
organization_ids=frozenset({"org-a", None}),
|
||||
)
|
||||
|
||||
assert presets.find_many.await_args.kwargs["where"]["AND"] == [
|
||||
{"OR": [{"organizationId": {"in": ["org-a"]}}, {"organizationId": None}]}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_external_save_cannot_version_a_graph_outside_its_organization(
|
||||
mocker,
|
||||
):
|
||||
mocker.patch(
|
||||
"backend.api.features.library.db.writable_organization_ids",
|
||||
new=AsyncMock(return_value=frozenset({"org-a"})),
|
||||
)
|
||||
mocker.patch(
|
||||
"backend.data.graph.get_graph_version_owners",
|
||||
new=AsyncMock(return_value={("test-user", None)}),
|
||||
)
|
||||
create = mocker.patch(
|
||||
"backend.api.features.library.db.graph_db.create_graph", new=AsyncMock()
|
||||
)
|
||||
|
||||
with pytest.raises(NotAuthorizedError):
|
||||
await db.update_graph_in_library(
|
||||
MagicMock(id="graph-id"), "test-user", organization_id="org-a"
|
||||
)
|
||||
|
||||
create.assert_not_awaited()
|
||||
@@ -34,7 +34,7 @@ logger = logging.getLogger(__name__)
|
||||
_background_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
|
||||
def _build_searchable_text(graph: graph_db.GraphModel) -> str:
|
||||
def _build_searchable_text(graph: graph_db.GraphModel | graph_db.AgentGraph) -> str:
|
||||
parts = [
|
||||
graph.name or "",
|
||||
graph.description or "",
|
||||
@@ -44,7 +44,9 @@ def _build_searchable_text(graph: graph_db.GraphModel) -> str:
|
||||
|
||||
|
||||
async def _run_embedding(
|
||||
library_agent_id: str, user_id: str, graph: graph_db.GraphModel
|
||||
library_agent_id: str,
|
||||
user_id: str,
|
||||
graph: graph_db.GraphModel | graph_db.AgentGraph,
|
||||
) -> None:
|
||||
try:
|
||||
searchable_text = _build_searchable_text(graph)
|
||||
@@ -83,7 +85,9 @@ async def _run_embedding(
|
||||
|
||||
|
||||
def schedule_library_agent_embedding(
|
||||
library_agent_id: str, user_id: str, graph: graph_db.GraphModel
|
||||
library_agent_id: str,
|
||||
user_id: str,
|
||||
graph: graph_db.GraphModel | graph_db.AgentGraph,
|
||||
) -> asyncio.Task[None]:
|
||||
"""Schedule a background (re-)embed. No-ops cheaply when the existing
|
||||
embedding's searchableText is unchanged. Failures are logged, not raised."""
|
||||
|
||||
@@ -8,3 +8,12 @@ class FolderAlreadyExistsError(FolderValidationError):
|
||||
"""Raised when a folder with the same name already exists in the location."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class LibraryAgentInAnotherOrganizationError(Exception):
|
||||
"""The user already has this agent in their library, in another organization.
|
||||
|
||||
A library entry is unique per user, graph and version, so a second one
|
||||
can't be made, and the existing one isn't this organization's to restore
|
||||
or return.
|
||||
"""
|
||||
@@ -0,0 +1,153 @@
|
||||
"""`get_library_agent` reports the marketplace listing of the agent's graph."""
|
||||
|
||||
import uuid
|
||||
|
||||
import prisma.models
|
||||
import pytest
|
||||
from autogpt_libs.auth.models import RequestContext
|
||||
|
||||
import backend.api.features.store.model as store_model
|
||||
from backend.api.features.graphs import routes as graphs_routes
|
||||
from backend.api.features.library import db as library_db
|
||||
from backend.api.model import CreateGraph
|
||||
from backend.blocks.io import AgentInputBlock
|
||||
from backend.data.graph import Graph, Node
|
||||
from backend.data.user import get_or_create_user
|
||||
from backend.util.test import SpinTestServer
|
||||
|
||||
_seeded_user_ids: list[str] = []
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def _clean_seeded_rows():
|
||||
yield
|
||||
user_ids = list(_seeded_user_ids)
|
||||
_seeded_user_ids.clear()
|
||||
if not user_ids:
|
||||
return
|
||||
# A listing pointing at a seed user's graph blocks the cascade from User.
|
||||
await prisma.models.StoreListing.prisma().delete_many(
|
||||
where={"owningUserId": {"in": user_ids}}
|
||||
)
|
||||
await prisma.models.LibraryAgent.prisma().delete_many(
|
||||
where={"userId": {"in": user_ids}}
|
||||
)
|
||||
await prisma.models.User.prisma().delete_many(where={"id": {"in": user_ids}})
|
||||
|
||||
|
||||
async def test_a_published_agent_reports_its_listing(server: SpinTestServer):
|
||||
owner, graph, listing_name = await _publish_graph(server)
|
||||
|
||||
agent = await library_db.get_library_agent(
|
||||
await _library_agent_id(owner.id, graph.id), owner.id
|
||||
)
|
||||
|
||||
listing = agent.marketplace_listing
|
||||
assert listing is not None
|
||||
assert listing.name == listing_name
|
||||
assert listing.slug == graph.id
|
||||
profile = await prisma.models.Profile.prisma().find_unique(
|
||||
where={"userId": owner.id}
|
||||
)
|
||||
assert profile is not None
|
||||
assert listing.creator.slug == profile.username
|
||||
|
||||
|
||||
async def test_a_version_saved_after_publishing_still_reports_the_listing(
|
||||
server: SpinTestServer,
|
||||
):
|
||||
"""Saving moves the owner's library agent to a version no listing names."""
|
||||
owner, graph, listing_name = await _publish_graph(server)
|
||||
library_agent_id = await _library_agent_id(owner.id, graph.id)
|
||||
|
||||
await graphs_routes.update_graph(
|
||||
graph.id,
|
||||
Graph(id=graph.id, name=graph.name, description="Edited", nodes=_nodes()),
|
||||
owner.id,
|
||||
_ctx(owner.id),
|
||||
)
|
||||
|
||||
agent = await library_db.get_library_agent(library_agent_id, owner.id)
|
||||
assert agent.graph_version == graph.version + 1
|
||||
assert agent.marketplace_listing is not None
|
||||
assert agent.marketplace_listing.name == listing_name
|
||||
|
||||
|
||||
async def _publish_graph(server: SpinTestServer):
|
||||
owner = await _create_user()
|
||||
admin = await _create_user()
|
||||
graph = await server.agent_server.test_create_graph(
|
||||
CreateGraph(
|
||||
graph=Graph(
|
||||
name=f"Listed graph {uuid.uuid4().hex[:8]}",
|
||||
description="Published to the marketplace",
|
||||
nodes=_nodes(),
|
||||
links=[],
|
||||
)
|
||||
),
|
||||
owner.id,
|
||||
)
|
||||
listing_name = f"Listing {uuid.uuid4().hex[:8]}"
|
||||
submission = await server.agent_server.test_create_store_listing(
|
||||
store_model.StoreSubmissionRequest(
|
||||
graph_id=graph.id,
|
||||
graph_version=graph.version,
|
||||
slug=graph.id,
|
||||
name=listing_name,
|
||||
sub_heading="Sub heading",
|
||||
video_url=None,
|
||||
image_urls=[],
|
||||
description="Description",
|
||||
categories=["operations"],
|
||||
),
|
||||
owner.id,
|
||||
)
|
||||
assert submission.listing_version_id is not None
|
||||
await server.agent_server.test_review_store_listing(
|
||||
store_model.ReviewSubmissionRequest(
|
||||
store_listing_version_id=submission.listing_version_id,
|
||||
is_approved=True,
|
||||
comments="ok",
|
||||
),
|
||||
user_id=admin.id,
|
||||
)
|
||||
return owner, graph, listing_name
|
||||
|
||||
|
||||
async def _library_agent_id(user_id: str, graph_id: str) -> str:
|
||||
library_agent = await prisma.models.LibraryAgent.prisma().find_first(
|
||||
where={"userId": user_id, "agentGraphId": graph_id}
|
||||
)
|
||||
assert library_agent is not None
|
||||
return library_agent.id
|
||||
|
||||
|
||||
async def _create_user():
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
user = await get_or_create_user(
|
||||
{
|
||||
"sub": str(uuid.uuid4()),
|
||||
"email": f"library-listing-{suffix}@example.com",
|
||||
"name": "Listing Owner",
|
||||
}
|
||||
)
|
||||
_seeded_user_ids.append(user.id)
|
||||
return user
|
||||
|
||||
|
||||
def _nodes() -> list[Node]:
|
||||
return [Node(block_id=AgentInputBlock().id, input_default={"name": "input_1"})]
|
||||
|
||||
|
||||
def _ctx(user_id: str) -> RequestContext:
|
||||
return RequestContext(
|
||||
user_id=user_id,
|
||||
org_id=f"test-org-{user_id}",
|
||||
team_id=None,
|
||||
is_org_owner=True,
|
||||
is_org_admin=True,
|
||||
is_org_billing_manager=False,
|
||||
is_team_admin=True,
|
||||
is_team_billing_manager=False,
|
||||
seat_status="ACTIVE",
|
||||
)
|
||||
@@ -42,6 +42,8 @@ class LibraryFolder(pydantic.BaseModel):
|
||||
updated_at: datetime.datetime
|
||||
agent_count: int = 0 # Direct agents in folder
|
||||
subfolder_count: int = 0 # Direct child folders
|
||||
organization_id: str | None = None
|
||||
team_id: str | None = None
|
||||
|
||||
@staticmethod
|
||||
def from_db(
|
||||
@@ -61,6 +63,8 @@ class LibraryFolder(pydantic.BaseModel):
|
||||
updated_at=folder.updatedAt,
|
||||
agent_count=agent_count,
|
||||
subfolder_count=subfolder_count,
|
||||
organization_id=folder.organizationId,
|
||||
team_id=folder.teamId,
|
||||
)
|
||||
|
||||
|
||||
@@ -244,6 +248,8 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
description="ISO 8601 timestamp of the next scheduled run, if any",
|
||||
)
|
||||
settings: GraphSettings = pydantic.Field(default_factory=GraphSettings)
|
||||
organization_id: str | None = None
|
||||
team_id: str | None = None
|
||||
marketplace_listing: Optional["MarketplaceListing"] = None
|
||||
|
||||
@staticmethod
|
||||
@@ -251,7 +257,6 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
agent: prisma.models.LibraryAgent,
|
||||
sub_graphs: Optional[list[prisma.models.AgentGraph]] = None,
|
||||
store_listing: Optional[prisma.models.StoreListing] = None,
|
||||
profile: Optional[prisma.models.Profile] = None,
|
||||
execution_count_override: Optional[int] = None,
|
||||
schedule_info: Optional[dict[str, str]] = None,
|
||||
) -> "LibraryAgent":
|
||||
@@ -349,19 +354,25 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
can_access_graph = agent.AgentGraph.userId == agent.userId
|
||||
is_latest_version = True
|
||||
|
||||
marketplace_listing_data = None
|
||||
if store_listing and store_listing.ActiveVersion and profile:
|
||||
creator_data = MarketplaceListingCreator(
|
||||
name=profile.name,
|
||||
id=profile.id,
|
||||
slug=profile.username,
|
||||
)
|
||||
marketplace_listing_data = MarketplaceListing(
|
||||
active_listing = store_listing.ActiveVersion if store_listing else None
|
||||
creator_profile = store_listing.CreatorProfile if store_listing else None
|
||||
marketplace_listing_info = (
|
||||
MarketplaceListing(
|
||||
id=store_listing.id,
|
||||
name=store_listing.ActiveVersion.name,
|
||||
name=active_listing.name,
|
||||
slug=store_listing.slug,
|
||||
creator=creator_data,
|
||||
creator=MarketplaceListingCreator(
|
||||
name=creator_profile.name,
|
||||
id=creator_profile.id,
|
||||
slug=creator_profile.username,
|
||||
),
|
||||
)
|
||||
if store_listing
|
||||
and active_listing
|
||||
and creator_profile
|
||||
and not store_listing.isDeleted
|
||||
else None
|
||||
)
|
||||
|
||||
return LibraryAgent(
|
||||
id=agent.id,
|
||||
@@ -404,6 +415,8 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
is_favorite=agent.isFavorite,
|
||||
is_hidden=agent.isHidden,
|
||||
folder_id=agent.folderId,
|
||||
organization_id=agent.organizationId,
|
||||
team_id=agent.teamId,
|
||||
folder_name=agent.Folder.name if agent.Folder else None,
|
||||
recommended_schedule_cron=agent.AgentGraph.recommendedScheduleCron,
|
||||
is_scheduled=bool(schedule_info and agent.agentGraphId in schedule_info),
|
||||
@@ -411,7 +424,7 @@ class LibraryAgent(pydantic.BaseModel):
|
||||
schedule_info.get(agent.agentGraphId) if schedule_info else None
|
||||
),
|
||||
settings=_parse_settings(agent.settings),
|
||||
marketplace_listing=marketplace_listing_data,
|
||||
marketplace_listing=marketplace_listing_info,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Loaded 100 of 247 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user