mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-21 09:47:13 +00:00
Compare commits
62 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0135ae31b8 | |||
| c222321f2d | |||
| 5084f2391d | |||
| b121c66f18 | |||
| 536534865e | |||
| ff528e8e0b | |||
| f7914a8900 | |||
| c18fc9dfe3 | |||
| bc40418bda | |||
| 45ed6efceb | |||
| c67a503532 | |||
| 8b14d2ab8e | |||
| 97428310b9 | |||
| 338ee733cb | |||
| c3299fd1a6 | |||
| be8478e940 | |||
| 96a535827f | |||
| 9b81a0b606 | |||
| bbc912d0ef | |||
| 20710df9cb | |||
| 90f3d880e5 | |||
| e37987215e | |||
| 22c389edc1 | |||
| 78068db9c8 | |||
| 7dc9dafb7c | |||
| 4bd899e77b | |||
| ddb6dbf593 | |||
| f59343fd5b | |||
| 211710e24c | |||
| cdd5c6589c | |||
| 3b4698463c | |||
| edd6cad449 | |||
| d78546967c | |||
| c08bfc8ced | |||
| 7820949d3a | |||
| e263a5d1d7 | |||
| 6bad7bcffc | |||
| f0b2c103c1 | |||
| 3101c9be6a | |||
| 1e6e4c0ca7 | |||
| 408c8031d4 | |||
| 0e6cca4690 | |||
| a7a7218afe | |||
| a67728c163 | |||
| e9c9e896c6 | |||
| e2331c4967 | |||
| 0ccbcd5f5f | |||
| c5aada494d | |||
| e36e3aaea8 | |||
| d64278ab3f | |||
| 161ea9b3eb | |||
| c7d14676fc | |||
| 2456bf1350 | |||
| 05a941ff16 | |||
| 0330788d14 | |||
| d8ab0ba567 | |||
| e3832ca536 | |||
| 5d9fd15671 | |||
| 404e3466d9 | |||
| 9df021eb8f | |||
| 98d0dba6d4 | |||
| 5ec2371879 |
@@ -1,78 +0,0 @@
|
|||||||
name: Build and deploy production
|
|
||||||
|
|
||||||
on:
|
|
||||||
push:
|
|
||||||
branches: [deploy/prod]
|
|
||||||
workflow_dispatch:
|
|
||||||
|
|
||||||
permissions:
|
|
||||||
contents: read
|
|
||||||
|
|
||||||
concurrency:
|
|
||||||
group: langbot-production
|
|
||||||
cancel-in-progress: false
|
|
||||||
|
|
||||||
env:
|
|
||||||
CORE_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot
|
|
||||||
CLOUD_IMAGE: ${{ secrets.DOCKER_USERNAME }}/langbot-cloud-core
|
|
||||||
SPACE_REF: e1b261dac45e886efc667b1096a4ec493c6a6111
|
|
||||||
|
|
||||||
jobs:
|
|
||||||
build-and-deploy:
|
|
||||||
runs-on: ubuntu-latest
|
|
||||||
environment: production
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
- uses: docker/setup-buildx-action@v3
|
|
||||||
- uses: docker/login-action@v3
|
|
||||||
with:
|
|
||||||
username: ${{ secrets.DOCKER_USERNAME }}
|
|
||||||
password: ${{ secrets.DOCKER_PASSWORD }}
|
|
||||||
- name: Build exact Core image
|
|
||||||
uses: docker/build-push-action@v6
|
|
||||||
with:
|
|
||||||
context: .
|
|
||||||
push: true
|
|
||||||
tags: |
|
|
||||||
${{ env.CORE_IMAGE }}:prod-${{ github.sha }}
|
|
||||||
${{ env.CORE_IMAGE }}:deploy-prod
|
|
||||||
cache-from: type=gha,scope=core-prod
|
|
||||||
cache-to: type=gha,mode=max,scope=core-prod
|
|
||||||
- name: Checkout production Cloud adapter
|
|
||||||
uses: actions/checkout@v4
|
|
||||||
with:
|
|
||||||
repository: langbot-app/langbot-space
|
|
||||||
ref: ${{ env.SPACE_REF }}
|
|
||||||
token: ${{ secrets.CLA_PAT }}
|
|
||||||
path: .space
|
|
||||||
- name: Build exact Cloud Core image
|
|
||||||
uses: docker/build-push-action@v6
|
|
||||||
with:
|
|
||||||
context: .space
|
|
||||||
file: .space/Dockerfile.cloud
|
|
||||||
push: true
|
|
||||||
build-args: LANGBOT_CORE_IMAGE=${{ env.CORE_IMAGE }}:prod-${{ github.sha }}
|
|
||||||
tags: |
|
|
||||||
${{ env.CLOUD_IMAGE }}:prod-${{ github.sha }}
|
|
||||||
${{ env.CLOUD_IMAGE }}:deploy-prod
|
|
||||||
cache-from: type=gha,scope=cloud-core-prod
|
|
||||||
cache-to: type=gha,mode=max,scope=cloud-core-prod
|
|
||||||
- name: Configure SSH
|
|
||||||
env:
|
|
||||||
SSH_KEY: ${{ secrets.JP09_SSH_KEY }}
|
|
||||||
KNOWN_HOSTS: ${{ secrets.JP09_KNOWN_HOSTS }}
|
|
||||||
run: |
|
|
||||||
install -m 700 -d ~/.ssh
|
|
||||||
install -m 600 /dev/null ~/.ssh/id_ed25519
|
|
||||||
printf '%s\n' "$SSH_KEY" > ~/.ssh/id_ed25519
|
|
||||||
printf '%s\n' "$KNOWN_HOSTS" > ~/.ssh/known_hosts
|
|
||||||
- name: Upload release manifest and deploy
|
|
||||||
env:
|
|
||||||
HOST: ${{ secrets.JP09_HOST }}
|
|
||||||
USER: ${{ secrets.JP09_USER }}
|
|
||||||
PORT: ${{ secrets.JP09_PORT }}
|
|
||||||
run: |
|
|
||||||
remote="$USER@$HOST"
|
|
||||||
ssh -p "$PORT" "$remote" 'install -d -m 700 /opt/langbot-cloud-prod'
|
|
||||||
scp -P "$PORT" deploy/prod/docker-compose.yml deploy/prod/deploy.sh "$remote:/opt/langbot-cloud-prod/"
|
|
||||||
ssh -p "$PORT" "$remote" "chmod 700 /opt/langbot-cloud-prod/deploy.sh && /opt/langbot-cloud-prod/deploy.sh prod-${GITHUB_SHA}"
|
|
||||||
@@ -1,92 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -Eeuo pipefail
|
|
||||||
|
|
||||||
cd /opt/langbot-cloud-prod
|
|
||||||
TAG=${1:?usage: deploy.sh prod-<40-char-sha>}
|
|
||||||
[[ "$TAG" =~ ^prod-[0-9a-f]{40}$ ]] || { echo 'invalid immutable image tag' >&2; exit 2; }
|
|
||||||
[[ -s .env ]] || { echo '/opt/langbot-cloud-prod/.env is missing' >&2; exit 3; }
|
|
||||||
|
|
||||||
rendered_compose=$(docker compose config)
|
|
||||||
grep -Fq 'LANGBOT_SPACE_CONTROL_PLANE_URL: https://space.langbot.app' <<<"$rendered_compose" || {
|
|
||||||
echo 'Cloud control-plane URL must be https://space.langbot.app' >&2
|
|
||||||
exit 4
|
|
||||||
}
|
|
||||||
grep -Fq 'SPACE__URL: https://space.langbot.app' <<<"$rendered_compose" || {
|
|
||||||
echo 'Cloud user-facing Space URL must be https://space.langbot.app' >&2
|
|
||||||
exit 5
|
|
||||||
}
|
|
||||||
|
|
||||||
update_env() {
|
|
||||||
local key=$1 value=$2
|
|
||||||
python3 - "$key" "$value" <<'PY'
|
|
||||||
from pathlib import Path
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
|
|
||||||
path = Path('.env')
|
|
||||||
key, value = sys.argv[1:]
|
|
||||||
lines = path.read_text().splitlines()
|
|
||||||
updated = False
|
|
||||||
for index, line in enumerate(lines):
|
|
||||||
if line.startswith(f'{key}='):
|
|
||||||
lines[index] = f'{key}={value}'
|
|
||||||
updated = True
|
|
||||||
break
|
|
||||||
if not updated:
|
|
||||||
lines.append(f'{key}={value}')
|
|
||||||
temporary = Path('.env.tmp')
|
|
||||||
temporary.write_text('\n'.join(lines) + '\n')
|
|
||||||
os.chmod(temporary, 0o600)
|
|
||||||
temporary.replace(path)
|
|
||||||
PY
|
|
||||||
}
|
|
||||||
update_env LANGBOT_IMAGE_TAG "$TAG"
|
|
||||||
set -a
|
|
||||||
. ./.env
|
|
||||||
set +a
|
|
||||||
|
|
||||||
for attempt in 1 2 3 4 5; do
|
|
||||||
if docker compose pull postgres redis migrate plugin-runtime core; then
|
|
||||||
break
|
|
||||||
fi
|
|
||||||
if [ "$attempt" -eq 5 ]; then
|
|
||||||
echo "docker compose pull failed after $attempt attempts" >&2
|
|
||||||
exit 1
|
|
||||||
fi
|
|
||||||
delay=$((attempt * 10))
|
|
||||||
echo "docker compose pull failed (attempt $attempt/5); retrying in ${delay}s" >&2
|
|
||||||
sleep "$delay"
|
|
||||||
done
|
|
||||||
docker compose up -d postgres redis
|
|
||||||
for _ in $(seq 1 60); do
|
|
||||||
if docker compose exec -T postgres pg_isready -U langbot_operator -d langbot >/dev/null 2>&1; then break; fi
|
|
||||||
sleep 2
|
|
||||||
done
|
|
||||||
docker compose exec -T postgres pg_isready -U langbot_operator -d langbot >/dev/null
|
|
||||||
|
|
||||||
docker compose exec -T postgres psql -v ON_ERROR_STOP=1 -U langbot_operator -d langbot \
|
|
||||||
-v runtime_password="$POSTGRES_RUNTIME_PASSWORD" <<'SQL'
|
|
||||||
SELECT format('CREATE ROLE langbot_runtime LOGIN PASSWORD %L', :'runtime_password')
|
|
||||||
WHERE NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'langbot_runtime')\gexec
|
|
||||||
ALTER ROLE langbot_runtime PASSWORD :'runtime_password';
|
|
||||||
GRANT CONNECT ON DATABASE langbot TO langbot_runtime;
|
|
||||||
REVOKE CREATE ON SCHEMA public FROM PUBLIC, langbot_runtime;
|
|
||||||
REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM langbot_runtime;
|
|
||||||
REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM langbot_runtime;
|
|
||||||
ALTER DEFAULT PRIVILEGES FOR ROLE langbot_operator IN SCHEMA public REVOKE ALL ON TABLES FROM langbot_runtime;
|
|
||||||
ALTER DEFAULT PRIVILEGES FOR ROLE langbot_operator IN SCHEMA public REVOKE ALL ON SEQUENCES FROM langbot_runtime;
|
|
||||||
GRANT USAGE ON SCHEMA public TO langbot_runtime;
|
|
||||||
SQL
|
|
||||||
|
|
||||||
docker compose --profile tools run --rm migrate
|
|
||||||
|
|
||||||
docker compose up -d --remove-orphans plugin-runtime core
|
|
||||||
for _ in $(seq 1 90); do
|
|
||||||
if docker compose exec -T core python -c 'import urllib.request; urllib.request.urlopen("http://127.0.0.1:5300/healthz", timeout=3)' >/dev/null 2>&1; then
|
|
||||||
docker compose ps
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
sleep 2
|
|
||||||
done
|
|
||||||
docker compose logs --tail=200 core plugin-runtime >&2
|
|
||||||
exit 1
|
|
||||||
@@ -1,161 +0,0 @@
|
|||||||
services:
|
|
||||||
postgres:
|
|
||||||
image: pgvector/pgvector:pg17
|
|
||||||
container_name: langbot-cloud-postgres
|
|
||||||
restart: unless-stopped
|
|
||||||
environment:
|
|
||||||
POSTGRES_DB: langbot
|
|
||||||
POSTGRES_USER: langbot_operator
|
|
||||||
POSTGRES_PASSWORD: ${POSTGRES_OPERATOR_PASSWORD}
|
|
||||||
volumes:
|
|
||||||
- postgres-data:/var/lib/postgresql/data
|
|
||||||
healthcheck:
|
|
||||||
test: [CMD-SHELL, "pg_isready -U langbot_operator -d langbot"]
|
|
||||||
interval: 5s
|
|
||||||
timeout: 5s
|
|
||||||
retries: 30
|
|
||||||
networks: [internal]
|
|
||||||
|
|
||||||
redis:
|
|
||||||
image: redis:7.4-alpine
|
|
||||||
container_name: langbot-cloud-redis
|
|
||||||
restart: unless-stopped
|
|
||||||
command: [redis-server, --appendonly, "yes", --requirepass, "${REDIS_PASSWORD}"]
|
|
||||||
volumes:
|
|
||||||
- redis-data:/data
|
|
||||||
healthcheck:
|
|
||||||
test: [CMD-SHELL, "redis-cli -a \"$${REDIS_PASSWORD}\" ping | grep PONG"]
|
|
||||||
interval: 5s
|
|
||||||
timeout: 5s
|
|
||||||
retries: 20
|
|
||||||
environment:
|
|
||||||
REDIS_PASSWORD: ${REDIS_PASSWORD}
|
|
||||||
networks: [internal]
|
|
||||||
|
|
||||||
migrate:
|
|
||||||
image: rockchin/langbot-cloud-core:${LANGBOT_IMAGE_TAG}
|
|
||||||
profiles: [tools]
|
|
||||||
command: [uv, run, langbot, migrate, --cloud]
|
|
||||||
environment: &core-env
|
|
||||||
TZ: Asia/Shanghai
|
|
||||||
SYSTEM__INSTANCE_ID: ${CLOUD_V2_INSTANCE_UUID}
|
|
||||||
SYSTEM__EDITION: cloud
|
|
||||||
SYSTEM__RECOVERY_KEY: ${SYSTEM_RECOVERY_KEY}
|
|
||||||
SYSTEM__JWT__SECRET: ${JWT_SECRET}
|
|
||||||
SYSTEM__LIMITATION__MAX_BOTS: "2"
|
|
||||||
SYSTEM__LIMITATION__MAX_PIPELINES: "3"
|
|
||||||
SYSTEM__LIMITATION__MAX_EXTENSIONS: "3"
|
|
||||||
SYSTEM__LIMITATION__MAX_KNOWLEDGE_BASES: "2"
|
|
||||||
API__WEBHOOK_PREFIX: https://cloud.langbot.app
|
|
||||||
API__WEBUI_URL: https://cloud.langbot.app
|
|
||||||
WORKSPACE__INVITATIONS__PUBLIC_WEB_URL: https://cloud.langbot.app
|
|
||||||
DATABASE__USE: postgresql
|
|
||||||
DATABASE__POSTGRESQL__URL: postgresql+asyncpg://langbot_runtime:${POSTGRES_RUNTIME_PASSWORD}@postgres:5432/langbot
|
|
||||||
DATABASE__CLOUD_MIGRATION__OPERATOR_DSN_ENV: LANGBOT_CLOUD_MIGRATION_DSN
|
|
||||||
LANGBOT_CLOUD_MIGRATION_DSN: postgresql://langbot_operator:${POSTGRES_OPERATOR_PASSWORD}@postgres:5432/langbot
|
|
||||||
VDB__USE: pgvector
|
|
||||||
VDB__PGVECTOR__USE_BUSINESS_DATABASE: "true"
|
|
||||||
VDB__PGVECTOR__ALLOWED_DIMENSIONS: "384,512,768,1024,1536"
|
|
||||||
PLUGIN__ENABLE: "true"
|
|
||||||
PLUGIN__RUNTIME_WS_URL: ws://plugin-runtime:5400/control/ws
|
|
||||||
PLUGIN__DISPLAY_PLUGIN_DEBUG_URL: wss://cloud.langbot.app/plugin/debug/ws
|
|
||||||
PLUGIN__WORKER__MAX_CPUS: "0.25"
|
|
||||||
PLUGIN__WORKER__MAX_MEMORY_MB: "256"
|
|
||||||
PLUGIN__WORKER__MAX_PIDS: "128"
|
|
||||||
PLUGIN__WORKER__MAX_WORKERS: "16"
|
|
||||||
PLUGIN__WORKER__MAX_TOTAL_CPUS: "4.0"
|
|
||||||
PLUGIN__WORKER__MAX_TOTAL_MEMORY_MB: "4096"
|
|
||||||
PLUGIN__WORKER__REQUIRE_HARD_LIMITS: "true"
|
|
||||||
LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN: ${PLUGIN_RUNTIME_CONTROL_TOKEN}
|
|
||||||
# Cloud v2 currently grants no managed Box capability. Keep the shared
|
|
||||||
# runtime deployed but disable Core integration until a hard-quota-capable
|
|
||||||
# backend can satisfy the fail-closed Cloud readiness contract.
|
|
||||||
BOX__ENABLED: "false"
|
|
||||||
BOX__BACKEND: nsjail
|
|
||||||
BOX__RUNTIME__ENDPOINT: ws://box:5410
|
|
||||||
BOX__ADMISSION__REQUIRED: "true"
|
|
||||||
BOX__ADMISSION__LOGICAL_SESSION_ID: global
|
|
||||||
BOX__ADMISSION__REQUIRED_BACKEND: nsjail
|
|
||||||
BOX__ADMISSION__MAX_SESSIONS: "1"
|
|
||||||
BOX__ADMISSION__MAX_MANAGED_PROCESSES: "0"
|
|
||||||
BOX__ADMISSION__CPUS: "0.25"
|
|
||||||
BOX__ADMISSION__MEMORY_MB: "256"
|
|
||||||
BOX__ADMISSION__WORKSPACE_QUOTA_MB: "256"
|
|
||||||
BOX__LOCAL__HOST_ROOT: /app/data/box
|
|
||||||
BOX__LOCAL__DEFAULT_WORKSPACE: /app/data/box
|
|
||||||
BOX__LOCAL__ALLOWED_MOUNT_ROOTS: /app/data/box
|
|
||||||
LANGBOT_BOX_CONTROL_TOKEN: ${BOX_CONTROL_TOKEN}
|
|
||||||
MCP__STDIO__ENABLED: "false"
|
|
||||||
LANGBOT_SPACE_CONTROL_PLANE_URL: https://space.langbot.app
|
|
||||||
LANGBOT_SPACE_CONTROL_PLANE_TOKEN: ${CLOUD_V2_CONTROL_PLANE_TOKEN}
|
|
||||||
LANGBOT_SPACE_CONTROL_PLANE_PUBLIC_KEY: ${CLOUD_V2_MANIFEST_PUBLIC_KEY}
|
|
||||||
LANGBOT_SPACE_CONTROL_PLANE_KEY_ID: ${CLOUD_V2_MANIFEST_KEY_ID}
|
|
||||||
SPACE__URL: https://space.langbot.app
|
|
||||||
depends_on:
|
|
||||||
postgres: {condition: service_healthy}
|
|
||||||
networks: [internal]
|
|
||||||
|
|
||||||
plugin-runtime:
|
|
||||||
image: rockchin/langbot:${LANGBOT_IMAGE_TAG}
|
|
||||||
container_name: langbot-cloud-plugin-runtime
|
|
||||||
restart: unless-stopped
|
|
||||||
command: [uv, run, python, -m, langbot_plugin.cli.__init__, rt]
|
|
||||||
environment:
|
|
||||||
LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN: ${PLUGIN_RUNTIME_CONTROL_TOKEN}
|
|
||||||
volumes:
|
|
||||||
- plugin-data:/app/data
|
|
||||||
- /sys/fs/cgroup:/sys/fs/cgroup:rw
|
|
||||||
cgroup: host
|
|
||||||
privileged: true
|
|
||||||
expose: ["5400"]
|
|
||||||
networks: [internal]
|
|
||||||
|
|
||||||
box:
|
|
||||||
image: rockchin/langbot:${LANGBOT_IMAGE_TAG}
|
|
||||||
container_name: langbot-cloud-box
|
|
||||||
restart: unless-stopped
|
|
||||||
command: [uv, run, lbp, box, --host, 0.0.0.0, --ws-control-port, "5410"]
|
|
||||||
environment:
|
|
||||||
LANGBOT_BOX_CONTROL_TOKEN: ${BOX_CONTROL_TOKEN}
|
|
||||||
LANGBOT_BOX_ROOT: /app/data/box
|
|
||||||
volumes:
|
|
||||||
- box-data:/app/data/box
|
|
||||||
- /sys/fs/cgroup:/sys/fs/cgroup:rw
|
|
||||||
cgroup: host
|
|
||||||
privileged: true
|
|
||||||
expose: ["5410"]
|
|
||||||
networks: [internal]
|
|
||||||
|
|
||||||
core:
|
|
||||||
image: rockchin/langbot-cloud-core:${LANGBOT_IMAGE_TAG}
|
|
||||||
container_name: langbot-cloud-core
|
|
||||||
restart: unless-stopped
|
|
||||||
environment: *core-env
|
|
||||||
volumes:
|
|
||||||
- core-data:/app/data
|
|
||||||
- box-data:/app/data/box
|
|
||||||
depends_on:
|
|
||||||
postgres: {condition: service_healthy}
|
|
||||||
redis: {condition: service_healthy}
|
|
||||||
plugin-runtime: {condition: service_started}
|
|
||||||
box: {condition: service_started}
|
|
||||||
expose: ["5300"]
|
|
||||||
healthcheck:
|
|
||||||
test: [CMD-SHELL, "python -c 'import urllib.request; urllib.request.urlopen(\"http://127.0.0.1:5300/healthz\", timeout=3)'" ]
|
|
||||||
interval: 10s
|
|
||||||
timeout: 5s
|
|
||||||
retries: 30
|
|
||||||
start_period: 30s
|
|
||||||
networks: [internal, shared-network]
|
|
||||||
|
|
||||||
networks:
|
|
||||||
internal:
|
|
||||||
shared-network:
|
|
||||||
external: true
|
|
||||||
|
|
||||||
volumes:
|
|
||||||
postgres-data:
|
|
||||||
redis-data:
|
|
||||||
plugin-data:
|
|
||||||
box-data:
|
|
||||||
core-data:
|
|
||||||
@@ -14,8 +14,8 @@ services:
|
|||||||
restart: on-failure
|
restart: on-failure
|
||||||
environment:
|
environment:
|
||||||
- TZ=Asia/Shanghai
|
- TZ=Asia/Shanghai
|
||||||
# Shared with the langbot service and sent only as a WebSocket handshake
|
# Optional. Leave unset on both OSS services, or set the same value on
|
||||||
# header. Generate with: openssl rand -hex 32
|
# both to protect the control WebSocket. Generate with: openssl rand -hex 32
|
||||||
- LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}
|
- LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}
|
||||||
# Process-wide admission for every asyncio.to_thread() call.
|
# Process-wide admission for every asyncio.to_thread() call.
|
||||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS=${LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS:-8}
|
- LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS=${LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS:-8}
|
||||||
@@ -77,8 +77,7 @@ services:
|
|||||||
restart: on-failure
|
restart: on-failure
|
||||||
environment:
|
environment:
|
||||||
- TZ=Asia/Shanghai
|
- TZ=Asia/Shanghai
|
||||||
# Must match langbot_plugin_runtime. Empty/missing values make the
|
# Optional. Leave unset on both OSS services, or match plugin Runtime.
|
||||||
# external control channel fail closed.
|
|
||||||
- LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}
|
- LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}
|
||||||
# Must match the value supplied to langbot_box. The token is sent only
|
# Must match the value supplied to langbot_box. The token is sent only
|
||||||
# in WebSocket handshake headers, never in URLs or action payloads.
|
# in WebSocket handshake headers, never in URLs or action payloads.
|
||||||
|
|||||||
@@ -81,10 +81,10 @@ This log records implementation choices made while delivering the Workspace arch
|
|||||||
- Decision: The MCP ASGI mount authenticates the API key once, binds an immutable per-request `RequestContext`, and every tool checks a fixed permission before calling tenant services with that same context.
|
- Decision: The MCP ASGI mount authenticates the API key once, binds an immutable per-request `RequestContext`, and every tool checks a fixed permission before calling tenant services with that same context.
|
||||||
- Reason: Authenticating the transport without propagating Workspace identity into tool calls would leave the direct service path globally scoped.
|
- Reason: Authenticating the transport without propagating Workspace identity into tool calls would leave the direct service path globally scoped.
|
||||||
|
|
||||||
### Released SDK protocol is pinned from PyPI
|
### Unreleased SDK protocol is pinned reproducibly without publishing
|
||||||
|
|
||||||
- Decision: The SDK tenancy protocol is released as `langbot-plugin==0.5.0` and LangBot pins that exact registry version.
|
- Decision: The SDK tenancy protocol is versioned as 0.4.18. This task does not create a GitHub release or publish PyPI because the user authorized pushing code, not a package release. After the SDK feature branch is final, LangBot's feature branch temporarily pins the exact pushed SDK Git commit. Before merging to master, the release gate is to publish `langbot-plugin==0.4.18` and replace the Git pin with the registry pin.
|
||||||
- Reason: The final PyPI release contains the complete tenant action context and shared Runtime hardening, while the exact version pin keeps production installs reproducible.
|
- Reason: The current registry release does not contain the complete tenant action context and shared Runtime hardening. An exact Git commit is reproducible and keeps the feature branch testable without expanding release authority.
|
||||||
|
|
||||||
### Cloud directory writes stay outside Core
|
### Cloud directory writes stay outside Core
|
||||||
|
|
||||||
@@ -103,11 +103,11 @@ This log records implementation choices made while delivering the Workspace arch
|
|||||||
- Decision: New Core JWTs require `iss=langbot-core`, an audience derived from the immutable instance UUID, and an expiry. Legacy community tokens are accepted only when they have the historical issuer, carry no audience, and the active policy is the OSS singleton policy.
|
- Decision: New Core JWTs require `iss=langbot-core`, an audience derived from the immutable instance UUID, and an expiry. Legacy community tokens are accepted only when they have the historical issuer, carry no audience, and the active policy is the OSS singleton policy.
|
||||||
- Reason: A token issued by one instance must not authenticate against another instance that happens to share a secret, and a compatibility decoder must not become an alternate path around the SaaS trust boundary.
|
- Reason: A token issued by one instance must not authenticate against another instance that happens to share a secret, and a compatibility decoder must not become an alternate path around the SaaS trust boundary.
|
||||||
|
|
||||||
### Runtime control transports authenticate before protocol dispatch
|
### Runtime control transports support opt-in shared-secret authentication
|
||||||
|
|
||||||
- Decision: External Plugin Runtime and Box WebSocket control channels require independent strong shared secrets in handshake headers. Locally managed child processes receive ephemeral secrets through their environment; secrets are not placed in URLs, process arguments, request payloads, or logs. Box additionally binds the first authenticated control channel to one trusted instance. Plugin Runtime debug and control credentials remain separate.
|
- Decision: OSS external Plugin Runtime and Box WebSocket control channels preserve tokenless standalone compatibility when the corresponding control token is unset. When a Runtime configures a token, it validates the independent shared secret in the handshake before protocol dispatch. Locally managed child processes still receive ephemeral secrets through their environment; secrets are not placed in URLs, process arguments, request payloads, or logs. Box additionally pins the first control channel to one declared instance identity. Plugin Runtime debug and control credentials remain separate.
|
||||||
- Reason: Workspace context inside an RPC payload is not trustworthy until the transport peer itself is authenticated. Separating control and debug credentials also limits accidental privilege reuse.
|
- Reason: Local OSS development must remain backward compatible, while exposed or shared Runtime endpoints can opt into transport authentication. Separating control and debug credentials also limits accidental privilege reuse.
|
||||||
- Deployment consequence: Docker Compose and Kubernetes wire one shared secret to each host/runtime pair. An empty external-runtime secret fails startup instead of silently exposing an unauthenticated socket.
|
- Deployment consequence: Docker Compose and Kubernetes should wire one strong shared secret to each host/runtime pair. Both sides must use the same value for protection to be effective; a Runtime configured with a token rejects clients that omit it or send a different value.
|
||||||
|
|
||||||
### Dashboard WebSocket sessions are tenant runtime objects
|
### Dashboard WebSocket sessions are tenant runtime objects
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,476 @@
|
|||||||
|
# 模型思考控制设计方案
|
||||||
|
|
||||||
|
> 日期:2026-07-31
|
||||||
|
> 状态:Phase 1 已审核并实现
|
||||||
|
> 范围:LangBot 主仓库的模型配置、LiteLLM 请求层、Local Agent、Web 管理面板、监控与测试
|
||||||
|
|
||||||
|
## 1. 结论
|
||||||
|
|
||||||
|
建议为 LangBot 增加一套与厂商参数解耦的“思考策略”模型,并明确区分三个概念:
|
||||||
|
|
||||||
|
1. **思考能力**:模型是否支持思考,以及支持开关、档位还是 token 预算。
|
||||||
|
2. **思考策略**:一次请求选择厂商默认、关闭、开启或指定思考档位。
|
||||||
|
3. **思考展示**:是否把模型返回的思考内容展示给最终用户。
|
||||||
|
|
||||||
|
现有 `remove-think` 只属于第 3 类。它会过滤输出,但不会阻止模型思考,也不会降低思考 token、费用或延迟。新能力不应复用或改写这个字段。
|
||||||
|
|
||||||
|
推荐实现原则:
|
||||||
|
|
||||||
|
- 默认值为 `provider_default`,不向上游增加任何新参数,现有模型行为完全不变。
|
||||||
|
- 用户显式选择的策略必须被准确执行;无法准确执行时返回明确错误,不静默降级。
|
||||||
|
- LangBot 内部只保存统一策略,Provider 请求层负责翻译成各厂商参数。
|
||||||
|
- `extra_args` 保留为高级逃生口,但不能成为主 UI 的思考配置方式。
|
||||||
|
- 模型页只管理并展示能力;可写策略归属于 Local Agent 流水线,同一模型可在不同业务中使用不同思考量。
|
||||||
|
- 原始 reasoning 数据与展示文本分开保存,保证多轮对话、工具调用和签名字段不丢失。
|
||||||
|
|
||||||
|
## 2. 调研结论
|
||||||
|
|
||||||
|
### 2.1 可验证资料
|
||||||
|
|
||||||
|
本次结论基于以下可验证来源:
|
||||||
|
|
||||||
|
- OpenAI 官方 Reasoning Guide:`reasoning.effort` 的可选值由模型决定,可包括 `none`、`minimal`、`low`、`medium`、`high`、`xhigh`、`max`;低档位偏向低延迟和低 token,高档位偏向质量。
|
||||||
|
- https://developers.openai.com/api/docs/guides/reasoning#reasoning-effort
|
||||||
|
- LangBot 锁定的 LiteLLM `1.88.1` 实现。`uv.lock` 已锁定该版本,本地缓存中的适配代码可以确认 LangBot 实际依赖所支持的翻译行为。
|
||||||
|
- LangBot 当前实现:模型级 `extra_args` 会在 `LiteLLMRequester._build_completion_args()` 中直接合并到 `acompletion()` 参数。
|
||||||
|
|
||||||
|
Anthropic、Google 和 LiteLLM 的官方文档域名在本次环境中被浏览器策略禁止访问,因此下表中这些厂商的结论以 LiteLLM `1.88.1` 实际适配代码为准。实施前应再用对应厂商官方文档做一次参数范围核验,尤其是模型代际和允许值。
|
||||||
|
|
||||||
|
### 2.2 厂商差异矩阵
|
||||||
|
|
||||||
|
| Provider / 生态 | 可控制能力 | LiteLLM 1.88.1 统一入口 | 关键限制 | 建议支持级别 |
|
||||||
|
| --- | --- | --- | --- | --- |
|
||||||
|
| OpenAI | 思考档位,部分模型支持 `none` | `reasoning_effort` | 每个模型支持的档位不同,不能把 `none` 当成通用能力 | 首批完整支持 |
|
||||||
|
| Anthropic | 旧模型使用 extended thinking + token budget;新模型可用 adaptive thinking + effort | `reasoning_effort` 或 `thinking` | `none` 表示不发送 thinking;新旧模型的映射不同 | 首批完整支持 |
|
||||||
|
| Gemini | 2.x 主要映射为 `thinkingBudget`;3.x 主要映射为 `thinkingLevel` | `reasoning_effort` 或 `thinking` | Gemini 3 的 `none` 可能只能降到最低档,不能保证真正关闭 | 首批支持,但严格限制关闭语义 |
|
||||||
|
| DeepSeek | 开启/关闭;当前适配不支持预算档位 | `thinking={type: enabled}`;非 `none` effort 会映射成开启 | 多轮思考模式要求回传 `reasoning_content` | 首批开关支持 |
|
||||||
|
| xAI | 思考档位 | `reasoning_effort` | 仅 reasoning-capable 模型接受 | 首批完整支持 |
|
||||||
|
| Ollama | `think` 布尔值;部分模型接受 low/medium/high | `reasoning_effort` | 非 gpt-oss 模型的档位可能退化为布尔开关 | 首批支持,按模型能力裁剪 UI |
|
||||||
|
| OpenRouter | 聚合多厂商的 reasoning 参数 | `reasoning_effort`、`thinking` | 实际能力由路由后的模型决定 | 首批支持,能力未知时要求测试 |
|
||||||
|
| Volcengine / Doubao | `thinking.type` 支持 enabled/disabled/auto | LiteLLM `volcengine` 适配器支持 `thinking` | LangBot 当前 manifest 使用 `openai`,不会进入该适配器 | 第二批,先修正路由并回归 |
|
||||||
|
| Bailian / Qwen | 厂商兼容接口有独立思考开关/预算 | LiteLLM `dashscope` 适配器目前未提供统一 reasoning 映射 | LangBot 当前 manifest 使用 `openai`,只能通过高级参数透传 | 第二批,实施前核对官方字段 |
|
||||||
|
| 其他 OpenAI-compatible 网关 | 取决于网关 | 尝试标准 `reasoning_effort` | 不能仅凭模型名推断完整能力 | 保守支持,默认不自动开启 |
|
||||||
|
|
||||||
|
### 2.3 对 LangBot 的直接含义
|
||||||
|
|
||||||
|
不能把这个功能实现成单一 `enable_thinking: bool`,原因如下:
|
||||||
|
|
||||||
|
- 有的模型只有开关,有的模型只有档位,有的模型允许精确 token 预算。
|
||||||
|
- 有的模型本身始终推理,只能降低思考量,无法真正关闭。
|
||||||
|
- 同一个通用档位在不同厂商会映射成不同的实际预算。
|
||||||
|
- 聚合网关和自定义 OpenAI-compatible 服务无法可靠地通过模型名识别能力。
|
||||||
|
- “不展示思考内容”不等于“关闭思考”。
|
||||||
|
|
||||||
|
## 3. 当前项目现状
|
||||||
|
|
||||||
|
### 3.1 已有能力
|
||||||
|
|
||||||
|
- `LLMModel.extra_args` 是 JSON 字段,Web 端已有通用高级参数编辑器。
|
||||||
|
- `LiteLLMRequester` 会按“模型级 `extra_args`,再调用级 `extra_args`”的顺序合并参数。
|
||||||
|
- LiteLLM 已统一处理多个 Provider 的 `reasoning_effort`、`thinking` 和返回的 `reasoning_content`。
|
||||||
|
- `LocalAgentRunner` 的非流式、流式、工具调用和 fallback 路径都经过 `RuntimeProvider.invoke_llm*()`。
|
||||||
|
- `remove-think` 已能控制 `<think>` 或独立 reasoning 内容是否进入展示文本。
|
||||||
|
- Gemini 工具调用所需的 `provider_specific_fields` / thought signature 已有保留逻辑和单元测试。
|
||||||
|
|
||||||
|
### 3.2 现有缺口
|
||||||
|
|
||||||
|
- 管理员只能手写 `extra_args`,没有统一语义、能力提示和校验。
|
||||||
|
- `remove-think` 名称容易被误解为关闭模型思考。
|
||||||
|
- 模型扫描只识别 `vision` 和 `func_call`,没有 reasoning 能力。
|
||||||
|
- 当前返回处理会把 `reasoning_content` 拼进 `<think>` 文本后删除原字段,可能损失多轮思考所需的结构化数据。
|
||||||
|
- DeepSeek 思考模式需要在后续轮次回传 `reasoning_content`,当前链路不能保证完整保留。
|
||||||
|
- Pipeline 只能选择模型,不能针对业务覆盖模型的思考策略。
|
||||||
|
- 监控只记录总输入/输出 token,没有单独展示 reasoning token。
|
||||||
|
- 部分 Provider manifest 仍声明为通用 `openai`,导致 LiteLLM 的厂商专用翻译器不会生效。
|
||||||
|
|
||||||
|
### 3.3 预计改动地图
|
||||||
|
|
||||||
|
| 层 | 主要文件 | 责任 |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| 持久化 | `src/langbot/pkg/entity/persistence/model.py`、`src/langbot/pkg/persistence/alembic/versions/` | 新增 `reasoning_config` JSON 列和 Alembic 迁移 |
|
||||||
|
| 模型服务 | `src/langbot/pkg/api/http/service/model.py` | CRUD 校验、冲突检测、测试模型时使用统一策略 |
|
||||||
|
| HTTP 控制器 | `src/langbot/pkg/api/http/controller/groups/provider/models.py` | 继续复用现有模型路由,不新增平行 API |
|
||||||
|
| 模型管理 | `src/langbot/pkg/provider/modelmgr/modelmgr.py` | 临时模型、数据库模型与扫描结果加载新字段 |
|
||||||
|
| 请求抽象 | `src/langbot/pkg/provider/modelmgr/requester.py` | 定义能力查询和 reasoning 参数构建接口 |
|
||||||
|
| LiteLLM 适配 | `src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py` | 能力识别、策略翻译、参数合并、reasoning 返回保留 |
|
||||||
|
| Provider manifest | `src/langbot/pkg/provider/modelmgr/requesters/*.yaml` | 必要时修正 Provider 路由;相关变更放到独立阶段 |
|
||||||
|
| Agent 调用 | `src/langbot/pkg/provider/runners/localagent.py` | 所有非流式、流式、工具调用、fallback 路径传递统一策略 |
|
||||||
|
| Pipeline 元数据 | `src/langbot/templates/metadata/pipeline/ai.yaml` | 第二阶段加入 Pipeline 级覆盖 |
|
||||||
|
| 输出配置 | `src/langbot/templates/metadata/pipeline/output.yaml` | 保留键名,澄清 `remove-think` 只控制展示 |
|
||||||
|
| Web 类型/API | `web/src/app/infra/entities/api/index.ts`、`web/src/app/infra/http/BackendClient.ts` | 增加配置与能力响应类型 |
|
||||||
|
| 模型 UI | `web/src/app/home/components/models-dialog/` | 能力标记、策略控件、校验、模型测试 |
|
||||||
|
| i18n | `web/src/i18n/locales/` | 至少补齐英文、简体中文及项目已有覆盖语言 |
|
||||||
|
| 测试 | `tests/unit_tests/provider/`、`web/tests/` | 翻译、服务、流式 round-trip、前端状态测试 |
|
||||||
|
|
||||||
|
Phase 1 不修改 `langbot-plugin-sdk` 的公共实体或运行时协议。现有 `provider_message.Message.provider_specific_fields` 已可承载 Provider 原始 reasoning 数据;只有后续要把 reasoning 升级为跨插件公开实体时,才需要跨仓库 SDK 变更。
|
||||||
|
|
||||||
|
## 4. 领域模型
|
||||||
|
|
||||||
|
### 4.1 统一策略
|
||||||
|
|
||||||
|
新增 `ReasoningConfig`,保存于 LLM 模型,Pipeline 可提供同结构覆盖。产品层只暴露一个离散档位:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"level": "provider_default"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
字段定义:
|
||||||
|
|
||||||
|
| 字段 | 类型 | 含义 |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| `level` | `provider_default \| disabled \| enabled \| minimal \| low \| medium \| high \| xhigh \| max` | 同时表达开关和思考强度 |
|
||||||
|
|
||||||
|
校验规则:
|
||||||
|
|
||||||
|
- `provider_default`:不发送任何 reasoning 参数,保持厂商和模型默认行为。
|
||||||
|
- `disabled`:明确关闭;仅当模型可真正关闭时允许保存/运行。
|
||||||
|
- `enabled`:明确开启,但由 Provider 决定具体强度,适用于只有开关的模型。
|
||||||
|
- `minimal` 到 `max`:明确开启,并指定强度;仅允许选择模型实际支持的档位。
|
||||||
|
- 厂商的 `auto` 统一映射为 `provider_default`,不再增加一个重复状态。
|
||||||
|
- 精确 token 预算不进入主数据结构。少数需要预算的场景继续通过高级参数配置,并由模型测试接口校验。
|
||||||
|
|
||||||
|
### 4.2 能力描述
|
||||||
|
|
||||||
|
沿用现有 `LLMModel.abilities`,新增 `reasoning` 能力标记。同时由后端在 API 返回中计算只读的 `reasoning_capabilities`:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"supported": true,
|
||||||
|
"controls": ["toggle", "effort"],
|
||||||
|
"efforts": ["none", "low", "medium", "high"],
|
||||||
|
"can_disable": true,
|
||||||
|
"source": "litellm"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
设计约束:
|
||||||
|
|
||||||
|
- `abilities` 仍是用户可编辑的粗粒度能力,符合现有 `vision`、`func_call` 模式。
|
||||||
|
- `reasoning_capabilities` 不持久化,优先从 LiteLLM 模型元数据计算,避免模型升级后数据库残留过期能力。
|
||||||
|
- 无法识别的自定义模型返回 `supported: null`、`source: unknown`,不猜测。
|
||||||
|
- 用户可手动添加 `reasoning` ability,但未知能力模型必须先通过“测试模型”验证显式策略。
|
||||||
|
- UI 只展示后端声明可用的控件;未知模型保留 Provider Default 和高级参数入口。
|
||||||
|
|
||||||
|
### 4.3 持久化
|
||||||
|
|
||||||
|
在 `llm_models` 表新增 JSON 列:
|
||||||
|
|
||||||
|
```text
|
||||||
|
reasoning_config JSON NOT NULL DEFAULT {"level":"provider_default"}
|
||||||
|
```
|
||||||
|
|
||||||
|
使用 Alembic 新迁移,不修改冻结的 legacy migration。
|
||||||
|
|
||||||
|
该列作为已实现版本的兼容字段保留;新的模型页不再提供写入口,Local Agent 请求以流水线中按模型 UUID 保存的策略为准。
|
||||||
|
|
||||||
|
不建议把内部策略塞进 `extra_args`,原因是当前 `extra_args` 会原样发送给 LiteLLM;使用保留键会让内部元数据泄漏到上游,并使高级参数与产品配置难以区分。
|
||||||
|
|
||||||
|
## 5. 配置优先级与请求流程
|
||||||
|
|
||||||
|
### 5.1 优先级
|
||||||
|
|
||||||
|
```text
|
||||||
|
Pipeline 当前候选模型策略
|
||||||
|
↓ 缺少配置时固定为 provider_default
|
||||||
|
Provider / 模型默认行为
|
||||||
|
```
|
||||||
|
|
||||||
|
请求参数合并顺序:
|
||||||
|
|
||||||
|
```text
|
||||||
|
基础参数
|
||||||
|
-> 模型 extra_args
|
||||||
|
-> 调用级 extra_args
|
||||||
|
-> 统一 reasoning 策略翻译结果(最后应用)
|
||||||
|
```
|
||||||
|
|
||||||
|
统一策略最后应用,可以确保流水线行为不受模型页历史设置影响。为了避免用户困惑,保存和测试时要检测 `extra_args` 中的冲突字段;当 `level != provider_default` 时,发现以下字段应直接报错:
|
||||||
|
|
||||||
|
- `reasoning_effort`
|
||||||
|
- `thinking`
|
||||||
|
- `reasoning`
|
||||||
|
- `extra_body` 内已知的 `thinking`、`enable_thinking`、`thinking_budget` 等字段
|
||||||
|
|
||||||
|
当 `level == provider_default` 时继续允许这些高级参数,保证旧配置兼容。
|
||||||
|
|
||||||
|
### 5.2 翻译层
|
||||||
|
|
||||||
|
在 `pkg/provider/modelmgr/` 内新增独立的 reasoning 规范化模块,职责是:
|
||||||
|
|
||||||
|
1. 读取当前流水线候选模型的请求级策略。
|
||||||
|
2. 查询 `ProviderAPIRequester.get_reasoning_capabilities(model)`。
|
||||||
|
3. 严格校验策略是否可以准确执行。
|
||||||
|
4. 生成 LiteLLM 参数,不直接发 HTTP。
|
||||||
|
5. 返回可观测的“最终生效策略”供日志和测试使用。
|
||||||
|
|
||||||
|
建议接口:
|
||||||
|
|
||||||
|
```python
|
||||||
|
class ProviderAPIRequester:
|
||||||
|
def get_reasoning_capabilities(self, model: RuntimeLLMModel) -> ReasoningCapabilities: ...
|
||||||
|
|
||||||
|
def build_reasoning_args(
|
||||||
|
self,
|
||||||
|
model: RuntimeLLMModel,
|
||||||
|
config: ReasoningConfig,
|
||||||
|
) -> dict[str, Any]: ...
|
||||||
|
```
|
||||||
|
|
||||||
|
LiteLLMRequester 默认优先生成统一参数:
|
||||||
|
|
||||||
|
- 强度档位:`reasoning_effort=<level>`
|
||||||
|
- 仅开启:`thinking={"type":"enabled"}` 或 Provider 等价参数
|
||||||
|
- 关闭:优先 `reasoning_effort="none"`
|
||||||
|
- 高级参数中的精确预算:`thinking={"type":"enabled","budget_tokens":N}`
|
||||||
|
|
||||||
|
Provider 特例只放在 requester 翻译层,不进入 Pipeline 或平台适配器。
|
||||||
|
|
||||||
|
### 5.3 Provider 特例
|
||||||
|
|
||||||
|
- **Gemini 3**:如果 LiteLLM 能力表不能确认真正关闭,`disabled` 必须报“不支持关闭,可选择 Provider Default 或最低档”,不能把 `none` 静默映射成 low/minimal。
|
||||||
|
- **DeepSeek**:所有非 `none` 档位最终都只是开启。能力 API 只返回 `toggle`,UI 不显示档位;多轮必须保存并回传 `reasoning_content`。
|
||||||
|
- **Ollama**:仅对明确支持等级的模型展示 effort;其他模型只展示开关。
|
||||||
|
- **OpenRouter**:以路由后的模型能力为准。模型未知时允许 Provider Default,显式策略必须通过测试接口。
|
||||||
|
- **Volcengine**:使用 `thinking.type=enabled/disabled/auto`。应先让该 requester 进入 LiteLLM `volcengine` 适配器,或增加等价的明确翻译,不能依赖模型名。
|
||||||
|
- **Bailian/Qwen**:作为第二批 Provider 专用翻译。实施前核对官方字段、模型范围、预算上下限和流式返回结构,不凭经验写接口。
|
||||||
|
|
||||||
|
## 6. 返回数据与思考展示
|
||||||
|
|
||||||
|
### 6.1 保留原始 reasoning
|
||||||
|
|
||||||
|
当前 `LiteLLMRequester` 会读取 `reasoning_content`,将其拼接成 `<think>` 文本,再删除原字段。建议改为:
|
||||||
|
|
||||||
|
```text
|
||||||
|
上游 reasoning_content
|
||||||
|
├─ 原样保存在 Message.provider_specific_fields.reasoning_content
|
||||||
|
└─ 根据 remove-think 决定是否渲染为 <think>...</think>
|
||||||
|
```
|
||||||
|
|
||||||
|
流式路径需要在 accumulator 中分别累计 `content` 与 `reasoning_content`,最终消息必须携带结构化 reasoning。不能只依赖已经渲染的 `<think>` 文本反向解析。
|
||||||
|
|
||||||
|
这样可以同时满足:
|
||||||
|
|
||||||
|
- `remove-think=true` 时用户看不到思考内容,但多轮协议仍能回传必要数据。
|
||||||
|
- `remove-think=false` 时保持当前用户体验。
|
||||||
|
- DeepSeek 多轮 thinking 不丢上下文。
|
||||||
|
- Gemini thought signature、Anthropic thinking block 等 Provider 字段可以继续按结构化方式 round-trip。
|
||||||
|
|
||||||
|
### 6.2 现有字段处理
|
||||||
|
|
||||||
|
保留数据库和 Pipeline 配置键 `remove-think`,避免破坏兼容。Web 文案改为更准确的:
|
||||||
|
|
||||||
|
- 中文:`向用户展示思考过程`
|
||||||
|
- 英文:`Show reasoning process`
|
||||||
|
|
||||||
|
UI 使用正向开关,保存时转换回 `remove-think = !showReasoning`。文案必须强调它只影响展示,不影响模型是否思考、token 或费用。
|
||||||
|
|
||||||
|
## 7. Web 管理面板
|
||||||
|
|
||||||
|
### 7.1 模型编辑
|
||||||
|
|
||||||
|
模型页只承担能力管理和只读展示:
|
||||||
|
|
||||||
|
1. `Reasoning` ability 复选框与 Vision、Function Calling 并列,供无法自动识别的自定义模型手动声明能力。
|
||||||
|
2. 模型卡片使用简短图标或 badge 标识 reasoning 能力。
|
||||||
|
3. 模型页不提供可写思考挡位,避免模型默认值与流水线策略形成两个控制源。
|
||||||
|
|
||||||
|
### 7.2 Local Agent 流水线策略
|
||||||
|
|
||||||
|
在 Local Agent 的主模型和每一个 fallback 模型下分别显示紧凑离散滑杆:
|
||||||
|
|
||||||
|
1. `Provider 默认` 始终为首个选项;选择它时不向上游增加任何思考参数。
|
||||||
|
2. 完整档位顺序为:`Provider 默认 / 关闭 / 开启 / 最低 / 低 / 中 / 高 / 极高 / 最大`。
|
||||||
|
3. 前端只渲染后端为该模型返回的可用档位;仅开关模型显示 `Provider 默认 / 关闭 / 开启`。
|
||||||
|
4. 模型不能真正关闭时不提供 `关闭`;能力未知时只显示不可调的 `Provider 默认`。
|
||||||
|
5. 主模型和 fallback 分别保存策略,切换候选模型时不会把一个模型的挡位错误应用到另一个模型。
|
||||||
|
6. Dify、Coze、Langflow、n8n 等外部 Runner 不显示该控件,因为 LangBot 不直接发起其内部模型请求。
|
||||||
|
|
||||||
|
流水线配置保持旧格式兼容,并在模型选择对象中增加按 UUID 保存的映射:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"model": {
|
||||||
|
"primary": "primary-model-uuid",
|
||||||
|
"fallbacks": ["fallback-model-uuid"],
|
||||||
|
"reasoning": {
|
||||||
|
"primary-model-uuid": "high"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
`provider_default` 不写入映射;缺少 `reasoning` 的旧流水线天然等价于全部使用 Provider 默认。
|
||||||
|
|
||||||
|
滑杆交互要求:轨道使用现有主色和中性灰,不使用渐变;当前档位同时显示文字;支持键盘方向键和正确的 ARIA value text;窄屏下不溢出。
|
||||||
|
|
||||||
|
### 7.3 i18n
|
||||||
|
|
||||||
|
新增文案至少覆盖 `en_US`、`zh_Hans`;`ja_JP` 在模型面板现有同类字段已覆盖时同步补齐。不要把厂商参数名直接作为用户文案。
|
||||||
|
|
||||||
|
## 8. API、MCP 与 Skill
|
||||||
|
|
||||||
|
### 8.1 HTTP API
|
||||||
|
|
||||||
|
模型 CRUD 增加:
|
||||||
|
|
||||||
|
- 请求字段:`reasoning_config`
|
||||||
|
- 响应字段:`reasoning_config`
|
||||||
|
- 只读字段:`reasoning_capabilities`
|
||||||
|
|
||||||
|
模型测试接口必须使用与真实请求完全相同的规范化和翻译逻辑,并在失败时返回可操作错误,例如:
|
||||||
|
|
||||||
|
```text
|
||||||
|
Model gemini-3-... cannot disable reasoning.
|
||||||
|
Supported controls: effort=[low, medium, high].
|
||||||
|
```
|
||||||
|
|
||||||
|
可选增加只读调试信息,仅在测试接口返回:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"effective_reasoning": {
|
||||||
|
"level": "low",
|
||||||
|
"translated_keys": ["reasoning_effort"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
不得返回 API key、完整请求正文或原始思考内容。
|
||||||
|
|
||||||
|
### 8.2 MCP 与技能
|
||||||
|
|
||||||
|
当前 MCP 仅列出模型 Provider,没有完整模型 CRUD 工具。如果本次不新增 agent-accessible HTTP 操作,则无需强行新增 MCP 工具。
|
||||||
|
|
||||||
|
如果后续让 Agent 修改模型思考策略,则必须同一提交更新:
|
||||||
|
|
||||||
|
- `src/langbot/pkg/api/mcp/server.py`
|
||||||
|
- 对应的 `skills/` 文档
|
||||||
|
- 参数 schema 和安全说明
|
||||||
|
|
||||||
|
## 9. 监控与可观测性
|
||||||
|
|
||||||
|
控制思考量后,管理员需要判断质量、延迟和成本是否值得。建议第二阶段增加:
|
||||||
|
|
||||||
|
- `reasoning_tokens`:从 `completion_tokens_details.reasoning_tokens` 或 Provider 等价字段提取。
|
||||||
|
- `effective_reasoning_level`:记录规范化后的生效档位,不记录原始思考内容。
|
||||||
|
- 模型监控页展示输入 token、可见输出 token、reasoning token、总延迟。
|
||||||
|
- Provider 不返回细分 token 时显示未知,不推算。
|
||||||
|
|
||||||
|
安全要求:日志、监控、debug API 默认都不得记录 reasoning 原文。思考内容可能包含敏感信息或系统提示,不应因为新增配置而扩大持久化范围。
|
||||||
|
|
||||||
|
## 10. 兼容与迁移
|
||||||
|
|
||||||
|
### 10.1 数据迁移
|
||||||
|
|
||||||
|
- 所有现有 LLM 记录迁移为 `{"level":"provider_default"}`。
|
||||||
|
- 不自动解析或迁移现有 `extra_args` 中的 reasoning 参数,避免误判嵌套结构和 Provider 语义。
|
||||||
|
- UI 检测到旧 `extra_args` reasoning 字段时显示“由高级参数控制”,统一策略保持 Provider Default。
|
||||||
|
- 用户主动改成统一策略时,要求先移除冲突高级参数。
|
||||||
|
|
||||||
|
### 10.2 运行时兼容
|
||||||
|
|
||||||
|
- `provider_default` 不产生任何新增请求参数。
|
||||||
|
- 不改变现有 `remove-think` 的存储键和默认值。
|
||||||
|
- 不改变已有 Provider 的 `litellm_provider`,除非该 Provider 在专项回归后单独切换。
|
||||||
|
- `drop_params` 不能用于掩盖显式 reasoning 配置错误;显式策略被丢弃应视为失败。
|
||||||
|
- 自托管和 toB 环境中的自定义兼容接口保持可用,未知能力不阻止 Provider Default 请求。
|
||||||
|
|
||||||
|
## 11. 实施拆分
|
||||||
|
|
||||||
|
### Phase 1:统一基础设施与主流 Provider
|
||||||
|
|
||||||
|
- Alembic 增加 `llm_models.reasoning_config`。
|
||||||
|
- Backend 模型实体、CRUD、测试接口支持统一配置。
|
||||||
|
- LiteLLMRequester 增加能力查询、严格校验和参数翻译。
|
||||||
|
- 支持 OpenAI、Anthropic、Gemini、DeepSeek、xAI、Ollama、OpenRouter 的已验证 LiteLLM 路径。
|
||||||
|
- 修复结构化 reasoning 的非流式/流式保留。
|
||||||
|
- 模型面板增加 reasoning ability 与只读能力标识。
|
||||||
|
- Local Agent 主模型和每个 fallback 增加独立的请求级策略。
|
||||||
|
|
||||||
|
### Phase 2:国内 Provider
|
||||||
|
|
||||||
|
- 专项核对并支持 Volcengine/Doubao、Bailian/Qwen。
|
||||||
|
- 对相关 requester 的 `litellm_provider` 变更做独立回归,避免把 reasoning 功能和通用请求行为回归混在一起。
|
||||||
|
- 补齐扫描结果中的 reasoning capability。
|
||||||
|
|
||||||
|
### Phase 3:监控与评估
|
||||||
|
|
||||||
|
- 持久化 reasoning token 和生效策略。
|
||||||
|
- 监控页增加 reasoning 成本/延迟指标。
|
||||||
|
- 建立不同 effort 的离线质量、首 token 延迟、总耗时和 token 对比基线。
|
||||||
|
|
||||||
|
## 12. 测试方案
|
||||||
|
|
||||||
|
### 12.1 单元测试
|
||||||
|
|
||||||
|
- `ReasoningConfig` 所有合法/非法组合。
|
||||||
|
- `provider_default` 不产生任何新增参数。
|
||||||
|
- 显式配置覆盖模型/调用 `extra_args` 的顺序。
|
||||||
|
- reasoning 配置与高级参数冲突时拒绝。
|
||||||
|
- OpenAI 档位原样映射。
|
||||||
|
- Anthropic 档位映射,以及高级参数预算兼容。
|
||||||
|
- Gemini 2 budget、Gemini 3 level,以及不支持真正关闭时拒绝。
|
||||||
|
- DeepSeek 只显示/接受 toggle,非 `none` effort 不伪装成不同档位。
|
||||||
|
- Ollama 布尔与分级模型差异。
|
||||||
|
- Volcengine enabled/disabled/auto 翻译。
|
||||||
|
- 未知 Provider 只允许 Provider Default,或在显式测试后使用标准参数。
|
||||||
|
- 非流式 `reasoning_content` 保存到 `provider_specific_fields`。
|
||||||
|
- 流式 reasoning 分片累计后仍能 round-trip。
|
||||||
|
- Gemini thought signature 和工具调用现有测试不能回归。
|
||||||
|
|
||||||
|
### 12.2 服务与持久化测试
|
||||||
|
|
||||||
|
- 新建、读取、更新模型的 `reasoning_config`。
|
||||||
|
- Alembic 从当前 head 升级后默认值正确。
|
||||||
|
- 模型测试接口与真实 Local Agent 使用同一翻译函数。
|
||||||
|
- 旧模型、旧 `extra_args` 和 `remove-think` 行为不变。
|
||||||
|
|
||||||
|
### 12.3 前端测试
|
||||||
|
|
||||||
|
- 能力不同的模型显示正确控件。
|
||||||
|
- 离散滑杆只能停在后端返回的可用档位。
|
||||||
|
- 当前档位文字、键盘操作和 ARIA value text 正确。
|
||||||
|
- 仅开关模型、不可关闭模型、完整档位模型分别显示正确刻度。
|
||||||
|
- fallback 能力不兼容时阻止保存并给出明确提示。
|
||||||
|
- 中英文文案完整,移动端 Popover 不溢出。
|
||||||
|
|
||||||
|
### 12.4 Provider 冒烟测试
|
||||||
|
|
||||||
|
至少选取以下真实或可控 mock:
|
||||||
|
|
||||||
|
- 一个支持 `none` 的 OpenAI reasoning 模型。
|
||||||
|
- 一个不支持 `none` 的 reasoning 模型。
|
||||||
|
- 一个 Anthropic adaptive thinking 模型。
|
||||||
|
- 一个 Gemini 2.x 与一个 Gemini 3.x 模型。
|
||||||
|
- 一个 DeepSeek hybrid thinking 模型,执行两轮含工具调用对话。
|
||||||
|
- 一个 Ollama 本地 reasoning 模型。
|
||||||
|
- 一个 OpenAI-compatible 自定义网关,验证 Provider Default 完全不变。
|
||||||
|
|
||||||
|
每个模型比较 Provider Default、最低档、中档、高档或关闭,记录成功率、首 token 延迟、总耗时、总 token 和 reasoning token(若可用)。
|
||||||
|
|
||||||
|
## 13. 风险与控制
|
||||||
|
|
||||||
|
| 风险 | 影响 | 控制措施 |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| 将“最低思考”误当成“关闭” | 用户以为节省了成本,实际仍在推理 | `can_disable` 严格校验,不静默降级 |
|
||||||
|
| 模型能力表过期 | 新模型无法配置或旧模型报错 | 能力未知时保守;允许测试;升级 LiteLLM 时回归 |
|
||||||
|
| 高 effort 导致延迟/费用陡增 | 用户体验和预算风险 | 默认 Provider Default;UI 提示;后续监控 reasoning token |
|
||||||
|
| `extra_args` 与统一配置冲突 | 实际生效值不可预测 | 保存/测试时拒绝冲突;统一策略最后应用 |
|
||||||
|
| reasoning 原文进入日志 | 敏感信息泄露 | 不记录原文,只记录策略和 token |
|
||||||
|
| 多轮 reasoning 丢失 | 工具调用或后续轮次失败/降质 | 结构化保存并 round-trip;流式专项测试 |
|
||||||
|
| 修改 Provider 路由造成通用回归 | 非 reasoning 请求也受影响 | 国内 Provider 路由放第二阶段,独立提交和回归 |
|
||||||
|
|
||||||
|
## 14. 需要审核确认的决策
|
||||||
|
|
||||||
|
1. **是否同意三层分离**:能力、策略、展示互不替代,保留 `remove-think` 仅控制展示。
|
||||||
|
2. **是否同意严格语义**:显式关闭无法准确执行时直接报错,不自动降为最低思考。
|
||||||
|
3. **是否同意请求级配置**:流水线按模型 UUID 保存挡位,不把产品配置塞进 `extra_args`。
|
||||||
|
4. **是否同意 Runner 边界**:仅 Local Agent 展示控制项,外部 Runner 由其外部系统管理模型策略。
|
||||||
|
5. **是否同意保守默认**:所有现有模型迁移为 Provider Default,不自动开启、关闭或迁移旧高级参数。
|
||||||
|
6. **是否把结构化 reasoning 保留纳入第一阶段**:这是 DeepSeek 多轮和工具调用正确性的必要条件,建议必须纳入。
|
||||||
|
|
||||||
|
## 15. 推荐审核结果
|
||||||
|
|
||||||
|
建议按以上 6 项全部通过,并将 Phase 1 作为一个完整功能单元实施。不要只增加前端开关或只在 `extra_args` 中写 `reasoning_effort`;那样虽然改动小,但会继续混淆展示与推理、无法处理 Provider 差异,也无法保证多轮对话正确性。
|
||||||
+3
-3
@@ -1,6 +1,6 @@
|
|||||||
[project]
|
[project]
|
||||||
name = "langbot"
|
name = "langbot"
|
||||||
version = "4.10.6"
|
version = "4.10.7"
|
||||||
description = "Production-grade platform for building agentic IM bots"
|
description = "Production-grade platform for building agentic IM bots"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
license-files = ["LICENSE"]
|
license-files = ["LICENSE"]
|
||||||
@@ -23,7 +23,7 @@ dependencies = [
|
|||||||
"pynacl>=1.5.0", # Required for Discord voice support
|
"pynacl>=1.5.0", # Required for Discord voice support
|
||||||
"gewechat-client>=0.1.5",
|
"gewechat-client>=0.1.5",
|
||||||
"lark-oapi>=1.5.5",
|
"lark-oapi>=1.5.5",
|
||||||
"mcp>=1.25.0",
|
"mcp>=1.25.0,<2.0.0",
|
||||||
"nakuru-project-idk>=0.0.2.1",
|
"nakuru-project-idk>=0.0.2.1",
|
||||||
"ollama>=0.4.8",
|
"ollama>=0.4.8",
|
||||||
"openai>1.0.0",
|
"openai>1.0.0",
|
||||||
@@ -71,7 +71,7 @@ dependencies = [
|
|||||||
"chromadb>=1.0.0,<2.0.0",
|
"chromadb>=1.0.0,<2.0.0",
|
||||||
"qdrant-client (>=1.15.1,<2.0.0)",
|
"qdrant-client (>=1.15.1,<2.0.0)",
|
||||||
"pyseekdb==1.1.0.post3",
|
"pyseekdb==1.1.0.post3",
|
||||||
"langbot-plugin==0.5.0",
|
"langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@555a58e5db3de28e977b08dd4cd116b332848a19",
|
||||||
"asyncpg>=0.30.0",
|
"asyncpg>=0.30.0",
|
||||||
"line-bot-sdk>=3.19.0",
|
"line-bot-sdk>=3.19.0",
|
||||||
"matrix-nio>=0.25.2",
|
"matrix-nio>=0.25.2",
|
||||||
|
|||||||
@@ -27,17 +27,19 @@ The `all` / `box` profile starts three services:
|
|||||||
- `langbot_box` — Box sandbox runtime (`:5410`). Uses the host Docker socket to
|
- `langbot_box` — Box sandbox runtime (`:5410`). Uses the host Docker socket to
|
||||||
spawn sandbox containers, so the **Box root host path and in-container path
|
spawn sandbox containers, so the **Box root host path and in-container path
|
||||||
must be identical** (`BOX__LOCAL__HOST_ROOT=${LANGBOT_BOX_ROOT:-${PWD}/data/box}`).
|
must be identical** (`BOX__LOCAL__HOST_ROOT=${LANGBOT_BOX_ROOT:-${PWD}/data/box}`).
|
||||||
Its RPC and managed-process relay require a shared
|
OSS allows its RPC and managed-process relay to run without a token when both
|
||||||
`LANGBOT_BOX_CONTROL_TOKEN` (at least 32 non-whitespace characters) in both
|
sides leave `LANGBOT_BOX_CONTROL_TOKEN` unset. For an exposed endpoint, set
|
||||||
the LangBot and Box containers. Generate it once with `openssl rand -hex 32`;
|
the same value of at least 32 non-whitespace characters in both the LangBot
|
||||||
never put it in `box.runtime.endpoint` or commit it to config.
|
and Box containers. Generate it once with `openssl rand -hex 32`; never put
|
||||||
|
it in `box.runtime.endpoint` or commit it to config.
|
||||||
|
|
||||||
Every Compose deployment also needs one
|
A Compose deployment may optionally set
|
||||||
`LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN` shared by `langbot` and
|
`LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN` on both `langbot` and
|
||||||
`langbot_plugin_runtime`. Generate it with `openssl rand -hex 32` and export it
|
`langbot_plugin_runtime` when port 5400 needs shared-secret protection. OSS
|
||||||
before `docker compose up`; the external Plugin Runtime fails closed when the
|
defaults to leaving it unset on both sides. If enabled, generate one value with
|
||||||
token is empty or weak. Kubernetes uses the `langbot-plugin-runtime-control`
|
`openssl rand -hex 32`; configuring only one side causes the control connection
|
||||||
Secret shown in `docker/kubernetes.yaml`.
|
to fail. Kubernetes may use the `langbot-plugin-runtime-control` Secret shown in
|
||||||
|
`docker/kubernetes.yaml`.
|
||||||
|
|
||||||
With Box off, the dashboard/skills list stays visible (read-only) but sandbox
|
With Box off, the dashboard/skills list stays visible (read-only) but sandbox
|
||||||
tools, skill add/edit, and stdio MCP are disabled. Set `box.enabled: false`
|
tools, skill add/edit, and stdio MCP are disabled. Set `box.enabled: false`
|
||||||
|
|||||||
@@ -6,7 +6,8 @@ description: Browse and search the LangBot Space marketplaces (plugins, MCP serv
|
|||||||
# LangBot Space MCP Operations
|
# LangBot Space MCP Operations
|
||||||
|
|
||||||
LangBot Space (space.langbot.app) exposes an **MCP server** so user-facing AI
|
LangBot Space (space.langbot.app) exposes an **MCP server** so user-facing AI
|
||||||
agents can browse and search the marketplaces (plugins, MCP servers, skills).
|
agents can browse and search the marketplaces (plugins, MCP servers, skills) and
|
||||||
|
rank live models for automated setup.
|
||||||
|
|
||||||
## Endpoint
|
## Endpoint
|
||||||
|
|
||||||
@@ -46,10 +47,12 @@ Authorization: Bearer lbpat_...uests without a valid PAT get `401 Unauthorized`.
|
|||||||
| `list_plugins` / `search_plugins` / `get_plugin` | Plugin marketplace |
|
| `list_plugins` / `search_plugins` / `get_plugin` | Plugin marketplace |
|
||||||
| `list_mcp_servers` / `search_mcp_servers` / `get_mcp_server` | MCP-server marketplace |
|
| `list_mcp_servers` / `search_mcp_servers` / `get_mcp_server` | MCP-server marketplace |
|
||||||
| `list_skills` / `search_skills` / `get_skill` | Skill marketplace |
|
| `list_skills` / `search_skills` / `get_skill` | Skill marketplace |
|
||||||
|
| `select_models` | Live best-first model list for setup wizards; optional `category` filter |
|
||||||
|
|
||||||
`list_*` and `search_*` are paged (`page`, `page_size`). `get_*` takes
|
`list_*` and `search_*` are paged (`page`, `page_size`). `get_*` takes
|
||||||
`author` + `name`. The tool surface mirrors the REST endpoints under
|
`author` + `name`. The tool surface mirrors the REST endpoints under
|
||||||
`/api/v1/marketplace/*` and is read/browse only.
|
`/api/v1/marketplace/*`; `select_models` mirrors `/api/v1/models/selection`.
|
||||||
|
All tools are read-only.
|
||||||
|
|
||||||
## How to use
|
## How to use
|
||||||
|
|
||||||
@@ -58,12 +61,16 @@ Authorization: Bearer lbpat_...uests without a valid PAT get `401 Unauthorized`.
|
|||||||
3. Use `search_plugins` / `search_mcp_servers` / `search_skills` to find items,
|
3. Use `search_plugins` / `search_mcp_servers` / `search_skills` to find items,
|
||||||
then `get_*` for details (e.g. to obtain author/name for installation in
|
then `get_*` for details (e.g. to obtain author/name for installation in
|
||||||
LangBot itself).
|
LangBot itself).
|
||||||
|
4. For automatic local-agent setup, call `select_models` (optionally with
|
||||||
|
`category`) and choose the first compatible item. Ordering is latest probe
|
||||||
|
state (available, unprobed, unavailable), then Space recommendation. Each
|
||||||
|
item includes `availability.up`, `last_probed_at`, latency, and HTTP status.
|
||||||
|
|
||||||
## Implementation & maintenance (for Space developers)
|
## Implementation & maintenance (for Space developers)
|
||||||
|
|
||||||
- Server: `internal/controller/mcp/server.go` (official Go MCP SDK
|
- Server: `internal/controller/mcp/server.go` (official Go MCP SDK
|
||||||
`github.com/modelcontextprotocol/go-sdk`). Tools call the service layer
|
`github.com/modelcontextprotocol/go-sdk`). Tools call the service layer
|
||||||
(`PluginService`, `MCPService`, `SkillService`) directly.
|
(`PluginService`, `MCPService`, `SkillService`, `ModelStatusService`) directly.
|
||||||
- Mount: `internal/controller/api.go` at `/mcp` and `/mcp/*any`.
|
- Mount: `internal/controller/api.go` at `/mcp` and `/mcp/*any`.
|
||||||
- Auth: PAT via `AccountService.ValidatePersonalAccessToken`.
|
- Auth: PAT via `AccountService.ValidatePersonalAccessToken`.
|
||||||
- Docs: `docs/MCP_SERVER.md`.
|
- Docs: `docs/MCP_SERVER.md`.
|
||||||
|
|||||||
@@ -19,7 +19,6 @@ class Permission(enum.StrEnum):
|
|||||||
WORKSPACE_VIEW = 'workspace.view'
|
WORKSPACE_VIEW = 'workspace.view'
|
||||||
WORKSPACE_UPDATE = 'workspace.update'
|
WORKSPACE_UPDATE = 'workspace.update'
|
||||||
WORKSPACE_DELETE = 'workspace.delete'
|
WORKSPACE_DELETE = 'workspace.delete'
|
||||||
OWNER_TRANSFER = 'owner.transfer'
|
|
||||||
MEMBER_VIEW = 'member.view'
|
MEMBER_VIEW = 'member.view'
|
||||||
MEMBER_INVITE = 'member.invite'
|
MEMBER_INVITE = 'member.invite'
|
||||||
MEMBER_UPDATE_ROLE = 'member.update_role'
|
MEMBER_UPDATE_ROLE = 'member.update_role'
|
||||||
@@ -49,7 +48,6 @@ _ROLE_PERMISSIONS: typing.Final = types.MappingProxyType(
|
|||||||
if permission
|
if permission
|
||||||
not in {
|
not in {
|
||||||
Permission.WORKSPACE_DELETE,
|
Permission.WORKSPACE_DELETE,
|
||||||
Permission.OWNER_TRANSFER,
|
|
||||||
Permission.BILLING_LINK_MANAGE,
|
Permission.BILLING_LINK_MANAGE,
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
@@ -74,6 +72,11 @@ class AuthorizationError(Exception):
|
|||||||
error_code = 'forbidden'
|
error_code = 'forbidden'
|
||||||
|
|
||||||
|
|
||||||
|
class AuthenticationDeniedError(AuthorizationError):
|
||||||
|
status_code = 401
|
||||||
|
error_code = 'invalid_authentication'
|
||||||
|
|
||||||
|
|
||||||
class WorkspaceRequiredError(AuthorizationError):
|
class WorkspaceRequiredError(AuthorizationError):
|
||||||
status_code = 400
|
status_code = 400
|
||||||
error_code = 'workspace_required'
|
error_code = 'workspace_required'
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ class PrincipalType(enum.StrEnum):
|
|||||||
|
|
||||||
ACCOUNT = 'account'
|
ACCOUNT = 'account'
|
||||||
API_KEY = 'api_key'
|
API_KEY = 'api_key'
|
||||||
|
SUPPORT_ADMIN = 'support_admin'
|
||||||
SYSTEM = 'system'
|
SYSTEM = 'system'
|
||||||
PUBLIC_BOT = 'public_bot'
|
PUBLIC_BOT = 'public_bot'
|
||||||
|
|
||||||
@@ -19,7 +20,9 @@ class PrincipalContext:
|
|||||||
|
|
||||||
principal_type: PrincipalType
|
principal_type: PrincipalType
|
||||||
account_uuid: str | None = None
|
account_uuid: str | None = None
|
||||||
|
actor_account_uuid: str | None = None
|
||||||
api_key_uuid: str | None = None
|
api_key_uuid: str | None = None
|
||||||
|
support_session_id: str | None = None
|
||||||
|
|
||||||
|
|
||||||
@dataclasses.dataclass(frozen=True, slots=True)
|
@dataclasses.dataclass(frozen=True, slots=True)
|
||||||
|
|||||||
@@ -14,10 +14,18 @@ from ....utils import bounded_executor
|
|||||||
from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError
|
from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError
|
||||||
from ....workspace.errors import WorkspaceNotFoundError
|
from ....workspace.errors import WorkspaceNotFoundError
|
||||||
from ....cloud.entitlements import EntitlementUnavailableError
|
from ....cloud.entitlements import EntitlementUnavailableError
|
||||||
from ....cloud.quotas import WorkspaceQuotaExceededError
|
|
||||||
from ....core.errors import TaskCapacityError
|
from ....core.errors import TaskCapacityError
|
||||||
from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission
|
from ..authz import (
|
||||||
|
AuthenticationDeniedError,
|
||||||
|
AuthorizationError,
|
||||||
|
Permission,
|
||||||
|
PermissionDeniedError,
|
||||||
|
WorkspaceRequiredError,
|
||||||
|
permissions_for_role,
|
||||||
|
require_permission,
|
||||||
|
)
|
||||||
from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
|
from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
|
||||||
|
from ....cloud.support_admin import SupportAdminSessionError
|
||||||
|
|
||||||
if typing.TYPE_CHECKING:
|
if typing.TYPE_CHECKING:
|
||||||
from ....core.app import Application
|
from ....core.app import Application
|
||||||
@@ -52,6 +60,16 @@ class AuthType(enum.Enum):
|
|||||||
USER_TOKEN_OR_API_KEY = 'user-token-or-api-key'
|
USER_TOKEN_OR_API_KEY = 'user-token-or-api-key'
|
||||||
|
|
||||||
|
|
||||||
|
_SUPPORT_ADMIN_DENIED_PERMISSIONS = frozenset(
|
||||||
|
{
|
||||||
|
Permission.MEMBER_VIEW.value,
|
||||||
|
Permission.MEMBER_INVITE.value,
|
||||||
|
Permission.MEMBER_UPDATE_ROLE.value,
|
||||||
|
Permission.MEMBER_REMOVE.value,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class RouterGroup(abc.ABC):
|
class RouterGroup(abc.ABC):
|
||||||
name: str
|
name: str
|
||||||
|
|
||||||
@@ -96,6 +114,10 @@ class RouterGroup(abc.ABC):
|
|||||||
return self.http_status(401, -1, 'No valid user token provided')
|
return self.http_status(401, -1, 'No valid user token provided')
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
if self._is_support_admin_token(token):
|
||||||
|
raise AuthenticationDeniedError(
|
||||||
|
'Support admin tokens cannot be refreshed or used on account endpoints'
|
||||||
|
)
|
||||||
account, user_email = await self._authenticate_account(token)
|
account, user_email = await self._authenticate_account(token)
|
||||||
# Account-token routes deliberately stop before Workspace
|
# Account-token routes deliberately stop before Workspace
|
||||||
# selection. They may bootstrap a selector, but cannot
|
# selection. They may bootstrap a selector, but cannot
|
||||||
@@ -112,8 +134,13 @@ class RouterGroup(abc.ABC):
|
|||||||
return self.http_status(401, -1, 'No valid user token provided')
|
return self.http_status(401, -1, 'No valid user token provided')
|
||||||
|
|
||||||
try:
|
try:
|
||||||
account, user_email = await self._authenticate_account(token)
|
request_context = await self._authenticate_support_admin(token, auth_type)
|
||||||
request_context = await self._resolve_account_context(account, auth_type)
|
if request_context is not None:
|
||||||
|
self._require_support_admin_route_allowed(rule, f, permission)
|
||||||
|
user_email = None
|
||||||
|
else:
|
||||||
|
account, user_email = await self._authenticate_account(token)
|
||||||
|
request_context = await self._resolve_account_context(account, auth_type)
|
||||||
if permission is not None:
|
if permission is not None:
|
||||||
if request_context is None:
|
if request_context is None:
|
||||||
raise AuthorizationError('Workspace authorization is unavailable')
|
raise AuthorizationError('Workspace authorization is unavailable')
|
||||||
@@ -142,10 +169,20 @@ class RouterGroup(abc.ABC):
|
|||||||
return self._auth_error_response(e)
|
return self._auth_error_response(e)
|
||||||
|
|
||||||
elif auth_type == AuthType.USER_TOKEN_OR_API_KEY:
|
elif auth_type == AuthType.USER_TOKEN_OR_API_KEY:
|
||||||
|
token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')
|
||||||
|
if token and self._is_support_admin_token(token):
|
||||||
|
try:
|
||||||
|
request_context = await self._authenticate_support_admin(token, auth_type)
|
||||||
|
if request_context is None:
|
||||||
|
raise AuthenticationDeniedError('Invalid support admin token')
|
||||||
|
self._require_support_admin_route_allowed(rule, f, permission)
|
||||||
|
if permission is not None:
|
||||||
|
require_permission(request_context, permission)
|
||||||
|
self._inject_handler_context(f, kwargs, None, request_context)
|
||||||
|
except Exception as e:
|
||||||
|
return self._auth_error_response(e)
|
||||||
# Try API key first (check X-API-Key header)
|
# Try API key first (check X-API-Key header)
|
||||||
api_key = quart.request.headers.get('X-API-Key', '')
|
elif api_key := quart.request.headers.get('X-API-Key', ''):
|
||||||
|
|
||||||
if api_key:
|
|
||||||
# API key authentication
|
# API key authentication
|
||||||
try:
|
try:
|
||||||
request_context = await self._authenticate_api_key(api_key, auth_type)
|
request_context = await self._authenticate_api_key(api_key, auth_type)
|
||||||
@@ -156,8 +193,6 @@ class RouterGroup(abc.ABC):
|
|||||||
return self._auth_error_response(e)
|
return self._auth_error_response(e)
|
||||||
else:
|
else:
|
||||||
# Try user token authentication (Authorization header)
|
# Try user token authentication (Authorization header)
|
||||||
token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')
|
|
||||||
|
|
||||||
if not token:
|
if not token:
|
||||||
return self.http_status(
|
return self.http_status(
|
||||||
401, -1, 'No valid authentication provided (user token or API key required)'
|
401, -1, 'No valid authentication provided (user token or API key required)'
|
||||||
@@ -220,8 +255,6 @@ class RouterGroup(abc.ABC):
|
|||||||
return self.http_status(403, e.code, str(e))
|
return self.http_status(403, e.code, str(e))
|
||||||
if isinstance(e, WorkspaceCollaborationError):
|
if isinstance(e, WorkspaceCollaborationError):
|
||||||
return self.http_status(400, e.code, str(e))
|
return self.http_status(400, e.code, str(e))
|
||||||
if isinstance(e, WorkspaceQuotaExceededError):
|
|
||||||
return self.http_status(409, e.error_code, str(e))
|
|
||||||
if isinstance(e, TaskCapacityError):
|
if isinstance(e, TaskCapacityError):
|
||||||
return self.http_status(429, 'task_capacity_exceeded', str(e))
|
return self.http_status(429, 'task_capacity_exceeded', str(e))
|
||||||
if isinstance(
|
if isinstance(
|
||||||
@@ -271,10 +304,83 @@ class RouterGroup(abc.ABC):
|
|||||||
raise ValueError('User not found')
|
raise ValueError('User not found')
|
||||||
return account, account.user
|
return account, account.user
|
||||||
|
|
||||||
|
def _is_support_admin_token(self, token: str) -> bool:
|
||||||
|
service = getattr(self.ap, 'support_admin_session_service', None)
|
||||||
|
detector = getattr(service, 'is_support_admin_token', None)
|
||||||
|
return callable(detector) and detector(token) is True
|
||||||
|
|
||||||
|
async def _authenticate_support_admin(
|
||||||
|
self,
|
||||||
|
token: str,
|
||||||
|
auth_type: AuthType,
|
||||||
|
*,
|
||||||
|
workspace_uuid: str | None = None,
|
||||||
|
request_id: str | None = None,
|
||||||
|
) -> RequestContext | None:
|
||||||
|
service = getattr(self.ap, 'support_admin_session_service', None)
|
||||||
|
detector = getattr(service, 'is_support_admin_token', None)
|
||||||
|
if service is None or not callable(detector) or detector(token) is not True:
|
||||||
|
return None
|
||||||
|
|
||||||
|
requested_workspace_uuid = (
|
||||||
|
workspace_uuid if workspace_uuid is not None else quart.request.headers.get('X-Workspace-Id')
|
||||||
|
)
|
||||||
|
if not requested_workspace_uuid:
|
||||||
|
raise WorkspaceRequiredError('Support admin token requires an explicit Workspace selector')
|
||||||
|
try:
|
||||||
|
identity = await service.authenticate_token(
|
||||||
|
token,
|
||||||
|
requested_workspace_uuid=requested_workspace_uuid,
|
||||||
|
)
|
||||||
|
except SupportAdminSessionError as exc:
|
||||||
|
raise AuthenticationDeniedError(str(exc)) from exc
|
||||||
|
|
||||||
|
entitlement_revision = await self._resolve_entitlement_revision(
|
||||||
|
identity.instance_uuid,
|
||||||
|
identity.workspace_uuid,
|
||||||
|
)
|
||||||
|
request_context = RequestContext(
|
||||||
|
instance_uuid=identity.instance_uuid,
|
||||||
|
placement_generation=identity.placement_generation,
|
||||||
|
request_id=request_id or self.request_id(),
|
||||||
|
auth_type=auth_type.value,
|
||||||
|
principal=PrincipalContext(
|
||||||
|
principal_type=PrincipalType.SUPPORT_ADMIN,
|
||||||
|
actor_account_uuid=identity.actor_account_uuid,
|
||||||
|
support_session_id=identity.grant_jti_hash,
|
||||||
|
),
|
||||||
|
workspace=WorkspaceContext(
|
||||||
|
workspace_uuid=identity.workspace_uuid,
|
||||||
|
membership_uuid=None,
|
||||||
|
role='owner',
|
||||||
|
permissions=permissions_for_role('owner') - _SUPPORT_ADMIN_DENIED_PERMISSIONS,
|
||||||
|
membership_revision=0,
|
||||||
|
),
|
||||||
|
entitlement_revision=entitlement_revision,
|
||||||
|
)
|
||||||
|
quart.g.request_context = request_context
|
||||||
|
quart.g.workspace_membership = None
|
||||||
|
return request_context
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _require_support_admin_route_allowed(
|
||||||
|
rule: str,
|
||||||
|
handler: RouteCallable,
|
||||||
|
permission: Permission | str | None,
|
||||||
|
) -> None:
|
||||||
|
parameters = inspect.signature(handler).parameters
|
||||||
|
if rule.startswith('/api/v1/user/') or 'account' in parameters or 'user_email' in parameters:
|
||||||
|
raise AuthenticationDeniedError('Support admin tokens are not permitted on account endpoints')
|
||||||
|
permission_value = permission.value if isinstance(permission, Permission) else permission
|
||||||
|
if permission_value in _SUPPORT_ADMIN_DENIED_PERMISSIONS:
|
||||||
|
raise PermissionDeniedError(permission_value)
|
||||||
|
|
||||||
async def _resolve_account_context(
|
async def _resolve_account_context(
|
||||||
self,
|
self,
|
||||||
account: typing.Any,
|
account: typing.Any,
|
||||||
auth_type: AuthType,
|
auth_type: AuthType,
|
||||||
|
*,
|
||||||
|
token: str | None = None,
|
||||||
) -> RequestContext | None:
|
) -> RequestContext | None:
|
||||||
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
||||||
account_uuid = getattr(account, 'uuid', None)
|
account_uuid = getattr(account, 'uuid', None)
|
||||||
|
|||||||
@@ -97,6 +97,16 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
|||||||
if not token or not workspace_uuid:
|
if not token or not workspace_uuid:
|
||||||
raise ValueError('Authentication is required')
|
raise ValueError('Authentication is required')
|
||||||
|
|
||||||
|
support_context = await self._authenticate_support_admin(
|
||||||
|
token,
|
||||||
|
group.AuthType.USER_TOKEN,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
request_id=quart.websocket.headers.get('X-Request-Id') or str(uuid.uuid4()),
|
||||||
|
)
|
||||||
|
if support_context is not None:
|
||||||
|
require_permission(support_context, Permission.RUNTIME_OPERATE)
|
||||||
|
return support_context, token
|
||||||
|
|
||||||
account, _ = await self._authenticate_account(token)
|
account, _ = await self._authenticate_account(token)
|
||||||
account_uuid = getattr(account, 'uuid', None)
|
account_uuid = getattr(account, 'uuid', None)
|
||||||
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
||||||
@@ -131,6 +141,23 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
|||||||
) -> RequestContext:
|
) -> RequestContext:
|
||||||
"""Recheck revocable account, membership, permission, and placement state."""
|
"""Recheck revocable account, membership, permission, and placement state."""
|
||||||
|
|
||||||
|
if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN:
|
||||||
|
current_context = await self._authenticate_support_admin(
|
||||||
|
token,
|
||||||
|
group.AuthType.USER_TOKEN,
|
||||||
|
workspace_uuid=request_context.workspace_uuid,
|
||||||
|
request_id=request_context.request_id,
|
||||||
|
)
|
||||||
|
if current_context is None or current_context.principal != request_context.principal:
|
||||||
|
raise ValueError('WebSocket support admin session changed')
|
||||||
|
if (
|
||||||
|
current_context.instance_uuid != request_context.instance_uuid
|
||||||
|
or current_context.placement_generation != request_context.placement_generation
|
||||||
|
):
|
||||||
|
raise ValueError('WebSocket authorization changed')
|
||||||
|
require_permission(current_context, Permission.RUNTIME_OPERATE)
|
||||||
|
return current_context
|
||||||
|
|
||||||
account, _ = await self._authenticate_account(token)
|
account, _ = await self._authenticate_account(token)
|
||||||
account_uuid = getattr(account, 'uuid', None)
|
account_uuid = getattr(account, 'uuid', None)
|
||||||
if account_uuid != request_context.account_uuid:
|
if account_uuid != request_context.account_uuid:
|
||||||
@@ -211,6 +238,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
|||||||
scope=WebSocketScope.from_context(request_context),
|
scope=WebSocketScope.from_context(request_context),
|
||||||
pipeline_uuid=pipeline_uuid,
|
pipeline_uuid=pipeline_uuid,
|
||||||
session_type=session_type,
|
session_type=session_type,
|
||||||
|
trigger_principal=request_context.principal,
|
||||||
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
|
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
|
||||||
send_queue_size=(
|
send_queue_size=(
|
||||||
self.ap.instance_config.data.get('system', {})
|
self.ap.instance_config.data.get('system', {})
|
||||||
@@ -391,7 +419,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
|||||||
)
|
)
|
||||||
elif message_type == 'message':
|
elif message_type == 'message':
|
||||||
try:
|
try:
|
||||||
await self._revalidate_websocket_authorization(request_context, token)
|
request_context = await self._revalidate_websocket_authorization(request_context, token)
|
||||||
except Exception:
|
except Exception:
|
||||||
await connection.send_queue.put({'type': 'error', 'message': 'Unauthorized'})
|
await connection.send_queue.put({'type': 'error', 'message': 'Unauthorized'})
|
||||||
break
|
break
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ class _AdapterSessionScope:
|
|||||||
principal_type: str
|
principal_type: str
|
||||||
account_uuid: str | None
|
account_uuid: str | None
|
||||||
api_key_uuid: str | None
|
api_key_uuid: str | None
|
||||||
|
support_session_id: str | None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_request_context(cls, request_context: RequestContext) -> '_AdapterSessionScope':
|
def from_request_context(cls, request_context: RequestContext) -> '_AdapterSessionScope':
|
||||||
@@ -33,6 +34,7 @@ class _AdapterSessionScope:
|
|||||||
principal_type=principal.principal_type.value,
|
principal_type=principal.principal_type.value,
|
||||||
account_uuid=principal.account_uuid,
|
account_uuid=principal.account_uuid,
|
||||||
api_key_uuid=principal.api_key_uuid,
|
api_key_uuid=principal.api_key_uuid,
|
||||||
|
support_session_id=principal.support_session_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
def matches(self, request_context: RequestContext) -> bool:
|
def matches(self, request_context: RequestContext) -> bool:
|
||||||
|
|||||||
@@ -15,7 +15,6 @@ import posixpath
|
|||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
|
|
||||||
from .....core import taskmgr
|
from .....core import taskmgr
|
||||||
from .....core.task_boundary import run_in_workspace_uow
|
|
||||||
from .....entity.persistence import plugin as persistence_plugin
|
from .....entity.persistence import plugin as persistence_plugin
|
||||||
from ...authz import Permission
|
from ...authz import Permission
|
||||||
from ...context import ExecutionContext, RequestContext
|
from ...context import ExecutionContext, RequestContext
|
||||||
@@ -311,11 +310,13 @@ class PluginsRouterGroup(group.RouterGroup):
|
|||||||
):
|
):
|
||||||
"""Revalidate a captured task context immediately before Runtime I/O."""
|
"""Revalidate a captured task context immediately before Runtime I/O."""
|
||||||
|
|
||||||
await run_in_workspace_uow(
|
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
|
||||||
self.ap,
|
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
|
||||||
execution_context.workspace_uuid,
|
if callable(tenant_scope):
|
||||||
lambda: self.ap.plugin_connector.require_workspace_context(execution_context),
|
async with tenant_scope(execution_context.workspace_uuid):
|
||||||
)
|
await self.ap.plugin_connector.require_workspace_context(execution_context)
|
||||||
|
return await operation()
|
||||||
|
await self.ap.plugin_connector.require_workspace_context(execution_context)
|
||||||
return await operation()
|
return await operation()
|
||||||
|
|
||||||
async def _require_authenticated_plugin_runtime_context(
|
async def _require_authenticated_plugin_runtime_context(
|
||||||
@@ -392,17 +393,27 @@ class PluginsRouterGroup(group.RouterGroup):
|
|||||||
)
|
)
|
||||||
async def _(request_context: RequestContext) -> str:
|
async def _(request_context: RequestContext) -> str:
|
||||||
"""Get plugin debug information including debug URL and key"""
|
"""Get plugin debug information including debug URL and key"""
|
||||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
execution_context = await self._require_authenticated_plugin_runtime_context(request_context)
|
||||||
debug_info = await self.ap.plugin_connector.get_debug_info()
|
debug_info = await self.ap.plugin_connector.get_debug_info(execution_context)
|
||||||
|
|
||||||
# Get debug URL from config
|
# Get debug URL from config
|
||||||
plugin_config = self.ap.instance_config.data.get('plugin', {})
|
plugin_config = self.ap.instance_config.data.get('plugin', {})
|
||||||
debug_url = plugin_config.get('display_plugin_debug_url', 'http://localhost:5401')
|
debug_url = plugin_config.get(
|
||||||
|
'display_plugin_debug_url',
|
||||||
|
'ws://localhost:5401/plugin/debug/ws',
|
||||||
|
)
|
||||||
|
parsed_debug_url = urlparse(debug_url)
|
||||||
|
if parsed_debug_url.scheme in {'http', 'https'}:
|
||||||
|
debug_url = parsed_debug_url._replace(
|
||||||
|
scheme='wss' if parsed_debug_url.scheme == 'https' else 'ws',
|
||||||
|
path=parsed_debug_url.path or '/plugin/debug/ws',
|
||||||
|
).geturl()
|
||||||
|
|
||||||
return self.success(
|
return self.success(
|
||||||
data={
|
data={
|
||||||
'debug_url': debug_url,
|
'debug_url': debug_url,
|
||||||
'plugin_debug_key': debug_info.get('plugin_debug_key', ''),
|
'plugin_debug_key': debug_info.get('plugin_debug_key', ''),
|
||||||
|
'expires_at': debug_info.get('expires_at', ''),
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
import quart
|
import quart
|
||||||
import argon2
|
import argon2
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import datetime
|
||||||
import uuid
|
import uuid
|
||||||
from urllib.parse import parse_qs, urlsplit
|
from urllib.parse import parse_qs, urlsplit
|
||||||
|
|
||||||
@@ -13,13 +14,6 @@ from ...service.user import ControlPlaneDirectoryRequiredError, PublicRegistrati
|
|||||||
|
|
||||||
@group.group_class('user', '/api/v1/user')
|
@group.group_class('user', '/api/v1/user')
|
||||||
class UserRouterGroup(group.RouterGroup):
|
class UserRouterGroup(group.RouterGroup):
|
||||||
@staticmethod
|
|
||||||
def _origin(value: str) -> tuple[str, str, int | None] | None:
|
|
||||||
parsed = urlsplit(value)
|
|
||||||
if parsed.scheme not in {'http', 'https'} or not parsed.hostname:
|
|
||||||
return None
|
|
||||||
return parsed.scheme, parsed.hostname.casefold(), parsed.port
|
|
||||||
|
|
||||||
def _validate_space_redirect_uri(self, redirect_uri: str, *, bind: bool) -> str:
|
def _validate_space_redirect_uri(self, redirect_uri: str, *, bind: bool) -> str:
|
||||||
parsed = urlsplit(redirect_uri)
|
parsed = urlsplit(redirect_uri)
|
||||||
if (
|
if (
|
||||||
@@ -37,17 +31,8 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
if query != {'mode': ['bind']}:
|
if query != {'mode': ['bind']}:
|
||||||
raise ValueError('Invalid Space binding redirect_uri')
|
raise ValueError('Invalid Space binding redirect_uri')
|
||||||
elif query:
|
elif query:
|
||||||
raise ValueError('Invalid Space login redirect_uri')
|
raise ValueError('Invalid LangBot Account login redirect_uri')
|
||||||
|
|
||||||
redirect_origin = self._origin(redirect_uri)
|
|
||||||
api_config = self.ap.instance_config.data.get('api', {})
|
|
||||||
trusted_origins = {
|
|
||||||
self._origin(str(api_config.get(config_key, '') or '').strip())
|
|
||||||
for config_key in ('webui_url', 'webhook_prefix')
|
|
||||||
}
|
|
||||||
trusted_origins.discard(None)
|
|
||||||
if redirect_origin not in trusted_origins:
|
|
||||||
raise ValueError('Untrusted redirect_uri origin')
|
|
||||||
return redirect_uri
|
return redirect_uri
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
@@ -218,7 +203,22 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
try:
|
try:
|
||||||
consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login')
|
consumed_state = await self.ap.user_service.consume_space_oauth_state_details(state, 'login')
|
||||||
# Exchange code for tokens
|
# Exchange code for tokens
|
||||||
token_data = await self.ap.space_service.exchange_oauth_code(code)
|
launch_workspace_uuid = consumed_state.launch_workspace_uuid
|
||||||
|
workspace_uuids = [launch_workspace_uuid] if launch_workspace_uuid else []
|
||||||
|
workspace_created_ats: dict[str, int] = {}
|
||||||
|
if not workspace_uuids and getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') != 'cloud':
|
||||||
|
binding = await self.ap.workspace_service.get_execution_binding()
|
||||||
|
workspace_uuids = [binding.workspace_uuid]
|
||||||
|
workspace_created_at = binding.workspace_created_at
|
||||||
|
if workspace_created_at is not None:
|
||||||
|
if workspace_created_at.tzinfo is None:
|
||||||
|
workspace_created_at = workspace_created_at.replace(tzinfo=datetime.UTC)
|
||||||
|
workspace_created_ats[binding.workspace_uuid] = int(workspace_created_at.timestamp())
|
||||||
|
token_data = await self.ap.space_service.exchange_oauth_code(
|
||||||
|
code,
|
||||||
|
workspace_uuids,
|
||||||
|
workspace_created_ats,
|
||||||
|
)
|
||||||
access_token = token_data.get('access_token')
|
access_token = token_data.get('access_token')
|
||||||
refresh_token = token_data.get('refresh_token')
|
refresh_token = token_data.get('refresh_token')
|
||||||
expires_in = token_data.get('expires_in', 0)
|
expires_in = token_data.get('expires_in', 0)
|
||||||
@@ -231,7 +231,6 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
access_token, refresh_token, expires_in
|
access_token, refresh_token, expires_in
|
||||||
)
|
)
|
||||||
|
|
||||||
launch_workspace_uuid = consumed_state.launch_workspace_uuid
|
|
||||||
if launch_workspace_uuid:
|
if launch_workspace_uuid:
|
||||||
try:
|
try:
|
||||||
access = await self.ap.workspace_collaboration_service.resolve_account_workspace(
|
access = await self.ap.workspace_collaboration_service.resolve_account_workspace(
|
||||||
@@ -285,8 +284,25 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
request_context.workspace_uuid,
|
request_context.workspace_uuid,
|
||||||
)
|
)
|
||||||
owner = await self.ap.user_service.get_workspace_owner(access.workspace.uuid)
|
owner = await self.ap.user_service.get_workspace_owner(access.workspace.uuid)
|
||||||
owner_space_bound = bool(owner and owner.space_account_uuid)
|
cloud_mode = getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud'
|
||||||
credits = await self.ap.space_service.get_credits(owner.user) if owner_space_bound else None
|
owner_has_local_space_credentials = bool(owner and owner.space_account_uuid)
|
||||||
|
# Cloud Accounts authenticate through LangBot Account, so every projected
|
||||||
|
# Workspace owner is already bound even when this Core has no local OAuth
|
||||||
|
# token row (model billing uses the owner's control-plane API key).
|
||||||
|
owner_space_bound = cloud_mode or owner_has_local_space_credentials
|
||||||
|
if cloud_mode:
|
||||||
|
catalog_service = getattr(self.ap, 'cloud_model_catalog_service', None)
|
||||||
|
credits = (
|
||||||
|
catalog_service.get_workspace_credits(access.workspace.uuid)
|
||||||
|
if catalog_service is not None
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
credits = (
|
||||||
|
await self.ap.space_service.get_credits(owner.user)
|
||||||
|
if owner is not None and owner.space_account_uuid
|
||||||
|
else None
|
||||||
|
)
|
||||||
return self.success(
|
return self.success(
|
||||||
data={
|
data={
|
||||||
'credits': credits,
|
'credits': credits,
|
||||||
@@ -302,8 +318,10 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
return self.success(data={'initialized': False})
|
return self.success(data={'initialized': False})
|
||||||
|
|
||||||
capabilities = await self.ap.user_service.get_login_capabilities()
|
capabilities = await self.ap.user_service.get_login_capabilities()
|
||||||
if getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud':
|
cloud_mode = getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud'
|
||||||
|
if cloud_mode:
|
||||||
capabilities['password_login_enabled'] = False
|
capabilities['password_login_enabled'] = False
|
||||||
|
capabilities['authenticated_invitation_acceptance_enabled'] = cloud_mode
|
||||||
return self.success(data={'initialized': True, **capabilities})
|
return self.success(data={'initialized': True, **capabilities})
|
||||||
|
|
||||||
@self.route('/set-password', methods=['POST'], auth_type=group.AuthType.USER_TOKEN)
|
@self.route('/set-password', methods=['POST'], auth_type=group.AuthType.USER_TOKEN)
|
||||||
@@ -382,7 +400,7 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
'Bind the LangBot Account with the same email as this local Account',
|
'Bind the LangBot Account with the same email as this local Account',
|
||||||
)
|
)
|
||||||
except ValueError:
|
except ValueError:
|
||||||
return self.http_status(400, -1, 'Space account binding failed')
|
return self.http_status(400, -1, 'LangBot Account binding failed')
|
||||||
except Exception:
|
except Exception:
|
||||||
raise
|
raise
|
||||||
|
|
||||||
@@ -396,6 +414,19 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
launch_assertion,
|
launch_assertion,
|
||||||
expected_workspace_uuid=workspace_uuid,
|
expected_workspace_uuid=workspace_uuid,
|
||||||
)
|
)
|
||||||
|
if launch.get('launch_mode') == 'support_admin':
|
||||||
|
token = launch.get('support_admin_token')
|
||||||
|
if not token:
|
||||||
|
raise SpaceLaunchError('Support admin launch session was not issued')
|
||||||
|
return self.success(
|
||||||
|
data={
|
||||||
|
'token': token,
|
||||||
|
'workspace_uuid': launch['workspace_uuid'],
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': launch['actor_account_uuid'],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid'])
|
account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid'])
|
||||||
if account is None:
|
if account is None:
|
||||||
raise SpaceLaunchError('Launch Account is not projected into Core')
|
raise SpaceLaunchError('Launch Account is not projected into Core')
|
||||||
@@ -410,7 +441,6 @@ class UserRouterGroup(group.RouterGroup):
|
|||||||
'token': token,
|
'token': token,
|
||||||
'user': account.user,
|
'user': account.user,
|
||||||
'workspace_uuid': access.workspace.uuid,
|
'workspace_uuid': access.workspace.uuid,
|
||||||
'return_path': launch.get('return_path', '/home'),
|
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
except SpaceLaunchError:
|
except SpaceLaunchError:
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ import typing
|
|||||||
import quart
|
import quart
|
||||||
|
|
||||||
from ...authz import Permission, permissions_for_role
|
from ...authz import Permission, permissions_for_role
|
||||||
from ...context import RequestContext
|
from ...context import PrincipalType, RequestContext
|
||||||
from ...service.user import AccountExistsLoginRequiredError, ControlPlaneDirectoryRequiredError
|
from ...service.user import AccountExistsLoginRequiredError, ControlPlaneDirectoryRequiredError
|
||||||
from .....entity.persistence.workspace import Workspace, WorkspaceInvitation, WorkspaceMembership
|
from .....entity.persistence.workspace import Workspace, WorkspaceInvitation, WorkspaceMembership
|
||||||
from .....entity.persistence.workspace import WorkspaceSource
|
from .....entity.persistence.workspace import WorkspaceSource
|
||||||
@@ -30,12 +30,14 @@ def _workspace_payload(workspace: Workspace) -> dict[str, typing.Any]:
|
|||||||
def _membership_payload(
|
def _membership_payload(
|
||||||
membership: WorkspaceMembership,
|
membership: WorkspaceMembership,
|
||||||
*,
|
*,
|
||||||
|
display_name: str,
|
||||||
email: str,
|
email: str,
|
||||||
) -> dict[str, typing.Any]:
|
) -> dict[str, typing.Any]:
|
||||||
return {
|
return {
|
||||||
'uuid': membership.uuid,
|
'uuid': membership.uuid,
|
||||||
'workspace_uuid': membership.workspace_uuid,
|
'workspace_uuid': membership.workspace_uuid,
|
||||||
'account_uuid': membership.account_uuid,
|
'account_uuid': membership.account_uuid,
|
||||||
|
'display_name': display_name,
|
||||||
'email': email,
|
'email': email,
|
||||||
'role': membership.role,
|
'role': membership.role,
|
||||||
'status': membership.status,
|
'status': membership.status,
|
||||||
@@ -94,7 +96,11 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
|||||||
workspaces.append(
|
workspaces.append(
|
||||||
{
|
{
|
||||||
'workspace': _workspace_payload(access.workspace),
|
'workspace': _workspace_payload(access.workspace),
|
||||||
'membership': _membership_payload(access.membership, email=account.user),
|
'membership': _membership_payload(
|
||||||
|
access.membership,
|
||||||
|
display_name=account.user,
|
||||||
|
email=account.normalized_email,
|
||||||
|
),
|
||||||
'permissions': sorted(permissions_for_role(access.membership.role)),
|
'permissions': sorted(permissions_for_role(access.membership.role)),
|
||||||
'placement_generation': access.execution.placement_generation,
|
'placement_generation': access.execution.placement_generation,
|
||||||
'plan_name': plan_name,
|
'plan_name': plan_name,
|
||||||
@@ -120,9 +126,6 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
|||||||
@self.route('/current', methods=['GET'], permission=Permission.WORKSPACE_VIEW)
|
@self.route('/current', methods=['GET'], permission=Permission.WORKSPACE_VIEW)
|
||||||
async def _(request_context: RequestContext) -> typing.Any:
|
async def _(request_context: RequestContext) -> typing.Any:
|
||||||
membership = quart.g.workspace_membership
|
membership = quart.g.workspace_membership
|
||||||
account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid)
|
|
||||||
if account is None:
|
|
||||||
return self.http_status(401, 'invalid_authentication', 'Account not found')
|
|
||||||
workspace = await self.ap.workspace_service.get_workspace(request_context.workspace_uuid)
|
workspace = await self.ap.workspace_service.get_workspace(request_context.workspace_uuid)
|
||||||
plan_name: str | None = None
|
plan_name: str | None = None
|
||||||
resolver = getattr(self.ap, 'entitlement_resolver', None)
|
resolver = getattr(self.ap, 'entitlement_resolver', None)
|
||||||
@@ -132,10 +135,37 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
|||||||
minimum_revision=request_context.entitlement_revision,
|
minimum_revision=request_context.entitlement_revision,
|
||||||
)
|
)
|
||||||
plan_name = entitlement.plan_name
|
plan_name = entitlement.plan_name
|
||||||
|
if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN:
|
||||||
|
return self.success(
|
||||||
|
data={
|
||||||
|
'workspace': _workspace_payload(workspace),
|
||||||
|
'membership': {
|
||||||
|
'uuid': None,
|
||||||
|
'workspace_uuid': request_context.workspace_uuid,
|
||||||
|
'account_uuid': None,
|
||||||
|
'display_name': None,
|
||||||
|
'email': None,
|
||||||
|
'role': 'owner',
|
||||||
|
'status': 'active',
|
||||||
|
'joined_at': None,
|
||||||
|
'created_at': None,
|
||||||
|
},
|
||||||
|
'permissions': sorted(request_context.workspace.permissions),
|
||||||
|
'placement_generation': request_context.placement_generation,
|
||||||
|
'plan_name': plan_name,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid)
|
||||||
|
if account is None:
|
||||||
|
return self.http_status(401, 'invalid_authentication', 'Account not found')
|
||||||
return self.success(
|
return self.success(
|
||||||
data={
|
data={
|
||||||
'workspace': _workspace_payload(workspace),
|
'workspace': _workspace_payload(workspace),
|
||||||
'membership': _membership_payload(membership, email=account.user),
|
'membership': _membership_payload(
|
||||||
|
membership,
|
||||||
|
display_name=account.user,
|
||||||
|
email=account.normalized_email,
|
||||||
|
),
|
||||||
'permissions': sorted(request_context.workspace.permissions),
|
'permissions': sorted(request_context.workspace.permissions),
|
||||||
'placement_generation': request_context.placement_generation,
|
'placement_generation': request_context.placement_generation,
|
||||||
'plan_name': plan_name,
|
'plan_name': plan_name,
|
||||||
@@ -264,7 +294,8 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
|||||||
data={
|
data={
|
||||||
'member': _membership_payload(
|
'member': _membership_payload(
|
||||||
member,
|
member,
|
||||||
email=account.user if account is not None else '',
|
display_name=account.user if account is not None else '',
|
||||||
|
email=account.normalized_email if account is not None else '',
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
@@ -283,7 +314,11 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
|||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _member_view_payload(view: WorkspaceMemberView) -> dict[str, typing.Any]:
|
def _member_view_payload(view: WorkspaceMemberView) -> dict[str, typing.Any]:
|
||||||
return _membership_payload(view.membership, email=view.email)
|
return _membership_payload(
|
||||||
|
view.membership,
|
||||||
|
display_name=view.display_name,
|
||||||
|
email=view.email,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@group.group_class('invitations', '/api/v1/invitations')
|
@group.group_class('invitations', '/api/v1/invitations')
|
||||||
|
|||||||
@@ -4,7 +4,6 @@ import uuid
|
|||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
|
|
||||||
from ....core import app
|
from ....core import app
|
||||||
from ....cloud.quotas import require_resource_capacity, resolve_workspace_quota
|
|
||||||
from ....entity.persistence import bot as persistence_bot
|
from ....entity.persistence import bot as persistence_bot
|
||||||
from ....entity.persistence import pipeline as persistence_pipeline
|
from ....entity.persistence import pipeline as persistence_pipeline
|
||||||
from ....workspace.errors import WorkspaceNotFoundError
|
from ....workspace.errors import WorkspaceNotFoundError
|
||||||
@@ -102,21 +101,20 @@ class BotService:
|
|||||||
async def create_bot(self, context: TenantContext, bot_data: dict) -> str:
|
async def create_bot(self, context: TenantContext, bot_data: dict) -> str:
|
||||||
"""Create bot"""
|
"""Create bot"""
|
||||||
workspace_uuid = require_workspace_uuid(context)
|
workspace_uuid = require_workspace_uuid(context)
|
||||||
|
# Check limitation
|
||||||
limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {})
|
limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {})
|
||||||
quota = await resolve_workspace_quota(
|
max_bots = limitation.get('max_bots', -1)
|
||||||
self.ap,
|
if max_bots >= 0:
|
||||||
workspace_uuid,
|
existing_bots = await self.get_bots(context)
|
||||||
'bots.max',
|
if len(existing_bots) >= max_bots:
|
||||||
fallback=limitation.get('max_bots', -1),
|
raise ValueError(f'Maximum number of bots ({max_bots}) reached')
|
||||||
)
|
|
||||||
|
|
||||||
# TODO: 检查配置信息格式
|
# TODO: 检查配置信息格式
|
||||||
bot_data = bot_data.copy()
|
bot_data = bot_data.copy()
|
||||||
bot_data['uuid'] = str(uuid.uuid4())
|
bot_data['uuid'] = str(uuid.uuid4())
|
||||||
bot_data['workspace_uuid'] = workspace_uuid
|
bot_data['workspace_uuid'] = workspace_uuid
|
||||||
|
|
||||||
# Preserve the legacy flat-row result shape for this optional lookup;
|
# bind the most recently updated pipeline if any exist
|
||||||
# quota admission and insertion below still share one transaction.
|
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
scope_statement(
|
scope_statement(
|
||||||
sqlalchemy.select(persistence_pipeline.LegacyPipeline),
|
sqlalchemy.select(persistence_pipeline.LegacyPipeline),
|
||||||
@@ -131,25 +129,7 @@ class BotService:
|
|||||||
bot_data['use_pipeline_uuid'] = pipeline.uuid
|
bot_data['use_pipeline_uuid'] = pipeline.uuid
|
||||||
bot_data['use_pipeline_name'] = pipeline.name
|
bot_data['use_pipeline_name'] = pipeline.name
|
||||||
|
|
||||||
async def persist(execute) -> None:
|
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_bot.Bot).values(bot_data))
|
||||||
await require_resource_capacity(
|
|
||||||
execute,
|
|
||||||
workspace_uuid=workspace_uuid,
|
|
||||||
model=persistence_bot.Bot,
|
|
||||||
quota=quota,
|
|
||||||
resource_name='bots',
|
|
||||||
)
|
|
||||||
|
|
||||||
await execute(sqlalchemy.insert(persistence_bot.Bot).values(bot_data))
|
|
||||||
|
|
||||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
|
||||||
if quota.requires_transaction_lock:
|
|
||||||
if not callable(tenant_uow):
|
|
||||||
raise RuntimeError('Cloud bot quota enforcement requires transactional persistence')
|
|
||||||
async with tenant_uow(workspace_uuid) as uow:
|
|
||||||
await persist(uow.execute)
|
|
||||||
else:
|
|
||||||
await persist(self.ap.persistence_mgr.execute_async)
|
|
||||||
|
|
||||||
bot = await self.get_bot(context, bot_data['uuid'], include_secret=True)
|
bot = await self.get_bot(context, bot_data['uuid'], include_secret=True)
|
||||||
|
|
||||||
|
|||||||
@@ -5,10 +5,12 @@ import uuid
|
|||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
from langbot_plugin.api.entities.builtin.provider import message as provider_message
|
from langbot_plugin.api.entities.builtin.provider import message as provider_message
|
||||||
|
|
||||||
|
from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER
|
||||||
from ....core import app
|
from ....core import app
|
||||||
from ....entity.persistence import model as persistence_model
|
from ....entity.persistence import model as persistence_model
|
||||||
from ....entity.persistence import pipeline as persistence_pipeline
|
from ....entity.persistence import pipeline as persistence_pipeline
|
||||||
from ....provider.modelmgr import requester as model_requester
|
from ....provider.modelmgr import requester as model_requester
|
||||||
|
from ....provider.modelmgr import reasoning as model_reasoning
|
||||||
from ....workspace.errors import WorkspaceNotFoundError
|
from ....workspace.errors import WorkspaceNotFoundError
|
||||||
from .secrets import mask_secret_value, redact_secrets, restore_secret_placeholders
|
from .secrets import mask_secret_value, redact_secrets, restore_secret_placeholders
|
||||||
from .tenant import TenantContext, require_workspace_uuid, scope_statement
|
from .tenant import TenantContext, require_workspace_uuid, scope_statement
|
||||||
@@ -54,6 +56,53 @@ def _redact_model_secrets(model_data: dict) -> dict:
|
|||||||
return redacted
|
return redacted
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize_llm_reasoning(model_data: dict) -> None:
|
||||||
|
model_data['reasoning_config'] = model_reasoning.validate_reasoning_config(
|
||||||
|
model_data.get('reasoning_config'),
|
||||||
|
model_data.get('abilities'),
|
||||||
|
model_data.get('extra_args'),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_llm_reasoning_capability(
|
||||||
|
model_entity: persistence_model.LLMModel,
|
||||||
|
runtime_provider: model_requester.RuntimeProvider,
|
||||||
|
) -> None:
|
||||||
|
config = model_reasoning.normalize_reasoning_config(model_entity.reasoning_config)
|
||||||
|
if config['level'] == 'provider_default':
|
||||||
|
return
|
||||||
|
|
||||||
|
runtime_model = model_requester.RuntimeLLMModel(
|
||||||
|
execution_context=runtime_provider.execution_context,
|
||||||
|
model_entity=model_entity,
|
||||||
|
provider=runtime_provider,
|
||||||
|
)
|
||||||
|
capabilities = runtime_provider.requester.get_reasoning_capabilities(runtime_model)
|
||||||
|
model_reasoning.validate_reasoning_capabilities(config, capabilities, model_entity.name)
|
||||||
|
|
||||||
|
|
||||||
|
def _reasoning_capabilities(ap: app.Application, model: persistence_model.LLMModel) -> dict:
|
||||||
|
model_mgr = getattr(ap, 'model_mgr', None)
|
||||||
|
runtime_models = getattr(model_mgr, 'llm_model_dict', {}) if model_mgr is not None else {}
|
||||||
|
for runtime_model in runtime_models.values():
|
||||||
|
if (
|
||||||
|
runtime_model.model_entity.uuid == model.uuid
|
||||||
|
and runtime_model.model_entity.workspace_uuid == model.workspace_uuid
|
||||||
|
):
|
||||||
|
return runtime_model.provider.requester.get_reasoning_capabilities(runtime_model)
|
||||||
|
return model_reasoning.default_reasoning_capabilities(
|
||||||
|
supported='reasoning' in (model.abilities or []),
|
||||||
|
source='manual' if 'reasoning' in (model.abilities or []) else 'unknown',
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _serialize_llm_model(ap: app.Application, model: persistence_model.LLMModel) -> dict:
|
||||||
|
model_dict = ap.persistence_mgr.serialize_model(persistence_model.LLMModel, model)
|
||||||
|
model_dict['reasoning_config'] = model_reasoning.normalize_reasoning_config(model_dict.get('reasoning_config'))
|
||||||
|
model_dict['reasoning_capabilities'] = _reasoning_capabilities(ap, model)
|
||||||
|
return model_dict
|
||||||
|
|
||||||
|
|
||||||
async def _validate_provider_supports(
|
async def _validate_provider_supports(
|
||||||
ap: app.Application,
|
ap: app.Application,
|
||||||
context: TenantContext,
|
context: TenantContext,
|
||||||
@@ -113,6 +162,23 @@ async def _require_workspace_provider(
|
|||||||
return provider
|
return provider
|
||||||
|
|
||||||
|
|
||||||
|
def _is_cloud_runtime(ap: app.Application) -> bool:
|
||||||
|
mode = getattr(ap.persistence_mgr, 'mode', None)
|
||||||
|
return getattr(mode, 'value', None) == 'cloud_runtime'
|
||||||
|
|
||||||
|
|
||||||
|
async def _assert_cloud_managed_provider_mutable(
|
||||||
|
ap: app.Application,
|
||||||
|
context: TenantContext,
|
||||||
|
provider_uuid: str,
|
||||||
|
) -> None:
|
||||||
|
if not _is_cloud_runtime(ap):
|
||||||
|
return
|
||||||
|
provider = await _require_workspace_provider(ap, context, provider_uuid)
|
||||||
|
if provider.get('requester') == LANGBOT_MODELS_PROVIDER_REQUESTER:
|
||||||
|
raise ValueError('LangBot Models is managed by Cloud and cannot be modified')
|
||||||
|
|
||||||
|
|
||||||
async def _require_runtime_provider(
|
async def _require_runtime_provider(
|
||||||
ap: app.Application,
|
ap: app.Application,
|
||||||
context: TenantContext,
|
context: TenantContext,
|
||||||
@@ -147,7 +213,7 @@ class LLMModelsService:
|
|||||||
|
|
||||||
models_list = []
|
models_list = []
|
||||||
for model in models:
|
for model in models:
|
||||||
model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, model)
|
model_dict = _serialize_llm_model(self.ap, model)
|
||||||
provider = providers.get(model.provider_uuid)
|
provider = providers.get(model.provider_uuid)
|
||||||
if provider:
|
if provider:
|
||||||
provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider)
|
provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider)
|
||||||
@@ -178,7 +244,7 @@ class LLMModelsService:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
models = result.all()
|
models = result.all()
|
||||||
serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, m) for m in models]
|
serialized = [_serialize_llm_model(self.ap, model) for model in models]
|
||||||
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
||||||
|
|
||||||
async def create_llm_model(
|
async def create_llm_model(
|
||||||
@@ -213,14 +279,19 @@ class LLMModelsService:
|
|||||||
model_data['provider_uuid'] = provider_uuid
|
model_data['provider_uuid'] = provider_uuid
|
||||||
|
|
||||||
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid'])
|
||||||
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'llm')
|
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'llm')
|
||||||
|
_normalize_llm_reasoning(model_data)
|
||||||
|
|
||||||
|
runtime_provider = await _require_runtime_provider(self.ap, context, model_data['provider_uuid'])
|
||||||
|
model_entity = persistence_model.LLMModel(**model_data)
|
||||||
|
_validate_llm_reasoning_capability(model_entity, runtime_provider)
|
||||||
|
|
||||||
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_model.LLMModel).values(**model_data))
|
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_model.LLMModel).values(**model_data))
|
||||||
|
|
||||||
runtime_provider = await _require_runtime_provider(self.ap, context, model_data['provider_uuid'])
|
|
||||||
runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
|
runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
|
||||||
context,
|
context,
|
||||||
persistence_model.LLMModel(**model_data),
|
model_entity,
|
||||||
runtime_provider,
|
runtime_provider,
|
||||||
)
|
)
|
||||||
await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model)
|
await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model)
|
||||||
@@ -268,7 +339,7 @@ class LLMModelsService:
|
|||||||
if model is None:
|
if model is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, model)
|
model_dict = _serialize_llm_model(self.ap, model)
|
||||||
|
|
||||||
# Get provider
|
# Get provider
|
||||||
provider_result = await self.ap.persistence_mgr.execute_async(
|
provider_result = await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -291,11 +362,17 @@ class LLMModelsService:
|
|||||||
|
|
||||||
return model_dict
|
return model_dict
|
||||||
|
|
||||||
async def update_llm_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None:
|
async def update_llm_model(
|
||||||
|
self,
|
||||||
|
context: TenantContext,
|
||||||
|
model_uuid: str,
|
||||||
|
model_data: dict,
|
||||||
|
) -> None:
|
||||||
"""Update an existing LLM model"""
|
"""Update an existing LLM model"""
|
||||||
existing_model = await self.get_llm_model(context, model_uuid, include_secret=True)
|
existing_model = await self.get_llm_model(context, model_uuid, include_secret=True)
|
||||||
if existing_model is None:
|
if existing_model is None:
|
||||||
raise WorkspaceNotFoundError('Model not found')
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
model_data = model_data.copy()
|
model_data = model_data.copy()
|
||||||
model_data.pop('uuid', None)
|
model_data.pop('uuid', None)
|
||||||
model_data.pop('workspace_uuid', None)
|
model_data.pop('workspace_uuid', None)
|
||||||
@@ -321,8 +398,21 @@ class LLMModelsService:
|
|||||||
|
|
||||||
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
||||||
await _require_workspace_provider(self.ap, context, provider_uuid)
|
await _require_workspace_provider(self.ap, context, provider_uuid)
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid)
|
||||||
await _validate_provider_supports(self.ap, context, provider_uuid, 'llm')
|
await _validate_provider_supports(self.ap, context, provider_uuid, 'llm')
|
||||||
|
|
||||||
|
merged_model_data = {
|
||||||
|
key: value
|
||||||
|
for key, value in {**existing_model, **model_data, 'provider_uuid': provider_uuid}.items()
|
||||||
|
if key not in {'provider', 'created_at', 'updated_at', 'reasoning_capabilities'}
|
||||||
|
}
|
||||||
|
_normalize_llm_reasoning(merged_model_data)
|
||||||
|
model_data['reasoning_config'] = merged_model_data['reasoning_config']
|
||||||
|
|
||||||
|
runtime_provider = await _require_runtime_provider(self.ap, context, provider_uuid)
|
||||||
|
model_entity = persistence_model.LLMModel(**_runtime_model_data(model_uuid, merged_model_data))
|
||||||
|
_validate_llm_reasoning_capability(model_entity, runtime_provider)
|
||||||
|
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
scope_statement(
|
scope_statement(
|
||||||
sqlalchemy.update(persistence_model.LLMModel)
|
sqlalchemy.update(persistence_model.LLMModel)
|
||||||
@@ -336,25 +426,20 @@ class LLMModelsService:
|
|||||||
raise WorkspaceNotFoundError('Model not found')
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
|
||||||
await self.ap.model_mgr.remove_llm_model(context, model_uuid)
|
await self.ap.model_mgr.remove_llm_model(context, model_uuid)
|
||||||
runtime_provider = await _require_runtime_provider(self.ap, context, provider_uuid)
|
|
||||||
runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
|
runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider(
|
||||||
context,
|
context,
|
||||||
persistence_model.LLMModel(
|
model_entity,
|
||||||
**_runtime_model_data(
|
|
||||||
model_uuid,
|
|
||||||
{
|
|
||||||
key: value
|
|
||||||
for key, value in {**existing_model, **model_data, 'provider_uuid': provider_uuid}.items()
|
|
||||||
if key not in {'provider', 'created_at', 'updated_at'}
|
|
||||||
},
|
|
||||||
)
|
|
||||||
),
|
|
||||||
runtime_provider,
|
runtime_provider,
|
||||||
)
|
)
|
||||||
await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model)
|
await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model)
|
||||||
|
|
||||||
async def delete_llm_model(self, context: TenantContext, model_uuid: str) -> None:
|
async def delete_llm_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||||
"""Delete an LLM model"""
|
"""Delete an LLM model"""
|
||||||
|
if _is_cloud_runtime(self.ap):
|
||||||
|
existing_model = await self.get_llm_model(context, model_uuid, include_secret=True)
|
||||||
|
if existing_model is None:
|
||||||
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
scope_statement(
|
scope_statement(
|
||||||
sqlalchemy.delete(persistence_model.LLMModel).where(persistence_model.LLMModel.uuid == model_uuid),
|
sqlalchemy.delete(persistence_model.LLMModel).where(persistence_model.LLMModel.uuid == model_uuid),
|
||||||
@@ -376,6 +461,7 @@ class LLMModelsService:
|
|||||||
raise WorkspaceNotFoundError('Model not found')
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
runtime_llm_model = await self.ap.model_mgr.get_model_by_uuid(context, model_uuid)
|
runtime_llm_model = await self.ap.model_mgr.get_model_by_uuid(context, model_uuid)
|
||||||
else:
|
else:
|
||||||
|
_normalize_llm_reasoning(model_data)
|
||||||
runtime_llm_model = await self.ap.model_mgr.init_temporary_runtime_llm_model(context, model_data)
|
runtime_llm_model = await self.ap.model_mgr.init_temporary_runtime_llm_model(context, model_data)
|
||||||
|
|
||||||
extra_args = model_data.get('extra_args', {})
|
extra_args = model_data.get('extra_args', {})
|
||||||
@@ -448,7 +534,10 @@ class EmbeddingModelsService:
|
|||||||
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
||||||
|
|
||||||
async def create_embedding_model(
|
async def create_embedding_model(
|
||||||
self, context: TenantContext, model_data: dict, preserve_uuid: bool = False
|
self,
|
||||||
|
context: TenantContext,
|
||||||
|
model_data: dict,
|
||||||
|
preserve_uuid: bool = False,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Create a new embedding model"""
|
"""Create a new embedding model"""
|
||||||
model_data = model_data.copy()
|
model_data = model_data.copy()
|
||||||
@@ -472,6 +561,7 @@ class EmbeddingModelsService:
|
|||||||
model_data['provider_uuid'] = provider_uuid
|
model_data['provider_uuid'] = provider_uuid
|
||||||
|
|
||||||
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid'])
|
||||||
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'text-embedding')
|
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'text-embedding')
|
||||||
|
|
||||||
await self.ap.persistence_mgr.execute_async(
|
await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -530,11 +620,17 @@ class EmbeddingModelsService:
|
|||||||
|
|
||||||
return model_dict
|
return model_dict
|
||||||
|
|
||||||
async def update_embedding_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None:
|
async def update_embedding_model(
|
||||||
|
self,
|
||||||
|
context: TenantContext,
|
||||||
|
model_uuid: str,
|
||||||
|
model_data: dict,
|
||||||
|
) -> None:
|
||||||
"""Update an existing embedding model"""
|
"""Update an existing embedding model"""
|
||||||
existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True)
|
existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True)
|
||||||
if existing_model is None:
|
if existing_model is None:
|
||||||
raise WorkspaceNotFoundError('Model not found')
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
model_data = model_data.copy()
|
model_data = model_data.copy()
|
||||||
model_data.pop('uuid', None)
|
model_data.pop('uuid', None)
|
||||||
model_data.pop('workspace_uuid', None)
|
model_data.pop('workspace_uuid', None)
|
||||||
@@ -559,6 +655,7 @@ class EmbeddingModelsService:
|
|||||||
|
|
||||||
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
||||||
await _require_workspace_provider(self.ap, context, provider_uuid)
|
await _require_workspace_provider(self.ap, context, provider_uuid)
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid)
|
||||||
await _validate_provider_supports(self.ap, context, provider_uuid, 'text-embedding')
|
await _validate_provider_supports(self.ap, context, provider_uuid, 'text-embedding')
|
||||||
|
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -593,6 +690,11 @@ class EmbeddingModelsService:
|
|||||||
|
|
||||||
async def delete_embedding_model(self, context: TenantContext, model_uuid: str) -> None:
|
async def delete_embedding_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||||
"""Delete an embedding model"""
|
"""Delete an embedding model"""
|
||||||
|
if _is_cloud_runtime(self.ap):
|
||||||
|
existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True)
|
||||||
|
if existing_model is None:
|
||||||
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
scope_statement(
|
scope_statement(
|
||||||
sqlalchemy.delete(persistence_model.EmbeddingModel).where(
|
sqlalchemy.delete(persistence_model.EmbeddingModel).where(
|
||||||
@@ -685,7 +787,12 @@ class RerankModelsService:
|
|||||||
serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, m) for m in models]
|
serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, m) for m in models]
|
||||||
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
return serialized if include_secret else [_redact_model_secrets(model) for model in serialized]
|
||||||
|
|
||||||
async def create_rerank_model(self, context: TenantContext, model_data: dict, preserve_uuid: bool = False) -> str:
|
async def create_rerank_model(
|
||||||
|
self,
|
||||||
|
context: TenantContext,
|
||||||
|
model_data: dict,
|
||||||
|
preserve_uuid: bool = False,
|
||||||
|
) -> str:
|
||||||
"""Create a new rerank model"""
|
"""Create a new rerank model"""
|
||||||
model_data = model_data.copy()
|
model_data = model_data.copy()
|
||||||
if not preserve_uuid:
|
if not preserve_uuid:
|
||||||
@@ -708,6 +815,7 @@ class RerankModelsService:
|
|||||||
model_data['provider_uuid'] = provider_uuid
|
model_data['provider_uuid'] = provider_uuid
|
||||||
|
|
||||||
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
await _require_workspace_provider(self.ap, context, model_data['provider_uuid'])
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid'])
|
||||||
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'rerank')
|
await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'rerank')
|
||||||
|
|
||||||
await self.ap.persistence_mgr.execute_async(
|
await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -766,11 +874,17 @@ class RerankModelsService:
|
|||||||
|
|
||||||
return model_dict
|
return model_dict
|
||||||
|
|
||||||
async def update_rerank_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None:
|
async def update_rerank_model(
|
||||||
|
self,
|
||||||
|
context: TenantContext,
|
||||||
|
model_uuid: str,
|
||||||
|
model_data: dict,
|
||||||
|
) -> None:
|
||||||
"""Update an existing rerank model"""
|
"""Update an existing rerank model"""
|
||||||
existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True)
|
existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True)
|
||||||
if existing_model is None:
|
if existing_model is None:
|
||||||
raise WorkspaceNotFoundError('Model not found')
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
model_data = model_data.copy()
|
model_data = model_data.copy()
|
||||||
model_data.pop('uuid', None)
|
model_data.pop('uuid', None)
|
||||||
model_data.pop('workspace_uuid', None)
|
model_data.pop('workspace_uuid', None)
|
||||||
@@ -795,6 +909,7 @@ class RerankModelsService:
|
|||||||
|
|
||||||
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid'])
|
||||||
await _require_workspace_provider(self.ap, context, provider_uuid)
|
await _require_workspace_provider(self.ap, context, provider_uuid)
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid)
|
||||||
await _validate_provider_supports(self.ap, context, provider_uuid, 'rerank')
|
await _validate_provider_supports(self.ap, context, provider_uuid, 'rerank')
|
||||||
|
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -829,6 +944,11 @@ class RerankModelsService:
|
|||||||
|
|
||||||
async def delete_rerank_model(self, context: TenantContext, model_uuid: str) -> None:
|
async def delete_rerank_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||||
"""Delete a rerank model"""
|
"""Delete a rerank model"""
|
||||||
|
if _is_cloud_runtime(self.ap):
|
||||||
|
existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True)
|
||||||
|
if existing_model is None:
|
||||||
|
raise WorkspaceNotFoundError('Model not found')
|
||||||
|
await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid'])
|
||||||
result = await self.ap.persistence_mgr.execute_async(
|
result = await self.ap.persistence_mgr.execute_async(
|
||||||
scope_statement(
|
scope_statement(
|
||||||
sqlalchemy.delete(persistence_model.RerankModel).where(
|
sqlalchemy.delete(persistence_model.RerankModel).where(
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import traceback
|
|||||||
|
|
||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
|
|
||||||
|
from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER
|
||||||
from ....core import app
|
from ....core import app
|
||||||
from ....entity.persistence import model as persistence_model
|
from ....entity.persistence import model as persistence_model
|
||||||
from ....workspace.errors import WorkspaceNotFoundError
|
from ....workspace.errors import WorkspaceNotFoundError
|
||||||
@@ -20,6 +21,20 @@ class ModelProviderService:
|
|||||||
def __init__(self, ap: app.Application) -> None:
|
def __init__(self, ap: app.Application) -> None:
|
||||||
self.ap = ap
|
self.ap = ap
|
||||||
|
|
||||||
|
def _is_cloud_runtime(self) -> bool:
|
||||||
|
mode = getattr(self.ap.persistence_mgr, 'mode', None)
|
||||||
|
return getattr(mode, 'value', None) == 'cloud_runtime'
|
||||||
|
|
||||||
|
def _system_requester_is_reserved(self, requester: object) -> bool:
|
||||||
|
return self._is_cloud_runtime() and requester == LANGBOT_MODELS_PROVIDER_REQUESTER
|
||||||
|
|
||||||
|
async def _assert_provider_mutable(self, context: TenantContext, provider_uuid: str) -> None:
|
||||||
|
if not self._is_cloud_runtime():
|
||||||
|
return
|
||||||
|
provider = await self.get_provider(context, provider_uuid)
|
||||||
|
if provider is not None and self._system_requester_is_reserved(provider.get('requester')):
|
||||||
|
raise ValueError('LangBot Models is managed by Cloud and cannot be modified')
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _normalize_api_keys(api_keys: str | list[str] | tuple[str, ...] | None) -> list[str]:
|
def _normalize_api_keys(api_keys: str | list[str] | tuple[str, ...] | None) -> list[str]:
|
||||||
if api_keys is None:
|
if api_keys is None:
|
||||||
@@ -99,6 +114,8 @@ class ModelProviderService:
|
|||||||
async def create_provider(self, context: TenantContext, provider_data: dict) -> str:
|
async def create_provider(self, context: TenantContext, provider_data: dict) -> str:
|
||||||
"""Create a new provider"""
|
"""Create a new provider"""
|
||||||
provider_data = provider_data.copy()
|
provider_data = provider_data.copy()
|
||||||
|
if self._system_requester_is_reserved(provider_data.get('requester')):
|
||||||
|
raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider')
|
||||||
provider_data['uuid'] = str(uuid.uuid4())
|
provider_data['uuid'] = str(uuid.uuid4())
|
||||||
provider_data['workspace_uuid'] = require_workspace_uuid(context)
|
provider_data['workspace_uuid'] = require_workspace_uuid(context)
|
||||||
provider_data['api_keys'] = self._normalize_api_keys(
|
provider_data['api_keys'] = self._normalize_api_keys(
|
||||||
@@ -115,7 +132,10 @@ class ModelProviderService:
|
|||||||
|
|
||||||
async def update_provider(self, context: TenantContext, provider_uuid: str, provider_data: dict) -> None:
|
async def update_provider(self, context: TenantContext, provider_uuid: str, provider_data: dict) -> None:
|
||||||
"""Update an existing provider"""
|
"""Update an existing provider"""
|
||||||
|
await self._assert_provider_mutable(context, provider_uuid)
|
||||||
provider_data = provider_data.copy()
|
provider_data = provider_data.copy()
|
||||||
|
if self._system_requester_is_reserved(provider_data.get('requester')):
|
||||||
|
raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider')
|
||||||
provider_data.pop('uuid', None)
|
provider_data.pop('uuid', None)
|
||||||
provider_data.pop('workspace_uuid', None)
|
provider_data.pop('workspace_uuid', None)
|
||||||
if 'api_keys' in provider_data:
|
if 'api_keys' in provider_data:
|
||||||
@@ -145,6 +165,7 @@ class ModelProviderService:
|
|||||||
|
|
||||||
async def delete_provider(self, context: TenantContext, provider_uuid: str) -> None:
|
async def delete_provider(self, context: TenantContext, provider_uuid: str) -> None:
|
||||||
"""Delete a provider (only if no models reference it)"""
|
"""Delete a provider (only if no models reference it)"""
|
||||||
|
await self._assert_provider_mutable(context, provider_uuid)
|
||||||
workspace_uuid = require_workspace_uuid(context)
|
workspace_uuid = require_workspace_uuid(context)
|
||||||
# Check if any models use this provider
|
# Check if any models use this provider
|
||||||
llm_result = await self.ap.persistence_mgr.execute_async(
|
llm_result = await self.ap.persistence_mgr.execute_async(
|
||||||
@@ -245,6 +266,8 @@ class ModelProviderService:
|
|||||||
api_keys: list,
|
api_keys: list,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Find existing provider or create new one"""
|
"""Find existing provider or create new one"""
|
||||||
|
if self._system_requester_is_reserved(requester):
|
||||||
|
raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider')
|
||||||
workspace_uuid = require_workspace_uuid(context)
|
workspace_uuid = require_workspace_uuid(context)
|
||||||
api_keys = self._normalize_api_keys(restore_secret_placeholders(api_keys, sensitive=True))
|
api_keys = self._normalize_api_keys(restore_secret_placeholders(api_keys, sensitive=True))
|
||||||
|
|
||||||
|
|||||||
@@ -59,6 +59,10 @@ class SpaceService:
|
|||||||
result_list = result.all()
|
result_list = result.all()
|
||||||
return result_list[0] if result_list else None
|
return result_list[0] if result_list else None
|
||||||
|
|
||||||
|
async def get_valid_access_token(self, user_email: str) -> str | None:
|
||||||
|
"""Return a current Space bearer, refreshing and persisting it when needed."""
|
||||||
|
return await self._ensure_valid_token(user_email)
|
||||||
|
|
||||||
async def _ensure_valid_token(self, user_email: str) -> str | None:
|
async def _ensure_valid_token(self, user_email: str) -> str | None:
|
||||||
"""Ensure access token is valid, refresh if expired. Returns valid access_token or None."""
|
"""Ensure access token is valid, refresh if expired. Returns valid access_token or None."""
|
||||||
user_obj = await self._get_user_by_email(user_email)
|
user_obj = await self._get_user_by_email(user_email)
|
||||||
@@ -117,7 +121,12 @@ class SpaceService:
|
|||||||
params['state'] = state
|
params['state'] = state
|
||||||
return f'{authorize_url}?{urlencode(params)}'
|
return f'{authorize_url}?{urlencode(params)}'
|
||||||
|
|
||||||
async def exchange_oauth_code(self, code: str) -> typing.Dict:
|
async def exchange_oauth_code(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
workspace_uuids: list[str] | None = None,
|
||||||
|
workspace_created_ats: dict[str, int] | None = None,
|
||||||
|
) -> typing.Dict:
|
||||||
"""Exchange OAuth authorization code for tokens"""
|
"""Exchange OAuth authorization code for tokens"""
|
||||||
from langbot.pkg.utils import constants
|
from langbot.pkg.utils import constants
|
||||||
|
|
||||||
@@ -127,7 +136,14 @@ class SpaceService:
|
|||||||
session = httpclient.get_session()
|
session = httpclient.get_session()
|
||||||
async with session.post(
|
async with session.post(
|
||||||
f'{space_url}/api/v1/accounts/oauth/token',
|
f'{space_url}/api/v1/accounts/oauth/token',
|
||||||
json={'code': code, 'instance_id': constants.instance_id},
|
json={
|
||||||
|
'code': code,
|
||||||
|
'instance_id': constants.instance_id,
|
||||||
|
# Sending an explicit empty list tells new Space servers not to
|
||||||
|
# synthesize a legacy instance-derived Workspace binding.
|
||||||
|
'workspace_uuids': workspace_uuids if workspace_uuids is not None else [],
|
||||||
|
'workspace_created_ats': workspace_created_ats or {},
|
||||||
|
},
|
||||||
) as response:
|
) as response:
|
||||||
if response.status != 200:
|
if response.status != 200:
|
||||||
error = await httpclient.read_text_limited(response)
|
error = await httpclient.read_text_limited(response)
|
||||||
|
|||||||
@@ -114,7 +114,7 @@ class UserService:
|
|||||||
if purpose == 'login' and account_uuid is not None:
|
if purpose == 'login' and account_uuid is not None:
|
||||||
raise ValueError('Login state cannot be bound to an Account')
|
raise ValueError('Login state cannot be bound to an Account')
|
||||||
if purpose != 'login' and launch_workspace_uuid is not None:
|
if purpose != 'login' and launch_workspace_uuid is not None:
|
||||||
raise ValueError('Launch Workspace state is only valid for Space login')
|
raise ValueError('Launch Workspace state is only valid for LangBot Account login')
|
||||||
if ttl_seconds <= 0:
|
if ttl_seconds <= 0:
|
||||||
raise ValueError('OAuth state lifetime must be positive')
|
raise ValueError('OAuth state lifetime must be positive')
|
||||||
|
|
||||||
@@ -327,7 +327,7 @@ class UserService:
|
|||||||
normalized_email = normalize_email(user_email)
|
normalized_email = normalize_email(user_email)
|
||||||
if self._uses_control_plane_directory():
|
if self._uses_control_plane_directory():
|
||||||
raise ControlPlaneDirectoryRequiredError(
|
raise ControlPlaneDirectoryRequiredError(
|
||||||
'Cloud invitation registration must use a Space account to preserve control-plane identity'
|
'Cloud invitation registration must use a LangBot Account to preserve control-plane identity'
|
||||||
)
|
)
|
||||||
invitation, _ = await self.ap.workspace_collaboration_service.inspect_invitation(invitation_token)
|
invitation, _ = await self.ap.workspace_collaboration_service.inspect_invitation(invitation_token)
|
||||||
if invitation.normalized_email != normalized_email:
|
if invitation.normalized_email != normalized_email:
|
||||||
@@ -394,13 +394,16 @@ class UserService:
|
|||||||
|
|
||||||
# Check if this user has a local password set
|
# Check if this user has a local password set
|
||||||
if not user_obj.password:
|
if not user_obj.password:
|
||||||
raise ValueError('请使用 Space 账户登录')
|
raise ValueError('请使用 LangBot 账号登录')
|
||||||
|
|
||||||
await self._verify_password(user_obj.password, password)
|
await self._verify_password(user_obj.password, password)
|
||||||
|
|
||||||
return await self.generate_jwt_token(user_obj)
|
return await self.generate_jwt_token(user_obj)
|
||||||
|
|
||||||
async def generate_jwt_token(self, account: user.User | str) -> str:
|
async def generate_jwt_token(
|
||||||
|
self,
|
||||||
|
account: user.User | str,
|
||||||
|
) -> str:
|
||||||
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||||
jwt_expire = self.ap.instance_config.data['system']['jwt']['expire']
|
jwt_expire = self.ap.instance_config.data['system']['jwt']['expire']
|
||||||
|
|
||||||
@@ -413,7 +416,7 @@ class UserService:
|
|||||||
# Lightweight unit-test and bootstrap callers may not have persistence wired.
|
# Lightweight unit-test and bootstrap callers may not have persistence wired.
|
||||||
account_obj = None
|
account_obj = None
|
||||||
|
|
||||||
payload = {
|
payload: dict[str, typing.Any] = {
|
||||||
'user': user_email,
|
'user': user_email,
|
||||||
'iss': self._jwt_identity()[0],
|
'iss': self._jwt_identity()[0],
|
||||||
'aud': self._jwt_identity()[1],
|
'aud': self._jwt_identity()[1],
|
||||||
@@ -776,8 +779,27 @@ class UserService:
|
|||||||
local_account = await self.get_user_by_email(user_email)
|
local_account = await self.get_user_by_email(user_email)
|
||||||
if local_account is None:
|
if local_account is None:
|
||||||
raise ValueError('User not found')
|
raise ValueError('User not found')
|
||||||
# Exchange code for tokens
|
# Exchange code for tokens and bind both installation and the active
|
||||||
token_data = await self.ap.space_service.exchange_oauth_code(code)
|
# OSS Workspace as independent identities.
|
||||||
|
workspace_service = getattr(self.ap, 'workspace_service', None)
|
||||||
|
if workspace_service is not None:
|
||||||
|
binding = await workspace_service.get_execution_binding()
|
||||||
|
created_at = binding.workspace_created_at
|
||||||
|
created_ts = (
|
||||||
|
int(created_at.replace(tzinfo=datetime.timezone.utc).timestamp())
|
||||||
|
if created_at.tzinfo is None
|
||||||
|
else int(created_at.timestamp())
|
||||||
|
)
|
||||||
|
token_data = await self.ap.space_service.exchange_oauth_code(
|
||||||
|
code,
|
||||||
|
[binding.workspace_uuid],
|
||||||
|
{binding.workspace_uuid: created_ts},
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Compatibility for early/bootstrap call sites that have not wired
|
||||||
|
# WorkspaceService yet; old Space servers still derive the legacy
|
||||||
|
# Workspace identity from instance_id when the field is omitted.
|
||||||
|
token_data = await self.ap.space_service.exchange_oauth_code(code)
|
||||||
access_token = token_data.get('access_token')
|
access_token = token_data.get('access_token')
|
||||||
refresh_token = token_data.get('refresh_token')
|
refresh_token = token_data.get('refresh_token')
|
||||||
expires_in = token_data.get('expires_in', 0)
|
expires_in = token_data.get('expires_in', 0)
|
||||||
@@ -803,7 +825,7 @@ class UserService:
|
|||||||
# Check if this Space account is already bound to another user
|
# Check if this Space account is already bound to another user
|
||||||
existing_space_user = await self.get_user_by_space_account_uuid(space_account_uuid)
|
existing_space_user = await self.get_user_by_space_account_uuid(space_account_uuid)
|
||||||
if existing_space_user and existing_space_user.normalized_email != normalize_email(user_email):
|
if existing_space_user and existing_space_user.normalized_email != normalize_email(user_email):
|
||||||
raise ValueError('This Space account is already bound to another user')
|
raise ValueError('This LangBot Account is already bound to another user')
|
||||||
|
|
||||||
# Update local account to Space account
|
# Update local account to Space account
|
||||||
normalized_email = normalize_email(user_email)
|
normalized_email = normalize_email(user_email)
|
||||||
|
|||||||
@@ -367,6 +367,12 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
def _ensure_control_token(self, *, allow_generate: bool) -> str:
|
def _ensure_control_token(self, *, allow_generate: bool) -> str:
|
||||||
if not self._control_token and allow_generate:
|
if not self._control_token and allow_generate:
|
||||||
self._control_token = secrets.token_urlsafe(48)
|
self._control_token = secrets.token_urlsafe(48)
|
||||||
|
if not self._control_token:
|
||||||
|
if getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'cloud':
|
||||||
|
raise BoxRuntimeUnavailableError(
|
||||||
|
f'{BOX_CONTROL_TOKEN_ENV} must be configured with a strong shared secret for a Cloud Box runtime'
|
||||||
|
)
|
||||||
|
return ''
|
||||||
try:
|
try:
|
||||||
self._control_token = validate_control_token(self._control_token)
|
self._control_token = validate_control_token(self._control_token)
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
@@ -376,19 +382,19 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
return self._control_token
|
return self._control_token
|
||||||
|
|
||||||
def get_control_headers(self) -> dict[str, str]:
|
def get_control_headers(self) -> dict[str, str]:
|
||||||
"""Headers for the instance-authenticated RPC control handshake."""
|
"""Return instance-scoped RPC headers and the optional shared secret."""
|
||||||
|
|
||||||
self._ensure_control_token(allow_generate=False)
|
self._ensure_control_token(allow_generate=False)
|
||||||
return {
|
headers = {BOX_INSTANCE_HEADER: self._trusted_instance_uuid}
|
||||||
BOX_CONTROL_TOKEN_HEADER: self._control_token,
|
if self._control_token:
|
||||||
BOX_INSTANCE_HEADER: self._trusted_instance_uuid,
|
headers[BOX_CONTROL_TOKEN_HEADER] = self._control_token
|
||||||
}
|
return headers
|
||||||
|
|
||||||
def get_relay_headers(
|
def get_relay_headers(
|
||||||
self,
|
self,
|
||||||
action_context: ActionContext,
|
action_context: ActionContext,
|
||||||
) -> dict[str, str]:
|
) -> dict[str, str]:
|
||||||
"""Return authenticated, placement-scoped relay handshake headers."""
|
"""Return instance- and placement-scoped relay handshake headers."""
|
||||||
|
|
||||||
context = ActionContext.model_validate(action_context).without_installation()
|
context = ActionContext.model_validate(action_context).without_installation()
|
||||||
if context.instance_uuid != self._trusted_instance_uuid:
|
if context.instance_uuid != self._trusted_instance_uuid:
|
||||||
|
|||||||
@@ -13,11 +13,12 @@ from typing import Any, Protocol, runtime_checkable
|
|||||||
from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
||||||
from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config
|
from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config
|
||||||
from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
|
from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
|
||||||
|
from .model_catalog import CloudModelCatalogProvider
|
||||||
|
|
||||||
|
|
||||||
CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap'
|
CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap'
|
||||||
REQUIRED_TENANT_ISOLATION_VERSION = 2
|
REQUIRED_TENANT_ISOLATION_VERSION = 2
|
||||||
SUPPORTED_PGVECTOR_DIMENSIONS = frozenset({384, 512, 768, 1024, 1536})
|
SUPPORTED_PGVECTOR_DIMENSIONS = frozenset({384, 512, 768, 1024, 1536, 3072})
|
||||||
|
|
||||||
|
|
||||||
class CloudBootstrapError(RuntimeError):
|
class CloudBootstrapError(RuntimeError):
|
||||||
@@ -50,6 +51,7 @@ class OpenSourceDeployment:
|
|||||||
)
|
)
|
||||||
directory_provider: None = None
|
directory_provider: None = None
|
||||||
manifest_provider: None = None
|
manifest_provider: None = None
|
||||||
|
model_catalog_provider: None = None
|
||||||
persistence_mode: str = 'oss_compat'
|
persistence_mode: str = 'oss_compat'
|
||||||
required_vector_backend: str | None = None
|
required_vector_backend: str | None = None
|
||||||
|
|
||||||
@@ -80,6 +82,7 @@ class VerifiedCloudDeployment:
|
|||||||
entitlement_provider: EntitlementProvider
|
entitlement_provider: EntitlementProvider
|
||||||
directory_provider: DirectoryProjectionProvider
|
directory_provider: DirectoryProjectionProvider
|
||||||
manifest_provider: CloudManifestProvider
|
manifest_provider: CloudManifestProvider
|
||||||
|
model_catalog_provider: CloudModelCatalogProvider
|
||||||
verification_key_id: str
|
verification_key_id: str
|
||||||
mode: str = dataclasses.field(default='cloud', init=False)
|
mode: str = dataclasses.field(default='cloud', init=False)
|
||||||
workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False)
|
workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False)
|
||||||
@@ -110,6 +113,8 @@ class VerifiedCloudDeployment:
|
|||||||
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter')
|
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter')
|
||||||
if not isinstance(self.manifest_provider, CloudManifestProvider):
|
if not isinstance(self.manifest_provider, CloudManifestProvider):
|
||||||
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter')
|
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter')
|
||||||
|
if not isinstance(self.model_catalog_provider, CloudModelCatalogProvider):
|
||||||
|
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a model catalog adapter')
|
||||||
|
|
||||||
def validate_instance_config(self, config: dict[str, Any]) -> None:
|
def validate_instance_config(self, config: dict[str, Any]) -> None:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ from ..entity.persistence.cloud_directory import DirectoryProjectionInbox, Direc
|
|||||||
from ..entity.persistence.user import AccountSource, AccountStatus, User
|
from ..entity.persistence.user import AccountSource, AccountStatus, User
|
||||||
from ..entity.persistence.workspace import (
|
from ..entity.persistence.workspace import (
|
||||||
MembershipRole,
|
MembershipRole,
|
||||||
|
MembershipSource,
|
||||||
MembershipStatus,
|
MembershipStatus,
|
||||||
Workspace,
|
Workspace,
|
||||||
WorkspaceExecutionSource,
|
WorkspaceExecutionSource,
|
||||||
@@ -358,6 +359,7 @@ class DirectoryProjectionService:
|
|||||||
|
|
||||||
await self._reconcile_entitlement_snapshot_set(snapshot)
|
await self._reconcile_entitlement_snapshot_set(snapshot)
|
||||||
self._publish_runtime_execution_projection(snapshot.workspaces)
|
self._publish_runtime_execution_projection(snapshot.workspaces)
|
||||||
|
self._request_model_catalog_sync()
|
||||||
self._record_batch_cardinality(
|
self._record_batch_cardinality(
|
||||||
active_workspaces=active_workspace_count,
|
active_workspaces=active_workspace_count,
|
||||||
workspaces=workspace_count,
|
workspaces=workspace_count,
|
||||||
@@ -466,6 +468,7 @@ class DirectoryProjectionService:
|
|||||||
returned.values(),
|
returned.values(),
|
||||||
affected_workspace_uuids=requested,
|
affected_workspace_uuids=requested,
|
||||||
)
|
)
|
||||||
|
self._request_model_catalog_sync()
|
||||||
self._record_batch_cardinality(
|
self._record_batch_cardinality(
|
||||||
active_workspaces=active_workspace_count,
|
active_workspaces=active_workspace_count,
|
||||||
workspaces=workspace_count,
|
workspaces=workspace_count,
|
||||||
@@ -475,6 +478,14 @@ class DirectoryProjectionService:
|
|||||||
self._record_success()
|
self._record_success()
|
||||||
self._consumer_cursor = batch.cursor
|
self._consumer_cursor = batch.cursor
|
||||||
|
|
||||||
|
def _request_model_catalog_sync(self) -> None:
|
||||||
|
"""Wake model provisioning after a committed directory change."""
|
||||||
|
|
||||||
|
service = getattr(self.ap, 'cloud_model_catalog_service', None)
|
||||||
|
request_sync = getattr(service, 'request_sync', None)
|
||||||
|
if callable(request_sync):
|
||||||
|
request_sync()
|
||||||
|
|
||||||
def _publish_runtime_execution_projection(
|
def _publish_runtime_execution_projection(
|
||||||
self,
|
self,
|
||||||
workspaces: Iterable[DirectoryWorkspace],
|
workspaces: Iterable[DirectoryWorkspace],
|
||||||
@@ -876,15 +887,15 @@ class DirectoryProjectionService:
|
|||||||
account_uuid=member.account_uuid,
|
account_uuid=member.account_uuid,
|
||||||
role=role,
|
role=role,
|
||||||
status=status,
|
status=status,
|
||||||
|
source=MembershipSource.CLOUD_PROJECTION.value,
|
||||||
joined_at=joined_at,
|
joined_at=joined_at,
|
||||||
projection_revision=member.projection_revision,
|
projection_revision=member.projection_revision,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
continue
|
continue
|
||||||
if membership.projection_revision == 0:
|
if membership.source != MembershipSource.CLOUD_PROJECTION.value:
|
||||||
# Revision zero is Core-owned collaboration state. Directory
|
# Core-owned collaboration state is never adopted based on
|
||||||
# projection seeds memberships, but must not overwrite later
|
# account provenance, revision, or matching account identity.
|
||||||
# invitation, role, or removal decisions made by Core.
|
|
||||||
continue
|
continue
|
||||||
if membership.uuid != member.membership_uuid:
|
if membership.uuid != member.membership_uuid:
|
||||||
raise DirectoryProjectionUnavailableError('Directory membership UUID changed for one account')
|
raise DirectoryProjectionUnavailableError('Directory membership UUID changed for one account')
|
||||||
@@ -896,11 +907,12 @@ class DirectoryProjectionService:
|
|||||||
raise DirectoryProjectionUnavailableError('Directory membership revision has conflicting contents')
|
raise DirectoryProjectionUnavailableError('Directory membership revision has conflicting contents')
|
||||||
membership.role = role
|
membership.role = role
|
||||||
membership.status = status
|
membership.status = status
|
||||||
|
membership.source = MembershipSource.CLOUD_PROJECTION.value
|
||||||
membership.joined_at = joined_at
|
membership.joined_at = joined_at
|
||||||
membership.projection_revision = member.projection_revision
|
membership.projection_revision = member.projection_revision
|
||||||
|
|
||||||
for account_uuid, membership in existing.items():
|
for account_uuid, membership in existing.items():
|
||||||
if account_uuid not in included_accounts and membership.projection_revision != 0:
|
if account_uuid not in included_accounts and membership.source == MembershipSource.CLOUD_PROJECTION.value:
|
||||||
membership.status = MembershipStatus.REMOVED.value
|
membership.status = MembershipStatus.REMOVED.value
|
||||||
membership.projection_revision = max(
|
membership.projection_revision = max(
|
||||||
int(membership.projection_revision),
|
int(membership.projection_revision),
|
||||||
|
|||||||
@@ -3,11 +3,9 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import base64
|
import base64
|
||||||
import binascii
|
import binascii
|
||||||
import datetime
|
|
||||||
import hashlib
|
import hashlib
|
||||||
import heapq
|
import heapq
|
||||||
import json
|
import json
|
||||||
import math
|
|
||||||
import os
|
import os
|
||||||
import time
|
import time
|
||||||
import typing
|
import typing
|
||||||
@@ -16,10 +14,8 @@ from collections.abc import Callable, Iterable
|
|||||||
from cryptography.exceptions import InvalidSignature
|
from cryptography.exceptions import InvalidSignature
|
||||||
from cryptography.hazmat.primitives import serialization
|
from cryptography.hazmat.primitives import serialization
|
||||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
|
||||||
import sqlalchemy
|
|
||||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
||||||
|
|
||||||
from ..entity.persistence.cloud_directory import SpaceLaunchAssertionConsumption
|
from .support_admin import SupportAdminReplayError, SupportAdminSessionError, hash_grant_jti
|
||||||
|
|
||||||
if typing.TYPE_CHECKING:
|
if typing.TYPE_CHECKING:
|
||||||
from ..core.app import Application
|
from ..core.app import Application
|
||||||
@@ -27,6 +23,7 @@ if typing.TYPE_CHECKING:
|
|||||||
|
|
||||||
CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
|
CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
|
||||||
LAUNCH_KIND = 'workspace.launch'
|
LAUNCH_KIND = 'workspace.launch'
|
||||||
|
SUPPORT_ADMIN_LAUNCH_KIND = 'workspace.support_admin_launch'
|
||||||
EXPECTED_ISSUER = 'langbot-space'
|
EXPECTED_ISSUER = 'langbot-space'
|
||||||
EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
|
EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
|
||||||
_CONSUMED_JTI_MAX_ENTRIES = 4096
|
_CONSUMED_JTI_MAX_ENTRIES = 4096
|
||||||
@@ -125,30 +122,66 @@ class SpaceLaunchService:
|
|||||||
*,
|
*,
|
||||||
expected_workspace_uuid: str | None = None,
|
expected_workspace_uuid: str | None = None,
|
||||||
) -> dict[str, str]:
|
) -> dict[str, str]:
|
||||||
claims, clock_skew_seconds = self._verify_assertion(assertion)
|
claims = self._verify_assertion(assertion)
|
||||||
payload = claims.get('payload')
|
payload = claims.get('payload')
|
||||||
if not isinstance(payload, dict):
|
if not isinstance(payload, dict):
|
||||||
raise SpaceLaunchError('Launch assertion payload must be a JSON object')
|
raise SpaceLaunchError('Launch assertion payload must be a JSON object')
|
||||||
account_uuid = _required_string(payload, 'account_uuid')
|
kind = _required_string(claims, 'kind')
|
||||||
workspace_uuid = _required_string(payload, 'workspace_uuid')
|
workspace_uuid = _required_string(payload, 'workspace_uuid')
|
||||||
return_path = _required_string(payload, 'return_path')
|
|
||||||
if (
|
|
||||||
not return_path.startswith('/')
|
|
||||||
or return_path.startswith('//')
|
|
||||||
or any(character in return_path for character in ('\\', '\r', '\n', '\t'))
|
|
||||||
):
|
|
||||||
raise SpaceLaunchError('Launch assertion return path is invalid')
|
|
||||||
if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
|
if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
|
||||||
raise SpaceLaunchError('Launch assertion targets another Workspace')
|
raise SpaceLaunchError('Launch assertion targets another Workspace')
|
||||||
replay_retention_expires_at = _required_int(claims, 'exp', minimum=1) + math.ceil(clock_skew_seconds)
|
if kind == SUPPORT_ADMIN_LAUNCH_KIND:
|
||||||
await self._consume_jti(_required_string(claims, 'jti'), replay_retention_expires_at)
|
if 'account_uuid' in payload:
|
||||||
return {
|
raise SpaceLaunchError('Admin launch assertion must not identify a customer Account')
|
||||||
|
if payload.get('launch_mode') != 'support_admin' or payload.get('principal_type') != 'support_admin':
|
||||||
|
raise SpaceLaunchError('Admin launch principal must be support_admin')
|
||||||
|
actor_account_uuid = _required_string(payload, 'actor_account_uuid')
|
||||||
|
if _required_string(payload, 'effective_role') != 'owner':
|
||||||
|
raise SpaceLaunchError('Admin launch effective role must be owner')
|
||||||
|
issued_at = _required_int(claims, 'iat')
|
||||||
|
expires_at = _required_int(claims, 'exp', minimum=1)
|
||||||
|
if expires_at - issued_at > 90:
|
||||||
|
raise SpaceLaunchError('Admin launch assertion lifetime exceeds 90 seconds')
|
||||||
|
grant_jti_hash = hash_grant_jti(_required_string(claims, 'jti'))
|
||||||
|
result = {
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'actor_account_uuid': actor_account_uuid,
|
||||||
|
'effective_role': 'owner',
|
||||||
|
'grant_jti_hash': grant_jti_hash,
|
||||||
|
}
|
||||||
|
support_service = getattr(self.ap, 'support_admin_session_service', None)
|
||||||
|
if support_service is None or not callable(getattr(support_service, 'consume_launch_grant', None)):
|
||||||
|
raise SpaceLaunchError('Durable support admin session service is unavailable')
|
||||||
|
try:
|
||||||
|
support_session = await support_service.consume_launch_grant(
|
||||||
|
grant_jti_hash=grant_jti_hash,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
actor_account_uuid=actor_account_uuid,
|
||||||
|
)
|
||||||
|
except SupportAdminReplayError as exc:
|
||||||
|
raise SpaceLaunchError('Launch assertion has already been consumed') from exc
|
||||||
|
except SupportAdminSessionError as exc:
|
||||||
|
raise SpaceLaunchError(str(exc)) from exc
|
||||||
|
result['support_admin_token'] = support_session.token
|
||||||
|
self.ap.logger.info(
|
||||||
|
'cloud_support_admin_launch_consumed actor_account_uuid=%s workspace_uuid=%s',
|
||||||
|
result['actor_account_uuid'],
|
||||||
|
workspace_uuid,
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
|
||||||
|
if payload.get('launch_mode') is not None:
|
||||||
|
raise SpaceLaunchError('Launch assertion mode is unsupported')
|
||||||
|
account_uuid = _required_string(payload, 'account_uuid')
|
||||||
|
result = {
|
||||||
'account_uuid': account_uuid,
|
'account_uuid': account_uuid,
|
||||||
'workspace_uuid': workspace_uuid,
|
'workspace_uuid': workspace_uuid,
|
||||||
'return_path': return_path,
|
|
||||||
}
|
}
|
||||||
|
await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
|
||||||
|
return result
|
||||||
|
|
||||||
def _verify_assertion(self, token: str) -> tuple[dict[str, typing.Any], float]:
|
def _verify_assertion(self, token: str) -> dict[str, typing.Any]:
|
||||||
if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
|
if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
|
||||||
raise SpaceLaunchError('Space direct launch requires verified Cloud mode')
|
raise SpaceLaunchError('Space direct launch requires verified Cloud mode')
|
||||||
public_key, key_id, clock_skew_seconds = self._trust_config()
|
public_key, key_id, clock_skew_seconds = self._trust_config()
|
||||||
@@ -184,8 +217,9 @@ class SpaceLaunchService:
|
|||||||
raise SpaceLaunchError('Launch assertion subject targets another instance')
|
raise SpaceLaunchError('Launch assertion subject targets another instance')
|
||||||
if _required_string(claims, 'instance_uuid') != instance_uuid:
|
if _required_string(claims, 'instance_uuid') != instance_uuid:
|
||||||
raise SpaceLaunchError('Launch assertion instance UUID does not match this Core')
|
raise SpaceLaunchError('Launch assertion instance UUID does not match this Core')
|
||||||
if _required_string(claims, 'kind') != LAUNCH_KIND:
|
kind = _required_string(claims, 'kind')
|
||||||
raise SpaceLaunchError('Launch assertion kind is not workspace.launch')
|
if kind not in {LAUNCH_KIND, SUPPORT_ADMIN_LAUNCH_KIND}:
|
||||||
|
raise SpaceLaunchError('Launch assertion kind is not supported')
|
||||||
|
|
||||||
issued_at = _required_int(claims, 'iat')
|
issued_at = _required_int(claims, 'iat')
|
||||||
not_before = _required_int(claims, 'nbf')
|
not_before = _required_int(claims, 'nbf')
|
||||||
@@ -199,7 +233,7 @@ class SpaceLaunchService:
|
|||||||
raise SpaceLaunchError('Launch assertion is expired')
|
raise SpaceLaunchError('Launch assertion is expired')
|
||||||
if expires_at <= max(issued_at, not_before):
|
if expires_at <= max(issued_at, not_before):
|
||||||
raise SpaceLaunchError('Launch assertion expiry must follow issue time')
|
raise SpaceLaunchError('Launch assertion expiry must follow issue time')
|
||||||
return claims, clock_skew_seconds
|
return claims
|
||||||
|
|
||||||
def _trust_config(self) -> tuple[Ed25519PublicKey, str, float]:
|
def _trust_config(self) -> tuple[Ed25519PublicKey, str, float]:
|
||||||
data = getattr(getattr(self.ap, 'instance_config', None), 'data', {}) or {}
|
data = getattr(getattr(self.ap, 'instance_config', None), 'data', {}) or {}
|
||||||
@@ -229,39 +263,19 @@ class SpaceLaunchService:
|
|||||||
async def _consume_jti(self, jti: str, expires_at: int) -> None:
|
async def _consume_jti(self, jti: str, expires_at: int) -> None:
|
||||||
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
|
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
|
||||||
now = int(self._wall_time())
|
now = int(self._wall_time())
|
||||||
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
|
|
||||||
instance_uuid = str(self.ap.workspace_service.instance_uuid)
|
|
||||||
if persistence_mgr is not None:
|
|
||||||
expires_at_datetime = datetime.datetime.fromtimestamp(expires_at, tz=datetime.timezone.utc)
|
|
||||||
now_datetime = datetime.datetime.fromtimestamp(now, tz=datetime.timezone.utc)
|
|
||||||
async with persistence_mgr.directory_projection_uow(instance_uuid) as uow:
|
|
||||||
await uow.session.execute(
|
|
||||||
sqlalchemy.delete(SpaceLaunchAssertionConsumption).where(
|
|
||||||
SpaceLaunchAssertionConsumption.instance_uuid == instance_uuid,
|
|
||||||
SpaceLaunchAssertionConsumption.expires_at < now_datetime,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
statement = (
|
|
||||||
pg_insert(SpaceLaunchAssertionConsumption)
|
|
||||||
.values(instance_uuid=instance_uuid, jti=digest, expires_at=expires_at_datetime)
|
|
||||||
.on_conflict_do_nothing(index_elements=['instance_uuid', 'jti'])
|
|
||||||
.returning(SpaceLaunchAssertionConsumption.jti)
|
|
||||||
)
|
|
||||||
result = await uow.session.execute(statement)
|
|
||||||
if result.scalar_one_or_none() is None:
|
|
||||||
raise SpaceLaunchError('Launch assertion has already been consumed')
|
|
||||||
return
|
|
||||||
|
|
||||||
# Lightweight unit-test and OSS compatibility fallback. Verified Cloud
|
|
||||||
# runtime always supplies the durable PostgreSQL persistence manager.
|
|
||||||
async with self._replay_lock:
|
async with self._replay_lock:
|
||||||
self._prune_consumed_jtis(now)
|
self._prune_consumed_jtis(now)
|
||||||
if digest in self._consumed_jtis:
|
if digest in self._consumed_jtis:
|
||||||
raise SpaceLaunchError('Launch assertion has already been consumed')
|
raise SpaceLaunchError('Launch assertion has already been consumed')
|
||||||
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES:
|
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES:
|
||||||
|
# Evicting a still-valid digest would make a signed launch
|
||||||
|
# assertion replayable. Bound memory by failing closed instead.
|
||||||
raise SpaceLaunchError('Launch assertion replay cache capacity reached')
|
raise SpaceLaunchError('Launch assertion replay cache capacity reached')
|
||||||
self._consumed_jtis[digest] = expires_at
|
self._consumed_jtis[digest] = expires_at
|
||||||
heapq.heappush(self._consumed_jti_expiry_heap, (expires_at, digest))
|
heapq.heappush(
|
||||||
|
self._consumed_jti_expiry_heap,
|
||||||
|
(expires_at, digest),
|
||||||
|
)
|
||||||
|
|
||||||
def _prune_consumed_jtis(self, now: int) -> None:
|
def _prune_consumed_jtis(self, now: int) -> None:
|
||||||
while self._consumed_jti_expiry_heap:
|
while self._consumed_jti_expiry_heap:
|
||||||
|
|||||||
@@ -0,0 +1,337 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import Any, Literal, Protocol, runtime_checkable
|
||||||
|
|
||||||
|
import sqlalchemy
|
||||||
|
from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
|
||||||
|
|
||||||
|
from ..entity.persistence import model as persistence_model
|
||||||
|
|
||||||
|
|
||||||
|
LANGBOT_MODELS_PROVIDER_REQUESTER = 'space-chat-completions'
|
||||||
|
LANGBOT_MODELS_PROVIDER_NAME = 'LangBot Models'
|
||||||
|
_MODEL_RESOURCE_NAMESPACE = uuid.UUID('94c703ca-1df5-4e91-bcd3-74ac65cb7921')
|
||||||
|
_SUPPORTED_CATEGORIES = {'chat', 'embedding', 'rerank'}
|
||||||
|
_MODEL_TABLES = (
|
||||||
|
persistence_model.LLMModel,
|
||||||
|
persistence_model.EmbeddingModel,
|
||||||
|
persistence_model.RerankModel,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class CloudModelCatalogItem(BaseModel):
|
||||||
|
model_config = ConfigDict(extra='forbid', frozen=True)
|
||||||
|
|
||||||
|
uuid: str = Field(min_length=1, max_length=255)
|
||||||
|
model_id: str = Field(min_length=1, max_length=255)
|
||||||
|
category: Literal['chat', 'embedding', 'rerank']
|
||||||
|
llm_abilities: tuple[str, ...] = ()
|
||||||
|
is_featured: bool = False
|
||||||
|
featured_order: int = 0
|
||||||
|
|
||||||
|
@field_validator('llm_abilities', mode='before')
|
||||||
|
@classmethod
|
||||||
|
def normalize_missing_abilities(cls, value: Any) -> Any:
|
||||||
|
return () if value is None else value
|
||||||
|
|
||||||
|
@field_validator('llm_abilities')
|
||||||
|
@classmethod
|
||||||
|
def validate_abilities(cls, value: tuple[str, ...]) -> tuple[str, ...]:
|
||||||
|
if any(not item.strip() or len(item) > 64 for item in value):
|
||||||
|
raise ValueError('Model abilities must be non-empty strings of at most 64 characters')
|
||||||
|
if len(set(value)) != len(value):
|
||||||
|
raise ValueError('Model abilities must be unique')
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
class CloudWorkspaceModelBilling(BaseModel):
|
||||||
|
model_config = ConfigDict(extra='forbid', frozen=True)
|
||||||
|
|
||||||
|
workspace_uuid: str = Field(min_length=36, max_length=36)
|
||||||
|
owner_account_uuid: str | None = Field(default=None, min_length=36, max_length=36)
|
||||||
|
api_key: SecretStr | None = None
|
||||||
|
credits: int | None = None
|
||||||
|
|
||||||
|
@field_validator('workspace_uuid')
|
||||||
|
@classmethod
|
||||||
|
def validate_uuid(cls, value: str) -> str:
|
||||||
|
return str(uuid.UUID(value))
|
||||||
|
|
||||||
|
@field_validator('owner_account_uuid')
|
||||||
|
@classmethod
|
||||||
|
def validate_optional_uuid(cls, value: str | None) -> str | None:
|
||||||
|
return None if value is None else str(uuid.UUID(value))
|
||||||
|
|
||||||
|
|
||||||
|
class CloudModelCatalogSnapshot(BaseModel):
|
||||||
|
model_config = ConfigDict(extra='forbid', frozen=True)
|
||||||
|
|
||||||
|
instance_uuid: str = Field(min_length=1, max_length=255)
|
||||||
|
generated_at: datetime
|
||||||
|
base_url: str = Field(min_length=1, max_length=512)
|
||||||
|
models: tuple[CloudModelCatalogItem, ...]
|
||||||
|
workspaces: tuple[CloudWorkspaceModelBilling, ...]
|
||||||
|
|
||||||
|
@field_validator('base_url')
|
||||||
|
@classmethod
|
||||||
|
def validate_base_url(cls, value: str) -> str:
|
||||||
|
normalized = value.rstrip('/')
|
||||||
|
if not normalized.startswith('https://'):
|
||||||
|
raise ValueError('Cloud model gateway base URL must use HTTPS')
|
||||||
|
return normalized
|
||||||
|
|
||||||
|
@field_validator('models')
|
||||||
|
@classmethod
|
||||||
|
def validate_models(cls, value: tuple[CloudModelCatalogItem, ...]) -> tuple[CloudModelCatalogItem, ...]:
|
||||||
|
if len(value) > 500:
|
||||||
|
raise ValueError('Cloud model catalog exceeds 500 models')
|
||||||
|
identities = {(item.category, item.uuid) for item in value}
|
||||||
|
if len(identities) != len(value):
|
||||||
|
raise ValueError('Cloud model catalog contains duplicate model identities')
|
||||||
|
return value
|
||||||
|
|
||||||
|
@field_validator('workspaces')
|
||||||
|
@classmethod
|
||||||
|
def validate_workspaces(
|
||||||
|
cls, value: tuple[CloudWorkspaceModelBilling, ...]
|
||||||
|
) -> tuple[CloudWorkspaceModelBilling, ...]:
|
||||||
|
if len(value) > 10_000:
|
||||||
|
raise ValueError('Cloud model catalog exceeds 10000 Workspaces')
|
||||||
|
identities = {item.workspace_uuid for item in value}
|
||||||
|
if len(identities) != len(value):
|
||||||
|
raise ValueError('Cloud model catalog contains duplicate Workspaces')
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
@runtime_checkable
|
||||||
|
class CloudModelCatalogProvider(Protocol):
|
||||||
|
async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot:
|
||||||
|
"""Fetch and verify the complete model catalog and Workspace billing projection."""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
def system_provider_uuid(workspace_uuid: str) -> str:
|
||||||
|
workspace = str(uuid.UUID(workspace_uuid))
|
||||||
|
return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:provider:{LANGBOT_MODELS_PROVIDER_REQUESTER}'))
|
||||||
|
|
||||||
|
|
||||||
|
def system_model_uuid(workspace_uuid: str, category: str, upstream_uuid: str) -> str:
|
||||||
|
workspace = str(uuid.UUID(workspace_uuid))
|
||||||
|
if category not in _SUPPORTED_CATEGORIES:
|
||||||
|
raise ValueError(f'Unsupported model category: {category}')
|
||||||
|
if not upstream_uuid:
|
||||||
|
raise ValueError('Upstream model UUID is required')
|
||||||
|
return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:model:{category}:{upstream_uuid}'))
|
||||||
|
|
||||||
|
|
||||||
|
class CloudModelCatalogSyncService:
|
||||||
|
"""Reconcile Space-owned model catalog and Owner billing tokens into every Cloud Workspace."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
ap: Any,
|
||||||
|
provider: CloudModelCatalogProvider,
|
||||||
|
instance_uuid: str,
|
||||||
|
*,
|
||||||
|
sync_interval_seconds: float = 3600.0,
|
||||||
|
) -> None:
|
||||||
|
if not isinstance(provider, CloudModelCatalogProvider):
|
||||||
|
raise TypeError('Cloud model catalog sync requires a CloudModelCatalogProvider')
|
||||||
|
if sync_interval_seconds < 10:
|
||||||
|
raise ValueError('Cloud model catalog sync interval must be at least 10 seconds')
|
||||||
|
self.ap = ap
|
||||||
|
self.provider = provider
|
||||||
|
self.instance_uuid = instance_uuid
|
||||||
|
self.sync_interval_seconds = float(sync_interval_seconds)
|
||||||
|
# A tenant UoW commits one Workspace at a time. Keep a durable in-memory
|
||||||
|
# convergence marker so a failed runtime reload is retried even when the
|
||||||
|
# following database reconciliation is a no-op.
|
||||||
|
self._runtime_reload_pending = False
|
||||||
|
self._workspace_credits: dict[str, int | None] = {}
|
||||||
|
self._sync_requested = asyncio.Event()
|
||||||
|
|
||||||
|
def get_workspace_credits(self, workspace_uuid: str) -> int | None:
|
||||||
|
"""Return the latest signed owner-credit projection for a Workspace."""
|
||||||
|
return self._workspace_credits.get(str(uuid.UUID(workspace_uuid)))
|
||||||
|
|
||||||
|
async def initialize(self) -> None:
|
||||||
|
await self.sync_once(reload_runtime=False)
|
||||||
|
|
||||||
|
def request_sync(self) -> None:
|
||||||
|
"""Wake the catalog loop after a directory Workspace change."""
|
||||||
|
|
||||||
|
self._sync_requested.set()
|
||||||
|
|
||||||
|
async def run(self) -> None:
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(self._sync_requested.wait(), timeout=self.sync_interval_seconds)
|
||||||
|
except TimeoutError:
|
||||||
|
pass
|
||||||
|
self._sync_requested.clear()
|
||||||
|
try:
|
||||||
|
await self.sync_once(reload_runtime=True)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception as exc:
|
||||||
|
# Exception messages can contain rendered SQL bound values,
|
||||||
|
# including provider API keys. Log only the exception class.
|
||||||
|
self.ap.logger.warning(f'Cloud model catalog synchronization failed ({type(exc).__name__})')
|
||||||
|
|
||||||
|
async def sync_once(self, *, reload_runtime: bool = True) -> dict[str, int]:
|
||||||
|
summary = {'workspaces': 0, 'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
snapshot: CloudModelCatalogSnapshot | None = None
|
||||||
|
sync_error: Exception | None = None
|
||||||
|
reload_error: Exception | None = None
|
||||||
|
try:
|
||||||
|
snapshot = await self.provider.fetch_model_catalog(self.instance_uuid)
|
||||||
|
if snapshot.instance_uuid != self.instance_uuid:
|
||||||
|
raise ValueError('Cloud model catalog targets another LangBot instance')
|
||||||
|
|
||||||
|
bindings = await self.ap.workspace_service.list_active_execution_bindings()
|
||||||
|
billing_by_workspace = {item.workspace_uuid: item for item in snapshot.workspaces}
|
||||||
|
missing = sorted(
|
||||||
|
binding.workspace_uuid for binding in bindings if binding.workspace_uuid not in billing_by_workspace
|
||||||
|
)
|
||||||
|
if missing:
|
||||||
|
raise ValueError(
|
||||||
|
f'Cloud model catalog is missing billing projections for {len(missing)} active Workspaces'
|
||||||
|
)
|
||||||
|
|
||||||
|
for binding in bindings:
|
||||||
|
counts = await self._sync_workspace(
|
||||||
|
binding.workspace_uuid,
|
||||||
|
snapshot,
|
||||||
|
billing_by_workspace[binding.workspace_uuid],
|
||||||
|
)
|
||||||
|
summary['workspaces'] += 1
|
||||||
|
workspace_changed = any(counts[key] > 0 for key in ('created', 'updated', 'deleted'))
|
||||||
|
if workspace_changed:
|
||||||
|
# _sync_workspace returns only after its tenant UoW commits.
|
||||||
|
self._runtime_reload_pending = True
|
||||||
|
for key in ('created', 'updated', 'deleted'):
|
||||||
|
summary[key] += counts[key]
|
||||||
|
self._workspace_credits[binding.workspace_uuid] = billing_by_workspace[binding.workspace_uuid].credits
|
||||||
|
except Exception as exc:
|
||||||
|
sync_error = exc
|
||||||
|
finally:
|
||||||
|
model_mgr = getattr(self.ap, 'model_mgr', None)
|
||||||
|
if reload_runtime and self._runtime_reload_pending and model_mgr is not None:
|
||||||
|
try:
|
||||||
|
await model_mgr.load_models_from_db()
|
||||||
|
except Exception as exc:
|
||||||
|
reload_error = exc
|
||||||
|
else:
|
||||||
|
self._runtime_reload_pending = False
|
||||||
|
|
||||||
|
if sync_error is not None:
|
||||||
|
if reload_error is not None:
|
||||||
|
raise sync_error from reload_error
|
||||||
|
raise sync_error
|
||||||
|
if reload_error is not None:
|
||||||
|
raise reload_error
|
||||||
|
|
||||||
|
changed = any(summary[key] > 0 for key in ('created', 'updated', 'deleted'))
|
||||||
|
if changed and snapshot is not None:
|
||||||
|
self.ap.logger.info(
|
||||||
|
'Cloud model catalog synchronized '
|
||||||
|
f'({summary["workspaces"]} Workspaces, {len(snapshot.models)} models, '
|
||||||
|
f'created={summary["created"]}, updated={summary["updated"]}, deleted={summary["deleted"]})'
|
||||||
|
)
|
||||||
|
return summary
|
||||||
|
|
||||||
|
async def _sync_workspace(
|
||||||
|
self,
|
||||||
|
workspace_uuid: str,
|
||||||
|
snapshot: CloudModelCatalogSnapshot,
|
||||||
|
billing: CloudWorkspaceModelBilling,
|
||||||
|
) -> dict[str, int]:
|
||||||
|
counts = {'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
provider_uuid = system_provider_uuid(workspace_uuid)
|
||||||
|
desired_keys = [billing.api_key.get_secret_value()] if billing.api_key is not None else []
|
||||||
|
|
||||||
|
async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow:
|
||||||
|
provider = await uow.session.scalar(
|
||||||
|
sqlalchemy.select(persistence_model.ModelProvider).where(
|
||||||
|
persistence_model.ModelProvider.uuid == provider_uuid
|
||||||
|
)
|
||||||
|
)
|
||||||
|
provider_values = {
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'name': LANGBOT_MODELS_PROVIDER_NAME,
|
||||||
|
'requester': LANGBOT_MODELS_PROVIDER_REQUESTER,
|
||||||
|
'base_url': snapshot.base_url,
|
||||||
|
'api_keys': desired_keys,
|
||||||
|
}
|
||||||
|
if provider is None:
|
||||||
|
provider = persistence_model.ModelProvider(uuid=provider_uuid, **provider_values)
|
||||||
|
uow.session.add(provider)
|
||||||
|
await uow.session.flush()
|
||||||
|
counts['created'] += 1
|
||||||
|
elif self._update_entity(provider, provider_values):
|
||||||
|
counts['updated'] += 1
|
||||||
|
|
||||||
|
existing_by_table: dict[type, dict[str, Any]] = {}
|
||||||
|
for table in _MODEL_TABLES:
|
||||||
|
rows = (
|
||||||
|
await uow.session.scalars(sqlalchemy.select(table).where(table.provider_uuid == provider_uuid))
|
||||||
|
).all()
|
||||||
|
existing_by_table[table] = {row.uuid: row for row in rows}
|
||||||
|
|
||||||
|
desired_ids: dict[type, set[str]] = {table: set() for table in _MODEL_TABLES}
|
||||||
|
for item in snapshot.models:
|
||||||
|
table, values = self._model_values(workspace_uuid, provider_uuid, item)
|
||||||
|
model_uuid = system_model_uuid(workspace_uuid, item.category, item.uuid)
|
||||||
|
desired_ids[table].add(model_uuid)
|
||||||
|
existing = existing_by_table[table].get(model_uuid)
|
||||||
|
if existing is None:
|
||||||
|
uow.session.add(table(uuid=model_uuid, **values))
|
||||||
|
counts['created'] += 1
|
||||||
|
elif self._update_entity(existing, values):
|
||||||
|
counts['updated'] += 1
|
||||||
|
|
||||||
|
for table, entities in existing_by_table.items():
|
||||||
|
for model_uuid, entity in entities.items():
|
||||||
|
if model_uuid not in desired_ids[table]:
|
||||||
|
await uow.session.delete(entity)
|
||||||
|
counts['deleted'] += 1
|
||||||
|
|
||||||
|
return counts
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _update_entity(entity: Any, values: dict[str, Any]) -> bool:
|
||||||
|
changed = False
|
||||||
|
for key, value in values.items():
|
||||||
|
if getattr(entity, key) != value:
|
||||||
|
setattr(entity, key, value)
|
||||||
|
changed = True
|
||||||
|
return changed
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _model_values(
|
||||||
|
workspace_uuid: str,
|
||||||
|
provider_uuid: str,
|
||||||
|
item: CloudModelCatalogItem,
|
||||||
|
) -> tuple[type, dict[str, Any]]:
|
||||||
|
ranking = 100 - item.featured_order if item.is_featured else 0
|
||||||
|
common = {
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'name': item.model_id,
|
||||||
|
'provider_uuid': provider_uuid,
|
||||||
|
'extra_args': {},
|
||||||
|
'prefered_ranking': ranking,
|
||||||
|
}
|
||||||
|
if item.category == 'chat':
|
||||||
|
return persistence_model.LLMModel, {
|
||||||
|
**common,
|
||||||
|
'abilities': list(item.llm_abilities),
|
||||||
|
'context_length': None,
|
||||||
|
}
|
||||||
|
if item.category == 'embedding':
|
||||||
|
return persistence_model.EmbeddingModel, common
|
||||||
|
if item.category == 'rerank':
|
||||||
|
return persistence_model.RerankModel, common
|
||||||
|
raise ValueError(f'Unsupported model category: {item.category}')
|
||||||
@@ -1,82 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from dataclasses import dataclass
|
|
||||||
from typing import Any, Awaitable, Callable
|
|
||||||
|
|
||||||
import sqlalchemy
|
|
||||||
|
|
||||||
from ..entity.persistence import workspace as persistence_workspace
|
|
||||||
from .entitlements import EntitlementResolver
|
|
||||||
|
|
||||||
|
|
||||||
Execute = Callable[[Any], Awaitable[Any]]
|
|
||||||
|
|
||||||
|
|
||||||
class WorkspaceQuotaExceededError(ValueError):
|
|
||||||
"""A stable business error raised when a workspace has no free slots."""
|
|
||||||
|
|
||||||
error_code = 'workspace_quota_exceeded'
|
|
||||||
|
|
||||||
def __init__(self, resource_name: str, limit: int) -> None:
|
|
||||||
self.resource_name = resource_name
|
|
||||||
self.limit = limit
|
|
||||||
super().__init__(f'Maximum number of {resource_name} ({limit}) reached')
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
|
||||||
class WorkspaceQuota:
|
|
||||||
limit: int
|
|
||||||
requires_transaction_lock: bool
|
|
||||||
|
|
||||||
|
|
||||||
async def resolve_workspace_quota(
|
|
||||||
ap: Any,
|
|
||||||
workspace_uuid: str,
|
|
||||||
limit_name: str,
|
|
||||||
*,
|
|
||||||
fallback: int = -1,
|
|
||||||
) -> WorkspaceQuota:
|
|
||||||
"""Resolve a plan-agnostic Cloud limit while preserving OSS configuration."""
|
|
||||||
|
|
||||||
resolver = getattr(ap, 'entitlement_resolver', None)
|
|
||||||
if isinstance(resolver, EntitlementResolver):
|
|
||||||
snapshot = await resolver.resolve(workspace_uuid)
|
|
||||||
return WorkspaceQuota(
|
|
||||||
limit=snapshot.limit(limit_name),
|
|
||||||
requires_transaction_lock=True,
|
|
||||||
)
|
|
||||||
return WorkspaceQuota(limit=fallback, requires_transaction_lock=False)
|
|
||||||
|
|
||||||
|
|
||||||
async def lock_workspace_for_quota(execute: Execute, workspace_uuid: str) -> None:
|
|
||||||
"""Serialize quota checks on the durable Workspace row within one transaction."""
|
|
||||||
|
|
||||||
result = await execute(
|
|
||||||
sqlalchemy.select(persistence_workspace.Workspace.uuid)
|
|
||||||
.where(persistence_workspace.Workspace.uuid == workspace_uuid)
|
|
||||||
.with_for_update()
|
|
||||||
)
|
|
||||||
if result.first() is None:
|
|
||||||
raise ValueError('Workspace does not exist')
|
|
||||||
|
|
||||||
|
|
||||||
async def require_resource_capacity(
|
|
||||||
execute: Execute,
|
|
||||||
*,
|
|
||||||
workspace_uuid: str,
|
|
||||||
model: type,
|
|
||||||
quota: WorkspaceQuota,
|
|
||||||
resource_name: str,
|
|
||||||
workspace_locked: bool = False,
|
|
||||||
) -> None:
|
|
||||||
if quota.limit < 0:
|
|
||||||
return
|
|
||||||
if quota.requires_transaction_lock and not workspace_locked:
|
|
||||||
await lock_workspace_for_quota(execute, workspace_uuid)
|
|
||||||
result = await execute(
|
|
||||||
sqlalchemy.select(sqlalchemy.func.count())
|
|
||||||
.select_from(model)
|
|
||||||
.where(model.workspace_uuid == workspace_uuid)
|
|
||||||
)
|
|
||||||
if int(result.scalar_one()) >= quota.limit:
|
|
||||||
raise WorkspaceQuotaExceededError(resource_name, quota.limit)
|
|
||||||
@@ -0,0 +1,248 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import dataclasses
|
||||||
|
import datetime
|
||||||
|
import hashlib
|
||||||
|
import re
|
||||||
|
import time
|
||||||
|
import typing
|
||||||
|
|
||||||
|
import jwt
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
|
||||||
|
from ..entity.persistence.support_admin import SupportAdminTemporarySession
|
||||||
|
from ..workspace.errors import WorkspaceError
|
||||||
|
|
||||||
|
if typing.TYPE_CHECKING:
|
||||||
|
from ..core.app import Application
|
||||||
|
|
||||||
|
|
||||||
|
SUPPORT_ADMIN_TOKEN_TYP = 'langbot-support-admin+jwt'
|
||||||
|
SUPPORT_ADMIN_TOKEN_KIND = 'support_admin.session'
|
||||||
|
SUPPORT_ADMIN_EFFECTIVE_ROLE = 'owner'
|
||||||
|
SUPPORT_ADMIN_MAX_TOKEN_SECONDS = 300
|
||||||
|
_SHA256_HEX = re.compile(r'^[0-9a-f]{64}$')
|
||||||
|
|
||||||
|
|
||||||
|
class SupportAdminSessionError(ValueError):
|
||||||
|
"""Raised when a support-admin session or token is not admissible."""
|
||||||
|
|
||||||
|
|
||||||
|
class SupportAdminReplayError(SupportAdminSessionError):
|
||||||
|
"""Raised when a launch grant JTI has already been consumed."""
|
||||||
|
|
||||||
|
|
||||||
|
@dataclasses.dataclass(frozen=True, slots=True)
|
||||||
|
class IssuedSupportAdminSession:
|
||||||
|
token: str
|
||||||
|
grant_jti_hash: str
|
||||||
|
workspace_uuid: str
|
||||||
|
actor_account_uuid: str
|
||||||
|
issued_at: datetime.datetime
|
||||||
|
expires_at: datetime.datetime
|
||||||
|
|
||||||
|
|
||||||
|
@dataclasses.dataclass(frozen=True, slots=True)
|
||||||
|
class SupportAdminSessionIdentity:
|
||||||
|
grant_jti_hash: str
|
||||||
|
workspace_uuid: str
|
||||||
|
actor_account_uuid: str
|
||||||
|
instance_uuid: str
|
||||||
|
placement_generation: int
|
||||||
|
|
||||||
|
|
||||||
|
def hash_grant_jti(jti: str) -> str:
|
||||||
|
return hashlib.sha256(jti.encode('utf-8')).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
class SupportAdminSessionService:
|
||||||
|
"""Issue and validate temporary Workspace-scoped support-admin sessions."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
ap: Application,
|
||||||
|
*,
|
||||||
|
wall_time: typing.Callable[[], float] = time.time,
|
||||||
|
) -> None:
|
||||||
|
self.ap = ap
|
||||||
|
self._wall_time = wall_time
|
||||||
|
|
||||||
|
async def consume_launch_grant(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
grant_jti_hash: str,
|
||||||
|
workspace_uuid: str,
|
||||||
|
actor_account_uuid: str,
|
||||||
|
) -> IssuedSupportAdminSession:
|
||||||
|
self._validate_grant_hash(grant_jti_hash)
|
||||||
|
if not workspace_uuid or not actor_account_uuid:
|
||||||
|
raise SupportAdminSessionError('Support admin session requires an actor and Workspace')
|
||||||
|
|
||||||
|
issued_at = self._utcnow()
|
||||||
|
expires_at = issued_at + datetime.timedelta(seconds=SUPPORT_ADMIN_MAX_TOKEN_SECONDS)
|
||||||
|
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||||
|
if not callable(tenant_uow):
|
||||||
|
raise SupportAdminSessionError('Support admin sessions require tenant persistence')
|
||||||
|
|
||||||
|
try:
|
||||||
|
async with tenant_uow(workspace_uuid) as uow:
|
||||||
|
await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
|
||||||
|
uow.session.add(
|
||||||
|
SupportAdminTemporarySession(
|
||||||
|
grant_jti_hash=grant_jti_hash,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
actor_account_uuid=actor_account_uuid,
|
||||||
|
issued_at=issued_at,
|
||||||
|
expires_at=expires_at,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await uow.session.flush()
|
||||||
|
except IntegrityError as exc:
|
||||||
|
raise SupportAdminReplayError('Launch assertion has already been consumed') from exc
|
||||||
|
except WorkspaceError as exc:
|
||||||
|
raise SupportAdminSessionError('Workspace is unavailable for support access') from exc
|
||||||
|
|
||||||
|
return IssuedSupportAdminSession(
|
||||||
|
token=self._encode_token(
|
||||||
|
grant_jti_hash=grant_jti_hash,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
actor_account_uuid=actor_account_uuid,
|
||||||
|
issued_at=issued_at,
|
||||||
|
expires_at=expires_at,
|
||||||
|
),
|
||||||
|
grant_jti_hash=grant_jti_hash,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
actor_account_uuid=actor_account_uuid,
|
||||||
|
issued_at=issued_at,
|
||||||
|
expires_at=expires_at,
|
||||||
|
)
|
||||||
|
|
||||||
|
def is_support_admin_token(self, token: str) -> bool:
|
||||||
|
"""Return True only for compact JWTs marked as support-admin tokens."""
|
||||||
|
|
||||||
|
if not isinstance(token, str) or token.count('.') != 2:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
header = jwt.get_unverified_header(token)
|
||||||
|
except jwt.PyJWTError:
|
||||||
|
return False
|
||||||
|
if header.get('typ') == SUPPORT_ADMIN_TOKEN_TYP:
|
||||||
|
return True
|
||||||
|
try:
|
||||||
|
payload = jwt.decode(token, options={'verify_signature': False})
|
||||||
|
except jwt.PyJWTError:
|
||||||
|
return False
|
||||||
|
return payload.get('kind') == SUPPORT_ADMIN_TOKEN_KIND
|
||||||
|
|
||||||
|
async def authenticate_token(
|
||||||
|
self,
|
||||||
|
token: str,
|
||||||
|
*,
|
||||||
|
requested_workspace_uuid: str | None,
|
||||||
|
) -> SupportAdminSessionIdentity:
|
||||||
|
if not self.is_support_admin_token(token):
|
||||||
|
raise SupportAdminSessionError('Not a support admin token')
|
||||||
|
workspace_uuid = (requested_workspace_uuid or '').strip()
|
||||||
|
if not workspace_uuid:
|
||||||
|
raise SupportAdminSessionError('Support admin token requires an explicit Workspace selector')
|
||||||
|
|
||||||
|
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||||
|
try:
|
||||||
|
payload = jwt.decode(
|
||||||
|
token,
|
||||||
|
jwt_secret,
|
||||||
|
algorithms=['HS256'],
|
||||||
|
issuer='langbot-core',
|
||||||
|
audience=self._audience(workspace_uuid),
|
||||||
|
options={'require': ['exp', 'iat', 'nbf', 'iss', 'aud']},
|
||||||
|
)
|
||||||
|
except jwt.PyJWTError as exc:
|
||||||
|
raise SupportAdminSessionError('Invalid support admin token') from exc
|
||||||
|
self._validate_payload(payload, workspace_uuid)
|
||||||
|
grant_jti_hash = payload['grant_jti_hash']
|
||||||
|
actor_account_uuid = payload['actor_account_uuid']
|
||||||
|
|
||||||
|
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||||
|
if not callable(tenant_uow):
|
||||||
|
raise SupportAdminSessionError('Support admin sessions require tenant persistence')
|
||||||
|
|
||||||
|
now = self._utcnow()
|
||||||
|
async with tenant_uow(workspace_uuid) as uow:
|
||||||
|
session = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
|
||||||
|
if (
|
||||||
|
session is None
|
||||||
|
or session.workspace_uuid != workspace_uuid
|
||||||
|
or session.actor_account_uuid != actor_account_uuid
|
||||||
|
or session.revoked_at is not None
|
||||||
|
or session.expires_at <= now
|
||||||
|
):
|
||||||
|
raise SupportAdminSessionError('Support admin session is inactive')
|
||||||
|
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
|
||||||
|
session.last_used_at = now
|
||||||
|
await uow.session.flush()
|
||||||
|
|
||||||
|
return SupportAdminSessionIdentity(
|
||||||
|
grant_jti_hash=grant_jti_hash,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
actor_account_uuid=actor_account_uuid,
|
||||||
|
instance_uuid=binding.instance_uuid,
|
||||||
|
placement_generation=binding.placement_generation,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def revoke_session(self, grant_jti_hash: str, workspace_uuid: str) -> None:
|
||||||
|
self._validate_grant_hash(grant_jti_hash)
|
||||||
|
now = self._utcnow()
|
||||||
|
async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow:
|
||||||
|
row = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
|
||||||
|
if row is not None and row.revoked_at is None:
|
||||||
|
row.revoked_at = now
|
||||||
|
|
||||||
|
def _encode_token(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
grant_jti_hash: str,
|
||||||
|
workspace_uuid: str,
|
||||||
|
actor_account_uuid: str,
|
||||||
|
issued_at: datetime.datetime,
|
||||||
|
expires_at: datetime.datetime,
|
||||||
|
) -> str:
|
||||||
|
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||||
|
payload: dict[str, typing.Any] = {
|
||||||
|
'kind': SUPPORT_ADMIN_TOKEN_KIND,
|
||||||
|
'iss': 'langbot-core',
|
||||||
|
'aud': self._audience(workspace_uuid),
|
||||||
|
'sub': f'support-admin:{actor_account_uuid}',
|
||||||
|
'iat': issued_at,
|
||||||
|
'nbf': issued_at,
|
||||||
|
'exp': expires_at,
|
||||||
|
'actor_account_uuid': actor_account_uuid,
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'effective_role': SUPPORT_ADMIN_EFFECTIVE_ROLE,
|
||||||
|
'grant_jti_hash': grant_jti_hash,
|
||||||
|
}
|
||||||
|
return jwt.encode(payload, jwt_secret, algorithm='HS256', headers={'typ': SUPPORT_ADMIN_TOKEN_TYP})
|
||||||
|
|
||||||
|
def _validate_payload(self, payload: dict[str, typing.Any], workspace_uuid: str) -> None:
|
||||||
|
if payload.get('kind') != SUPPORT_ADMIN_TOKEN_KIND:
|
||||||
|
raise SupportAdminSessionError('Invalid support admin token kind')
|
||||||
|
if payload.get('workspace_uuid') != workspace_uuid:
|
||||||
|
raise SupportAdminSessionError('Support admin session is scoped to another Workspace')
|
||||||
|
if payload.get('effective_role') != SUPPORT_ADMIN_EFFECTIVE_ROLE:
|
||||||
|
raise SupportAdminSessionError('Invalid support admin token role')
|
||||||
|
actor_account_uuid = payload.get('actor_account_uuid')
|
||||||
|
if not isinstance(actor_account_uuid, str) or not actor_account_uuid.strip():
|
||||||
|
raise SupportAdminSessionError('Invalid support admin actor')
|
||||||
|
grant_jti_hash = payload.get('grant_jti_hash')
|
||||||
|
if not isinstance(grant_jti_hash, str) or not _SHA256_HEX.match(grant_jti_hash):
|
||||||
|
raise SupportAdminSessionError('Invalid support admin grant')
|
||||||
|
|
||||||
|
def _audience(self, workspace_uuid: str) -> str:
|
||||||
|
return f'langbot-support-admin:{self.ap.workspace_service.instance_uuid}:{workspace_uuid}'
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _validate_grant_hash(grant_jti_hash: str) -> None:
|
||||||
|
if not _SHA256_HEX.match(grant_jti_hash):
|
||||||
|
raise SupportAdminSessionError('Invalid support admin grant')
|
||||||
|
|
||||||
|
def _utcnow(self) -> datetime.datetime:
|
||||||
|
return datetime.datetime.fromtimestamp(self._wall_time(), datetime.UTC).replace(tzinfo=None)
|
||||||
@@ -51,8 +51,10 @@ from ..workspace import collaboration as workspace_collaboration_module
|
|||||||
from ..workspace import invitation_delivery as invitation_delivery_module
|
from ..workspace import invitation_delivery as invitation_delivery_module
|
||||||
from ..cloud import bootstrap as cloud_bootstrap_module
|
from ..cloud import bootstrap as cloud_bootstrap_module
|
||||||
from ..cloud import launch as cloud_launch_module
|
from ..cloud import launch as cloud_launch_module
|
||||||
|
from ..cloud import support_admin as cloud_support_admin_module
|
||||||
from ..cloud import directory_projection as cloud_directory_projection_module
|
from ..cloud import directory_projection as cloud_directory_projection_module
|
||||||
from ..cloud import entitlements as cloud_entitlements_module
|
from ..cloud import entitlements as cloud_entitlements_module
|
||||||
|
from ..cloud import model_catalog as cloud_model_catalog_module
|
||||||
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType
|
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType
|
||||||
|
|
||||||
|
|
||||||
@@ -136,16 +138,17 @@ class Application:
|
|||||||
|
|
||||||
space_launch_service: cloud_launch_module.SpaceLaunchService = None
|
space_launch_service: cloud_launch_module.SpaceLaunchService = None
|
||||||
|
|
||||||
|
support_admin_session_service: cloud_support_admin_module.SupportAdminSessionService = None
|
||||||
|
|
||||||
deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None
|
deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None
|
||||||
|
|
||||||
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
|
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
|
||||||
|
directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None
|
||||||
|
cloud_model_catalog_service: cloud_model_catalog_module.CloudModelCatalogSyncService | None = None
|
||||||
manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None
|
manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None
|
||||||
|
|
||||||
entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None
|
entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None
|
||||||
|
|
||||||
directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None
|
|
||||||
|
|
||||||
vector_db_mgr: vectordb_mgr.VectorDBManager = None
|
vector_db_mgr: vectordb_mgr.VectorDBManager = None
|
||||||
|
|
||||||
http_ctrl: http_controller.HTTPController = None
|
http_ctrl: http_controller.HTTPController = None
|
||||||
@@ -246,6 +249,10 @@ class Application:
|
|||||||
{},
|
{},
|
||||||
)
|
)
|
||||||
),
|
),
|
||||||
|
'plugin_runtime_connected': bool(
|
||||||
|
self.plugin_connector is not None
|
||||||
|
and getattr(self.plugin_connector, '_runtime_available', lambda: False)()
|
||||||
|
),
|
||||||
}
|
}
|
||||||
mcp_loader = getattr(self.tool_mgr, 'mcp_tool_loader', None)
|
mcp_loader = getattr(self.tool_mgr, 'mcp_tool_loader', None)
|
||||||
runtime_stats.update(
|
runtime_stats.update(
|
||||||
@@ -303,6 +310,12 @@ class Application:
|
|||||||
name='cloud-directory-projection',
|
name='cloud-directory-projection',
|
||||||
scopes=[core_entities.LifecycleControlScope.APPLICATION],
|
scopes=[core_entities.LifecycleControlScope.APPLICATION],
|
||||||
)
|
)
|
||||||
|
if self.cloud_model_catalog_service is not None:
|
||||||
|
self.task_mgr.create_task(
|
||||||
|
self.cloud_model_catalog_service.run(),
|
||||||
|
name='cloud-model-catalog-sync',
|
||||||
|
scopes=[core_entities.LifecycleControlScope.APPLICATION],
|
||||||
|
)
|
||||||
if self.manifest_refresh_service is not None:
|
if self.manifest_refresh_service is not None:
|
||||||
self.task_mgr.create_task(
|
self.task_mgr.create_task(
|
||||||
self.manifest_refresh_service.run(),
|
self.manifest_refresh_service.run(),
|
||||||
|
|||||||
@@ -42,9 +42,11 @@ from ...workspace import collaboration as workspace_collaboration_module
|
|||||||
from ...workspace import invitation_delivery as invitation_delivery_module
|
from ...workspace import invitation_delivery as invitation_delivery_module
|
||||||
from ...cloud import bootstrap as cloud_bootstrap
|
from ...cloud import bootstrap as cloud_bootstrap
|
||||||
from ...cloud import launch as cloud_launch_module
|
from ...cloud import launch as cloud_launch_module
|
||||||
|
from ...cloud import support_admin as cloud_support_admin_module
|
||||||
from ...cloud.directory import directory_projection_limits_from_config
|
from ...cloud.directory import directory_projection_limits_from_config
|
||||||
from ...cloud.directory_projection import DirectoryProjectionService
|
from ...cloud.directory_projection import DirectoryProjectionService
|
||||||
from ...cloud.entitlements import EntitlementResolver
|
from ...cloud.entitlements import EntitlementResolver
|
||||||
|
from ...cloud.model_catalog import CloudModelCatalogSyncService
|
||||||
from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType
|
from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType
|
||||||
from ...api.http.authz import WorkspaceRequiredError
|
from ...api.http.authz import WorkspaceRequiredError
|
||||||
|
|
||||||
@@ -175,11 +177,22 @@ class BuildAppStage(stage.BootingStage):
|
|||||||
# of repeating tenant validation for every manager.
|
# of repeating tenant validation for every manager.
|
||||||
await workspace_service_inst.prime_startup_execution_bindings()
|
await workspace_service_inst.prime_startup_execution_bindings()
|
||||||
|
|
||||||
|
if not isinstance(deployment, cloud_bootstrap.VerifiedCloudDeployment):
|
||||||
|
raise RuntimeError('Multi-Workspace runtime requires a verified Cloud deployment')
|
||||||
|
cloud_model_catalog_service = CloudModelCatalogSyncService(
|
||||||
|
ap,
|
||||||
|
deployment.model_catalog_provider,
|
||||||
|
constants.instance_id,
|
||||||
|
)
|
||||||
|
await cloud_model_catalog_service.initialize()
|
||||||
|
ap.cloud_model_catalog_service = cloud_model_catalog_service
|
||||||
|
|
||||||
ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService(
|
ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService(
|
||||||
ap,
|
ap,
|
||||||
workspace_service_inst,
|
workspace_service_inst,
|
||||||
)
|
)
|
||||||
ap.invitation_delivery_service = invitation_delivery_module.InvitationDeliveryService(ap)
|
ap.invitation_delivery_service = invitation_delivery_module.InvitationDeliveryService(ap)
|
||||||
|
ap.support_admin_session_service = cloud_support_admin_module.SupportAdminSessionService(ap)
|
||||||
ap.space_launch_service = cloud_launch_module.SpaceLaunchService(ap)
|
ap.space_launch_service = cloud_launch_module.SpaceLaunchService(ap)
|
||||||
|
|
||||||
user_service_inst = user_service.UserService(ap)
|
user_service_inst = user_service.UserService(ap)
|
||||||
|
|||||||
@@ -41,6 +41,7 @@ _RUNTIME_POLICY_DEFAULTS = {
|
|||||||
}
|
}
|
||||||
},
|
},
|
||||||
'plugin': {
|
'plugin': {
|
||||||
|
'connect_timeout_seconds': 180.0,
|
||||||
'worker': {
|
'worker': {
|
||||||
'max_cpus': 1.0,
|
'max_cpus': 1.0,
|
||||||
'max_memory_mb': 512,
|
'max_memory_mb': 512,
|
||||||
@@ -56,7 +57,7 @@ _RUNTIME_POLICY_DEFAULTS = {
|
|||||||
'restart_failure_window_seconds': 30.0,
|
'restart_failure_window_seconds': 30.0,
|
||||||
'restart_circuit_open_seconds': 60.0,
|
'restart_circuit_open_seconds': 60.0,
|
||||||
'require_hard_limits': False,
|
'require_hard_limits': False,
|
||||||
}
|
},
|
||||||
},
|
},
|
||||||
'mcp': {'stdio': {'enabled': True}},
|
'mcp': {'stdio': {'enabled': True}},
|
||||||
'monitoring': {
|
'monitoring': {
|
||||||
|
|||||||
@@ -17,4 +17,4 @@ class SpaceAccountBindingRequiredError(AccountEmailMismatchError):
|
|||||||
code = 'space_account_binding_required'
|
code = 'space_account_binding_required'
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return 'This local Account must bind Space from Account settings before Space login'
|
return 'This local account must bind a LangBot Account from Account settings before LangBot Account login'
|
||||||
|
|||||||
@@ -67,26 +67,3 @@ class DirectoryProjectionInbox(Base):
|
|||||||
name='ck_directory_projection_inbox_fingerprint',
|
name='ck_directory_projection_inbox_fingerprint',
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class SpaceLaunchAssertionConsumption(Base):
|
|
||||||
"""Durable, instance-scoped replay ledger for signed Space launch assertions."""
|
|
||||||
|
|
||||||
__tablename__ = 'space_launch_assertion_consumptions'
|
|
||||||
|
|
||||||
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
|
|
||||||
jti = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
|
|
||||||
expires_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=False)
|
|
||||||
consumed_at = sqlalchemy.Column(
|
|
||||||
sqlalchemy.DateTime(timezone=True),
|
|
||||||
nullable=False,
|
|
||||||
server_default=sqlalchemy.func.now(),
|
|
||||||
)
|
|
||||||
|
|
||||||
__table_args__ = (
|
|
||||||
sqlalchemy.Index(
|
|
||||||
'ix_space_launch_assertion_consumptions_expiry',
|
|
||||||
'instance_uuid',
|
|
||||||
'expires_at',
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|||||||
@@ -48,6 +48,12 @@ class LLMModel(Base):
|
|||||||
provider_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
|
provider_uuid = sqlalchemy.Column(sqlalchemy.String(255), nullable=False)
|
||||||
abilities = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default=[])
|
abilities = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default=[])
|
||||||
context_length = sqlalchemy.Column(sqlalchemy.Integer, nullable=True)
|
context_length = sqlalchemy.Column(sqlalchemy.Integer, nullable=True)
|
||||||
|
reasoning_config = sqlalchemy.Column(
|
||||||
|
sqlalchemy.JSON,
|
||||||
|
nullable=False,
|
||||||
|
default=lambda: {'level': 'provider_default'},
|
||||||
|
server_default=sqlalchemy.text('\'{"level":"provider_default"}\''),
|
||||||
|
)
|
||||||
extra_args = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default={})
|
extra_args = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default={})
|
||||||
prefered_ranking = sqlalchemy.Column(sqlalchemy.Integer, nullable=False, default=0)
|
prefered_ranking = sqlalchemy.Column(sqlalchemy.Integer, nullable=False, default=0)
|
||||||
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
|
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
|
||||||
|
|||||||
@@ -0,0 +1,35 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlalchemy
|
||||||
|
|
||||||
|
from .base import Base
|
||||||
|
|
||||||
|
|
||||||
|
class SupportAdminTemporarySession(Base):
|
||||||
|
"""Temporary support-admin Workspace access session."""
|
||||||
|
|
||||||
|
__tablename__ = 'support_admin_temporary_sessions'
|
||||||
|
|
||||||
|
grant_jti_hash = sqlalchemy.Column(sqlalchemy.String(64), primary_key=True)
|
||||||
|
workspace_uuid = sqlalchemy.Column(
|
||||||
|
sqlalchemy.String(36),
|
||||||
|
sqlalchemy.ForeignKey('workspaces.uuid', ondelete='CASCADE'),
|
||||||
|
nullable=False,
|
||||||
|
)
|
||||||
|
actor_account_uuid = sqlalchemy.Column(sqlalchemy.String(36), nullable=False)
|
||||||
|
issued_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||||
|
expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||||
|
revoked_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||||
|
last_used_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||||
|
|
||||||
|
__table_args__ = (
|
||||||
|
sqlalchemy.Index(
|
||||||
|
'ix_support_admin_sessions_workspace_expiry',
|
||||||
|
'workspace_uuid',
|
||||||
|
'expires_at',
|
||||||
|
),
|
||||||
|
sqlalchemy.CheckConstraint(
|
||||||
|
'length(grant_jti_hash) = 64',
|
||||||
|
name='ck_support_admin_sessions_grant_jti_hash',
|
||||||
|
),
|
||||||
|
)
|
||||||
@@ -40,6 +40,11 @@ class MembershipStatus(enum.StrEnum):
|
|||||||
REMOVED = 'removed'
|
REMOVED = 'removed'
|
||||||
|
|
||||||
|
|
||||||
|
class MembershipSource(enum.StrEnum):
|
||||||
|
LOCAL = 'local'
|
||||||
|
CLOUD_PROJECTION = 'cloud_projection'
|
||||||
|
|
||||||
|
|
||||||
class InvitationStatus(enum.StrEnum):
|
class InvitationStatus(enum.StrEnum):
|
||||||
PENDING = 'pending'
|
PENDING = 'pending'
|
||||||
ACCEPTED = 'accepted'
|
ACCEPTED = 'accepted'
|
||||||
@@ -151,6 +156,11 @@ class WorkspaceMembership(Base):
|
|||||||
nullable=True,
|
nullable=True,
|
||||||
)
|
)
|
||||||
joined_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
joined_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||||
|
source = sqlalchemy.Column(
|
||||||
|
sqlalchemy.String(32),
|
||||||
|
nullable=False,
|
||||||
|
server_default=MembershipSource.LOCAL.value,
|
||||||
|
)
|
||||||
projection_revision = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
|
projection_revision = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
|
||||||
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
|
created_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False, server_default=sqlalchemy.func.now())
|
||||||
updated_at = sqlalchemy.Column(
|
updated_at = sqlalchemy.Column(
|
||||||
@@ -163,6 +173,13 @@ class WorkspaceMembership(Base):
|
|||||||
__table_args__ = (
|
__table_args__ = (
|
||||||
sqlalchemy.UniqueConstraint('workspace_uuid', 'account_uuid', name='uq_workspace_membership_account'),
|
sqlalchemy.UniqueConstraint('workspace_uuid', 'account_uuid', name='uq_workspace_membership_account'),
|
||||||
sqlalchemy.Index('ix_workspace_memberships_account_status', 'account_uuid', 'status'),
|
sqlalchemy.Index('ix_workspace_memberships_account_status', 'account_uuid', 'status'),
|
||||||
|
sqlalchemy.Index(
|
||||||
|
'uq_workspace_memberships_one_active_owner',
|
||||||
|
'workspace_uuid',
|
||||||
|
unique=True,
|
||||||
|
sqlite_where=sqlalchemy.text("role = 'owner' AND status = 'active'"),
|
||||||
|
postgresql_where=sqlalchemy.text("role = 'owner' AND status = 'active'"),
|
||||||
|
),
|
||||||
sqlalchemy.CheckConstraint(
|
sqlalchemy.CheckConstraint(
|
||||||
"role IN ('owner', 'admin', 'developer', 'operator', 'viewer')",
|
"role IN ('owner', 'admin', 'developer', 'operator', 'viewer')",
|
||||||
name='ck_workspace_memberships_role',
|
name='ck_workspace_memberships_role',
|
||||||
@@ -171,6 +188,10 @@ class WorkspaceMembership(Base):
|
|||||||
"status IN ('active', 'disabled', 'removed')",
|
"status IN ('active', 'disabled', 'removed')",
|
||||||
name='ck_workspace_memberships_status',
|
name='ck_workspace_memberships_status',
|
||||||
),
|
),
|
||||||
|
sqlalchemy.CheckConstraint(
|
||||||
|
"source IN ('local', 'cloud_projection')",
|
||||||
|
name='ck_workspace_memberships_source',
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -18,6 +18,17 @@ down_revision = '0008_mcp_resource_prefs'
|
|||||||
branch_labels = None
|
branch_labels = None
|
||||||
depends_on = None
|
depends_on = None
|
||||||
|
|
||||||
|
_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID('8ea04f29-8528-4cc3-bb28-30a838c89d76')
|
||||||
|
|
||||||
|
|
||||||
|
def _workspace_uuid_from_instance_id(instance_id: str) -> str:
|
||||||
|
value = instance_id.strip()
|
||||||
|
candidate = value[len('instance_') :] if value.startswith('instance_') else value
|
||||||
|
try:
|
||||||
|
return str(uuid.UUID(candidate))
|
||||||
|
except ValueError:
|
||||||
|
return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value))
|
||||||
|
|
||||||
|
|
||||||
def _table_names(conn: sa.Connection) -> set[str]:
|
def _table_names(conn: sa.Connection) -> set[str]:
|
||||||
return set(sa.inspect(conn).get_table_names())
|
return set(sa.inspect(conn).get_table_names())
|
||||||
@@ -403,7 +414,7 @@ def _bootstrap_default_workspace(conn: sa.Connection) -> None:
|
|||||||
.values(created_by_account_uuid=owner_account_uuid)
|
.values(created_by_account_uuid=owner_account_uuid)
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
workspace_uuid = str(uuid.uuid4())
|
workspace_uuid = _workspace_uuid_from_instance_id(instance_uuid)
|
||||||
conn.execute(
|
conn.execute(
|
||||||
workspaces.insert().values(
|
workspaces.insert().values(
|
||||||
uuid=workspace_uuid,
|
uuid=workspace_uuid,
|
||||||
|
|||||||
@@ -0,0 +1,88 @@
|
|||||||
|
"""add temporary support-admin sessions
|
||||||
|
|
||||||
|
Revision ID: 0016_support_admin_sessions
|
||||||
|
Revises: 0015_cloud_core_collab
|
||||||
|
Create Date: 2026-07-31
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
|
||||||
|
revision = '0016_support_admin_sessions'
|
||||||
|
down_revision = '0015_cloud_core_collab'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
|
||||||
|
_TABLE_NAME = 'support_admin_temporary_sessions'
|
||||||
|
_POLICY_NAME = 'langbot_workspace_isolation'
|
||||||
|
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||||
|
|
||||||
|
|
||||||
|
def _setting(name: str) -> str:
|
||||||
|
return f"NULLIF(current_setting('{name}', true), '')"
|
||||||
|
|
||||||
|
|
||||||
|
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||||
|
return conn.dialect.identifier_preparer.quote(identifier)
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||||
|
if _TABLE_NAME not in existing_tables:
|
||||||
|
op.create_table(
|
||||||
|
_TABLE_NAME,
|
||||||
|
sa.Column('grant_jti_hash', sa.String(64), nullable=False),
|
||||||
|
sa.Column(
|
||||||
|
'workspace_uuid',
|
||||||
|
sa.String(36),
|
||||||
|
sa.ForeignKey('workspaces.uuid', ondelete='CASCADE'),
|
||||||
|
nullable=False,
|
||||||
|
),
|
||||||
|
sa.Column('actor_account_uuid', sa.String(36), nullable=False),
|
||||||
|
sa.Column('issued_at', sa.DateTime(), nullable=False),
|
||||||
|
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||||
|
sa.Column('revoked_at', sa.DateTime(), nullable=True),
|
||||||
|
sa.Column('last_used_at', sa.DateTime(), nullable=True),
|
||||||
|
sa.CheckConstraint(
|
||||||
|
'length(grant_jti_hash) = 64',
|
||||||
|
name='ck_support_admin_sessions_grant_jti_hash',
|
||||||
|
),
|
||||||
|
sa.PrimaryKeyConstraint('grant_jti_hash'),
|
||||||
|
)
|
||||||
|
op.create_index(
|
||||||
|
'ix_support_admin_sessions_workspace_expiry',
|
||||||
|
_TABLE_NAME,
|
||||||
|
['workspace_uuid', 'expires_at'],
|
||||||
|
unique=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
if conn.dialect.name != 'postgresql':
|
||||||
|
return
|
||||||
|
|
||||||
|
table = _quote(conn, _TABLE_NAME)
|
||||||
|
policy = _quote(conn, _POLICY_NAME)
|
||||||
|
expression = f'workspace_uuid::text = {_setting(_TENANT_SETTING)}'
|
||||||
|
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||||
|
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||||
|
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||||
|
op.execute(
|
||||||
|
sa.text(
|
||||||
|
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
|
||||||
|
f'USING ({expression}) WITH CHECK ({expression})'
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
if conn.dialect.name == 'postgresql':
|
||||||
|
table = _quote(conn, _TABLE_NAME)
|
||||||
|
policy = _quote(conn, _POLICY_NAME)
|
||||||
|
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||||
|
op.drop_index('ix_support_admin_sessions_workspace_expiry', table_name=_TABLE_NAME)
|
||||||
|
op.drop_table(_TABLE_NAME)
|
||||||
@@ -0,0 +1,167 @@
|
|||||||
|
"""align the OSS Workspace UUID with the persisted instance identity
|
||||||
|
|
||||||
|
Revision ID: 0017_oss_workspace_identity
|
||||||
|
Revises: 0016_support_admin_sessions
|
||||||
|
Create Date: 2026-07-31
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
|
||||||
|
revision = '0017_oss_workspace_identity'
|
||||||
|
down_revision = '0016_support_admin_sessions'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID('8ea04f29-8528-4cc3-bb28-30a838c89d76')
|
||||||
|
_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
|
||||||
|
|
||||||
|
|
||||||
|
def _workspace_uuid_from_instance_id(instance_id: str) -> str:
|
||||||
|
value = instance_id.strip()
|
||||||
|
candidate = value[len('instance_') :] if value.startswith('instance_') else value
|
||||||
|
try:
|
||||||
|
return str(uuid.UUID(candidate))
|
||||||
|
except ValueError:
|
||||||
|
return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value))
|
||||||
|
|
||||||
|
|
||||||
|
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||||
|
return conn.dialect.identifier_preparer.quote(identifier)
|
||||||
|
|
||||||
|
|
||||||
|
def _defer_foreign_keys(conn: sa.Connection, inspector: sa.Inspector, table_names: list[str]) -> None:
|
||||||
|
"""Allow the transaction to re-key a connected tenant graph atomically."""
|
||||||
|
|
||||||
|
if conn.dialect.name == 'sqlite':
|
||||||
|
conn.execute(sa.text('PRAGMA defer_foreign_keys = ON'))
|
||||||
|
return
|
||||||
|
if conn.dialect.name != 'postgresql':
|
||||||
|
raise RuntimeError(f'Unsupported Workspace identity migration dialect: {conn.dialect.name}')
|
||||||
|
|
||||||
|
for table_name in table_names:
|
||||||
|
for foreign_key in inspector.get_foreign_keys(table_name):
|
||||||
|
constraint_name = foreign_key.get('name')
|
||||||
|
if not constraint_name:
|
||||||
|
continue
|
||||||
|
conn.execute(
|
||||||
|
sa.text(
|
||||||
|
f'ALTER TABLE {_quote(conn, table_name)} '
|
||||||
|
f'ALTER CONSTRAINT {_quote(conn, constraint_name)} DEFERRABLE INITIALLY DEFERRED'
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _suspend_postgres_rls(
|
||||||
|
conn: sa.Connection,
|
||||||
|
table_names: list[str],
|
||||||
|
) -> dict[str, tuple[bool, bool]]:
|
||||||
|
if conn.dialect.name != 'postgresql':
|
||||||
|
return {}
|
||||||
|
|
||||||
|
states: dict[str, tuple[bool, bool]] = {}
|
||||||
|
for table_name in table_names:
|
||||||
|
row = conn.execute(
|
||||||
|
sa.text('SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE oid = to_regclass(:table_name)'),
|
||||||
|
{'table_name': table_name},
|
||||||
|
).one()
|
||||||
|
enabled, forced = bool(row.relrowsecurity), bool(row.relforcerowsecurity)
|
||||||
|
states[table_name] = (enabled, forced)
|
||||||
|
table = _quote(conn, table_name)
|
||||||
|
if forced:
|
||||||
|
conn.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
|
||||||
|
if enabled:
|
||||||
|
conn.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
|
||||||
|
return states
|
||||||
|
|
||||||
|
|
||||||
|
def _restore_postgres_rls(conn: sa.Connection, states: dict[str, tuple[bool, bool]]) -> None:
|
||||||
|
for table_name, (enabled, forced) in states.items():
|
||||||
|
table = _quote(conn, table_name)
|
||||||
|
if enabled:
|
||||||
|
conn.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||||
|
if forced:
|
||||||
|
conn.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
table_names = inspector.get_table_names()
|
||||||
|
if 'workspaces' not in table_names:
|
||||||
|
return
|
||||||
|
|
||||||
|
metadata = sa.MetaData()
|
||||||
|
workspaces = sa.Table('workspaces', metadata, autoload_with=conn)
|
||||||
|
local_rows = conn.execute(sa.select(workspaces).where(workspaces.c.source == 'local')).mappings().all()
|
||||||
|
if not local_rows:
|
||||||
|
return
|
||||||
|
if len(local_rows) != 1:
|
||||||
|
raise RuntimeError('Cannot align OSS Workspace identity: expected exactly one local Workspace')
|
||||||
|
|
||||||
|
old_row = dict(local_rows[0])
|
||||||
|
old_uuid = old_row['uuid']
|
||||||
|
canonical_uuid = _workspace_uuid_from_instance_id(old_row['instance_uuid'])
|
||||||
|
if old_uuid == canonical_uuid:
|
||||||
|
return
|
||||||
|
if conn.execute(sa.select(workspaces.c.uuid).where(workspaces.c.uuid == canonical_uuid)).scalar_one_or_none():
|
||||||
|
raise RuntimeError(f'Cannot align OSS Workspace identity: target {canonical_uuid!r} already exists')
|
||||||
|
|
||||||
|
tenant_tables = [
|
||||||
|
table_name
|
||||||
|
for table_name in table_names
|
||||||
|
if table_name == 'workspaces'
|
||||||
|
or 'workspace_uuid' in {column['name'] for column in inspector.get_columns(table_name)}
|
||||||
|
]
|
||||||
|
rls_states = _suspend_postgres_rls(conn, tenant_tables)
|
||||||
|
try:
|
||||||
|
_defer_foreign_keys(conn, inspector, table_names)
|
||||||
|
|
||||||
|
# Release local source/slug uniqueness while the canonical parent exists
|
||||||
|
# alongside the old parent for the duration of this transaction.
|
||||||
|
temporary_slug = f'__workspace_rekey__{old_uuid}'
|
||||||
|
conn.execute(
|
||||||
|
workspaces.update()
|
||||||
|
.where(workspaces.c.uuid == old_uuid)
|
||||||
|
.values(source='cloud_projection', slug=temporary_slug)
|
||||||
|
)
|
||||||
|
new_row = dict(old_row)
|
||||||
|
new_row['uuid'] = canonical_uuid
|
||||||
|
conn.execute(workspaces.insert().values(**new_row))
|
||||||
|
|
||||||
|
for table_name in tenant_tables:
|
||||||
|
if table_name == 'workspaces':
|
||||||
|
continue
|
||||||
|
table = sa.Table(table_name, metadata, autoload_with=conn, extend_existing=True)
|
||||||
|
conn.execute(table.update().where(table.c.workspace_uuid == old_uuid).values(workspace_uuid=canonical_uuid))
|
||||||
|
|
||||||
|
if 'metadata' in table_names:
|
||||||
|
conn.execute(
|
||||||
|
sa.text('UPDATE metadata SET value = :canonical_uuid WHERE key = :key AND value = :old_uuid'),
|
||||||
|
{
|
||||||
|
'canonical_uuid': canonical_uuid,
|
||||||
|
'key': _OSS_WORKSPACE_METADATA_KEY,
|
||||||
|
'old_uuid': old_uuid,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
conn.execute(workspaces.delete().where(workspaces.c.uuid == old_uuid))
|
||||||
|
if conn.dialect.name == 'postgresql':
|
||||||
|
# Fire deferred FK triggers before ALTER TABLE restores RLS; PostgreSQL
|
||||||
|
# rejects ALTER TABLE while a relation has pending trigger events.
|
||||||
|
conn.execute(sa.text('SET CONSTRAINTS ALL IMMEDIATE'))
|
||||||
|
except Exception:
|
||||||
|
# Alembic owns the transaction. Rollback restores the transactional RLS DDL.
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
_restore_postgres_rls(conn, rls_states)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
# The previous random UUID is intentionally not recoverable. Keeping the
|
||||||
|
# canonical identity preserves every FK and is safe for older application code.
|
||||||
|
pass
|
||||||
@@ -0,0 +1,57 @@
|
|||||||
|
"""add llm reasoning config
|
||||||
|
|
||||||
|
Revision ID: 0018_llm_reasoning_config
|
||||||
|
Revises: 0017_oss_workspace_identity
|
||||||
|
Create Date: 2026-07-27
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision = '0018_llm_reasoning_config'
|
||||||
|
down_revision = '0017_oss_workspace_identity'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
|
||||||
|
_LLM_MODELS = sa.table(
|
||||||
|
'llm_models',
|
||||||
|
sa.column('reasoning_config', sa.JSON()),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'llm_models' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
|
||||||
|
columns = {column['name'] for column in inspector.get_columns('llm_models')}
|
||||||
|
if 'reasoning_config' in columns:
|
||||||
|
return
|
||||||
|
|
||||||
|
op.add_column(
|
||||||
|
'llm_models',
|
||||||
|
sa.Column(
|
||||||
|
'reasoning_config',
|
||||||
|
sa.JSON(),
|
||||||
|
nullable=True,
|
||||||
|
server_default=sa.text('\'{"level":"provider_default"}\''),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
conn.execute(_LLM_MODELS.update().values(reasoning_config={'level': 'provider_default'}))
|
||||||
|
with op.batch_alter_table('llm_models') as batch_op:
|
||||||
|
batch_op.alter_column('reasoning_config', existing_type=sa.JSON(), nullable=False)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'llm_models' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
columns = {column['name'] for column in inspector.get_columns('llm_models')}
|
||||||
|
if 'reasoning_config' in columns:
|
||||||
|
with op.batch_alter_table('llm_models') as batch_op:
|
||||||
|
batch_op.drop_column('reasoning_config')
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
"""merge the published Space launch replay and main migration branches
|
||||||
|
|
||||||
|
Revision ID: 0018_merge_launch_replay
|
||||||
|
Revises: 0016_space_launch_replay, 0017_oss_workspace_identity
|
||||||
|
Create Date: 2026-08-01
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
revision = '0018_merge_launch_replay'
|
||||||
|
down_revision = ('0016_space_launch_replay', '0017_oss_workspace_identity')
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
pass
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
"""enforce one active owner per Workspace
|
||||||
|
|
||||||
|
Revision ID: 0019_single_workspace_owner
|
||||||
|
Revises: 0018_merge_launch_replay
|
||||||
|
Create Date: 2026-08-02
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision = '0019_single_workspace_owner'
|
||||||
|
down_revision = '0018_merge_launch_replay'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
_INDEX_NAME = 'uq_workspace_memberships_one_active_owner'
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'workspace_memberships' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
|
||||||
|
# Ownership transfer used to promote a second member without demoting the
|
||||||
|
# original owner. Preserve the Workspace creator where possible and demote
|
||||||
|
# every historical extra owner before installing the database invariant.
|
||||||
|
op.execute(
|
||||||
|
sa.text(
|
||||||
|
"""
|
||||||
|
WITH ranked_owners AS (
|
||||||
|
SELECT membership.uuid,
|
||||||
|
ROW_NUMBER() OVER (
|
||||||
|
PARTITION BY membership.workspace_uuid
|
||||||
|
ORDER BY
|
||||||
|
CASE
|
||||||
|
WHEN membership.account_uuid = workspace.created_by_account_uuid THEN 0
|
||||||
|
ELSE 1
|
||||||
|
END,
|
||||||
|
COALESCE(membership.joined_at, membership.created_at),
|
||||||
|
membership.uuid
|
||||||
|
) AS owner_rank
|
||||||
|
FROM workspace_memberships AS membership
|
||||||
|
JOIN workspaces AS workspace
|
||||||
|
ON workspace.uuid = membership.workspace_uuid
|
||||||
|
WHERE membership.role = 'owner'
|
||||||
|
AND membership.status = 'active'
|
||||||
|
)
|
||||||
|
UPDATE workspace_memberships
|
||||||
|
SET role = 'admin'
|
||||||
|
WHERE uuid IN (
|
||||||
|
SELECT uuid
|
||||||
|
FROM ranked_owners
|
||||||
|
WHERE owner_rank > 1
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
# Fresh installations may already have this index because SQLAlchemy
|
||||||
|
# metadata is created before Alembic advances the revision marker.
|
||||||
|
op.execute(
|
||||||
|
sa.text(
|
||||||
|
'CREATE UNIQUE INDEX IF NOT EXISTS '
|
||||||
|
'uq_workspace_memberships_one_active_owner '
|
||||||
|
'ON workspace_memberships (workspace_uuid) '
|
||||||
|
"WHERE role = 'owner' AND status = 'active'"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'workspace_memberships' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
index_names = {index['name'] for index in inspector.get_indexes('workspace_memberships')}
|
||||||
|
if _INDEX_NAME in index_names:
|
||||||
|
op.drop_index(_INDEX_NAME, table_name='workspace_memberships')
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
"""enable 3072-dimensional pgvector embeddings
|
||||||
|
|
||||||
|
Revision ID: 001a_pgvector_dimension_3072
|
||||||
|
Revises: 0019_single_workspace_owner
|
||||||
|
Create Date: 2026-08-05
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision = '001a_pgvector_dimension_3072'
|
||||||
|
down_revision = '0019_single_workspace_owner'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
_TABLE = 'langbot_vectors'
|
||||||
|
_CHECK = 'ck_langbot_vectors_embedding_dimension_enabled'
|
||||||
|
_INDEX = 'ix_langbot_vectors_hnsw_cosine_3072'
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
if conn.dialect.name != 'postgresql' or _TABLE not in sa.inspect(conn).get_table_names():
|
||||||
|
return
|
||||||
|
op.drop_constraint(_CHECK, _TABLE, type_='check')
|
||||||
|
op.create_check_constraint(_CHECK, _TABLE, 'embedding_dimension IN (384, 512, 768, 1024, 1536, 3072)')
|
||||||
|
op.execute(
|
||||||
|
sa.text(
|
||||||
|
f'CREATE INDEX {_INDEX} ON {_TABLE} USING hnsw ((embedding::halfvec(3072)) halfvec_cosine_ops) WHERE embedding_dimension = 3072'
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
if conn.dialect.name != 'postgresql' or _TABLE not in sa.inspect(conn).get_table_names():
|
||||||
|
return
|
||||||
|
count = conn.scalar(sa.text(f'SELECT COUNT(*) FROM {_TABLE} WHERE embedding_dimension = 3072'))
|
||||||
|
if count:
|
||||||
|
raise RuntimeError('Cannot disable 3072-dimensional pgvector while matching embeddings exist')
|
||||||
|
op.drop_index(_INDEX, table_name=_TABLE)
|
||||||
|
op.drop_constraint(_CHECK, _TABLE, type_='check')
|
||||||
|
op.create_check_constraint(_CHECK, _TABLE, 'embedding_dimension IN (384, 512, 768, 1024, 1536)')
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
"""add explicit Workspace membership source
|
||||||
|
|
||||||
|
Revision ID: 0020_membership_source
|
||||||
|
Revises: 001a_pgvector_dimension_3072
|
||||||
|
Create Date: 2026-08-06
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from alembic import op
|
||||||
|
|
||||||
|
revision = '0020_membership_source'
|
||||||
|
down_revision = '001a_pgvector_dimension_3072'
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
_CONSTRAINT_NAME = 'ck_workspace_memberships_source'
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'workspace_memberships' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
if 'source' in {column['name'] for column in inspector.get_columns('workspace_memberships')}:
|
||||||
|
return
|
||||||
|
|
||||||
|
# No durable historical field distinguishes Directory-created revision-zero
|
||||||
|
# rows from Core invitations. Protect every existing row; production can
|
||||||
|
# reclassify separately after UUIDs have been verified against Space.
|
||||||
|
with op.batch_alter_table('workspace_memberships') as batch_op:
|
||||||
|
batch_op.add_column(sa.Column('source', sa.String(length=32), nullable=False, server_default='local'))
|
||||||
|
batch_op.create_check_constraint(
|
||||||
|
_CONSTRAINT_NAME,
|
||||||
|
"source IN ('local', 'cloud_projection')",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
conn = op.get_bind()
|
||||||
|
inspector = sa.inspect(conn)
|
||||||
|
if 'workspace_memberships' not in inspector.get_table_names():
|
||||||
|
return
|
||||||
|
if 'source' not in {column['name'] for column in inspector.get_columns('workspace_memberships')}:
|
||||||
|
return
|
||||||
|
with op.batch_alter_table('workspace_memberships') as batch_op:
|
||||||
|
batch_op.drop_constraint(_CONSTRAINT_NAME, type_='check')
|
||||||
|
batch_op.drop_column('source')
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
"""merge reasoning config with the main migration branch
|
||||||
|
|
||||||
|
Revision ID: 0021_merge_reasoning_config
|
||||||
|
Revises: 0020_membership_source, 0018_llm_reasoning_config
|
||||||
|
Create Date: 2026-08-09
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
revision = '0021_merge_reasoning_config'
|
||||||
|
down_revision = ('0020_membership_source', '0018_llm_reasoning_config')
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
pass
|
||||||
@@ -54,6 +54,7 @@ _ALEMBIC_TENANT_TABLES = {
|
|||||||
'workspace_memberships',
|
'workspace_memberships',
|
||||||
'workspace_invitations',
|
'workspace_invitations',
|
||||||
'workspace_execution_states',
|
'workspace_execution_states',
|
||||||
|
'support_admin_temporary_sessions',
|
||||||
'workspace_metadata',
|
'workspace_metadata',
|
||||||
'api_keys',
|
'api_keys',
|
||||||
'bots',
|
'bots',
|
||||||
@@ -97,7 +98,7 @@ _WORKSPACE_ALEMBIC_REVISION = '0009_workspace_tenancy'
|
|||||||
_RESOURCE_SCOPE_ALEMBIC_REVISION = '0010_scope_resources'
|
_RESOURCE_SCOPE_ALEMBIC_REVISION = '0010_scope_resources'
|
||||||
_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
|
_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
|
||||||
_RELEASE_MIGRATION_ADVISORY_LOCK_ID = 0x4C414E47424F5432
|
_RELEASE_MIGRATION_ADVISORY_LOCK_ID = 0x4C414E47424F5432
|
||||||
_PGVECTOR_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
|
_PGVECTOR_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536, 3072)
|
||||||
_RUNTIME_SCHEMA = 'public'
|
_RUNTIME_SCHEMA = 'public'
|
||||||
_ALEMBIC_RUNTIME_TABLE = 'alembic_version'
|
_ALEMBIC_RUNTIME_TABLE = 'alembic_version'
|
||||||
_RUNTIME_TABLE_PRIVILEGES = frozenset({'SELECT', 'INSERT', 'UPDATE', 'DELETE'})
|
_RUNTIME_TABLE_PRIVILEGES = frozenset({'SELECT', 'INSERT', 'UPDATE', 'DELETE'})
|
||||||
@@ -1355,14 +1356,16 @@ class PersistenceManager:
|
|||||||
index = by_index.get(index_name)
|
index = by_index.get(index_name)
|
||||||
index_definition = normalized(None if index is None else index['definition'])
|
index_definition = normalized(None if index is None else index['definition'])
|
||||||
predicate = normalized(None if index is None else index['predicate'])
|
predicate = normalized(None if index is None else index['predicate'])
|
||||||
|
vector_type = 'halfvec' if dimension > 2000 else 'vector'
|
||||||
|
operator_class = f'{vector_type}_cosine_ops'
|
||||||
if (
|
if (
|
||||||
index is None
|
index is None
|
||||||
or index['access_method'] != 'hnsw'
|
or index['access_method'] != 'hnsw'
|
||||||
or index['is_valid'] is not True
|
or index['is_valid'] is not True
|
||||||
or index['is_ready'] is not True
|
or index['is_ready'] is not True
|
||||||
or f'vector({dimension})' not in index_definition
|
or f'{vector_type}({dimension})' not in index_definition
|
||||||
or f'(embedding)::vector({dimension})' not in index_definition
|
or f'(embedding)::{vector_type}({dimension})' not in index_definition
|
||||||
or 'vector_cosine_ops' not in index_definition
|
or operator_class not in index_definition
|
||||||
or predicate.strip('() ') != f'embedding_dimension = {dimension}'
|
or predicate.strip('() ') != f'embedding_dimension = {dimension}'
|
||||||
):
|
):
|
||||||
raise RuntimeError(f'PostgreSQL pgvector ANN index {index_name!r} is invalid')
|
raise RuntimeError(f'PostgreSQL pgvector ANN index {index_name!r} is invalid')
|
||||||
|
|||||||
@@ -13,7 +13,7 @@ import typing
|
|||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||||
import sqlalchemy.orm as sqlalchemy_orm
|
import sqlalchemy.orm as sqlalchemy_orm
|
||||||
from pgvector.sqlalchemy import Vector
|
from pgvector.sqlalchemy import HALFVEC, Vector
|
||||||
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
|
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
|
||||||
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
|
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
|
||||||
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
|
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
|
||||||
@@ -43,6 +43,7 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
|
|||||||
'workspace_memberships': 'workspace_uuid',
|
'workspace_memberships': 'workspace_uuid',
|
||||||
'workspace_invitations': 'workspace_uuid',
|
'workspace_invitations': 'workspace_uuid',
|
||||||
'workspace_execution_states': 'workspace_uuid',
|
'workspace_execution_states': 'workspace_uuid',
|
||||||
|
'support_admin_temporary_sessions': 'workspace_uuid',
|
||||||
'workspace_metadata': 'workspace_uuid',
|
'workspace_metadata': 'workspace_uuid',
|
||||||
'api_keys': 'workspace_uuid',
|
'api_keys': 'workspace_uuid',
|
||||||
'bots': 'workspace_uuid',
|
'bots': 'workspace_uuid',
|
||||||
@@ -75,7 +76,6 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
|
|||||||
DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
|
DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
|
||||||
'directory_projection_states': 'instance_uuid',
|
'directory_projection_states': 'instance_uuid',
|
||||||
'directory_projection_inbox': 'instance_uuid',
|
'directory_projection_inbox': 'instance_uuid',
|
||||||
'space_launch_assertion_consumptions': 'instance_uuid',
|
|
||||||
}
|
}
|
||||||
|
|
||||||
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
|
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
|
||||||
@@ -209,7 +209,7 @@ _ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES = {
|
|||||||
'now': sqlalchemy.sql.functions.now,
|
'now': sqlalchemy.sql.functions.now,
|
||||||
'sum': sqlalchemy.sql.functions.sum,
|
'sum': sqlalchemy.sql.functions.sum,
|
||||||
}
|
}
|
||||||
_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'length', 'nullif'})
|
_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'date_trunc', 'length', 'nullif'})
|
||||||
_ALLOWED_SCOPED_CUSTOM_OPERATORS = frozenset({'<=>'})
|
_ALLOWED_SCOPED_CUSTOM_OPERATORS = frozenset({'<=>'})
|
||||||
_ALLOWED_SCOPED_STATEMENT_TYPES = (
|
_ALLOWED_SCOPED_STATEMENT_TYPES = (
|
||||||
sqlalchemy.sql.dml.UpdateBase,
|
sqlalchemy.sql.dml.UpdateBase,
|
||||||
@@ -281,7 +281,7 @@ def _validate_scoped_sql_type(
|
|||||||
return
|
return
|
||||||
seen.add(identity)
|
seen.add(identity)
|
||||||
|
|
||||||
if type(sql_type) is Vector:
|
if type(sql_type) in {Vector, HALFVEC}:
|
||||||
return
|
return
|
||||||
if not type(sql_type).__module__.startswith('sqlalchemy.'):
|
if not type(sql_type).__module__.startswith('sqlalchemy.'):
|
||||||
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements')
|
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements')
|
||||||
@@ -462,7 +462,7 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
|
|||||||
if isinstance(element, sqlalchemy.sql.elements.BindParameter) and element.literal_execute:
|
if isinstance(element, sqlalchemy.sql.elements.BindParameter) and element.literal_execute:
|
||||||
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow literal-execute SQL parameters')
|
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow literal-execute SQL parameters')
|
||||||
|
|
||||||
if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) is not Vector:
|
if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) not in {Vector, HALFVEC}:
|
||||||
raise ScopedSessionTransactionError(
|
raise ScopedSessionTransactionError(
|
||||||
'TenantUnitOfWork only allows the trusted pgvector cast used by tenant vector search'
|
'TenantUnitOfWork only allows the trusted pgvector cast used by tenant vector search'
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -132,9 +132,7 @@ class Controller:
|
|||||||
|
|
||||||
break
|
break
|
||||||
|
|
||||||
if selected_query: # 找到了
|
if not selected_query: # No query is runnable under the current session limits.
|
||||||
queries.remove(selected_query)
|
|
||||||
else: # 没找到 说明:没有请求 或者 所有query对应的session都已达到并发上限
|
|
||||||
await self.ap.query_pool.condition.wait()
|
await self.ap.query_pool.condition.wait()
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
|||||||
@@ -41,6 +41,29 @@ class PreProcessor(stage.PipelineStage):
|
|||||||
selected_tool_names = {tool for tool in selected_tools if isinstance(tool, str)}
|
selected_tool_names = {tool for tool in selected_tools if isinstance(tool, str)}
|
||||||
return [tool for tool in tools if tool.name in selected_tool_names]
|
return [tool for tool in tools if tool.name in selected_tool_names]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _append_to_system_prompt(
|
||||||
|
messages: list[provider_message.Message],
|
||||||
|
addition: str,
|
||||||
|
) -> None:
|
||||||
|
"""Append text to the first system message, creating one if none exists.
|
||||||
|
|
||||||
|
Handles both plain-string and content-element (list) message bodies.
|
||||||
|
"""
|
||||||
|
if messages and messages[0].role == 'system':
|
||||||
|
head = messages[0]
|
||||||
|
if isinstance(head.content, str):
|
||||||
|
head.content = head.content + addition
|
||||||
|
elif isinstance(head.content, list):
|
||||||
|
for ce in head.content:
|
||||||
|
if getattr(ce, 'type', None) == 'text':
|
||||||
|
ce.text = (ce.text or '') + addition
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
head.content.append(provider_message.ContentElement(type='text', text=addition))
|
||||||
|
else:
|
||||||
|
messages.insert(0, provider_message.Message(role='system', content=addition.strip()))
|
||||||
|
|
||||||
async def process(
|
async def process(
|
||||||
self,
|
self,
|
||||||
query: pipeline_query.Query,
|
query: pipeline_query.Query,
|
||||||
@@ -275,6 +298,23 @@ class PreProcessor(stage.PipelineStage):
|
|||||||
query.prompt.messages = event_ctx.event.default_prompt
|
query.prompt.messages = event_ctx.event.default_prompt
|
||||||
query.messages = event_ctx.event.prompt
|
query.messages = event_ctx.event.prompt
|
||||||
|
|
||||||
|
# =========== Current date grounding for the local-agent runner ===========
|
||||||
|
# local-agent system prompts are static strings with no template-variable
|
||||||
|
# support, so without an explicit anchor the LLM resolves relative time
|
||||||
|
# references (e.g. "this quarter", "latest", "currently") against whichever
|
||||||
|
# period is best represented in its training data instead of the real date,
|
||||||
|
# and won't reliably know to double check time-sensitive facts with a tool.
|
||||||
|
if selected_runner == 'local-agent':
|
||||||
|
date_addition = (
|
||||||
|
f'\n\nCurrent date: {datetime.datetime.now().strftime("%Y-%m-%d (%A)")}. '
|
||||||
|
'Resolve relative time references (e.g. "today", "this quarter", "latest", '
|
||||||
|
'"currently") based on this date, not your training cutoff. For anything '
|
||||||
|
'time-sensitive that may have changed since training — stock prices, '
|
||||||
|
'financial results, news, current events, exchange rates, or similar — '
|
||||||
|
'verify with a search tool if one is available rather than answering from memory.'
|
||||||
|
)
|
||||||
|
self._append_to_system_prompt(query.prompt.messages, date_addition)
|
||||||
|
|
||||||
# =========== Skill awareness for the local-agent runner ===========
|
# =========== Skill awareness for the local-agent runner ===========
|
||||||
# The actual activation goes through the ``activate`` Tool Call so the
|
# The actual activation goes through the ``activate`` Tool Call so the
|
||||||
# LLM doesn't see full SKILL.md instructions until it commits to a
|
# LLM doesn't see full SKILL.md instructions until it commits to a
|
||||||
@@ -310,27 +350,7 @@ class PreProcessor(stage.PipelineStage):
|
|||||||
bound_skills=bound_skills,
|
bound_skills=bound_skills,
|
||||||
)
|
)
|
||||||
if skill_addition:
|
if skill_addition:
|
||||||
# Append to the first system message; create one if the
|
self._append_to_system_prompt(query.prompt.messages, skill_addition)
|
||||||
# prompt has none. Handles both plain-string and
|
|
||||||
# content-element (list) message bodies.
|
|
||||||
if query.prompt.messages and query.prompt.messages[0].role == 'system':
|
|
||||||
head = query.prompt.messages[0]
|
|
||||||
if isinstance(head.content, str):
|
|
||||||
head.content = head.content + skill_addition
|
|
||||||
elif isinstance(head.content, list):
|
|
||||||
appended = False
|
|
||||||
for ce in head.content:
|
|
||||||
if getattr(ce, 'type', None) == 'text':
|
|
||||||
ce.text = (ce.text or '') + skill_addition
|
|
||||||
appended = True
|
|
||||||
break
|
|
||||||
if not appended:
|
|
||||||
head.content.append(provider_message.ContentElement(type='text', text=skill_addition))
|
|
||||||
else:
|
|
||||||
query.prompt.messages.insert(
|
|
||||||
0,
|
|
||||||
provider_message.Message(role='system', content=skill_addition.strip()),
|
|
||||||
)
|
|
||||||
self.ap.logger.debug(
|
self.ap.logger.debug(
|
||||||
f'Skill index injected into system prompt: '
|
f'Skill index injected into system prompt: '
|
||||||
f'pipeline={query.pipeline_uuid} '
|
f'pipeline={query.pipeline_uuid} '
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ from ....provider import runner as runner_module
|
|||||||
import langbot_plugin.api.entities.events as events
|
import langbot_plugin.api.entities.events as events
|
||||||
from ....utils import importutil, constants, runner as runner_utils
|
from ....utils import importutil, constants, runner as runner_utils
|
||||||
from ....telemetry import features as telemetry_features
|
from ....telemetry import features as telemetry_features
|
||||||
|
from ....telemetry.identity import workspace_identity
|
||||||
from ....provider import runners
|
from ....provider import runners
|
||||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
@@ -265,7 +266,8 @@ class ChatMessageHandler(handler.MessageHandler):
|
|||||||
'duration_ms': duration_ms,
|
'duration_ms': duration_ms,
|
||||||
'model_name': model_name,
|
'model_name': model_name,
|
||||||
'version': constants.semantic_version,
|
'version': constants.semantic_version,
|
||||||
'instance_id': constants.instance_id,
|
**workspace_identity(get_query_execution_context(query)),
|
||||||
|
'runtime_instance_id': constants.instance_id,
|
||||||
'edition': constants.edition,
|
'edition': constants.edition,
|
||||||
'pipeline_plugins': pipeline_plugins,
|
'pipeline_plugins': pipeline_plugins,
|
||||||
'features': features,
|
'features': features,
|
||||||
|
|||||||
@@ -179,10 +179,13 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
|||||||
)
|
)
|
||||||
file_format = 'image/jpeg'
|
file_format = 'image/jpeg'
|
||||||
|
|
||||||
|
# NOTE: Telegram's file.file_path is a full URL of the form
|
||||||
|
# https://api.telegram.org/file/bot<TOKEN>/<path> which embeds the
|
||||||
|
# bot token. Unlike the public CDN URLs used by other adapters, it
|
||||||
|
# cannot be exposed safely, so only base64 is stored here.
|
||||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||||
message_components.append(
|
message_components.append(
|
||||||
platform_message.Image(
|
platform_message.Image(
|
||||||
url=file.file_path,
|
|
||||||
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
|
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import contextvars
|
|||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
import typing
|
import typing
|
||||||
|
from dataclasses import dataclass
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
import pydantic
|
import pydantic
|
||||||
@@ -25,6 +26,15 @@ _current_pipeline_uuid: contextvars.ContextVar[str | None] = contextvars.Context
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class WebSocketReplyContext:
|
||||||
|
"""Trusted routing context retained when the originating socket reconnects."""
|
||||||
|
|
||||||
|
scope: WebSocketScope
|
||||||
|
pipeline_uuid: str
|
||||||
|
session_id: str | None
|
||||||
|
|
||||||
|
|
||||||
class WebSocketMessage(pydantic.BaseModel):
|
class WebSocketMessage(pydantic.BaseModel):
|
||||||
"""WebSocket消息格式"""
|
"""WebSocket消息格式"""
|
||||||
|
|
||||||
@@ -265,6 +275,11 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
|||||||
embed_target = self._parse_embed_target(sender_id)
|
embed_target = self._parse_embed_target(sender_id)
|
||||||
if embed_target is not None:
|
if embed_target is not None:
|
||||||
return embed_target
|
return embed_target
|
||||||
|
reply_context = getattr(message_source, '_websocket_reply_context', None)
|
||||||
|
if isinstance(reply_context, WebSocketReplyContext):
|
||||||
|
if reply_context.scope != self._scope():
|
||||||
|
raise ValueError('WebSocket reply context does not match this adapter scope')
|
||||||
|
return reply_context.pipeline_uuid, reply_context.session_id
|
||||||
raise ValueError('WebSocket reply target is not bound to this adapter scope')
|
raise ValueError('WebSocket reply target is not bound to this adapter scope')
|
||||||
|
|
||||||
async def send_message(
|
async def send_message(
|
||||||
@@ -685,6 +700,16 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
|||||||
|
|
||||||
# 异步触发事件处理
|
# 异步触发事件处理
|
||||||
# Use owner_bot's listeners if available, otherwise fall back to proxy bot
|
# Use owner_bot's listeners if available, otherwise fall back to proxy bot
|
||||||
|
object.__setattr__(
|
||||||
|
event,
|
||||||
|
'_websocket_reply_context',
|
||||||
|
WebSocketReplyContext(
|
||||||
|
scope=connection.scope,
|
||||||
|
pipeline_uuid=pipeline_uuid,
|
||||||
|
session_id=connection.session_id,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
listeners = (
|
listeners = (
|
||||||
owner_bot.adapter.listeners
|
owner_bot.adapter.listeners
|
||||||
if (owner_bot and hasattr(owner_bot.adapter, 'listeners') and owner_bot.adapter.listeners)
|
if (owner_bot and hasattr(owner_bot.adapter, 'listeners') and owner_bot.adapter.listeners)
|
||||||
@@ -707,28 +732,37 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
|||||||
if len(listener_tasks) >= 100:
|
if len(listener_tasks) >= 100:
|
||||||
await self.logger.warning('WebSocket inbound listener capacity reached; dropping message')
|
await self.logger.warning('WebSocket inbound listener capacity reached; dropping message')
|
||||||
return
|
return
|
||||||
token = _current_pipeline_uuid.set(pipeline_uuid)
|
listener = typing.cast(
|
||||||
try:
|
typing.Callable[[typing.Any, typing.Any], typing.Awaitable[None]],
|
||||||
task_manager = getattr(self.ap, 'task_mgr', None)
|
listeners[event.__class__],
|
||||||
if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list):
|
)
|
||||||
listener_task = asyncio.create_task(listeners[event.__class__](event, callback_adapter))
|
|
||||||
else:
|
async def run_listener():
|
||||||
listener_task = task_manager.create_task(
|
token = _current_pipeline_uuid.set(pipeline_uuid)
|
||||||
listeners[event.__class__](event, callback_adapter),
|
try:
|
||||||
kind='websocket-message',
|
await listener(event, callback_adapter)
|
||||||
name=f'websocket-message-{connection.connection_id}',
|
finally:
|
||||||
scopes=[
|
_current_pipeline_uuid.reset(token)
|
||||||
core_entities.LifecycleControlScope.APPLICATION,
|
|
||||||
core_entities.LifecycleControlScope.PLATFORM,
|
listener_coro = run_listener()
|
||||||
],
|
task_manager = getattr(self.ap, 'task_mgr', None)
|
||||||
instance_uuid=connection.instance_uuid,
|
if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list):
|
||||||
workspace_uuid=connection.workspace_uuid,
|
listener_task = asyncio.create_task(listener_coro)
|
||||||
placement_generation=connection.placement_generation,
|
else:
|
||||||
).task
|
listener_task = task_manager.create_task(
|
||||||
listener_tasks.add(listener_task)
|
listener_coro,
|
||||||
listener_task.add_done_callback(self._listener_task_done)
|
kind='websocket-message',
|
||||||
finally:
|
name=f'websocket-message-{connection.connection_id}',
|
||||||
_current_pipeline_uuid.reset(token)
|
scopes=[
|
||||||
|
core_entities.LifecycleControlScope.APPLICATION,
|
||||||
|
core_entities.LifecycleControlScope.PLATFORM,
|
||||||
|
],
|
||||||
|
instance_uuid=connection.instance_uuid,
|
||||||
|
workspace_uuid=connection.workspace_uuid,
|
||||||
|
placement_generation=connection.placement_generation,
|
||||||
|
).task
|
||||||
|
listener_tasks.add(listener_task)
|
||||||
|
listener_task.add_done_callback(self._listener_task_done)
|
||||||
|
|
||||||
def get_websocket_messages(
|
def get_websocket_messages(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ from datetime import datetime
|
|||||||
|
|
||||||
import pydantic
|
import pydantic
|
||||||
|
|
||||||
from ...api.http.context import ExecutionContext
|
from ...api.http.context import ExecutionContext, PrincipalContext
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
_SESSION_FILTER_UNSET = object()
|
_SESSION_FILTER_UNSET = object()
|
||||||
@@ -95,6 +95,9 @@ class WebSocketConnection(pydantic.BaseModel):
|
|||||||
metadata: dict = pydantic.Field(default_factory=dict)
|
metadata: dict = pydantic.Field(default_factory=dict)
|
||||||
"""连接元数据(可存储额外信息)"""
|
"""连接元数据(可存储额外信息)"""
|
||||||
|
|
||||||
|
trigger_principal: PrincipalContext | None = None
|
||||||
|
"""Authenticated principal that opened this dashboard connection."""
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def scope(self) -> WebSocketScope:
|
def scope(self) -> WebSocketScope:
|
||||||
return WebSocketScope(
|
return WebSocketScope(
|
||||||
@@ -112,6 +115,7 @@ class WebSocketConnection(pydantic.BaseModel):
|
|||||||
workspace_uuid=self.workspace_uuid,
|
workspace_uuid=self.workspace_uuid,
|
||||||
placement_generation=self.placement_generation,
|
placement_generation=self.placement_generation,
|
||||||
pipeline_uuid=self.pipeline_uuid,
|
pipeline_uuid=self.pipeline_uuid,
|
||||||
|
trigger_principal=self.trigger_principal,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -138,6 +142,7 @@ class WebSocketConnectionManager:
|
|||||||
pipeline_uuid: str,
|
pipeline_uuid: str,
|
||||||
session_type: str,
|
session_type: str,
|
||||||
metadata: dict | None = None,
|
metadata: dict | None = None,
|
||||||
|
trigger_principal: PrincipalContext | None = None,
|
||||||
session_id: str | None = None,
|
session_id: str | None = None,
|
||||||
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
|
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
|
||||||
max_connections: int = 1024,
|
max_connections: int = 1024,
|
||||||
@@ -174,6 +179,7 @@ class WebSocketConnectionManager:
|
|||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
websocket=websocket,
|
websocket=websocket,
|
||||||
metadata=metadata or {},
|
metadata=metadata or {},
|
||||||
|
trigger_principal=trigger_principal,
|
||||||
send_queue=asyncio.Queue(maxsize=send_queue_size),
|
send_queue=asyncio.Queue(maxsize=send_queue_size),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import contextlib
|
|||||||
import contextvars
|
import contextvars
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
|
import math
|
||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
from typing import Any
|
from typing import Any
|
||||||
@@ -19,11 +20,6 @@ from urllib.parse import urljoin, urlparse
|
|||||||
from langbot_plugin.api.entities.builtin.pipeline.query import provider_session
|
from langbot_plugin.api.entities.builtin.pipeline.query import provider_session
|
||||||
|
|
||||||
from ..core import app
|
from ..core import app
|
||||||
from ..cloud.quotas import (
|
|
||||||
lock_workspace_for_quota,
|
|
||||||
require_resource_capacity,
|
|
||||||
resolve_workspace_quota,
|
|
||||||
)
|
|
||||||
from . import handler
|
from . import handler
|
||||||
from .archive import inspect_plugin_archive_metadata
|
from .archive import inspect_plugin_archive_metadata
|
||||||
from .github import (
|
from .github import (
|
||||||
@@ -81,7 +77,7 @@ _GITHUB_ASSET_HOSTS = frozenset(
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
_HTTP_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
|
_HTTP_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
|
||||||
_CONNECT_TIMEOUT_SEC = 30.0
|
_DEFAULT_CONNECT_TIMEOUT_SECONDS = 180.0
|
||||||
_HEARTBEAT_INTERVAL_SEC = 20.0
|
_HEARTBEAT_INTERVAL_SEC = 20.0
|
||||||
_HEARTBEAT_FAILURE_THRESHOLD = 3
|
_HEARTBEAT_FAILURE_THRESHOLD = 3
|
||||||
_RECONNECT_MAX_DELAY_SEC = 60.0
|
_RECONNECT_MAX_DELAY_SEC = 60.0
|
||||||
@@ -211,6 +207,17 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
|
|
||||||
return f'{constants.instance_id}:plugin-runtime'
|
return f'{constants.instance_id}:plugin-runtime'
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _runtime_connect_timeout(plugin_config: dict[str, Any]) -> float:
|
||||||
|
value = plugin_config.get('connect_timeout_seconds', _DEFAULT_CONNECT_TIMEOUT_SECONDS)
|
||||||
|
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0:
|
||||||
|
raise ValueError('plugin.connect_timeout_seconds must be a positive number')
|
||||||
|
return float(value)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _runtime_connect_timeout_error(timeout_seconds: float) -> str:
|
||||||
|
return f'Plugin runtime did not become ready within {timeout_seconds:g} seconds'
|
||||||
|
|
||||||
def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
|
def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
|
||||||
runtime_handler = getattr(self, 'handler', None)
|
runtime_handler = getattr(self, 'handler', None)
|
||||||
if runtime_handler is None:
|
if runtime_handler is None:
|
||||||
@@ -256,6 +263,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
def _control_headers(self, *, allow_generate: bool) -> dict[str, str]:
|
def _control_headers(self, *, allow_generate: bool) -> dict[str, str]:
|
||||||
if not self._control_token and allow_generate:
|
if not self._control_token and allow_generate:
|
||||||
self._control_token = secrets.token_urlsafe(48)
|
self._control_token = secrets.token_urlsafe(48)
|
||||||
|
if not self._control_token:
|
||||||
|
if self.runtime_profile == 'shared':
|
||||||
|
raise PluginRuntimeNotConnectedError(
|
||||||
|
f'{PLUGIN_RUNTIME_CONTROL_TOKEN_ENV} must be configured with a strong shared secret '
|
||||||
|
'for a Cloud Plugin Runtime'
|
||||||
|
)
|
||||||
|
return {}
|
||||||
try:
|
try:
|
||||||
self._control_token = validate_runtime_secret(
|
self._control_token = validate_runtime_secret(
|
||||||
self._control_token,
|
self._control_token,
|
||||||
@@ -704,10 +718,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
runtime_handler = self._runtime_handler()
|
runtime_handler = self._runtime_handler()
|
||||||
|
started_at = time.monotonic()
|
||||||
async with self._state_lock:
|
async with self._state_lock:
|
||||||
all_states: dict[str, PluginInstallationDesiredState] = {}
|
all_states: dict[str, PluginInstallationDesiredState] = {}
|
||||||
workspace_installations: dict[str, set[str]] = {}
|
workspace_installations: dict[str, set[str]] = {}
|
||||||
|
workspace_count = 0
|
||||||
for context in contexts:
|
for context in contexts:
|
||||||
|
workspace_count += 1
|
||||||
execution_context = await self._validate_execution_context(context)
|
execution_context = await self._validate_execution_context(context)
|
||||||
states = await self._load_workspace_desired_states(execution_context)
|
states = await self._load_workspace_desired_states(execution_context)
|
||||||
installation_ids = {state.binding.installation_uuid for state in states}
|
installation_ids = {state.binding.installation_uuid for state in states}
|
||||||
@@ -727,6 +744,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
runtime_handler.unregister_installation_binding(previous.binding)
|
runtime_handler.unregister_installation_binding(previous.binding)
|
||||||
self._known_desired_states = all_states
|
self._known_desired_states = all_states
|
||||||
self._workspace_installations = workspace_installations
|
self._workspace_installations = workspace_installations
|
||||||
|
self.ap.logger.info(
|
||||||
|
'Shared plugin runtime reconcile completed: workspaces=%d desired_installations=%d '
|
||||||
|
'elapsed_seconds=%.3f',
|
||||||
|
workspace_count,
|
||||||
|
len(all_states),
|
||||||
|
time.monotonic() - started_at,
|
||||||
|
)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
async def _validate_execution_context(self, context: TenantContext) -> ExecutionContext:
|
async def _validate_execution_context(self, context: TenantContext) -> ExecutionContext:
|
||||||
@@ -822,6 +846,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
runtime_id=self._runtime_id,
|
runtime_id=self._runtime_id,
|
||||||
)
|
)
|
||||||
self.worker_policy = self._load_worker_policy()
|
self.worker_policy = self._load_worker_policy()
|
||||||
|
plugin_config = self.ap.instance_config.data.get('plugin', {})
|
||||||
|
connect_timeout_seconds = self._runtime_connect_timeout(plugin_config)
|
||||||
|
|
||||||
async with self._lifecycle_lock:
|
async with self._lifecycle_lock:
|
||||||
if self._closing:
|
if self._closing:
|
||||||
@@ -963,10 +989,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
|
|
||||||
self._transport_task = asyncio.create_task(task_coro)
|
self._transport_task = asyncio.create_task(task_coro)
|
||||||
try:
|
try:
|
||||||
await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC)
|
await asyncio.wait_for(self._connected.wait(), timeout=connect_timeout_seconds)
|
||||||
except asyncio.TimeoutError as exc:
|
except asyncio.TimeoutError as exc:
|
||||||
await self._stop_transport()
|
await self._stop_transport()
|
||||||
raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc
|
raise PluginRuntimeNotConnectedError(
|
||||||
|
self._runtime_connect_timeout_error(connect_timeout_seconds)
|
||||||
|
) from exc
|
||||||
if connect_errors:
|
if connect_errors:
|
||||||
await self._stop_transport()
|
await self._stop_transport()
|
||||||
raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
|
raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
|
||||||
@@ -1300,11 +1328,6 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
install_info: dict[str, Any],
|
install_info: dict[str, Any],
|
||||||
artifact_digest: str,
|
artifact_digest: str,
|
||||||
) -> tuple[InstallationBinding, str | None, bool]:
|
) -> tuple[InstallationBinding, str | None, bool]:
|
||||||
quota = await resolve_workspace_quota(
|
|
||||||
self.ap,
|
|
||||||
execution_context.workspace_uuid,
|
|
||||||
'plugins.max',
|
|
||||||
)
|
|
||||||
safe_install_info = {
|
safe_install_info = {
|
||||||
key: value
|
key: value
|
||||||
for key, value in install_info.items()
|
for key, value in install_info.items()
|
||||||
@@ -1326,19 +1349,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
)
|
)
|
||||||
|
|
||||||
async def persist(execute):
|
async def persist(execute):
|
||||||
if quota.requires_transaction_lock:
|
|
||||||
await lock_workspace_for_quota(execute, execution_context.workspace_uuid)
|
|
||||||
result = await execute(statement)
|
result = await execute(statement)
|
||||||
setting = result.first()
|
setting = result.first()
|
||||||
if setting is None:
|
if setting is None:
|
||||||
await require_resource_capacity(
|
|
||||||
execute,
|
|
||||||
workspace_uuid=execution_context.workspace_uuid,
|
|
||||||
model=persistence_plugin.PluginSetting,
|
|
||||||
quota=quota,
|
|
||||||
resource_name='plugins',
|
|
||||||
workspace_locked=quota.requires_transaction_lock,
|
|
||||||
)
|
|
||||||
installation_uuid = str(uuid.uuid4())
|
installation_uuid = str(uuid.uuid4())
|
||||||
runtime_revision = 1
|
runtime_revision = 1
|
||||||
previous_digest = None
|
previous_digest = None
|
||||||
@@ -1393,8 +1406,6 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
)
|
)
|
||||||
|
|
||||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||||
if quota.requires_transaction_lock and not callable(tenant_uow):
|
|
||||||
raise RuntimeError('Cloud plugin quota enforcement requires transactional persistence')
|
|
||||||
if callable(tenant_uow):
|
if callable(tenant_uow):
|
||||||
async with tenant_uow(execution_context.workspace_uuid) as uow:
|
async with tenant_uow(execution_context.workspace_uuid) as uow:
|
||||||
return await persist(uow.execute)
|
return await persist(uow.execute)
|
||||||
@@ -1990,11 +2001,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
|||||||
with runtime_handler.installation_scope(binding):
|
with runtime_handler.installation_scope(binding):
|
||||||
return await runtime_handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
|
return await runtime_handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
|
||||||
|
|
||||||
async def get_debug_info(self) -> dict[str, Any]:
|
async def get_debug_info(self, execution_context: ExecutionContext) -> dict[str, Any]:
|
||||||
"""Get debug information including debug key and WS URL"""
|
"""Get debug information including debug key and WS URL"""
|
||||||
if not self.is_enable_plugin or not self._runtime_available():
|
if not self.is_enable_plugin or not self._runtime_available():
|
||||||
return {}
|
return {}
|
||||||
return await self._runtime_handler().get_debug_info()
|
return await self._runtime_handler().get_debug_info(execution_context)
|
||||||
|
|
||||||
async def emit_event(
|
async def emit_event(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -1579,7 +1579,7 @@ class RuntimeConnectionHandler(handler.Handler):
|
|||||||
return await self.call_action(
|
return await self.call_action(
|
||||||
LangBotToRuntimeAction.RECONCILE_PLUGIN_INSTALLATIONS,
|
LangBotToRuntimeAction.RECONCILE_PLUGIN_INSTALLATIONS,
|
||||||
request.model_dump(),
|
request.model_dump(),
|
||||||
timeout=120,
|
timeout=300,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def apply_plugin_installation(
|
async def apply_plugin_installation(
|
||||||
@@ -1960,14 +1960,19 @@ class RuntimeConnectionHandler(handler.Handler):
|
|||||||
)
|
)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
async def get_debug_info(self) -> dict[str, Any]:
|
async def get_debug_info(self, execution_context: ExecutionContext) -> dict[str, Any]:
|
||||||
"""Get debug information including debug key and WS URL"""
|
"""Get debug information including debug key and WS URL"""
|
||||||
with self.installation_scope(None):
|
action_context = ActionContext(
|
||||||
result = await self.call_action(
|
instance_uuid=execution_context.instance_uuid,
|
||||||
LangBotToRuntimeAction.GET_DEBUG_INFO,
|
workspace_uuid=execution_context.workspace_uuid,
|
||||||
{},
|
placement_generation=execution_context.placement_generation,
|
||||||
timeout=10,
|
)
|
||||||
)
|
result = await self.call_action(
|
||||||
|
LangBotToRuntimeAction.GET_DEBUG_INFO,
|
||||||
|
{},
|
||||||
|
timeout=10,
|
||||||
|
action_context=action_context,
|
||||||
|
)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
# ================= RAG Capability Callers (LangBot -> Runtime) =================
|
# ================= RAG Capability Callers (LangBot -> Runtime) =================
|
||||||
|
|||||||
@@ -649,6 +649,7 @@ class ModelManager:
|
|||||||
provider_uuid=runtime_provider.provider_entity.uuid,
|
provider_uuid=runtime_provider.provider_entity.uuid,
|
||||||
abilities=model_info.get('abilities', []),
|
abilities=model_info.get('abilities', []),
|
||||||
context_length=model_info.get('context_length'),
|
context_length=model_info.get('context_length'),
|
||||||
|
reasoning_config=model_info.get('reasoning_config', {'level': 'provider_default'}),
|
||||||
extra_args=model_info.get('extra_args', {}),
|
extra_args=model_info.get('extra_args', {}),
|
||||||
)
|
)
|
||||||
return self._build_llm_model(execution_context, model_entity, runtime_provider)
|
return self._build_llm_model(execution_context, model_entity, runtime_provider)
|
||||||
@@ -717,7 +718,10 @@ class ModelManager:
|
|||||||
provider_entity = self._coerce_provider(provider_info, context)
|
provider_entity = self._coerce_provider(provider_info, context)
|
||||||
requester_manifest = self.get_available_requester_manifest_by_name(provider_entity.requester)
|
requester_manifest = self.get_available_requester_manifest_by_name(provider_entity.requester)
|
||||||
litellm_provider = self._get_litellm_provider_from_manifest(requester_manifest)
|
litellm_provider = self._get_litellm_provider_from_manifest(requester_manifest)
|
||||||
config = {'base_url': provider_entity.base_url}
|
config = {
|
||||||
|
'base_url': provider_entity.base_url,
|
||||||
|
'requester_name': provider_entity.requester,
|
||||||
|
}
|
||||||
|
|
||||||
if litellm_provider:
|
if litellm_provider:
|
||||||
from .requesters import litellmchat
|
from .requesters import litellmchat
|
||||||
|
|||||||
@@ -0,0 +1,125 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import typing
|
||||||
|
|
||||||
|
|
||||||
|
ReasoningLevel = typing.Literal[
|
||||||
|
'provider_default',
|
||||||
|
'disabled',
|
||||||
|
'enabled',
|
||||||
|
'minimal',
|
||||||
|
'low',
|
||||||
|
'medium',
|
||||||
|
'high',
|
||||||
|
'xhigh',
|
||||||
|
'max',
|
||||||
|
]
|
||||||
|
|
||||||
|
REASONING_LEVELS: tuple[str, ...] = (
|
||||||
|
'provider_default',
|
||||||
|
'disabled',
|
||||||
|
'enabled',
|
||||||
|
'minimal',
|
||||||
|
'low',
|
||||||
|
'medium',
|
||||||
|
'high',
|
||||||
|
'xhigh',
|
||||||
|
'max',
|
||||||
|
)
|
||||||
|
DEFAULT_REASONING_CONFIG: dict[str, str] = {'level': 'provider_default'}
|
||||||
|
|
||||||
|
_CONFLICTING_TOP_LEVEL_ARGS = {
|
||||||
|
'reasoning_effort',
|
||||||
|
'thinking',
|
||||||
|
'enable_thinking',
|
||||||
|
'thinking_budget',
|
||||||
|
'reasoning',
|
||||||
|
}
|
||||||
|
_CONFLICTING_EXTRA_BODY_ARGS = {
|
||||||
|
'reasoning_effort',
|
||||||
|
'thinking',
|
||||||
|
'enable_thinking',
|
||||||
|
'thinking_budget',
|
||||||
|
'reasoning',
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_reasoning_config(value: typing.Any) -> dict[str, str]:
|
||||||
|
"""Return the canonical model reasoning configuration."""
|
||||||
|
if value is None:
|
||||||
|
return dict(DEFAULT_REASONING_CONFIG)
|
||||||
|
if not isinstance(value, dict):
|
||||||
|
raise ValueError('reasoning_config must be an object')
|
||||||
|
|
||||||
|
unknown_fields = set(value) - {'level'}
|
||||||
|
if unknown_fields:
|
||||||
|
raise ValueError(f'Unsupported reasoning_config fields: {", ".join(sorted(unknown_fields))}')
|
||||||
|
|
||||||
|
level = value.get('level', 'provider_default')
|
||||||
|
if level not in REASONING_LEVELS:
|
||||||
|
raise ValueError(f'Unsupported reasoning level: {level}')
|
||||||
|
return {'level': typing.cast(str, level)}
|
||||||
|
|
||||||
|
|
||||||
|
def validate_reasoning_config(
|
||||||
|
value: typing.Any,
|
||||||
|
abilities: typing.Iterable[str] | None,
|
||||||
|
extra_args: typing.Any,
|
||||||
|
) -> dict[str, str]:
|
||||||
|
"""Validate a model-facing reasoning config and conflicting raw arguments."""
|
||||||
|
config = normalize_reasoning_config(value)
|
||||||
|
if config['level'] == 'provider_default':
|
||||||
|
return config
|
||||||
|
|
||||||
|
if 'reasoning' not in set(abilities or []):
|
||||||
|
raise ValueError('The reasoning ability must be enabled before selecting a reasoning level')
|
||||||
|
|
||||||
|
conflicts = find_reasoning_arg_conflicts(extra_args)
|
||||||
|
if conflicts:
|
||||||
|
raise ValueError('reasoning_config conflicts with advanced parameters: ' + ', '.join(conflicts))
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
def find_reasoning_arg_conflicts(extra_args: typing.Any) -> list[str]:
|
||||||
|
if not isinstance(extra_args, dict):
|
||||||
|
return []
|
||||||
|
|
||||||
|
conflicts = [key for key in sorted(_CONFLICTING_TOP_LEVEL_ARGS) if key in extra_args]
|
||||||
|
extra_body = extra_args.get('extra_body')
|
||||||
|
if isinstance(extra_body, dict):
|
||||||
|
conflicts.extend(f'extra_body.{key}' for key in sorted(_CONFLICTING_EXTRA_BODY_ARGS) if key in extra_body)
|
||||||
|
return conflicts
|
||||||
|
|
||||||
|
|
||||||
|
def validate_reasoning_capabilities(
|
||||||
|
config: typing.Any,
|
||||||
|
capabilities: typing.Mapping[str, typing.Any],
|
||||||
|
model_name: str,
|
||||||
|
) -> None:
|
||||||
|
"""Ensure an explicit reasoning level can be honored by the requester."""
|
||||||
|
level = normalize_reasoning_config(config)['level']
|
||||||
|
if level == 'provider_default':
|
||||||
|
return
|
||||||
|
|
||||||
|
available_levels = capabilities.get('levels')
|
||||||
|
if not isinstance(available_levels, list):
|
||||||
|
available_levels = []
|
||||||
|
legacy_levels = capabilities.get('legacy_levels')
|
||||||
|
if not isinstance(legacy_levels, list):
|
||||||
|
legacy_levels = []
|
||||||
|
if capabilities.get('supported') is not True or (level not in available_levels and level not in legacy_levels):
|
||||||
|
available_text = ', '.join(str(item) for item in available_levels) or 'provider_default'
|
||||||
|
raise ValueError(
|
||||||
|
f'Reasoning level "{level}" is not supported by model {model_name}. Available levels: {available_text}'
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def default_reasoning_capabilities(
|
||||||
|
supported: bool = False,
|
||||||
|
source: str = 'unknown',
|
||||||
|
) -> dict[str, typing.Any]:
|
||||||
|
return {
|
||||||
|
'supported': supported,
|
||||||
|
'levels': ['provider_default'],
|
||||||
|
'source': source,
|
||||||
|
}
|
||||||
@@ -10,6 +10,7 @@ from ...entity.persistence import model as persistence_model
|
|||||||
from ...workspace.errors import WorkspaceInvariantError
|
from ...workspace.errors import WorkspaceInvariantError
|
||||||
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
||||||
from . import token
|
from . import token
|
||||||
|
from . import reasoning
|
||||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
|
|
||||||
@@ -377,11 +378,15 @@ class RuntimeLLMModel:
|
|||||||
provider: RuntimeProvider
|
provider: RuntimeProvider
|
||||||
"""提供商实例"""
|
"""提供商实例"""
|
||||||
|
|
||||||
|
reasoning_config_override: dict[str, str] | None
|
||||||
|
"""Request-scoped reasoning policy supplied by the active pipeline."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
execution_context: ExecutionContext,
|
execution_context: ExecutionContext,
|
||||||
model_entity: persistence_model.LLMModel,
|
model_entity: persistence_model.LLMModel,
|
||||||
provider: RuntimeProvider,
|
provider: RuntimeProvider,
|
||||||
|
reasoning_config_override: dict[str, str] | None = None,
|
||||||
):
|
):
|
||||||
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='LLM model')
|
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='LLM model')
|
||||||
if model_entity.workspace_uuid != execution_context.workspace_uuid:
|
if model_entity.workspace_uuid != execution_context.workspace_uuid:
|
||||||
@@ -391,6 +396,7 @@ class RuntimeLLMModel:
|
|||||||
self.execution_context = execution_context
|
self.execution_context = execution_context
|
||||||
self.model_entity = model_entity
|
self.model_entity = model_entity
|
||||||
self.provider = provider
|
self.provider = provider
|
||||||
|
self.reasoning_config_override = reasoning_config_override
|
||||||
|
|
||||||
|
|
||||||
class RuntimeEmbeddingModel:
|
class RuntimeEmbeddingModel:
|
||||||
@@ -482,6 +488,13 @@ class ProviderAPIRequester(metaclass=abc.ABCMeta):
|
|||||||
"""
|
"""
|
||||||
raise NotImplementedError('This provider does not support model scanning')
|
raise NotImplementedError('This provider does not support model scanning')
|
||||||
|
|
||||||
|
def get_reasoning_capabilities(self, model: RuntimeLLMModel) -> dict[str, typing.Any]:
|
||||||
|
"""Return normalized reasoning controls supported by a model."""
|
||||||
|
return reasoning.default_reasoning_capabilities(
|
||||||
|
supported='reasoning' in (model.model_entity.abilities or []),
|
||||||
|
source='manual' if 'reasoning' in (model.model_entity.abilities or []) else 'unknown',
|
||||||
|
)
|
||||||
|
|
||||||
@abc.abstractmethod
|
@abc.abstractmethod
|
||||||
async def invoke_llm(
|
async def invoke_llm(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ import typing
|
|||||||
import litellm
|
import litellm
|
||||||
from litellm import acompletion, aembedding, arerank
|
from litellm import acompletion, aembedding, arerank
|
||||||
|
|
||||||
from .. import errors, requester
|
from .. import errors, reasoning, requester
|
||||||
from ....utils import httpclient
|
from ....utils import httpclient
|
||||||
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
||||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
@@ -164,6 +164,39 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
|
|
||||||
_EMBEDDING_MODEL_HINTS = ('embedding', 'embed', 'bge-', 'e5-', 'm3e', 'gte-', 'text-embedding')
|
_EMBEDDING_MODEL_HINTS = ('embedding', 'embed', 'bge-', 'e5-', 'm3e', 'gte-', 'text-embedding')
|
||||||
_RERANK_MODEL_HINTS = ('rerank', 're-rank', 're_rank')
|
_RERANK_MODEL_HINTS = ('rerank', 're-rank', 're_rank')
|
||||||
|
_QWEN_DEDICATED_THINKING_MODELS = frozenset(
|
||||||
|
{
|
||||||
|
'qwen3.7-max-preview',
|
||||||
|
'qwen3.7-max-2026-05-17',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
_QWEN_REASONING_BUDGETS = {
|
||||||
|
'low': 1024,
|
||||||
|
'medium': 4096,
|
||||||
|
'high': 8192,
|
||||||
|
}
|
||||||
|
_INFERRED_EFFORT_PROVIDERS = frozenset(
|
||||||
|
{
|
||||||
|
'anthropic',
|
||||||
|
'gemini',
|
||||||
|
'groq',
|
||||||
|
'mistral',
|
||||||
|
'openai',
|
||||||
|
'openrouter',
|
||||||
|
'together_ai',
|
||||||
|
'xai',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
_REQUESTER_REASONING_FAMILIES = {
|
||||||
|
'openai-chat-completions': 'openai',
|
||||||
|
'anthropic-messages': 'anthropic',
|
||||||
|
'deepseek-chat-completions': 'deepseek',
|
||||||
|
'moonshot-chat-completions': 'kimi',
|
||||||
|
'moonshot-cn-chat-completions': 'kimi',
|
||||||
|
'bailian-chat-completions': 'qwen',
|
||||||
|
'doubao-chat-completions': 'doubao',
|
||||||
|
'mimo-chat-completions': 'mimo',
|
||||||
|
}
|
||||||
|
|
||||||
default_config: dict[str, typing.Any] = {
|
default_config: dict[str, typing.Any] = {
|
||||||
'base_url': '',
|
'base_url': '',
|
||||||
@@ -172,6 +205,7 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
'drop_params': False,
|
'drop_params': False,
|
||||||
'num_retries': 0,
|
'num_retries': 0,
|
||||||
'api_version': '',
|
'api_version': '',
|
||||||
|
'requester_name': '',
|
||||||
}
|
}
|
||||||
|
|
||||||
async def initialize(self):
|
async def initialize(self):
|
||||||
@@ -201,7 +235,10 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
provider = self._get_custom_llm_provider()
|
provider = self._get_custom_llm_provider()
|
||||||
candidates: list[tuple[str, str | None]] = [(model_name, provider)]
|
candidates: list[tuple[str, str | None]] = [
|
||||||
|
(candidate, None) for candidate in self._metadata_model_candidates(model_name)
|
||||||
|
]
|
||||||
|
candidates.append((model_name, provider))
|
||||||
litellm_model_name = self._build_litellm_model_name(model_name)
|
litellm_model_name = self._build_litellm_model_name(model_name)
|
||||||
if litellm_model_name != model_name:
|
if litellm_model_name != model_name:
|
||||||
candidates.append((litellm_model_name, None))
|
candidates.append((litellm_model_name, None))
|
||||||
@@ -268,6 +305,14 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
deduped_candidates.append(candidate)
|
deduped_candidates.append(candidate)
|
||||||
return deduped_candidates
|
return deduped_candidates
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _metadata_model_candidates(model_name: str) -> list[str]:
|
||||||
|
"""Return known equivalent model IDs used only for LiteLLM metadata lookup."""
|
||||||
|
normalized_model_name = (model_name or '').lower()
|
||||||
|
if normalized_model_name.startswith('mimo-v2.5'):
|
||||||
|
return [f'openrouter/xiaomi/{normalized_model_name}']
|
||||||
|
return []
|
||||||
|
|
||||||
def _known_context_length_fallback(self, model_name: str) -> int | None:
|
def _known_context_length_fallback(self, model_name: str) -> int | None:
|
||||||
normalized_model_name = (model_name or '').lower()
|
normalized_model_name = (model_name or '').lower()
|
||||||
if normalized_model_name.startswith('deepseek-v4-'):
|
if normalized_model_name.startswith('deepseek-v4-'):
|
||||||
@@ -287,7 +332,8 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
if not callable(helper):
|
if not callable(helper):
|
||||||
return self._known_context_length_fallback(model_name)
|
return self._known_context_length_fallback(model_name)
|
||||||
|
|
||||||
candidates = [model_name]
|
candidates = self._metadata_model_candidates(model_name)
|
||||||
|
candidates.append(model_name)
|
||||||
litellm_model_name = self._build_litellm_model_name(model_name)
|
litellm_model_name = self._build_litellm_model_name(model_name)
|
||||||
if litellm_model_name != model_name:
|
if litellm_model_name != model_name:
|
||||||
candidates.append(litellm_model_name)
|
candidates.append(litellm_model_name)
|
||||||
@@ -314,6 +360,297 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
def _supports_vision(self, model_name: str) -> bool:
|
def _supports_vision(self, model_name: str) -> bool:
|
||||||
return self._safe_litellm_bool_helper('supports_vision', model_name)
|
return self._safe_litellm_bool_helper('supports_vision', model_name)
|
||||||
|
|
||||||
|
def _supports_reasoning(self, model_name: str) -> bool:
|
||||||
|
return self._safe_litellm_bool_helper('supports_reasoning', model_name)
|
||||||
|
|
||||||
|
def _requester_name(self, model: requester.RuntimeLLMModel | None = None) -> str:
|
||||||
|
if model is not None:
|
||||||
|
provider_entity = getattr(getattr(model, 'provider', None), 'provider_entity', None)
|
||||||
|
name = getattr(provider_entity, 'requester', None)
|
||||||
|
if isinstance(name, str) and name:
|
||||||
|
return name.lower()
|
||||||
|
return str(self.requester_cfg.get('requester_name') or '').lower()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _infer_reasoning_family_from_model_name(model_name: str) -> str:
|
||||||
|
normalized_name = (model_name or '').lower()
|
||||||
|
basename = normalized_name.rsplit('/', 1)[-1]
|
||||||
|
if basename.startswith(('gpt-', 'chatgpt-', 'o1', 'o3', 'o4')):
|
||||||
|
return 'openai'
|
||||||
|
if basename.startswith('claude-'):
|
||||||
|
return 'anthropic'
|
||||||
|
if basename.startswith('deepseek-'):
|
||||||
|
return 'deepseek'
|
||||||
|
if basename.startswith(('kimi-', 'moonshot-')):
|
||||||
|
return 'kimi'
|
||||||
|
if basename.startswith(('qwen-', 'qwen3', 'qwq')):
|
||||||
|
return 'qwen'
|
||||||
|
if basename.startswith(('doubao-', 'seed-')):
|
||||||
|
return 'doubao'
|
||||||
|
if basename.startswith('mimo-'):
|
||||||
|
return 'mimo'
|
||||||
|
return ''
|
||||||
|
|
||||||
|
def _reasoning_family(
|
||||||
|
self,
|
||||||
|
model_name: str,
|
||||||
|
model: requester.RuntimeLLMModel | None = None,
|
||||||
|
) -> str:
|
||||||
|
requester_name = self._requester_name(model)
|
||||||
|
if requester_name in {'new-api-chat-completions', 'volcark-chat-completions'}:
|
||||||
|
inferred_family = self._infer_reasoning_family_from_model_name(model_name)
|
||||||
|
if inferred_family:
|
||||||
|
return inferred_family
|
||||||
|
return 'volcengine' if requester_name == 'volcark-chat-completions' else ''
|
||||||
|
|
||||||
|
# Bailian's compatible endpoint also hosts Kimi models. Keep those
|
||||||
|
# models on Kimi's ``thinking`` protocol instead of Qwen's
|
||||||
|
# ``enable_thinking`` protocol.
|
||||||
|
if requester_name == 'bailian-chat-completions':
|
||||||
|
inferred_family = self._infer_reasoning_family_from_model_name(model_name)
|
||||||
|
if inferred_family == 'kimi':
|
||||||
|
return inferred_family
|
||||||
|
|
||||||
|
requester_family = self._REQUESTER_REASONING_FAMILIES.get(requester_name)
|
||||||
|
if requester_family:
|
||||||
|
return requester_family
|
||||||
|
|
||||||
|
inferred_family = self._infer_reasoning_family_from_model_name(model_name)
|
||||||
|
provider = (self._get_custom_llm_provider() or '').lower()
|
||||||
|
if provider == 'openai':
|
||||||
|
return inferred_family or ('openai' if requester_name in {'', 'openai'} else '')
|
||||||
|
if provider:
|
||||||
|
return provider
|
||||||
|
return inferred_family
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_anthropic_adaptive_model(model_name: str) -> bool:
|
||||||
|
basename = model_name.lower().rsplit('/', 1)[-1]
|
||||||
|
if 'mythos-preview' in basename:
|
||||||
|
return True
|
||||||
|
|
||||||
|
parts = basename.split('-')
|
||||||
|
if len(parts) < 3 or parts[0] != 'claude':
|
||||||
|
return False
|
||||||
|
model_families = {'opus', 'sonnet', 'fable', 'mythos'}
|
||||||
|
if parts[1] in model_families:
|
||||||
|
if parts[2] == '5':
|
||||||
|
return True
|
||||||
|
return len(parts) >= 4 and parts[2] == '4' and parts[3] in {'6', '7', '8'}
|
||||||
|
return parts[1] == '5' and parts[2] in model_families
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_anthropic_always_thinking_model(model_name: str) -> bool:
|
||||||
|
normalized_name = model_name.lower()
|
||||||
|
return any(marker in normalized_name for marker in ('fable-5', 'mythos-5', 'mythos-preview'))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_dedicated_qwen_thinking_model(model_name: str) -> bool:
|
||||||
|
normalized_name = model_name.lower().rsplit('/', 1)[-1]
|
||||||
|
return (
|
||||||
|
normalized_name in LiteLLMRequester._QWEN_DEDICATED_THINKING_MODELS
|
||||||
|
or normalized_name.startswith('qwq')
|
||||||
|
or '-thinking' in normalized_name
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _supports_qwen_thinking_budget(model_name: str) -> bool:
|
||||||
|
"""Return whether the documented Qwen3 family supports thinking_budget."""
|
||||||
|
normalized_name = model_name.lower().rsplit('/', 1)[-1]
|
||||||
|
return normalized_name.startswith('qwen3')
|
||||||
|
|
||||||
|
def _known_reasoning_levels(self, model_name: str, family: str) -> list[str] | None:
|
||||||
|
normalized_name = model_name.lower().rsplit('/', 1)[-1]
|
||||||
|
|
||||||
|
if family == 'deepseek' and normalized_name.startswith('deepseek-'):
|
||||||
|
if normalized_name.startswith('deepseek-v4-'):
|
||||||
|
return ['provider_default', 'disabled', 'low', 'high', 'xhigh', 'max']
|
||||||
|
if 'reasoner' in normalized_name or '-r1' in normalized_name:
|
||||||
|
return ['provider_default']
|
||||||
|
return ['provider_default', 'disabled', 'enabled']
|
||||||
|
|
||||||
|
if family == 'kimi':
|
||||||
|
if normalized_name.startswith('kimi-k3'):
|
||||||
|
return ['provider_default', 'low', 'high', 'max']
|
||||||
|
if normalized_name.startswith('kimi-k2.7-code'):
|
||||||
|
return ['provider_default']
|
||||||
|
if normalized_name.startswith(('kimi-k2.5', 'kimi-k2.6')):
|
||||||
|
return ['provider_default', 'disabled', 'enabled']
|
||||||
|
if 'thinking' in normalized_name:
|
||||||
|
return ['provider_default']
|
||||||
|
|
||||||
|
if family == 'qwen' and normalized_name.startswith(('qwen-', 'qwen3', 'qwq')):
|
||||||
|
if self._is_dedicated_qwen_thinking_model(normalized_name):
|
||||||
|
if self._supports_qwen_thinking_budget(normalized_name):
|
||||||
|
return ['provider_default', 'low', 'medium', 'high']
|
||||||
|
return ['provider_default']
|
||||||
|
if self._supports_qwen_thinking_budget(normalized_name):
|
||||||
|
return ['provider_default', 'disabled', 'low', 'medium', 'high']
|
||||||
|
return ['provider_default', 'disabled', 'enabled']
|
||||||
|
|
||||||
|
if family == 'doubao' and normalized_name.startswith(('doubao-', 'seed-')):
|
||||||
|
return ['provider_default', 'disabled', 'low', 'medium', 'high']
|
||||||
|
|
||||||
|
if family == 'mimo' and normalized_name.startswith(('mimo-v2.5',)):
|
||||||
|
return ['provider_default', 'disabled', 'enabled']
|
||||||
|
|
||||||
|
if family == 'anthropic' and normalized_name.startswith('claude-'):
|
||||||
|
levels = ['provider_default']
|
||||||
|
adaptive = self._is_anthropic_adaptive_model(normalized_name)
|
||||||
|
if adaptive and not self._is_anthropic_always_thinking_model(normalized_name):
|
||||||
|
levels.append('disabled')
|
||||||
|
levels.extend(['low', 'medium', 'high'])
|
||||||
|
if adaptive:
|
||||||
|
levels.extend(['xhigh', 'max'])
|
||||||
|
return levels
|
||||||
|
|
||||||
|
if family == 'openai' and normalized_name.startswith(('gpt-5', 'o1', 'o3', 'o4')):
|
||||||
|
return ['provider_default', 'low', 'medium', 'high']
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _openai_reasoning_levels(self, model_name: str) -> list[str]:
|
||||||
|
model_info = self._safe_model_info(model_name)
|
||||||
|
levels = ['provider_default']
|
||||||
|
if model_info.get('supports_none_reasoning_effort') is True:
|
||||||
|
levels.append('disabled')
|
||||||
|
if model_info.get('supports_minimal_reasoning_effort') is True:
|
||||||
|
levels.append('minimal')
|
||||||
|
for level in ('low', 'medium', 'high'):
|
||||||
|
if model_info.get(f'supports_{level}_reasoning_effort') is not False:
|
||||||
|
levels.append(level)
|
||||||
|
for level in ('xhigh', 'max'):
|
||||||
|
if model_info.get(f'supports_{level}_reasoning_effort') is True:
|
||||||
|
levels.append(level)
|
||||||
|
return levels
|
||||||
|
|
||||||
|
def _safe_model_info(self, model_name: str) -> dict[str, typing.Any]:
|
||||||
|
helper = getattr(litellm, 'get_model_info', None)
|
||||||
|
if not callable(helper):
|
||||||
|
return {}
|
||||||
|
|
||||||
|
candidates = [
|
||||||
|
*self._metadata_model_candidates(model_name),
|
||||||
|
model_name,
|
||||||
|
self._build_litellm_model_name(model_name),
|
||||||
|
]
|
||||||
|
for candidate in candidates:
|
||||||
|
try:
|
||||||
|
info = helper(candidate)
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
if isinstance(info, dict):
|
||||||
|
return info
|
||||||
|
model_dump = getattr(info, 'model_dump', None)
|
||||||
|
if callable(model_dump):
|
||||||
|
try:
|
||||||
|
dumped = model_dump()
|
||||||
|
if isinstance(dumped, dict):
|
||||||
|
return dumped
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def get_reasoning_capabilities(self, model: requester.RuntimeLLMModel) -> dict[str, typing.Any]:
|
||||||
|
model_name = model.model_entity.name
|
||||||
|
abilities = model.model_entity.abilities or []
|
||||||
|
detected = self._supports_reasoning(model_name)
|
||||||
|
declared = 'reasoning' in abilities
|
||||||
|
family = self._reasoning_family(model_name, model)
|
||||||
|
known_levels = self._known_reasoning_levels(model_name, family)
|
||||||
|
supported = detected or declared or known_levels is not None
|
||||||
|
if not supported:
|
||||||
|
return reasoning.default_reasoning_capabilities()
|
||||||
|
|
||||||
|
normalized_name = model_name.lower()
|
||||||
|
if family == 'openai':
|
||||||
|
levels = self._openai_reasoning_levels(model_name)
|
||||||
|
elif known_levels is not None:
|
||||||
|
levels = known_levels
|
||||||
|
elif family == 'anthropic':
|
||||||
|
levels = ['provider_default', 'low', 'medium', 'high']
|
||||||
|
elif family in {'deepseek', 'qwen', 'mimo', 'volcengine'}:
|
||||||
|
levels = ['provider_default', 'disabled', 'enabled']
|
||||||
|
elif family == 'doubao':
|
||||||
|
levels = ['provider_default', 'disabled', 'low', 'medium', 'high']
|
||||||
|
elif family == 'ollama':
|
||||||
|
levels = ['provider_default']
|
||||||
|
levels.append('disabled')
|
||||||
|
if normalized_name.startswith('gpt-oss') or '/gpt-oss' in normalized_name:
|
||||||
|
levels.extend(['low', 'medium', 'high'])
|
||||||
|
else:
|
||||||
|
levels.append('enabled')
|
||||||
|
elif family in self._INFERRED_EFFORT_PROVIDERS:
|
||||||
|
levels = ['provider_default', 'low', 'medium', 'high']
|
||||||
|
else:
|
||||||
|
levels = ['provider_default']
|
||||||
|
|
||||||
|
capabilities = {
|
||||||
|
'supported': True,
|
||||||
|
'levels': list(dict.fromkeys(levels)),
|
||||||
|
'source': 'litellm' if detected else ('provider' if known_levels is not None else 'manual'),
|
||||||
|
}
|
||||||
|
if family == 'qwen' and 'disabled' in capabilities['levels'] and 'enabled' not in capabilities['levels']:
|
||||||
|
capabilities['legacy_levels'] = ['enabled']
|
||||||
|
return capabilities
|
||||||
|
|
||||||
|
def _build_reasoning_args(self, model: requester.RuntimeLLMModel) -> dict[str, typing.Any]:
|
||||||
|
level = self._reasoning_level(model)
|
||||||
|
if level == 'provider_default':
|
||||||
|
return {}
|
||||||
|
|
||||||
|
config = {'level': level}
|
||||||
|
capabilities = self.get_reasoning_capabilities(model)
|
||||||
|
try:
|
||||||
|
reasoning.validate_reasoning_capabilities(config, capabilities, model.model_entity.name)
|
||||||
|
except ValueError as exc:
|
||||||
|
raise errors.RequesterError(str(exc)) from exc
|
||||||
|
|
||||||
|
family = self._reasoning_family(model.model_entity.name, model)
|
||||||
|
if level == 'disabled':
|
||||||
|
if family in {'deepseek', 'kimi', 'mimo', 'doubao'}:
|
||||||
|
return {'extra_body': {'thinking': {'type': 'disabled'}}}
|
||||||
|
if family == 'qwen':
|
||||||
|
return {'extra_body': {'enable_thinking': False}}
|
||||||
|
if family == 'volcengine':
|
||||||
|
return {'extra_body': {'thinking': {'type': 'disabled'}}}
|
||||||
|
if family == 'anthropic':
|
||||||
|
return {'thinking': {'type': 'disabled'}}
|
||||||
|
return {'reasoning_effort': 'none'}
|
||||||
|
if level == 'enabled':
|
||||||
|
if family in {'deepseek', 'kimi', 'mimo', 'volcengine'}:
|
||||||
|
return {'extra_body': {'thinking': {'type': 'enabled'}}}
|
||||||
|
if family == 'qwen':
|
||||||
|
return {'extra_body': {'enable_thinking': True}}
|
||||||
|
return {'reasoning_effort': 'low'}
|
||||||
|
if family == 'qwen' and level in self._QWEN_REASONING_BUDGETS:
|
||||||
|
return {
|
||||||
|
'extra_body': {
|
||||||
|
'enable_thinking': True,
|
||||||
|
'thinking_budget': self._QWEN_REASONING_BUDGETS[level],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if family == 'deepseek':
|
||||||
|
return {
|
||||||
|
'extra_body': {
|
||||||
|
'thinking': {'type': 'enabled'},
|
||||||
|
'reasoning_effort': level,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return {'reasoning_effort': level}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _reasoning_config_value(model: requester.RuntimeLLMModel) -> typing.Any:
|
||||||
|
raw_config = getattr(model, 'reasoning_config_override', None)
|
||||||
|
if raw_config is None:
|
||||||
|
raw_config = getattr(model.model_entity, 'reasoning_config', None)
|
||||||
|
if not isinstance(raw_config, dict):
|
||||||
|
return None
|
||||||
|
return raw_config
|
||||||
|
|
||||||
|
def _reasoning_level(self, model: requester.RuntimeLLMModel) -> str:
|
||||||
|
return reasoning.normalize_reasoning_config(self._reasoning_config_value(model))['level']
|
||||||
|
|
||||||
def _infer_model_type(self, model_id: str) -> str:
|
def _infer_model_type(self, model_id: str) -> str:
|
||||||
normalized_id = (model_id or '').lower()
|
normalized_id = (model_id or '').lower()
|
||||||
if any(kw in normalized_id for kw in self._RERANK_MODEL_HINTS):
|
if any(kw in normalized_id for kw in self._RERANK_MODEL_HINTS):
|
||||||
@@ -344,6 +681,13 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
)
|
)
|
||||||
if supports_provider_reported_vision or self._supports_vision(model_id):
|
if supports_provider_reported_vision or self._supports_vision(model_id):
|
||||||
abilities.append('vision')
|
abilities.append('vision')
|
||||||
|
supports_provider_reported_reasoning = bool(
|
||||||
|
model_payload and model_payload.get('supports_reasoning') is True
|
||||||
|
)
|
||||||
|
family = self._reasoning_family(model_id)
|
||||||
|
supports_known_reasoning = self._known_reasoning_levels(model_id, family) is not None
|
||||||
|
if supports_provider_reported_reasoning or supports_known_reasoning or self._supports_reasoning(model_id):
|
||||||
|
abilities.append('reasoning')
|
||||||
scanned_model['abilities'] = abilities
|
scanned_model['abilities'] = abilities
|
||||||
|
|
||||||
context_length = self._context_length_from_scan_payload(model_payload)
|
context_length = self._context_length_from_scan_payload(model_payload)
|
||||||
@@ -354,13 +698,51 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
|
|
||||||
return scanned_model
|
return scanned_model
|
||||||
|
|
||||||
def _convert_messages(self, messages: typing.List[provider_message.Message]) -> list[dict]:
|
def _convert_messages(
|
||||||
|
self,
|
||||||
|
messages: typing.List[provider_message.Message],
|
||||||
|
reasoning_family: str = '',
|
||||||
|
include_reasoning_context: bool = True,
|
||||||
|
) -> list[dict]:
|
||||||
"""Convert LangBot messages to LiteLLM/OpenAI format."""
|
"""Convert LangBot messages to LiteLLM/OpenAI format."""
|
||||||
req_messages = []
|
req_messages = []
|
||||||
for m in messages:
|
for m in messages:
|
||||||
msg_dict = m.dict(exclude_none=True)
|
msg_dict = m.dict(exclude_none=True)
|
||||||
content = msg_dict.get('content')
|
content = msg_dict.get('content')
|
||||||
|
|
||||||
|
if msg_dict.get('role') == 'assistant' and reasoning_family:
|
||||||
|
provider_fields = msg_dict.get('provider_specific_fields')
|
||||||
|
if isinstance(provider_fields, dict):
|
||||||
|
cleaned_provider_fields = dict(provider_fields)
|
||||||
|
reasoning_content = cleaned_provider_fields.pop('reasoning_content', None)
|
||||||
|
thinking_blocks = cleaned_provider_fields.pop('thinking_blocks', None)
|
||||||
|
|
||||||
|
# ``content`` is also used for the user-facing rendering.
|
||||||
|
# Do not replay that rendered <think> wrapper alongside the
|
||||||
|
# structured provider reasoning on the next request.
|
||||||
|
if reasoning_content or thinking_blocks:
|
||||||
|
content = msg_dict.get('content')
|
||||||
|
if isinstance(content, str):
|
||||||
|
msg_dict['content'] = self._strip_think(content)
|
||||||
|
|
||||||
|
if include_reasoning_context:
|
||||||
|
if reasoning_family == 'anthropic' and thinking_blocks:
|
||||||
|
msg_dict['thinking_blocks'] = thinking_blocks
|
||||||
|
elif reasoning_family in {
|
||||||
|
'deepseek',
|
||||||
|
'kimi',
|
||||||
|
'qwen',
|
||||||
|
'doubao',
|
||||||
|
'mimo',
|
||||||
|
'volcengine',
|
||||||
|
} and isinstance(reasoning_content, str):
|
||||||
|
msg_dict['reasoning_content'] = reasoning_content
|
||||||
|
|
||||||
|
if cleaned_provider_fields:
|
||||||
|
msg_dict['provider_specific_fields'] = cleaned_provider_fields
|
||||||
|
else:
|
||||||
|
msg_dict.pop('provider_specific_fields', None)
|
||||||
|
|
||||||
if isinstance(content, list):
|
if isinstance(content, list):
|
||||||
converted_parts = []
|
converted_parts = []
|
||||||
for part in content:
|
for part in content:
|
||||||
@@ -421,6 +803,52 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
|
|
||||||
return content or ''
|
return content or ''
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _thinking_blocks_text(thinking_blocks: typing.Any) -> str:
|
||||||
|
if not isinstance(thinking_blocks, list):
|
||||||
|
return ''
|
||||||
|
parts = []
|
||||||
|
for block in thinking_blocks:
|
||||||
|
if isinstance(block, dict):
|
||||||
|
text = block.get('thinking')
|
||||||
|
else:
|
||||||
|
text = getattr(block, 'thinking', None)
|
||||||
|
if isinstance(text, str) and text:
|
||||||
|
parts.append(text)
|
||||||
|
return ''.join(parts)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _merge_thinking_blocks(
|
||||||
|
cls,
|
||||||
|
current: list[dict[str, typing.Any]],
|
||||||
|
incoming: typing.Any,
|
||||||
|
) -> list[dict[str, typing.Any]]:
|
||||||
|
"""Merge Anthropic thinking block fragments emitted by a stream."""
|
||||||
|
if not isinstance(incoming, list):
|
||||||
|
return current
|
||||||
|
merged = [dict(block) for block in current]
|
||||||
|
for raw_block in incoming:
|
||||||
|
block = cls._as_dict(raw_block)
|
||||||
|
if not block:
|
||||||
|
continue
|
||||||
|
block_type = block.get('type')
|
||||||
|
if block_type == 'redacted_thinking':
|
||||||
|
merged.append(block)
|
||||||
|
continue
|
||||||
|
|
||||||
|
text = block.get('thinking') if isinstance(block.get('thinking'), str) else ''
|
||||||
|
signature = block.get('signature')
|
||||||
|
if merged and merged[-1].get('type') == 'thinking' and not merged[-1].get('signature'):
|
||||||
|
merged[-1]['thinking'] = f'{merged[-1].get("thinking", "")}{text}'
|
||||||
|
if signature:
|
||||||
|
merged[-1]['signature'] = signature
|
||||||
|
elif merged and signature and merged[-1].get('signature') == signature:
|
||||||
|
if text and text != merged[-1].get('thinking', ''):
|
||||||
|
merged[-1]['thinking'] = f'{merged[-1].get("thinking", "")}{text}'
|
||||||
|
else:
|
||||||
|
merged.append(block)
|
||||||
|
return merged
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _normalize_usage(usage: typing.Any) -> dict:
|
def _normalize_usage(usage: typing.Any) -> dict:
|
||||||
"""Normalize a LiteLLM/OpenAI usage object into a plain token dict.
|
"""Normalize a LiteLLM/OpenAI usage object into a plain token dict.
|
||||||
@@ -651,7 +1079,13 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
stream: bool = False,
|
stream: bool = False,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
"""Build common completion arguments for invoke_llm and invoke_llm_stream."""
|
"""Build common completion arguments for invoke_llm and invoke_llm_stream."""
|
||||||
req_messages = self._convert_messages(messages)
|
reasoning_family = self._reasoning_family(model.model_entity.name, model)
|
||||||
|
reasoning_level = self._reasoning_level(model)
|
||||||
|
req_messages = self._convert_messages(
|
||||||
|
messages,
|
||||||
|
reasoning_family=reasoning_family,
|
||||||
|
include_reasoning_context=reasoning_level != 'disabled',
|
||||||
|
)
|
||||||
model_name = self._build_litellm_model_name(model.model_entity.name)
|
model_name = self._build_litellm_model_name(model.model_entity.name)
|
||||||
api_key = model.provider.token_mgr.get_token()
|
api_key = model.provider.token_mgr.get_token()
|
||||||
|
|
||||||
@@ -670,6 +1104,29 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
args.update(model.model_entity.extra_args)
|
args.update(model.model_entity.extra_args)
|
||||||
args.update(extra_args)
|
args.update(extra_args)
|
||||||
|
|
||||||
|
reasoning_args = self._build_reasoning_args(model)
|
||||||
|
if reasoning_args:
|
||||||
|
conflicts = reasoning.find_reasoning_arg_conflicts(model.model_entity.extra_args)
|
||||||
|
conflicts.extend(reasoning.find_reasoning_arg_conflicts(extra_args))
|
||||||
|
if conflicts:
|
||||||
|
raise errors.RequesterError(
|
||||||
|
'reasoning_config conflicts with advanced parameters: ' + ', '.join(dict.fromkeys(conflicts))
|
||||||
|
)
|
||||||
|
reasoning_extra_body = reasoning_args.get('extra_body')
|
||||||
|
if isinstance(reasoning_extra_body, dict):
|
||||||
|
existing_extra_body = args.get('extra_body') or {}
|
||||||
|
if not isinstance(existing_extra_body, dict):
|
||||||
|
raise errors.RequesterError('extra_body must be an object')
|
||||||
|
args.update({key: value for key, value in reasoning_args.items() if key != 'extra_body'})
|
||||||
|
args['extra_body'] = {**existing_extra_body, **reasoning_extra_body}
|
||||||
|
else:
|
||||||
|
args.update(reasoning_args)
|
||||||
|
if 'reasoning_effort' in reasoning_args and self._get_custom_llm_provider() == 'openai':
|
||||||
|
allowed_openai_params = args.get('allowed_openai_params') or []
|
||||||
|
if not isinstance(allowed_openai_params, (list, tuple, set)):
|
||||||
|
raise errors.RequesterError('allowed_openai_params must be an array')
|
||||||
|
args['allowed_openai_params'] = list(dict.fromkeys([*allowed_openai_params, 'reasoning_effort']))
|
||||||
|
|
||||||
if funcs:
|
if funcs:
|
||||||
tools = await self.ap.tool_mgr.generate_tools_for_openai(funcs)
|
tools = await self.ap.tool_mgr.generate_tools_for_openai(funcs)
|
||||||
if tools:
|
if tools:
|
||||||
@@ -699,10 +1156,21 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
|
|
||||||
content = message_data.get('content', '')
|
content = message_data.get('content', '')
|
||||||
reasoning_content = message_data.get('reasoning_content', None)
|
reasoning_content = message_data.get('reasoning_content', None)
|
||||||
message_data['content'] = self._process_thinking_content(content, reasoning_content, remove_think)
|
thinking_blocks = message_data.get('thinking_blocks')
|
||||||
|
if reasoning_content or thinking_blocks:
|
||||||
|
provider_fields = dict(message_data.get('provider_specific_fields') or {})
|
||||||
|
if reasoning_content:
|
||||||
|
provider_fields['reasoning_content'] = reasoning_content
|
||||||
|
if thinking_blocks:
|
||||||
|
provider_fields['thinking_blocks'] = thinking_blocks
|
||||||
|
message_data['provider_specific_fields'] = provider_fields
|
||||||
|
display_reasoning = reasoning_content or self._thinking_blocks_text(thinking_blocks) or None
|
||||||
|
message_data['content'] = self._process_thinking_content(content, display_reasoning, remove_think)
|
||||||
|
|
||||||
if 'reasoning_content' in message_data:
|
if 'reasoning_content' in message_data:
|
||||||
del message_data['reasoning_content']
|
del message_data['reasoning_content']
|
||||||
|
if 'thinking_blocks' in message_data:
|
||||||
|
del message_data['thinking_blocks']
|
||||||
|
|
||||||
message = provider_message.Message(**message_data)
|
message = provider_message.Message(**message_data)
|
||||||
usage_info = self._extract_usage(response)
|
usage_info = self._extract_usage(response)
|
||||||
@@ -728,6 +1196,9 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
role = 'assistant'
|
role = 'assistant'
|
||||||
tool_call_state: dict[int, dict[str, typing.Any]] = {}
|
tool_call_state: dict[int, dict[str, typing.Any]] = {}
|
||||||
think_state = _ThinkStripState() if remove_think else None
|
think_state = _ThinkStripState() if remove_think else None
|
||||||
|
reasoning_started = False
|
||||||
|
reasoning_closed = False
|
||||||
|
thinking_blocks_state: list[dict[str, typing.Any]] = []
|
||||||
|
|
||||||
try:
|
try:
|
||||||
response = await acompletion(**args)
|
response = await acompletion(**args)
|
||||||
@@ -758,28 +1229,63 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
if 'role' in delta and delta['role']:
|
if 'role' in delta and delta['role']:
|
||||||
role = delta['role']
|
role = delta['role']
|
||||||
|
|
||||||
delta_content = delta.get('content', '')
|
delta_content = delta.get('content') or ''
|
||||||
reasoning_content = delta.get('reasoning_content', '')
|
reasoning_content = delta.get('reasoning_content') or ''
|
||||||
|
provider_fields = dict(delta.get('provider_specific_fields') or {})
|
||||||
|
raw_thinking_blocks = delta.get('thinking_blocks')
|
||||||
|
if raw_thinking_blocks:
|
||||||
|
thinking_blocks_state = self._merge_thinking_blocks(thinking_blocks_state, raw_thinking_blocks)
|
||||||
|
provider_fields['thinking_blocks'] = thinking_blocks_state
|
||||||
|
thinking_blocks_text = self._thinking_blocks_text(raw_thinking_blocks)
|
||||||
|
display_reasoning_content = reasoning_content or thinking_blocks_text
|
||||||
|
|
||||||
# Handle reasoning_content based on remove_think flag
|
# Handle reasoning_content based on remove_think flag
|
||||||
if reasoning_content:
|
if reasoning_content:
|
||||||
|
provider_fields['reasoning_content'] = reasoning_content
|
||||||
if remove_think:
|
if remove_think:
|
||||||
# Skip reasoning content when remove_think is True
|
delta_content = delta_content or None
|
||||||
chunk_idx += 1
|
|
||||||
continue
|
|
||||||
else:
|
else:
|
||||||
# Use reasoning_content as the displayed content
|
# Stream explicit markers so downstream adapters and
|
||||||
delta_content = reasoning_content
|
# the debug page see the same format as non-streaming
|
||||||
|
# responses.
|
||||||
|
if not reasoning_started:
|
||||||
|
delta_content = '<think>\n'
|
||||||
|
reasoning_started = True
|
||||||
|
else:
|
||||||
|
delta_content = ''
|
||||||
|
delta_content += display_reasoning_content
|
||||||
|
if delta.get('content'):
|
||||||
|
delta_content += f'\n</think>\n{delta.get("content")}'
|
||||||
|
reasoning_closed = True
|
||||||
|
|
||||||
|
elif display_reasoning_content:
|
||||||
|
if remove_think:
|
||||||
|
delta_content = delta_content or None
|
||||||
|
else:
|
||||||
|
if not reasoning_started:
|
||||||
|
delta_content = '<think>\n'
|
||||||
|
reasoning_started = True
|
||||||
|
else:
|
||||||
|
delta_content = ''
|
||||||
|
delta_content += display_reasoning_content
|
||||||
|
if delta.get('content'):
|
||||||
|
delta_content += f'\n</think>\n{delta.get("content")}'
|
||||||
|
reasoning_closed = True
|
||||||
|
|
||||||
|
elif delta_content and not remove_think and reasoning_started and not reasoning_closed:
|
||||||
|
delta_content = f'\n</think>\n{delta_content}'
|
||||||
|
reasoning_closed = True
|
||||||
|
|
||||||
|
if finish_reason and not remove_think and reasoning_started and not reasoning_closed:
|
||||||
|
delta_content = f'{delta_content}\n</think>\n'
|
||||||
|
reasoning_closed = True
|
||||||
|
|
||||||
if think_state is not None and delta_content:
|
if think_state is not None and delta_content:
|
||||||
delta_content = think_state.feed(delta_content)
|
delta_content = think_state.feed(delta_content)
|
||||||
if not delta_content:
|
|
||||||
chunk_idx += 1
|
|
||||||
continue
|
|
||||||
|
|
||||||
tool_calls = self._normalize_stream_tool_calls(delta.get('tool_calls'), tool_call_state)
|
tool_calls = self._normalize_stream_tool_calls(delta.get('tool_calls'), tool_call_state)
|
||||||
|
|
||||||
if chunk_idx == 0 and not delta_content and not tool_calls:
|
if not delta_content and not tool_calls and not provider_fields and not finish_reason:
|
||||||
chunk_idx += 1
|
chunk_idx += 1
|
||||||
continue
|
continue
|
||||||
|
|
||||||
@@ -791,13 +1297,20 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
|||||||
}
|
}
|
||||||
|
|
||||||
# Preserve provider_specific_fields from delta (e.g., Gemini thought_signatures)
|
# Preserve provider_specific_fields from delta (e.g., Gemini thought_signatures)
|
||||||
if delta.get('provider_specific_fields'):
|
if provider_fields:
|
||||||
chunk_data['provider_specific_fields'] = delta['provider_specific_fields']
|
chunk_data['provider_specific_fields'] = provider_fields
|
||||||
|
|
||||||
chunk_data = {k: v for k, v in chunk_data.items() if v is not None}
|
chunk_data = {k: v for k, v in chunk_data.items() if v is not None}
|
||||||
yield provider_message.MessageChunk(**chunk_data)
|
yield provider_message.MessageChunk(**chunk_data)
|
||||||
chunk_idx += 1
|
chunk_idx += 1
|
||||||
|
|
||||||
|
if reasoning_started and not reasoning_closed:
|
||||||
|
yield provider_message.MessageChunk(
|
||||||
|
role=role,
|
||||||
|
content='\n</think>\n',
|
||||||
|
is_final=True,
|
||||||
|
)
|
||||||
|
|
||||||
if think_state is not None:
|
if think_state is not None:
|
||||||
pending_content = think_state.flush()
|
pending_content = think_state.flush()
|
||||||
if pending_content:
|
if pending_content:
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import typing
|
|||||||
from .. import runner
|
from .. import runner
|
||||||
from ...telemetry import features as telemetry_features
|
from ...telemetry import features as telemetry_features
|
||||||
from ..modelmgr import requester as modelmgr_requester
|
from ..modelmgr import requester as modelmgr_requester
|
||||||
|
from ..modelmgr import reasoning as modelmgr_reasoning
|
||||||
from ..tools.loaders.native import EXEC_TOOL_NAME
|
from ..tools.loaders.native import EXEC_TOOL_NAME
|
||||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||||
@@ -60,6 +61,7 @@ class _StreamAccumulator:
|
|||||||
self.msg_idx = 0
|
self.msg_idx = 0
|
||||||
self.accumulated_content = initial_content or ''
|
self.accumulated_content = initial_content or ''
|
||||||
self.last_role = 'assistant'
|
self.last_role = 'assistant'
|
||||||
|
self.provider_specific_fields: dict[str, typing.Any] = {}
|
||||||
self.msg_sequence = msg_sequence
|
self.msg_sequence = msg_sequence
|
||||||
self.remove_think = remove_think
|
self.remove_think = remove_think
|
||||||
self._think_state = None
|
self._think_state = None
|
||||||
@@ -90,10 +92,27 @@ class _StreamAccumulator:
|
|||||||
name=tool_call.function.name if tool_call.function else '',
|
name=tool_call.function.name if tool_call.function else '',
|
||||||
arguments='',
|
arguments='',
|
||||||
),
|
),
|
||||||
|
provider_specific_fields=(
|
||||||
|
dict(tool_call.provider_specific_fields) if tool_call.provider_specific_fields else None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
elif tool_call.provider_specific_fields:
|
||||||
|
existing_fields = self.tool_calls_map[tool_call.id].provider_specific_fields or {}
|
||||||
|
self.tool_calls_map[tool_call.id].provider_specific_fields = {
|
||||||
|
**existing_fields,
|
||||||
|
**tool_call.provider_specific_fields,
|
||||||
|
}
|
||||||
if tool_call.function and tool_call.function.arguments:
|
if tool_call.function and tool_call.function.arguments:
|
||||||
self.tool_calls_map[tool_call.id].function.arguments += tool_call.function.arguments
|
self.tool_calls_map[tool_call.id].function.arguments += tool_call.function.arguments
|
||||||
|
|
||||||
|
if msg.provider_specific_fields:
|
||||||
|
for key, value in msg.provider_specific_fields.items():
|
||||||
|
if key == 'reasoning_content' and isinstance(value, str):
|
||||||
|
previous = self.provider_specific_fields.get(key, '')
|
||||||
|
self.provider_specific_fields[key] = f'{previous}{value}'
|
||||||
|
else:
|
||||||
|
self.provider_specific_fields[key] = value
|
||||||
|
|
||||||
if msg.is_final:
|
if msg.is_final:
|
||||||
self._flush_think_state()
|
self._flush_think_state()
|
||||||
|
|
||||||
@@ -103,6 +122,7 @@ class _StreamAccumulator:
|
|||||||
role=self.last_role,
|
role=self.last_role,
|
||||||
content=self._maybe_strip_think(self.accumulated_content),
|
content=self._maybe_strip_think(self.accumulated_content),
|
||||||
tool_calls=list(self.tool_calls_map.values()) if (self.tool_calls_map and msg.is_final) else None,
|
tool_calls=list(self.tool_calls_map.values()) if (self.tool_calls_map and msg.is_final) else None,
|
||||||
|
provider_specific_fields=(self.provider_specific_fields or None) if msg.is_final else None,
|
||||||
is_final=msg.is_final,
|
is_final=msg.is_final,
|
||||||
msg_sequence=self.msg_sequence,
|
msg_sequence=self.msg_sequence,
|
||||||
)
|
)
|
||||||
@@ -115,6 +135,7 @@ class _StreamAccumulator:
|
|||||||
role=self.last_role,
|
role=self.last_role,
|
||||||
content=self._maybe_strip_think(self.accumulated_content),
|
content=self._maybe_strip_think(self.accumulated_content),
|
||||||
tool_calls=list(self.tool_calls_map.values()) if self.tool_calls_map else None,
|
tool_calls=list(self.tool_calls_map.values()) if self.tool_calls_map else None,
|
||||||
|
provider_specific_fields=self.provider_specific_fields or None,
|
||||||
msg_sequence=self.msg_sequence,
|
msg_sequence=self.msg_sequence,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -233,9 +254,10 @@ class LocalAgentRunner(runner.RequestRunner):
|
|||||||
execution_context,
|
execution_context,
|
||||||
query.use_llm_model_uuid,
|
query.use_llm_model_uuid,
|
||||||
)
|
)
|
||||||
candidates.append(primary)
|
|
||||||
except ValueError:
|
except ValueError:
|
||||||
self.ap.logger.warning(f'Primary model {query.use_llm_model_uuid} not found')
|
self.ap.logger.warning(f'Primary model {query.use_llm_model_uuid} not found')
|
||||||
|
else:
|
||||||
|
candidates.append(LocalAgentRunner._apply_pipeline_reasoning_config(query, primary))
|
||||||
|
|
||||||
# Fallback models
|
# Fallback models
|
||||||
fallback_uuids = (query.variables or {}).get('_fallback_model_uuids', [])
|
fallback_uuids = (query.variables or {}).get('_fallback_model_uuids', [])
|
||||||
@@ -245,12 +267,31 @@ class LocalAgentRunner(runner.RequestRunner):
|
|||||||
execution_context,
|
execution_context,
|
||||||
fb_uuid,
|
fb_uuid,
|
||||||
)
|
)
|
||||||
candidates.append(fb_model)
|
|
||||||
except ValueError:
|
except ValueError:
|
||||||
self.ap.logger.warning(f'Fallback model {fb_uuid} not found, skipping')
|
self.ap.logger.warning(f'Fallback model {fb_uuid} not found, skipping')
|
||||||
|
else:
|
||||||
|
candidates.append(LocalAgentRunner._apply_pipeline_reasoning_config(query, fb_model))
|
||||||
|
|
||||||
return candidates
|
return candidates
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _apply_pipeline_reasoning_config(
|
||||||
|
query: pipeline_query.Query,
|
||||||
|
model: modelmgr_requester.RuntimeLLMModel,
|
||||||
|
) -> modelmgr_requester.RuntimeLLMModel:
|
||||||
|
local_agent_config = query.pipeline_config.get('ai', {}).get('local-agent', {})
|
||||||
|
model_config = local_agent_config.get('model', {})
|
||||||
|
reasoning_by_model = model_config.get('reasoning', {}) if isinstance(model_config, dict) else {}
|
||||||
|
level = (
|
||||||
|
reasoning_by_model.get(model.model_entity.uuid, 'provider_default')
|
||||||
|
if isinstance(reasoning_by_model, dict)
|
||||||
|
else 'provider_default'
|
||||||
|
)
|
||||||
|
reasoning_config = modelmgr_reasoning.normalize_reasoning_config({'level': level})
|
||||||
|
configured_model = copy.copy(model)
|
||||||
|
configured_model.reasoning_config_override = reasoning_config
|
||||||
|
return configured_model
|
||||||
|
|
||||||
async def _invoke_with_fallback(
|
async def _invoke_with_fallback(
|
||||||
self,
|
self,
|
||||||
query: pipeline_query.Query,
|
query: pipeline_query.Query,
|
||||||
|
|||||||
@@ -27,6 +27,19 @@ if typing.TYPE_CHECKING:
|
|||||||
HEARTBEAT_INTERVAL_SECONDS = 24 * 3600
|
HEARTBEAT_INTERVAL_SECONDS = 24 * 3600
|
||||||
|
|
||||||
|
|
||||||
|
class WorkspaceResourceSnapshot(typing.TypedDict):
|
||||||
|
workspace_uuid: str
|
||||||
|
bot_count: int
|
||||||
|
pipeline_count: int
|
||||||
|
knowledge_base_count: int
|
||||||
|
plugin_count: int
|
||||||
|
mcp_server_count: int
|
||||||
|
extension_count: int
|
||||||
|
skill_count: int
|
||||||
|
adapters: list[str]
|
||||||
|
execution_generation: int
|
||||||
|
|
||||||
|
|
||||||
async def _count(
|
async def _count(
|
||||||
ap: core_app.Application,
|
ap: core_app.Application,
|
||||||
table,
|
table,
|
||||||
@@ -52,14 +65,13 @@ async def _count(
|
|||||||
return -1
|
return -1
|
||||||
|
|
||||||
|
|
||||||
async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dict]:
|
async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) -> list[WorkspaceResourceSnapshot]:
|
||||||
"""Summarize already-loaded Cloud registries without per-tenant SQL."""
|
"""Summarize already-loaded Cloud registries without per-tenant SQL."""
|
||||||
persistence_mgr = ap.persistence_mgr
|
persistence_mgr = ap.persistence_mgr
|
||||||
if getattr(getattr(persistence_mgr, 'mode', None), 'value', None) != 'cloud_runtime':
|
if getattr(getattr(persistence_mgr, 'mode', None), 'value', None) != 'cloud_runtime':
|
||||||
return []
|
return []
|
||||||
|
|
||||||
bindings = await ap.workspace_service.list_active_execution_bindings()
|
counts: dict[str, WorkspaceResourceSnapshot] = {
|
||||||
counts = {
|
|
||||||
binding.workspace_uuid: {
|
binding.workspace_uuid: {
|
||||||
'workspace_uuid': binding.workspace_uuid,
|
'workspace_uuid': binding.workspace_uuid,
|
||||||
'bot_count': 0,
|
'bot_count': 0,
|
||||||
@@ -68,13 +80,20 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dic
|
|||||||
'plugin_count': 0,
|
'plugin_count': 0,
|
||||||
'mcp_server_count': 0,
|
'mcp_server_count': 0,
|
||||||
'extension_count': 0,
|
'extension_count': 0,
|
||||||
|
'skill_count': 0,
|
||||||
|
'adapters': [],
|
||||||
|
'execution_generation': binding.placement_generation,
|
||||||
}
|
}
|
||||||
for binding in bindings
|
for binding in bindings
|
||||||
}
|
}
|
||||||
|
|
||||||
for key in getattr(ap.platform_mgr, '_bots_by_key', {}):
|
adapter_sets: dict[str, set[str]] = {workspace_uuid: set() for workspace_uuid in counts}
|
||||||
|
for key, bot in getattr(ap.platform_mgr, '_bots_by_key', {}).items():
|
||||||
if len(key) >= 2 and key[1] in counts:
|
if len(key) >= 2 and key[1] in counts:
|
||||||
counts[key[1]]['bot_count'] += 1
|
counts[key[1]]['bot_count'] += 1
|
||||||
|
adapter = getattr(bot, 'adapter', None)
|
||||||
|
if adapter is not None and getattr(bot, 'enable', False):
|
||||||
|
adapter_sets[key[1]].add(adapter.__class__.__name__)
|
||||||
for key in getattr(ap.pipeline_mgr, '_pipelines_by_key', {}):
|
for key in getattr(ap.pipeline_mgr, '_pipelines_by_key', {}):
|
||||||
if len(key) >= 2 and key[1] in counts:
|
if len(key) >= 2 and key[1] in counts:
|
||||||
counts[key[1]]['pipeline_count'] += 1
|
counts[key[1]]['pipeline_count'] += 1
|
||||||
@@ -87,14 +106,24 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dic
|
|||||||
for workspace_uuid, installations in getattr(ap.plugin_connector, '_workspace_installations', {}).items():
|
for workspace_uuid, installations in getattr(ap.plugin_connector, '_workspace_installations', {}).items():
|
||||||
if workspace_uuid in counts:
|
if workspace_uuid in counts:
|
||||||
counts[workspace_uuid]['plugin_count'] = len(installations)
|
counts[workspace_uuid]['plugin_count'] = len(installations)
|
||||||
|
for key, skills in getattr(ap.skill_mgr, '_skills_by_scope', {}).items():
|
||||||
|
if len(key) >= 2 and key[1] in counts:
|
||||||
|
counts[key[1]]['skill_count'] += len(skills)
|
||||||
|
|
||||||
for resource in counts.values():
|
for workspace_uuid, resource in counts.items():
|
||||||
resource['extension_count'] = resource['plugin_count'] + resource['mcp_server_count']
|
resource['extension_count'] = resource['plugin_count'] + resource['mcp_server_count']
|
||||||
|
resource['adapters'] = sorted(adapter_sets[workspace_uuid])
|
||||||
return list(counts.values())
|
return list(counts.values())
|
||||||
|
|
||||||
|
|
||||||
async def build_heartbeat_payload(ap: core_app.Application) -> dict:
|
async def build_heartbeat_payload(
|
||||||
"""Collect the anonymous instance profile snapshot."""
|
ap: core_app.Application,
|
||||||
|
*,
|
||||||
|
workspace_uuid: str,
|
||||||
|
workspace_create_ts: int = 0,
|
||||||
|
workspace_resource: WorkspaceResourceSnapshot | None = None,
|
||||||
|
) -> dict:
|
||||||
|
"""Collect one anonymous Workspace profile snapshot."""
|
||||||
from ..entity.persistence import bot as persistence_bot
|
from ..entity.persistence import bot as persistence_bot
|
||||||
from ..entity.persistence import mcp as persistence_mcp
|
from ..entity.persistence import mcp as persistence_mcp
|
||||||
from ..entity.persistence import pipeline as persistence_pipeline
|
from ..entity.persistence import pipeline as persistence_pipeline
|
||||||
@@ -177,15 +206,16 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict:
|
|||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
workspace_resources = await _cloud_workspace_resource_counts(ap)
|
if workspace_resource is not None:
|
||||||
if workspace_resources:
|
features.update({key: value for key, value in workspace_resource.items() if key != 'workspace_uuid'})
|
||||||
features['workspace_resources'] = workspace_resources
|
|
||||||
|
|
||||||
return {
|
return {
|
||||||
'event_type': 'instance_heartbeat',
|
'event_type': 'instance_heartbeat',
|
||||||
'query_id': '',
|
'query_id': '',
|
||||||
'version': constants.semantic_version,
|
'version': constants.semantic_version,
|
||||||
'instance_id': constants.instance_id,
|
'instance_id': constants.instance_id,
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'workspace_create_ts': workspace_create_ts,
|
||||||
'instance_create_ts': constants.instance_create_ts,
|
'instance_create_ts': constants.instance_create_ts,
|
||||||
'edition': constants.edition,
|
'edition': constants.edition,
|
||||||
'features': features,
|
'features': features,
|
||||||
@@ -193,14 +223,49 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _workspace_created_timestamp(created_at: datetime | None) -> int:
|
||||||
|
if created_at is None:
|
||||||
|
return 0
|
||||||
|
if created_at.tzinfo is None:
|
||||||
|
# SQLAlchemy may return persisted UTC values without tzinfo. Never
|
||||||
|
# reinterpret them in the host's local timezone.
|
||||||
|
created_at = created_at.replace(tzinfo=timezone.utc)
|
||||||
|
return int(created_at.timestamp())
|
||||||
|
|
||||||
|
|
||||||
|
async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]:
|
||||||
|
"""Build one heartbeat per active Workspace."""
|
||||||
|
bindings = await ap.workspace_service.list_active_execution_bindings()
|
||||||
|
workspace_uuids = sorted({binding.workspace_uuid for binding in bindings})
|
||||||
|
workspace_create_ts = {
|
||||||
|
binding.workspace_uuid: _workspace_created_timestamp(getattr(binding, 'workspace_created_at', None))
|
||||||
|
for binding in bindings
|
||||||
|
}
|
||||||
|
resources = {
|
||||||
|
resource['workspace_uuid']: resource for resource in await _cloud_workspace_resource_counts(ap, bindings)
|
||||||
|
}
|
||||||
|
return [
|
||||||
|
await build_heartbeat_payload(
|
||||||
|
ap,
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
workspace_create_ts=workspace_create_ts.get(workspace_uuid, 0),
|
||||||
|
workspace_resource=resources.get(workspace_uuid),
|
||||||
|
)
|
||||||
|
for workspace_uuid in workspace_uuids
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
async def heartbeat_loop(ap: core_app.Application) -> None:
|
async def heartbeat_loop(ap: core_app.Application) -> None:
|
||||||
"""Send one heartbeat shortly after startup, then daily."""
|
"""Send one heartbeat shortly after startup, then daily."""
|
||||||
# Small delay so managers (platform, skills, plugins) finish loading first
|
# Small delay so managers (platform, skills, plugins) finish loading first
|
||||||
await asyncio.sleep(30)
|
await asyncio.sleep(30)
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
payload = await build_heartbeat_payload(ap)
|
for payload in await build_heartbeat_payloads(ap):
|
||||||
await ap.telemetry.start_send_task(payload)
|
# Heartbeats are a daily bounded batch, not best-effort query events.
|
||||||
|
# Await each send so the TelemetryManager's 8-task queue cannot drop
|
||||||
|
# Workspaces after the first batch.
|
||||||
|
await ap.telemetry.send(payload)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
try:
|
try:
|
||||||
ap.logger.debug(f'Telemetry heartbeat failed: {e}')
|
ap.logger.debug(f'Telemetry heartbeat failed: {e}')
|
||||||
|
|||||||
@@ -0,0 +1,22 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import typing
|
||||||
|
|
||||||
|
|
||||||
|
class WorkspaceExecutionContext(typing.Protocol):
|
||||||
|
@property
|
||||||
|
def instance_uuid(self) -> str: ...
|
||||||
|
|
||||||
|
@property
|
||||||
|
def workspace_uuid(self) -> str: ...
|
||||||
|
|
||||||
|
|
||||||
|
def workspace_identity(execution_context: WorkspaceExecutionContext) -> dict[str, str]:
|
||||||
|
"""Build both first-class telemetry identities for one execution."""
|
||||||
|
instance_id = execution_context.instance_uuid.strip()
|
||||||
|
workspace_uuid = execution_context.workspace_uuid.strip()
|
||||||
|
if not instance_id:
|
||||||
|
raise ValueError('Telemetry execution instance ID is empty')
|
||||||
|
if not workspace_uuid:
|
||||||
|
raise ValueError('Telemetry execution Workspace UUID is empty')
|
||||||
|
return {'instance_id': instance_id, 'workspace_uuid': workspace_uuid}
|
||||||
@@ -2,7 +2,11 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
|
import os
|
||||||
|
import typing
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
from ..core import app as core_app
|
from ..core import app as core_app
|
||||||
from ..utils import httpclient
|
from ..utils import httpclient
|
||||||
|
|
||||||
@@ -21,7 +25,7 @@ class TelemetryManager:
|
|||||||
def __init__(self, ap: core_app.Application):
|
def __init__(self, ap: core_app.Application):
|
||||||
self.ap = ap
|
self.ap = ap
|
||||||
|
|
||||||
self.telemetry_config = {}
|
self.telemetry_config: dict[str, typing.Any] = {}
|
||||||
self.send_tasks: list[asyncio.Task] = []
|
self.send_tasks: list[asyncio.Task] = []
|
||||||
self._client: httpx.AsyncClient | None = None
|
self._client: httpx.AsyncClient | None = None
|
||||||
|
|
||||||
@@ -131,7 +135,35 @@ class TelemetryManager:
|
|||||||
async with self._client_context() as client:
|
async with self._client_context() as client:
|
||||||
try:
|
try:
|
||||||
# Use asyncio.wait_for to ensure we always bound the total time
|
# Use asyncio.wait_for to ensure we always bound the total time
|
||||||
resp = await asyncio.wait_for(client.post(url, json=sanitized), timeout=10 + 1)
|
telemetry_token = os.getenv('LANGBOT_TELEMETRY_INGEST_TOKEN', '').strip()
|
||||||
|
headers: dict[str, str] = {}
|
||||||
|
if telemetry_token:
|
||||||
|
headers['X-LangBot-Telemetry-Token'] = telemetry_token
|
||||||
|
else:
|
||||||
|
workspace_uuid = str(sanitized.get('workspace_uuid', '')).strip()
|
||||||
|
user_service = getattr(self.ap, 'user_service', None)
|
||||||
|
if workspace_uuid and user_service is not None:
|
||||||
|
try:
|
||||||
|
owner = await user_service.get_workspace_owner(workspace_uuid)
|
||||||
|
owner_email = str(getattr(owner, 'user', '') or '').strip()
|
||||||
|
space_service = getattr(self.ap, 'space_service', None)
|
||||||
|
access_token = (
|
||||||
|
await space_service.get_valid_access_token(owner_email)
|
||||||
|
if owner_email and space_service is not None
|
||||||
|
else None
|
||||||
|
)
|
||||||
|
access_token = str(access_token or '').strip()
|
||||||
|
if access_token:
|
||||||
|
headers['Authorization'] = f'Bearer {access_token}'
|
||||||
|
except Exception:
|
||||||
|
self.ap.logger.debug(
|
||||||
|
'Could not resolve authenticated telemetry reporter', exc_info=True
|
||||||
|
)
|
||||||
|
if headers:
|
||||||
|
request = client.post(url, json=sanitized, headers=headers)
|
||||||
|
else:
|
||||||
|
request = client.post(url, json=sanitized)
|
||||||
|
resp = await asyncio.wait_for(request, timeout=10 + 1)
|
||||||
|
|
||||||
if resp.status_code >= 400:
|
if resp.status_code >= 400:
|
||||||
body = await httpclient.response_text(resp, max_chars=200)
|
body = await httpclient.response_text(resp, max_chars=200)
|
||||||
@@ -143,7 +175,8 @@ class TelemetryManager:
|
|||||||
app_err = False
|
app_err = False
|
||||||
try:
|
try:
|
||||||
j = await httpclient.parse_json_response(resp)
|
j = await httpclient.parse_json_response(resp)
|
||||||
if isinstance(j, dict) and j.get('code') is not None and int(j.get('code')) >= 400:
|
app_code = j.get('code') if isinstance(j, dict) else None
|
||||||
|
if app_code is not None and int(app_code) >= 400:
|
||||||
app_err = True
|
app_err = True
|
||||||
self.ap.logger.warning(
|
self.ap.logger.warning(
|
||||||
f'Telemetry post to {url} returned application error code {j.get("code")} - {j.get("msg")}'
|
f'Telemetry post to {url} returned application error code {j.get("code")} - {j.get("msg")}'
|
||||||
|
|||||||
@@ -67,7 +67,7 @@ class VectorDBManager:
|
|||||||
use_business_database = pgvector_config.get('use_business_database', False)
|
use_business_database = pgvector_config.get('use_business_database', False)
|
||||||
allowed_dimensions = pgvector_config.get(
|
allowed_dimensions = pgvector_config.get(
|
||||||
'allowed_dimensions',
|
'allowed_dimensions',
|
||||||
[384, 512, 768, 1024, 1536],
|
[384, 512, 768, 1024, 1536, 3072],
|
||||||
)
|
)
|
||||||
common_options = {
|
common_options = {
|
||||||
'use_business_database': use_business_database,
|
'use_business_database': use_business_database,
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ from collections.abc import AsyncIterator
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
from pgvector.sqlalchemy import Vector
|
from pgvector.sqlalchemy import HALFVEC, Vector
|
||||||
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
|
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||||
from sqlalchemy.orm import declarative_base
|
from sqlalchemy.orm import declarative_base
|
||||||
@@ -18,7 +18,7 @@ from langbot.pkg.vector.vdb import VectorDatabase
|
|||||||
|
|
||||||
Base = declarative_base()
|
Base = declarative_base()
|
||||||
|
|
||||||
DEFAULT_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
|
DEFAULT_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536, 3072)
|
||||||
|
|
||||||
# pgvector schema only stores these metadata fields.
|
# pgvector schema only stores these metadata fields.
|
||||||
_PG_SUPPORTED_FIELDS = {'text', 'file_id', 'chunk_uuid'}
|
_PG_SUPPORTED_FIELDS = {'text', 'file_id', 'chunk_uuid'}
|
||||||
@@ -321,7 +321,12 @@ class PgVectorDatabase(VectorDatabase):
|
|||||||
if len(query_embedding) != scope.embedding_dimension:
|
if len(query_embedding) != scope.embedding_dimension:
|
||||||
raise ValueError(f'Query embedding must have the selected dimension {scope.embedding_dimension}')
|
raise ValueError(f'Query embedding must have the selected dimension {scope.embedding_dimension}')
|
||||||
|
|
||||||
typed_embedding = sqlalchemy.cast(PgVectorEntry.embedding, Vector(scope.embedding_dimension))
|
typed_embedding = sqlalchemy.cast(
|
||||||
|
PgVectorEntry.embedding,
|
||||||
|
HALFVEC(scope.embedding_dimension)
|
||||||
|
if scope.embedding_dimension > 2000
|
||||||
|
else Vector(scope.embedding_dimension),
|
||||||
|
)
|
||||||
distance = typed_embedding.cosine_distance(query_embedding)
|
distance = typed_embedding.cosine_distance(query_embedding)
|
||||||
statement = (
|
statement = (
|
||||||
sqlalchemy.select(
|
sqlalchemy.select(
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from ..entity.persistence.user import AccountStatus, User
|
|||||||
from ..entity.persistence.workspace import (
|
from ..entity.persistence.workspace import (
|
||||||
InvitationStatus,
|
InvitationStatus,
|
||||||
MembershipRole,
|
MembershipRole,
|
||||||
|
MembershipSource,
|
||||||
MembershipStatus,
|
MembershipStatus,
|
||||||
Workspace,
|
Workspace,
|
||||||
WorkspaceInvitation,
|
WorkspaceInvitation,
|
||||||
@@ -88,6 +89,7 @@ class ResolvedWorkspaceAccess:
|
|||||||
@dataclasses.dataclass(frozen=True, slots=True)
|
@dataclasses.dataclass(frozen=True, slots=True)
|
||||||
class WorkspaceMemberView:
|
class WorkspaceMemberView:
|
||||||
membership: WorkspaceMembership
|
membership: WorkspaceMembership
|
||||||
|
display_name: str
|
||||||
email: str
|
email: str
|
||||||
|
|
||||||
|
|
||||||
@@ -294,7 +296,7 @@ class WorkspaceCollaborationService:
|
|||||||
async def operation(active_session: AsyncSession) -> list[WorkspaceMemberView]:
|
async def operation(active_session: AsyncSession) -> list[WorkspaceMemberView]:
|
||||||
await self._load_actor(active_session, workspace_uuid, actor)
|
await self._load_actor(active_session, workspace_uuid, actor)
|
||||||
statement = (
|
statement = (
|
||||||
sqlalchemy.select(WorkspaceMembership, User.user)
|
sqlalchemy.select(WorkspaceMembership, User.user, User.normalized_email)
|
||||||
.join(User, User.uuid == WorkspaceMembership.account_uuid)
|
.join(User, User.uuid == WorkspaceMembership.account_uuid)
|
||||||
.where(
|
.where(
|
||||||
WorkspaceMembership.workspace_uuid == workspace_uuid,
|
WorkspaceMembership.workspace_uuid == workspace_uuid,
|
||||||
@@ -304,8 +306,12 @@ class WorkspaceCollaborationService:
|
|||||||
.order_by(WorkspaceMembership.created_at, WorkspaceMembership.uuid)
|
.order_by(WorkspaceMembership.created_at, WorkspaceMembership.uuid)
|
||||||
)
|
)
|
||||||
return [
|
return [
|
||||||
WorkspaceMemberView(membership=membership, email=email)
|
WorkspaceMemberView(
|
||||||
for membership, email in (await active_session.execute(statement)).all()
|
membership=membership,
|
||||||
|
display_name=display_name,
|
||||||
|
email=email,
|
||||||
|
)
|
||||||
|
for membership, display_name, email in (await active_session.execute(statement)).all()
|
||||||
]
|
]
|
||||||
|
|
||||||
return await self._run(operation, session=session, read_only=True)
|
return await self._run(operation, session=session, read_only=True)
|
||||||
@@ -478,6 +484,7 @@ class WorkspaceCollaborationService:
|
|||||||
account_uuid=account_uuid,
|
account_uuid=account_uuid,
|
||||||
role=invitation.role,
|
role=invitation.role,
|
||||||
status=MembershipStatus.ACTIVE.value,
|
status=MembershipStatus.ACTIVE.value,
|
||||||
|
source=MembershipSource.LOCAL.value,
|
||||||
invited_by_account_uuid=invitation.created_by_account_uuid,
|
invited_by_account_uuid=invitation.created_by_account_uuid,
|
||||||
joined_at=now,
|
joined_at=now,
|
||||||
projection_revision=0,
|
projection_revision=0,
|
||||||
@@ -486,6 +493,7 @@ class WorkspaceCollaborationService:
|
|||||||
elif membership.status != MembershipStatus.ACTIVE.value:
|
elif membership.status != MembershipStatus.ACTIVE.value:
|
||||||
membership.role = invitation.role
|
membership.role = invitation.role
|
||||||
membership.status = MembershipStatus.ACTIVE.value
|
membership.status = MembershipStatus.ACTIVE.value
|
||||||
|
membership.source = MembershipSource.LOCAL.value
|
||||||
membership.invited_by_account_uuid = invitation.created_by_account_uuid
|
membership.invited_by_account_uuid = invitation.created_by_account_uuid
|
||||||
membership.joined_at = now
|
membership.joined_at = now
|
||||||
|
|
||||||
@@ -606,6 +614,8 @@ class WorkspaceCollaborationService:
|
|||||||
) -> WorkspaceMembership:
|
) -> WorkspaceMembership:
|
||||||
if role not in {item.value for item in MembershipRole}:
|
if role not in {item.value for item in MembershipRole}:
|
||||||
raise MembershipPermissionError('Unknown Workspace role')
|
raise MembershipPermissionError('Unknown Workspace role')
|
||||||
|
if role == MembershipRole.OWNER.value:
|
||||||
|
raise MembershipPermissionError('Workspace ownership cannot be transferred')
|
||||||
|
|
||||||
async def operation(active_session: AsyncSession) -> WorkspaceMembership:
|
async def operation(active_session: AsyncSession) -> WorkspaceMembership:
|
||||||
await self._require_active_workspace(active_session, workspace_uuid)
|
await self._require_active_workspace(active_session, workspace_uuid)
|
||||||
@@ -617,8 +627,8 @@ class WorkspaceCollaborationService:
|
|||||||
target_account_uuid,
|
target_account_uuid,
|
||||||
)
|
)
|
||||||
self._require_can_manage_target(persisted_actor, target, new_role=role)
|
self._require_can_manage_target(persisted_actor, target, new_role=role)
|
||||||
if target.role == MembershipRole.OWNER.value and role != MembershipRole.OWNER.value:
|
if target.role == MembershipRole.OWNER.value:
|
||||||
await self._require_another_owner(active_session, workspace_uuid, target.account_uuid)
|
raise LastOwnerError('The Workspace owner cannot be removed or demoted')
|
||||||
target.role = role
|
target.role = role
|
||||||
await active_session.flush()
|
await active_session.flush()
|
||||||
return target
|
return target
|
||||||
@@ -644,7 +654,7 @@ class WorkspaceCollaborationService:
|
|||||||
)
|
)
|
||||||
self._require_can_manage_target(persisted_actor, target)
|
self._require_can_manage_target(persisted_actor, target)
|
||||||
if target.role == MembershipRole.OWNER.value:
|
if target.role == MembershipRole.OWNER.value:
|
||||||
await self._require_another_owner(active_session, workspace_uuid, target.account_uuid)
|
raise LastOwnerError('The Workspace owner cannot be removed or demoted')
|
||||||
target.status = MembershipStatus.REMOVED.value
|
target.status = MembershipStatus.REMOVED.value
|
||||||
await active_session.flush()
|
await active_session.flush()
|
||||||
return target
|
return target
|
||||||
@@ -751,26 +761,6 @@ class WorkspaceCollaborationService:
|
|||||||
raise WorkspaceNotFoundError('Workspace not found')
|
raise WorkspaceNotFoundError('Workspace not found')
|
||||||
return persisted_actor
|
return persisted_actor
|
||||||
|
|
||||||
async def _require_another_owner(
|
|
||||||
self,
|
|
||||||
session: AsyncSession,
|
|
||||||
workspace_uuid: str,
|
|
||||||
excluded_account_uuid: str,
|
|
||||||
) -> None:
|
|
||||||
owners = (
|
|
||||||
await session.scalars(
|
|
||||||
sqlalchemy.select(WorkspaceMembership)
|
|
||||||
.where(
|
|
||||||
WorkspaceMembership.workspace_uuid == workspace_uuid,
|
|
||||||
WorkspaceMembership.status == MembershipStatus.ACTIVE.value,
|
|
||||||
WorkspaceMembership.role == MembershipRole.OWNER.value,
|
|
||||||
)
|
|
||||||
.with_for_update()
|
|
||||||
)
|
|
||||||
).all()
|
|
||||||
if not any(owner.account_uuid != excluded_account_uuid for owner in owners):
|
|
||||||
raise LastOwnerError('The last Workspace owner cannot be removed or demoted')
|
|
||||||
|
|
||||||
def _require_actor_workspace(self, actor: WorkspaceMembership, workspace_uuid: str) -> None:
|
def _require_actor_workspace(self, actor: WorkspaceMembership, workspace_uuid: str) -> None:
|
||||||
if actor.workspace_uuid != workspace_uuid or actor.status != MembershipStatus.ACTIVE.value:
|
if actor.workspace_uuid != workspace_uuid or actor.status != MembershipStatus.ACTIVE.value:
|
||||||
raise WorkspaceNotFoundError('Workspace not found')
|
raise WorkspaceNotFoundError('Workspace not found')
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import datetime
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
|
||||||
@@ -12,3 +13,4 @@ class WorkspaceExecutionBinding:
|
|||||||
placement_generation: int
|
placement_generation: int
|
||||||
write_fenced: bool
|
write_fenced: bool
|
||||||
state: str
|
state: str
|
||||||
|
workspace_created_at: datetime.datetime | None = None
|
||||||
|
|||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
|
||||||
|
_INSTANCE_PREFIX = 'instance_'
|
||||||
|
_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID('8ea04f29-8528-4cc3-bb28-30a838c89d76')
|
||||||
|
|
||||||
|
|
||||||
|
def workspace_uuid_from_instance_id(instance_id: str) -> str:
|
||||||
|
"""Return the stable OSS Workspace UUID for a persisted instance identity."""
|
||||||
|
value = instance_id.strip()
|
||||||
|
if not value:
|
||||||
|
raise ValueError('LangBot instance identity is empty')
|
||||||
|
|
||||||
|
candidate = value[len(_INSTANCE_PREFIX) :] if value.startswith(_INSTANCE_PREFIX) else value
|
||||||
|
try:
|
||||||
|
return str(uuid.UUID(candidate))
|
||||||
|
except ValueError:
|
||||||
|
return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value))
|
||||||
@@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|||||||
|
|
||||||
from ..entity.persistence.workspace import (
|
from ..entity.persistence.workspace import (
|
||||||
MembershipRole,
|
MembershipRole,
|
||||||
|
MembershipSource,
|
||||||
MembershipStatus,
|
MembershipStatus,
|
||||||
Workspace,
|
Workspace,
|
||||||
WorkspaceExecutionSource,
|
WorkspaceExecutionSource,
|
||||||
@@ -30,6 +31,7 @@ from .errors import (
|
|||||||
WorkspaceOwnerAlreadyExistsError,
|
WorkspaceOwnerAlreadyExistsError,
|
||||||
)
|
)
|
||||||
from .entities import WorkspaceExecutionBinding
|
from .entities import WorkspaceExecutionBinding
|
||||||
|
from .identity import workspace_uuid_from_instance_id
|
||||||
from .policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
from .policy import CloudWorkspacePolicy, SingleWorkspacePolicy
|
||||||
from .repository import WorkspaceRepository
|
from .repository import WorkspaceRepository
|
||||||
|
|
||||||
@@ -282,6 +284,7 @@ class WorkspaceService:
|
|||||||
placement_generation=execution_state.active_generation,
|
placement_generation=execution_state.active_generation,
|
||||||
write_fenced=execution_state.write_fenced,
|
write_fenced=execution_state.write_fenced,
|
||||||
state=execution_state.state,
|
state=execution_state.state,
|
||||||
|
workspace_created_at=workspace.created_at,
|
||||||
)
|
)
|
||||||
|
|
||||||
binding = await self._run(operation, session=session)
|
binding = await self._run(operation, session=session)
|
||||||
@@ -449,6 +452,7 @@ class WorkspaceService:
|
|||||||
account_uuid=account_uuid,
|
account_uuid=account_uuid,
|
||||||
role=MembershipRole.OWNER.value,
|
role=MembershipRole.OWNER.value,
|
||||||
status=MembershipStatus.ACTIVE.value,
|
status=MembershipStatus.ACTIVE.value,
|
||||||
|
source=MembershipSource.LOCAL.value,
|
||||||
joined_at=joined_at,
|
joined_at=joined_at,
|
||||||
projection_revision=0,
|
projection_revision=0,
|
||||||
)
|
)
|
||||||
@@ -456,6 +460,7 @@ class WorkspaceService:
|
|||||||
else:
|
else:
|
||||||
membership.role = MembershipRole.OWNER.value
|
membership.role = MembershipRole.OWNER.value
|
||||||
membership.status = MembershipStatus.ACTIVE.value
|
membership.status = MembershipStatus.ACTIVE.value
|
||||||
|
membership.source = MembershipSource.LOCAL.value
|
||||||
membership.joined_at = membership.joined_at or joined_at
|
membership.joined_at = membership.joined_at or joined_at
|
||||||
|
|
||||||
if workspace.created_by_account_uuid is None:
|
if workspace.created_by_account_uuid is None:
|
||||||
@@ -497,7 +502,7 @@ class WorkspaceService:
|
|||||||
created_by_account_uuid: str | None = None,
|
created_by_account_uuid: str | None = None,
|
||||||
) -> Workspace:
|
) -> Workspace:
|
||||||
return Workspace(
|
return Workspace(
|
||||||
uuid=str(uuid.uuid4()),
|
uuid=workspace_uuid_from_instance_id(self.instance_uuid),
|
||||||
instance_uuid=self.instance_uuid,
|
instance_uuid=self.instance_uuid,
|
||||||
name=name,
|
name=name,
|
||||||
slug=slug,
|
slug=slug,
|
||||||
|
|||||||
@@ -201,7 +201,7 @@ vdb:
|
|||||||
# keep this false when deliberately using an external pgvector DB.
|
# keep this false when deliberately using an external pgvector DB.
|
||||||
use_business_database: false
|
use_business_database: false
|
||||||
# Release migrations create one partial ANN index per enabled value.
|
# Release migrations create one partial ANN index per enabled value.
|
||||||
allowed_dimensions: [384, 512, 768, 1024, 1536]
|
allowed_dimensions: [384, 512, 768, 1024, 1536, 3072]
|
||||||
host: '127.0.0.1'
|
host: '127.0.0.1'
|
||||||
port: 5433
|
port: 5433
|
||||||
database: 'langbot'
|
database: 'langbot'
|
||||||
@@ -245,6 +245,8 @@ storage:
|
|||||||
max_concurrency: 16
|
max_concurrency: 16
|
||||||
plugin:
|
plugin:
|
||||||
enable: true
|
enable: true
|
||||||
|
# Maximum time for the Runtime transport, handshake, and desired-state replay.
|
||||||
|
connect_timeout_seconds: 180.0
|
||||||
runtime_ws_url: 'ws://langbot_plugin_runtime:5400/control/ws'
|
runtime_ws_url: 'ws://langbot_plugin_runtime:5400/control/ws'
|
||||||
enable_marketplace: true
|
enable_marketplace: true
|
||||||
display_plugin_debug_url: 'ws://localhost:5401/plugin/debug/ws'
|
display_plugin_debug_url: 'ws://localhost:5401/plugin/debug/ws'
|
||||||
@@ -326,8 +328,9 @@ box:
|
|||||||
enabled: true
|
enabled: true
|
||||||
backend: 'local' # 'local' (Docker/nsjail), 'docker', 'nsjail', or 'e2b'. Can be written via BOX__BACKEND.
|
backend: 'local' # 'local' (Docker/nsjail), 'docker', 'nsjail', or 'e2b'. Can be written via BOX__BACKEND.
|
||||||
runtime:
|
runtime:
|
||||||
# External WebSocket runtimes also require LANGBOT_BOX_CONTROL_TOKEN in
|
# LANGBOT_BOX_CONTROL_TOKEN is optional for OSS external WebSocket
|
||||||
# both LangBot and Box. Keep the shared secret out of this config file.
|
# runtimes. To protect an exposed endpoint, set the same strong secret
|
||||||
|
# in both LangBot and Box. Keep it out of this config file.
|
||||||
endpoint: '' # External Box Runtime base URL, e.g. 'ws://127.0.0.1:5410'. Leave empty for local auto-managed runtime.
|
endpoint: '' # External Box Runtime base URL, e.g. 'ws://127.0.0.1:5410'. Leave empty for local auto-managed runtime.
|
||||||
limits:
|
limits:
|
||||||
max_sessions: 64
|
max_sessions: 64
|
||||||
|
|||||||
@@ -92,6 +92,7 @@ stages:
|
|||||||
default:
|
default:
|
||||||
primary: ''
|
primary: ''
|
||||||
fallbacks: []
|
fallbacks: []
|
||||||
|
reasoning: {}
|
||||||
- name: max-round
|
- name: max-round
|
||||||
label:
|
label:
|
||||||
en_US: Max Round
|
en_US: Max Round
|
||||||
|
|||||||
@@ -106,7 +106,12 @@ async def plugin_security_api(plugin_module):
|
|||||||
application.plugin_connector.require_workspace_context = AsyncMock()
|
application.plugin_connector.require_workspace_context = AsyncMock()
|
||||||
application.plugin_connector.list_plugins = AsyncMock(return_value=[raw_plugin])
|
application.plugin_connector.list_plugins = AsyncMock(return_value=[raw_plugin])
|
||||||
application.plugin_connector.get_plugin_info = AsyncMock(return_value=raw_plugin)
|
application.plugin_connector.get_plugin_info = AsyncMock(return_value=raw_plugin)
|
||||||
application.plugin_connector.get_debug_info = AsyncMock(return_value={'plugin_debug_key': 'runtime-debug-secret'})
|
application.plugin_connector.get_debug_info = AsyncMock(
|
||||||
|
return_value={
|
||||||
|
'plugin_debug_key': 'runtime-debug-secret',
|
||||||
|
'expires_at': '2026-08-04T12:00:00Z',
|
||||||
|
}
|
||||||
|
)
|
||||||
application.plugin_connector.get_plugin_logs = AsyncMock(return_value=['private runtime line'])
|
application.plugin_connector.get_plugin_logs = AsyncMock(return_value=['private runtime line'])
|
||||||
application.plugin_connector.set_plugin_config = AsyncMock()
|
application.plugin_connector.set_plugin_config = AsyncMock()
|
||||||
|
|
||||||
@@ -230,10 +235,22 @@ async def test_debug_key_requires_resource_manage_permission(plugin_security_api
|
|||||||
assert operator_denied.status_code == 403
|
assert operator_denied.status_code == 403
|
||||||
assert allowed.status_code == 200
|
assert allowed.status_code == 200
|
||||||
assert (await allowed.get_json())['data'] == {
|
assert (await allowed.get_json())['data'] == {
|
||||||
'debug_url': 'http://localhost:5401',
|
'debug_url': 'ws://localhost:5401/plugin/debug/ws',
|
||||||
'plugin_debug_key': 'runtime-debug-secret',
|
'plugin_debug_key': 'runtime-debug-secret',
|
||||||
|
'expires_at': '2026-08-04T12:00:00Z',
|
||||||
}
|
}
|
||||||
application.plugin_connector.get_debug_info.assert_awaited_once_with()
|
application.plugin_connector.get_debug_info.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_debug_info_uses_websocket_endpoint_for_legacy_config(plugin_security_api):
|
||||||
|
application, client, _ = plugin_security_api
|
||||||
|
application.instance_config.data['plugin'].pop('display_plugin_debug_url')
|
||||||
|
|
||||||
|
response = await client.get('/api/v1/plugins/debug-info', headers=_headers('manager-token'))
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert (await response.get_json())['data']['debug_url'] == 'ws://localhost:5401/plugin/debug/ws'
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|||||||
@@ -9,6 +9,8 @@ Run: uv run pytest tests/integration/api/test_smoke.py -q
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from unittest.mock import MagicMock, AsyncMock, Mock
|
from unittest.mock import MagicMock, AsyncMock, Mock
|
||||||
|
|
||||||
@@ -304,12 +306,34 @@ class TestUserInitEndpoint:
|
|||||||
data = await response.get_json()
|
data = await response.get_json()
|
||||||
assert data['data'] == {
|
assert data['data'] == {
|
||||||
'initialized': True,
|
'initialized': True,
|
||||||
|
'authenticated_invitation_acceptance_enabled': False,
|
||||||
'password_login_enabled': True,
|
'password_login_enabled': True,
|
||||||
'space_login_enabled': False,
|
'space_login_enabled': False,
|
||||||
}
|
}
|
||||||
fake_api_app.user_service.get_login_capabilities.assert_awaited_once_with()
|
fake_api_app.user_service.get_login_capabilities.assert_awaited_once_with()
|
||||||
fake_api_app.user_service.get_first_user.assert_not_awaited()
|
fake_api_app.user_service.get_first_user.assert_not_awaited()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_account_info_enables_authenticated_invitation_acceptance_in_cloud(
|
||||||
|
self, quart_test_client, fake_api_app
|
||||||
|
):
|
||||||
|
fake_api_app.deployment = SimpleNamespace(mode='cloud')
|
||||||
|
fake_api_app.user_service.is_initialized.return_value = True
|
||||||
|
fake_api_app.user_service.get_login_capabilities = AsyncMock(
|
||||||
|
return_value={'password_login_enabled': True, 'space_login_enabled': True}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await quart_test_client.get('/api/v1/user/account-info')
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
data = await response.get_json()
|
||||||
|
assert data['data'] == {
|
||||||
|
'initialized': True,
|
||||||
|
'authenticated_invitation_acceptance_enabled': True,
|
||||||
|
'password_login_enabled': False,
|
||||||
|
'space_login_enabled': True,
|
||||||
|
}
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_recovery_key_resets_any_existing_account(self, quart_test_client, fake_api_app, monkeypatch):
|
async def test_recovery_key_resets_any_existing_account(self, quart_test_client, fake_api_app, monkeypatch):
|
||||||
fake_api_app.user_service.is_initialized.return_value = True
|
fake_api_app.user_service.is_initialized.return_value = True
|
||||||
|
|||||||
@@ -0,0 +1,452 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import base64
|
||||||
|
import datetime
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import sqlalchemy
|
||||||
|
from cryptography.hazmat.primitives import serialization
|
||||||
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
||||||
|
from quart import Quart
|
||||||
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||||
|
|
||||||
|
from langbot.pkg.api.http.authz import Permission
|
||||||
|
from langbot.pkg.api.http.context import PrincipalType, RequestContext
|
||||||
|
from langbot.pkg.api.http.controller import group
|
||||||
|
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
|
||||||
|
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
|
||||||
|
from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService
|
||||||
|
from langbot.pkg.cloud.support_admin import SupportAdminSessionService
|
||||||
|
from langbot.pkg.entity.persistence.base import Base
|
||||||
|
from langbot.pkg.entity.persistence.support_admin import SupportAdminTemporarySession
|
||||||
|
from langbot.pkg.entity.persistence.user import User
|
||||||
|
from langbot.pkg.entity.persistence.workspace import (
|
||||||
|
Workspace,
|
||||||
|
WorkspaceExecutionState,
|
||||||
|
WorkspaceMembership,
|
||||||
|
)
|
||||||
|
from langbot.pkg.workspace.service import WorkspaceService
|
||||||
|
|
||||||
|
|
||||||
|
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||||
|
|
||||||
|
|
||||||
|
INSTANCE_UUID = 'instance-support-admin'
|
||||||
|
WORKSPACE_UUID = '10000000-0000-4000-8000-000000000001'
|
||||||
|
OTHER_WORKSPACE_UUID = '10000000-0000-4000-8000-000000000002'
|
||||||
|
ACTOR_ACCOUNT_UUID = '20000000-0000-4000-8000-000000000001'
|
||||||
|
KEY_ID = 'support-admin-key-1'
|
||||||
|
|
||||||
|
|
||||||
|
def _base64url(raw: bytes) -> str:
|
||||||
|
return base64.urlsafe_b64encode(raw).rstrip(b'=').decode('ascii')
|
||||||
|
|
||||||
|
|
||||||
|
def _sign(private_key: Ed25519PrivateKey, claims: dict, *, key_id: str = KEY_ID) -> str:
|
||||||
|
header = {'alg': 'EdDSA', 'kid': key_id, 'typ': 'langbot-control-plane+jwt'}
|
||||||
|
encoded_header = _base64url(json.dumps(header, separators=(',', ':')).encode('utf-8'))
|
||||||
|
encoded_claims = _base64url(json.dumps(claims, separators=(',', ':')).encode('utf-8'))
|
||||||
|
signing_input = f'{encoded_header}.{encoded_claims}'
|
||||||
|
return f'{signing_input}.{_base64url(private_key.sign(signing_input.encode("ascii")))}'
|
||||||
|
|
||||||
|
|
||||||
|
def _admin_claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE_UUID) -> dict:
|
||||||
|
return {
|
||||||
|
'iss': 'langbot-space',
|
||||||
|
'aud': 'langbot-cloud-runtime',
|
||||||
|
'sub': f'langbot-instance:{INSTANCE_UUID}',
|
||||||
|
'jti': jti or str(uuid.uuid4()),
|
||||||
|
'iat': now,
|
||||||
|
'nbf': now - 5,
|
||||||
|
'exp': now + 90,
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'kind': 'workspace.support_admin_launch',
|
||||||
|
'payload': {
|
||||||
|
'workspace_uuid': workspace_uuid,
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': ACTOR_ACCOUNT_UUID,
|
||||||
|
'effective_role': 'owner',
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@group.group_class('support_admin_probe', '/api/v1/support-admin-probe')
|
||||||
|
class SupportAdminProbeGroup(group.RouterGroup):
|
||||||
|
async def initialize(self) -> None:
|
||||||
|
@self.route('/user-token', auth_type=group.AuthType.USER_TOKEN, permission=Permission.WORKSPACE_VIEW)
|
||||||
|
async def _(request_context: RequestContext) -> str:
|
||||||
|
return self.success(data=_context_payload(request_context))
|
||||||
|
|
||||||
|
@self.route(
|
||||||
|
'/member-operation',
|
||||||
|
auth_type=group.AuthType.USER_TOKEN,
|
||||||
|
permission=Permission.MEMBER_VIEW,
|
||||||
|
)
|
||||||
|
async def member_operation(request_context: RequestContext) -> str:
|
||||||
|
return self.success(data=_context_payload(request_context))
|
||||||
|
|
||||||
|
@self.route(
|
||||||
|
'/user-token-or-api-key',
|
||||||
|
auth_type=group.AuthType.USER_TOKEN_OR_API_KEY,
|
||||||
|
permission=Permission.WORKSPACE_VIEW,
|
||||||
|
)
|
||||||
|
async def _(request_context: RequestContext) -> str:
|
||||||
|
return self.success(data=_context_payload(request_context))
|
||||||
|
|
||||||
|
|
||||||
|
def _context_payload(request_context: RequestContext) -> dict:
|
||||||
|
return {
|
||||||
|
'principal_type': request_context.principal.principal_type.value,
|
||||||
|
'actor_account_uuid': request_context.principal.actor_account_uuid,
|
||||||
|
'account_uuid': request_context.principal.account_uuid,
|
||||||
|
'role': request_context.workspace.role,
|
||||||
|
'membership_uuid': request_context.workspace.membership_uuid,
|
||||||
|
'permissions': sorted(request_context.workspace.permissions),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class _TenantUow:
|
||||||
|
def __init__(self, engine):
|
||||||
|
self._engine = engine
|
||||||
|
self.session = None
|
||||||
|
self._transaction = None
|
||||||
|
|
||||||
|
async def __aenter__(self):
|
||||||
|
session_factory = async_sessionmaker(self._engine, expire_on_commit=False)
|
||||||
|
self.session = session_factory()
|
||||||
|
self._transaction = await self.session.begin()
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, exc_type, exc, traceback):
|
||||||
|
try:
|
||||||
|
if exc_type is None:
|
||||||
|
await self._transaction.commit()
|
||||||
|
else:
|
||||||
|
await self._transaction.rollback()
|
||||||
|
finally:
|
||||||
|
await self.session.close()
|
||||||
|
|
||||||
|
|
||||||
|
class _TenantScope:
|
||||||
|
async def __aenter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, exc_type, exc, traceback):
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
class _PersistenceManager:
|
||||||
|
def __init__(self, engine):
|
||||||
|
self._engine = engine
|
||||||
|
self.mode = SimpleNamespace(value='oss_compat')
|
||||||
|
|
||||||
|
def get_db_engine(self):
|
||||||
|
return self._engine
|
||||||
|
|
||||||
|
def tenant_uow(self, workspace_uuid: str):
|
||||||
|
del workspace_uuid
|
||||||
|
return _TenantUow(self._engine)
|
||||||
|
|
||||||
|
def tenant_scope(self, workspace_uuid: str):
|
||||||
|
del workspace_uuid
|
||||||
|
return _TenantScope()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
async def support_admin_api(tmp_path):
|
||||||
|
private_key = Ed25519PrivateKey.generate()
|
||||||
|
public_key = private_key.public_key().public_bytes(
|
||||||
|
encoding=serialization.Encoding.Raw,
|
||||||
|
format=serialization.PublicFormat.Raw,
|
||||||
|
)
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "support-admin.db"}')
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.run_sync(
|
||||||
|
Base.metadata.create_all,
|
||||||
|
tables=[
|
||||||
|
User.__table__,
|
||||||
|
Workspace.__table__,
|
||||||
|
WorkspaceExecutionState.__table__,
|
||||||
|
WorkspaceMembership.__table__,
|
||||||
|
SupportAdminTemporarySession.__table__,
|
||||||
|
],
|
||||||
|
)
|
||||||
|
for workspace_uuid, slug in (
|
||||||
|
(WORKSPACE_UUID, 'support-admin-a'),
|
||||||
|
(OTHER_WORKSPACE_UUID, 'support-admin-b'),
|
||||||
|
):
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(Workspace).values(
|
||||||
|
uuid=workspace_uuid,
|
||||||
|
instance_uuid=INSTANCE_UUID,
|
||||||
|
name=slug,
|
||||||
|
slug=slug,
|
||||||
|
type='team',
|
||||||
|
status='active',
|
||||||
|
source='cloud_projection',
|
||||||
|
projection_revision=1,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(WorkspaceExecutionState).values(
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
instance_uuid=INSTANCE_UUID,
|
||||||
|
active_generation=1,
|
||||||
|
state='active',
|
||||||
|
write_fenced=False,
|
||||||
|
source='cloud',
|
||||||
|
desired_state_revision=1,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
app = SimpleNamespace()
|
||||||
|
app.persistence_mgr = _PersistenceManager(engine)
|
||||||
|
app.instance_config = SimpleNamespace(
|
||||||
|
data={
|
||||||
|
'system': {
|
||||||
|
'jwt': {'secret': 'support-admin-secret', 'expire': 3600},
|
||||||
|
'websocket_retention': {},
|
||||||
|
},
|
||||||
|
'space': {
|
||||||
|
'launch': {
|
||||||
|
'control_plane_public_key': _base64url(public_key),
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'api': {'global_api_key': ''},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
app.logger = logging.getLogger('support-admin-test')
|
||||||
|
app.deployment = SimpleNamespace(mode='cloud', multi_workspace_enabled=True, verification_key_id=KEY_ID)
|
||||||
|
app.directory_projection_service = SimpleNamespace(require_ready=lambda: None)
|
||||||
|
app.workspace_service = WorkspaceService(app, instance_uuid=INSTANCE_UUID)
|
||||||
|
app.entitlement_resolver = SimpleNamespace(
|
||||||
|
instance_uuid=INSTANCE_UUID,
|
||||||
|
resolve=AsyncMock(return_value=SimpleNamespace(entitlement_revision=7)),
|
||||||
|
)
|
||||||
|
app.support_admin_session_service = SupportAdminSessionService(app)
|
||||||
|
app.space_launch_service = SpaceLaunchService(app)
|
||||||
|
app.user_service = SimpleNamespace()
|
||||||
|
app.user_service.get_authenticated_account = AsyncMock(side_effect=AssertionError('normal account auth used'))
|
||||||
|
app.user_service.verify_jwt_token = AsyncMock(side_effect=AssertionError('normal token verification used'))
|
||||||
|
app.user_service.get_user_by_email = AsyncMock(side_effect=AssertionError('user lookup used'))
|
||||||
|
app.apikey_service = SimpleNamespace()
|
||||||
|
app.apikey_service.authenticate_api_key = AsyncMock(
|
||||||
|
return_value=SimpleNamespace(
|
||||||
|
instance_uuid=INSTANCE_UUID,
|
||||||
|
workspace_uuid=OTHER_WORKSPACE_UUID,
|
||||||
|
placement_generation=1,
|
||||||
|
api_key_uuid='api-key',
|
||||||
|
permissions=frozenset(permission.value for permission in Permission),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
quart_app = Quart(__name__)
|
||||||
|
await UserRouterGroup(app, quart_app).initialize()
|
||||||
|
await SupportAdminProbeGroup(app, quart_app).initialize()
|
||||||
|
|
||||||
|
yield app, quart_app.test_client(), engine, private_key
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
async def _issue_support_token(app, private_key: Ed25519PrivateKey, *, jti: str | None = None) -> dict[str, str]:
|
||||||
|
launch = await app.space_launch_service.consume_assertion(
|
||||||
|
_sign(private_key, _admin_claims(now=int(time.time()), jti=jti)),
|
||||||
|
expected_workspace_uuid=WORKSPACE_UUID,
|
||||||
|
)
|
||||||
|
return launch
|
||||||
|
|
||||||
|
|
||||||
|
def _auth(token: str, workspace_uuid: str = WORKSPACE_UUID) -> dict[str, str]:
|
||||||
|
return {'Authorization': f'Bearer {token}', 'X-Workspace-Id': workspace_uuid}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_membership_only_routes_are_denied(support_admin_api):
|
||||||
|
app, client, _engine, private_key = support_admin_api
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
|
||||||
|
response = await client.get(
|
||||||
|
'/api/v1/support-admin-probe/member-operation',
|
||||||
|
headers=_auth(launch['support_admin_token']),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 403
|
||||||
|
assert (await response.get_json())['code'] == 'permission_denied'
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_check_token_is_rejected(support_admin_api):
|
||||||
|
app, client, _engine, private_key = support_admin_api
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
|
||||||
|
response = await client.get('/api/v1/user/check-token', headers=_auth(launch['support_admin_token']))
|
||||||
|
|
||||||
|
assert response.status_code == 401
|
||||||
|
assert (await response.get_json())['code'] == 'invalid_authentication'
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_cross_workspace_denied_for_user_token_and_or_api_key(support_admin_api):
|
||||||
|
app, client, _engine, private_key = support_admin_api
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
|
||||||
|
missing_selector = await client.get(
|
||||||
|
'/api/v1/support-admin-probe/user-token',
|
||||||
|
headers={'Authorization': f'Bearer {launch["support_admin_token"]}'},
|
||||||
|
)
|
||||||
|
user_response = await client.get(
|
||||||
|
'/api/v1/support-admin-probe/user-token',
|
||||||
|
headers=_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID),
|
||||||
|
)
|
||||||
|
either_response = await client.get(
|
||||||
|
'/api/v1/support-admin-probe/user-token-or-api-key',
|
||||||
|
headers={
|
||||||
|
**_auth(launch['support_admin_token'], OTHER_WORKSPACE_UUID),
|
||||||
|
'X-API-Key': 'valid-api-key',
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert missing_selector.status_code == 400
|
||||||
|
assert user_response.status_code == 401
|
||||||
|
assert either_response.status_code == 401
|
||||||
|
app.apikey_service.authenticate_api_key.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_request_context_has_actor_owner_and_no_membership(support_admin_api):
|
||||||
|
app, client, engine, private_key = support_admin_api
|
||||||
|
before_count = await _membership_count(engine)
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
|
||||||
|
response = await client.get(
|
||||||
|
'/api/v1/support-admin-probe/user-token',
|
||||||
|
headers=_auth(launch['support_admin_token']),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
data = (await response.get_json())['data']
|
||||||
|
permissions = set(data.pop('permissions'))
|
||||||
|
assert Permission.WORKSPACE_VIEW.value in permissions
|
||||||
|
assert Permission.RESOURCE_MANAGE.value in permissions
|
||||||
|
assert not permissions.intersection(
|
||||||
|
{
|
||||||
|
Permission.MEMBER_VIEW.value,
|
||||||
|
Permission.MEMBER_INVITE.value,
|
||||||
|
Permission.MEMBER_UPDATE_ROLE.value,
|
||||||
|
Permission.MEMBER_REMOVE.value,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
assert data == {
|
||||||
|
'principal_type': PrincipalType.SUPPORT_ADMIN.value,
|
||||||
|
'actor_account_uuid': ACTOR_ACCOUNT_UUID,
|
||||||
|
'account_uuid': None,
|
||||||
|
'role': 'owner',
|
||||||
|
'membership_uuid': None,
|
||||||
|
}
|
||||||
|
assert await _membership_count(engine) == before_count
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_missing_workspace_is_controlled_launch_failure(support_admin_api):
|
||||||
|
app, _client, engine, private_key = support_admin_api
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.delete(WorkspaceExecutionState).where(WorkspaceExecutionState.workspace_uuid == WORKSPACE_UUID)
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(SpaceLaunchError, match='unavailable'):
|
||||||
|
await _issue_support_token(app, private_key)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_launch_replay_is_durable_across_service_instances(support_admin_api):
|
||||||
|
app, _client, _engine, private_key = support_admin_api
|
||||||
|
jti = str(uuid.uuid4())
|
||||||
|
|
||||||
|
await _issue_support_token(app, private_key, jti=jti)
|
||||||
|
second_service = SpaceLaunchService(app)
|
||||||
|
|
||||||
|
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
||||||
|
await second_service.consume_assertion(
|
||||||
|
_sign(private_key, _admin_claims(now=int(time.time()), jti=jti)),
|
||||||
|
expected_workspace_uuid=WORKSPACE_UUID,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_persisted_expiry_and_revocation_are_enforced(support_admin_api):
|
||||||
|
app, client, engine, private_key = support_admin_api
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
token = launch['support_admin_token']
|
||||||
|
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.update(SupportAdminTemporarySession)
|
||||||
|
.where(SupportAdminTemporarySession.grant_jti_hash == launch['grant_jti_hash'])
|
||||||
|
.values(expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None) - datetime.timedelta(minutes=1))
|
||||||
|
)
|
||||||
|
expired = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(token))
|
||||||
|
assert expired.status_code == 401
|
||||||
|
|
||||||
|
second = await _issue_support_token(app, private_key)
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.update(SupportAdminTemporarySession)
|
||||||
|
.where(SupportAdminTemporarySession.grant_jti_hash == second['grant_jti_hash'])
|
||||||
|
.values(revoked_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None))
|
||||||
|
)
|
||||||
|
revoked = await client.get('/api/v1/support-admin-probe/user-token', headers=_auth(second['support_admin_token']))
|
||||||
|
assert revoked.status_code == 401
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_websocket_preserves_actor_and_revalidates(support_admin_api):
|
||||||
|
app, _client, _engine, private_key = support_admin_api
|
||||||
|
launch = await _issue_support_token(app, private_key)
|
||||||
|
captured_contexts = []
|
||||||
|
|
||||||
|
class Adapter:
|
||||||
|
async def handle_websocket_message(self, connection, data):
|
||||||
|
del data
|
||||||
|
captured_contexts.append(connection.execution_context)
|
||||||
|
await connection.send_queue.put({'type': 'handled'})
|
||||||
|
connection.is_active = False
|
||||||
|
|
||||||
|
app.pipeline_service = SimpleNamespace(get_pipeline=AsyncMock(return_value=SimpleNamespace(uuid='pipeline-1')))
|
||||||
|
app.platform_mgr = SimpleNamespace(
|
||||||
|
get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=Adapter()))
|
||||||
|
)
|
||||||
|
|
||||||
|
quart_app = Quart(__name__)
|
||||||
|
await WebSocketChatRouterGroup(app, quart_app).initialize()
|
||||||
|
|
||||||
|
async with quart_app.test_client().websocket('/api/v1/pipelines/pipeline-1/ws/connect') as websocket:
|
||||||
|
await websocket.send(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
'type': 'authenticate',
|
||||||
|
'token': launch['support_admin_token'],
|
||||||
|
'workspace_uuid': WORKSPACE_UUID,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
connected = json.loads(await websocket.receive())
|
||||||
|
assert connected['type'] == 'connected'
|
||||||
|
await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'hi'}]}))
|
||||||
|
handled = json.loads(await websocket.receive())
|
||||||
|
assert handled['type'] == 'handled'
|
||||||
|
|
||||||
|
assert captured_contexts
|
||||||
|
principal = captured_contexts[0].trigger_principal
|
||||||
|
assert principal is not None
|
||||||
|
assert principal.principal_type == PrincipalType.SUPPORT_ADMIN
|
||||||
|
assert principal.actor_account_uuid == ACTOR_ACCOUNT_UUID
|
||||||
|
|
||||||
|
|
||||||
|
async def _membership_count(engine) -> int:
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
return int(
|
||||||
|
await connection.scalar(
|
||||||
|
sqlalchemy.select(sqlalchemy.func.count()).select_from(WorkspaceMembership),
|
||||||
|
)
|
||||||
|
or 0
|
||||||
|
)
|
||||||
@@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import datetime
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from unittest.mock import AsyncMock, Mock
|
from unittest.mock import AsyncMock, Mock
|
||||||
from urllib.parse import parse_qs, urlsplit
|
from urllib.parse import parse_qs, urlsplit
|
||||||
@@ -14,6 +15,7 @@ from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
|
|||||||
|
|
||||||
pytestmark = pytest.mark.integration
|
pytestmark = pytest.mark.integration
|
||||||
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
|
WORKSPACE_UUID = '11111111-1111-4111-8111-111111111111'
|
||||||
|
WORKSPACE_CREATED_AT = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.UTC)
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
@@ -58,6 +60,12 @@ async def space_oauth_api():
|
|||||||
return_value={'account_uuid': 'account-a', 'workspace_uuid': WORKSPACE_UUID}
|
return_value={'account_uuid': 'account-a', 'workspace_uuid': WORKSPACE_UUID}
|
||||||
)
|
)
|
||||||
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access)
|
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(return_value=access)
|
||||||
|
application.workspace_service.get_execution_binding = AsyncMock(
|
||||||
|
return_value=SimpleNamespace(
|
||||||
|
workspace_uuid=WORKSPACE_UUID,
|
||||||
|
workspace_created_at=WORKSPACE_CREATED_AT,
|
||||||
|
)
|
||||||
|
)
|
||||||
application.space_service.get_oauth_authorize_url = Mock(
|
application.space_service.get_oauth_authorize_url = Mock(
|
||||||
side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}'
|
side_effect=lambda redirect_uri, state: f'https://space.example/authorize?state={state}'
|
||||||
)
|
)
|
||||||
@@ -157,34 +165,51 @@ async def test_bind_state_is_account_bound_and_requires_authentication(space_oau
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_redirect_origin_and_callback_path_are_restricted(space_oauth_api):
|
async def test_redirect_allows_any_http_or_https_origin(space_oauth_api):
|
||||||
_, client = space_oauth_api
|
_, client = space_oauth_api
|
||||||
|
|
||||||
wrong_origin = await client.get(
|
responses = [
|
||||||
'/api/v1/user/space/authorize-url',
|
await client.get(
|
||||||
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
|
'/api/v1/user/space/authorize-url',
|
||||||
headers={'Origin': 'http://localhost'},
|
query_string={'redirect_uri': redirect_uri},
|
||||||
)
|
headers={'Origin': 'https://irrelevant.example'},
|
||||||
wrong_path = await client.get(
|
)
|
||||||
'/api/v1/user/space/authorize-url',
|
for redirect_uri in (
|
||||||
query_string={'redirect_uri': 'http://localhost/arbitrary'},
|
'https://langbot.example/auth/space/callback',
|
||||||
headers={'Origin': 'http://localhost'},
|
'https://gateway.example:8443/auth/space/callback',
|
||||||
)
|
'https://192.0.2.10/auth/space/callback',
|
||||||
forged_origin = await client.get(
|
'http://localhost:5300/auth/space/callback',
|
||||||
'/api/v1/user/space/authorize-url',
|
'http://127.0.0.1:5300/auth/space/callback',
|
||||||
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
|
'http://[::1]:5300/auth/space/callback',
|
||||||
headers={'Origin': 'https://evil.example'},
|
'http://langbot.example/auth/space/callback',
|
||||||
)
|
'http://192.0.2.10:5300/auth/space/callback',
|
||||||
forged_host = await client.get(
|
)
|
||||||
'/api/v1/user/space/authorize-url',
|
]
|
||||||
query_string={'redirect_uri': 'https://evil.example/auth/space/callback'},
|
|
||||||
headers={'Host': 'evil.example'},
|
|
||||||
)
|
|
||||||
|
|
||||||
assert (await wrong_origin.get_json())['code'] == 1
|
assert all(response.status_code == 200 for response in responses)
|
||||||
assert (await wrong_path.get_json())['code'] == 1
|
payloads = [await response.get_json() for response in responses]
|
||||||
assert (await forged_origin.get_json())['code'] == 1
|
assert all(payload['code'] == 0 for payload in payloads)
|
||||||
assert (await forged_host.get_json())['code'] == 1
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_redirect_rejects_invalid_callback_shape(space_oauth_api):
|
||||||
|
_, client = space_oauth_api
|
||||||
|
|
||||||
|
responses = [
|
||||||
|
await client.get(
|
||||||
|
'/api/v1/user/space/authorize-url',
|
||||||
|
query_string={'redirect_uri': redirect_uri},
|
||||||
|
)
|
||||||
|
for redirect_uri in (
|
||||||
|
'https://langbot.example/arbitrary',
|
||||||
|
'https://langbot.example/auth/space/callback?next=https://evil.example',
|
||||||
|
'https://user@langbot.example/auth/space/callback',
|
||||||
|
'https://langbot.example/auth/space/callback#fragment',
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
payloads = [await response.get_json() for response in responses]
|
||||||
|
assert all(payload['code'] == 1 for payload in payloads)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -234,7 +259,11 @@ async def test_login_callback_requires_and_consumes_server_state(space_oauth_api
|
|||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert (await response.get_json())['data']['token'] == 'space-login-token'
|
assert (await response.get_json())['data']['token'] == 'space-login-token'
|
||||||
application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login')
|
application.user_service.consume_space_oauth_state_details.assert_awaited_once_with('opaque-login-state', 'login')
|
||||||
application.space_service.exchange_oauth_code.assert_awaited_once_with('oauth-code')
|
application.space_service.exchange_oauth_code.assert_awaited_once_with(
|
||||||
|
'oauth-code',
|
||||||
|
[WORKSPACE_UUID],
|
||||||
|
{WORKSPACE_UUID: int(WORKSPACE_CREATED_AT.timestamp())},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -270,11 +299,12 @@ async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api):
|
|||||||
|
|
||||||
response = await client.get(
|
response = await client.get(
|
||||||
'/api/v1/user/space-credits',
|
'/api/v1/user/space-credits',
|
||||||
headers={'Authorization': 'Bearer account-token', 'X-Workspace-UUID': WORKSPACE_UUID},
|
headers={'Authorization': 'Bearer account-token', 'X-Workspace-Id': WORKSPACE_UUID},
|
||||||
)
|
)
|
||||||
|
payload = await response.get_json()
|
||||||
|
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert (await response.get_json())['data'] == {
|
assert payload['data'] == {
|
||||||
'credits': 25000,
|
'credits': 25000,
|
||||||
'owner_space_bound': True,
|
'owner_space_bound': True,
|
||||||
'is_workspace_owner': True,
|
'is_workspace_owner': True,
|
||||||
@@ -282,6 +312,31 @@ async def test_space_credits_are_resolved_from_workspace_owner(space_oauth_api):
|
|||||||
application.space_service.get_credits.assert_awaited_once_with('owner@example.com')
|
application.space_service.get_credits.assert_awaited_once_with('owner@example.com')
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_cloud_workspace_owner_is_always_space_bound_after_login(space_oauth_api):
|
||||||
|
application, client = space_oauth_api
|
||||||
|
application.deployment.mode = 'cloud'
|
||||||
|
application.user_service.get_workspace_owner = AsyncMock(return_value=None)
|
||||||
|
application.space_service.get_credits = AsyncMock()
|
||||||
|
application.cloud_model_catalog_service = SimpleNamespace(
|
||||||
|
get_workspace_credits=lambda workspace_uuid: 25000 if workspace_uuid == WORKSPACE_UUID else None
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await client.get(
|
||||||
|
'/api/v1/user/space-credits',
|
||||||
|
headers={'Authorization': 'Bearer account-token', 'X-Workspace-Id': WORKSPACE_UUID},
|
||||||
|
)
|
||||||
|
payload = await response.get_json()
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert payload['data'] == {
|
||||||
|
'credits': 25000,
|
||||||
|
'owner_space_bound': True,
|
||||||
|
'is_workspace_owner': True,
|
||||||
|
}
|
||||||
|
application.space_service.get_credits.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_oauth_api):
|
async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_oauth_api):
|
||||||
application, client = space_oauth_api
|
application, client = space_oauth_api
|
||||||
|
|||||||
@@ -188,6 +188,7 @@ async def test_owner_invites_second_account_and_secret_is_not_persisted(workspac
|
|||||||
workspace_uuid = current['workspace']['uuid']
|
workspace_uuid = current['workspace']['uuid']
|
||||||
assert current['membership']['role'] == 'owner'
|
assert current['membership']['role'] == 'owner'
|
||||||
assert 'member.invite' in current['permissions']
|
assert 'member.invite' in current['permissions']
|
||||||
|
assert 'owner.transfer' not in current['permissions']
|
||||||
|
|
||||||
invite_response = await client.post(
|
invite_response = await client.post(
|
||||||
f'/api/v1/workspaces/{workspace_uuid}/invitations',
|
f'/api/v1/workspaces/{workspace_uuid}/invitations',
|
||||||
@@ -263,6 +264,14 @@ async def test_owner_invites_second_account_and_secret_is_not_persisted(workspac
|
|||||||
assert member_current['membership']['role'] == 'viewer'
|
assert member_current['membership']['role'] == 'viewer'
|
||||||
assert 'member.invite' not in member_current['permissions']
|
assert 'member.invite' not in member_current['permissions']
|
||||||
|
|
||||||
|
transfer_response = await client.patch(
|
||||||
|
f'/api/v1/workspaces/{workspace_uuid}/members/{member_current["membership"]["account_uuid"]}',
|
||||||
|
headers=_auth(owner_token, workspace_uuid),
|
||||||
|
json={'role': 'owner'},
|
||||||
|
)
|
||||||
|
assert transfer_response.status_code == 403
|
||||||
|
assert (await transfer_response.get_json())['code'] == 'permission_denied'
|
||||||
|
|
||||||
forbidden_invite = await client.post(
|
forbidden_invite = await client.post(
|
||||||
f'/api/v1/workspaces/{workspace_uuid}/invitations',
|
f'/api/v1/workspaces/{workspace_uuid}/invitations',
|
||||||
headers=_auth(member_token, workspace_uuid),
|
headers=_auth(member_token, workspace_uuid),
|
||||||
@@ -272,6 +281,31 @@ async def test_owner_invites_second_account_and_secret_is_not_persisted(workspac
|
|||||||
assert (await forbidden_invite.get_json())['code'] == 'permission_denied'
|
assert (await forbidden_invite.get_json())['code'] == 'permission_denied'
|
||||||
|
|
||||||
|
|
||||||
|
async def test_workspace_member_list_returns_display_name_and_email(workspace_api):
|
||||||
|
_, client, engine, owner_token = workspace_api
|
||||||
|
|
||||||
|
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
|
||||||
|
current = (await current_response.get_json())['data']
|
||||||
|
workspace_uuid = current['workspace']['uuid']
|
||||||
|
owner_uuid = current['membership']['account_uuid']
|
||||||
|
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.update(User).where(User.uuid == owner_uuid).values(user='Owner Display Name')
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await client.get(
|
||||||
|
f'/api/v1/workspaces/{workspace_uuid}/members',
|
||||||
|
headers=_auth(owner_token, workspace_uuid),
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
members = (await response.get_json())['data']['members']
|
||||||
|
assert len(members) == 1
|
||||||
|
assert members[0]['display_name'] == 'Owner Display Name'
|
||||||
|
assert members[0]['email'] == 'owner@example.com'
|
||||||
|
|
||||||
|
|
||||||
async def test_oss_invitation_accept_requires_logout_before_registration(workspace_api):
|
async def test_oss_invitation_accept_requires_logout_before_registration(workspace_api):
|
||||||
_, client, _, owner_token = workspace_api
|
_, client, _, owner_token = workspace_api
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,70 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
|
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_membership_source_migration_backfills_existing_rows_as_local_and_enforces_constraint(tmp_path):
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "membership-source.db"}')
|
||||||
|
try:
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sa.text(
|
||||||
|
"""
|
||||||
|
CREATE TABLE workspace_memberships (
|
||||||
|
uuid VARCHAR(36) PRIMARY KEY,
|
||||||
|
workspace_uuid VARCHAR(36) NOT NULL,
|
||||||
|
account_uuid VARCHAR(36) NOT NULL,
|
||||||
|
role VARCHAR(32) NOT NULL,
|
||||||
|
status VARCHAR(32) NOT NULL,
|
||||||
|
projection_revision BIGINT NOT NULL DEFAULT 0,
|
||||||
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await connection.execute(
|
||||||
|
sa.text(
|
||||||
|
"""
|
||||||
|
INSERT INTO workspace_memberships
|
||||||
|
(uuid, workspace_uuid, account_uuid, role, status, projection_revision)
|
||||||
|
VALUES
|
||||||
|
('00000000-0000-4000-8000-000000000001', 'workspace', 'local-account',
|
||||||
|
'viewer', 'active', 0),
|
||||||
|
('00000000-0000-4000-8000-000000000002', 'workspace', 'cloud-account',
|
||||||
|
'viewer', 'active', 0)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
await run_alembic_stamp(engine, '0019_single_workspace_owner')
|
||||||
|
await run_alembic_upgrade(engine, 'head')
|
||||||
|
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
rows = (
|
||||||
|
await connection.execute(sa.text('SELECT uuid, source FROM workspace_memberships ORDER BY uuid'))
|
||||||
|
).all()
|
||||||
|
columns = await connection.run_sync(
|
||||||
|
lambda sync_connection: {
|
||||||
|
column['name']: column
|
||||||
|
for column in sa.inspect(sync_connection).get_columns('workspace_memberships')
|
||||||
|
}
|
||||||
|
)
|
||||||
|
assert rows == [
|
||||||
|
('00000000-0000-4000-8000-000000000001', 'local'),
|
||||||
|
('00000000-0000-4000-8000-000000000002', 'local'),
|
||||||
|
]
|
||||||
|
assert columns['source']['nullable'] is False
|
||||||
|
|
||||||
|
with pytest.raises(sa.exc.IntegrityError):
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sa.text("UPDATE workspace_memberships SET source = 'guessed-from-user-source'")
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
@@ -9,8 +9,11 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
import sqlalchemy
|
import sqlalchemy
|
||||||
|
from sqlalchemy import text
|
||||||
from sqlalchemy.ext.asyncio import create_async_engine
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
from langbot.pkg.entity.persistence.base import Base
|
from langbot.pkg.entity.persistence.base import Base
|
||||||
@@ -95,6 +98,29 @@ class TestSQLiteMigrationBaseline:
|
|||||||
class TestSQLiteMigrationUpgrade:
|
class TestSQLiteMigrationUpgrade:
|
||||||
"""Tests for upgrade to head workflow."""
|
"""Tests for upgrade to head workflow."""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_upgrade_from_published_space_launch_head_to_merged_head(self, sqlite_engine):
|
||||||
|
"""A database released at the production-only 0016 head must remain upgradable."""
|
||||||
|
async with sqlite_engine.begin() as conn:
|
||||||
|
await conn.run_sync(Base.metadata.create_all)
|
||||||
|
|
||||||
|
await run_alembic_stamp(sqlite_engine, '0016_space_launch_replay')
|
||||||
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||||
|
|
||||||
|
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||||
|
assert _get_script_head() == '0021_merge_reasoning_config'
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_upgrade_from_reasoning_config_head_to_merged_head(self, sqlite_engine):
|
||||||
|
"""A database that already ran the feature migration must remain upgradable."""
|
||||||
|
async with sqlite_engine.begin() as conn:
|
||||||
|
await conn.run_sync(Base.metadata.create_all)
|
||||||
|
|
||||||
|
await run_alembic_stamp(sqlite_engine, '0018_llm_reasoning_config')
|
||||||
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||||
|
|
||||||
|
assert await get_alembic_current(sqlite_engine) == '0021_merge_reasoning_config'
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
|
async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
|
||||||
"""
|
"""
|
||||||
@@ -190,6 +216,66 @@ class TestSQLiteMigrationUpgrade:
|
|||||||
await run_alembic_upgrade(sqlite_engine, 'head')
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||||
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_reasoning_config_migrates_existing_models(self, sqlite_engine):
|
||||||
|
"""Upgrade from 0017 backfills reasoning config and keeps a database default."""
|
||||||
|
async with sqlite_engine.begin() as conn:
|
||||||
|
await conn.execute(
|
||||||
|
text(
|
||||||
|
"""
|
||||||
|
CREATE TABLE llm_models (
|
||||||
|
uuid VARCHAR(255) PRIMARY KEY,
|
||||||
|
name VARCHAR(255) NOT NULL,
|
||||||
|
provider_uuid VARCHAR(255) NOT NULL,
|
||||||
|
abilities JSON NOT NULL,
|
||||||
|
context_length INTEGER,
|
||||||
|
extra_args JSON NOT NULL,
|
||||||
|
prefered_ranking INTEGER NOT NULL DEFAULT 0
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await conn.execute(
|
||||||
|
text(
|
||||||
|
"""
|
||||||
|
INSERT INTO llm_models (
|
||||||
|
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
||||||
|
) VALUES (
|
||||||
|
'existing-model', 'Existing Model', 'provider', '[]', '{}', 0
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
await run_alembic_stamp(sqlite_engine, '0017_oss_workspace_identity')
|
||||||
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
||||||
|
|
||||||
|
async with sqlite_engine.begin() as conn:
|
||||||
|
columns = await conn.run_sync(lambda sync_conn: sqlalchemy.inspect(sync_conn).get_columns('llm_models'))
|
||||||
|
reasoning_column = next(column for column in columns if column['name'] == 'reasoning_config')
|
||||||
|
assert reasoning_column['nullable'] is False
|
||||||
|
|
||||||
|
existing_value = (
|
||||||
|
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'existing-model'"))
|
||||||
|
).scalar_one()
|
||||||
|
assert json.loads(existing_value) == {'level': 'provider_default'}
|
||||||
|
|
||||||
|
await conn.execute(
|
||||||
|
text(
|
||||||
|
"""
|
||||||
|
INSERT INTO llm_models (
|
||||||
|
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
||||||
|
) VALUES (
|
||||||
|
'new-model', 'New Model', 'provider', '[]', '{}', 0
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
)
|
||||||
|
new_value = (
|
||||||
|
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'new-model'"))
|
||||||
|
).scalar_one()
|
||||||
|
assert json.loads(new_value) == {'level': 'provider_default'}
|
||||||
|
|
||||||
|
|
||||||
class TestSQLiteMigrationFreshDatabase:
|
class TestSQLiteMigrationFreshDatabase:
|
||||||
"""Tests for fresh database workflow."""
|
"""Tests for fresh database workflow."""
|
||||||
|
|||||||
@@ -85,6 +85,32 @@ async def clean_database(postgres_engine: AsyncEngine):
|
|||||||
await clean()
|
await clean()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_upgrade_adds_3072_dimension_index_and_constraint(
|
||||||
|
postgres_engine: AsyncEngine,
|
||||||
|
clean_database,
|
||||||
|
) -> None:
|
||||||
|
async with postgres_engine.begin() as conn:
|
||||||
|
await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
|
||||||
|
await conn.run_sync(Base.metadata.create_all)
|
||||||
|
await run_alembic_stamp(postgres_engine, '0010_scope_resources')
|
||||||
|
await run_alembic_upgrade(postgres_engine, 'head')
|
||||||
|
|
||||||
|
async with postgres_engine.connect() as conn:
|
||||||
|
constraint = await conn.scalar(
|
||||||
|
text(
|
||||||
|
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
|
||||||
|
"WHERE conrelid = 'langbot_vectors'::regclass "
|
||||||
|
"AND conname = 'ck_langbot_vectors_embedding_dimension_enabled'"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
assert '3072' in constraint
|
||||||
|
index_definition = await conn.scalar(
|
||||||
|
text("SELECT indexdef FROM pg_indexes WHERE indexname = 'ix_langbot_vectors_hnsw_cosine_3072'")
|
||||||
|
)
|
||||||
|
assert 'halfvec(3072)' in index_definition
|
||||||
|
assert 'halfvec_cosine_ops' in index_definition
|
||||||
|
|
||||||
|
|
||||||
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
|
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
|
||||||
postgres_url: str,
|
postgres_url: str,
|
||||||
postgres_engine: AsyncEngine,
|
postgres_engine: AsyncEngine,
|
||||||
|
|||||||
@@ -92,7 +92,7 @@ def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_
|
|||||||
'use': 'pgvector',
|
'use': 'pgvector',
|
||||||
'pgvector': {
|
'pgvector': {
|
||||||
'use_business_database': True,
|
'use_business_database': True,
|
||||||
'allowed_dimensions': [384, 512, 768, 1024, 1536],
|
'allowed_dimensions': [384, 512, 768, 1024, 1536, 3072],
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,100 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import sqlalchemy as sa
|
||||||
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||||
|
|
||||||
|
from langbot.pkg.entity.persistence.base import Base
|
||||||
|
from langbot.pkg.entity.persistence.user import User
|
||||||
|
from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceMembership
|
||||||
|
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_single_owner_migration_demotes_historical_extra_owner_and_installs_unique_index(tmp_path):
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "single-owner.db"}')
|
||||||
|
try:
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.run_sync(Base.metadata.create_all)
|
||||||
|
await connection.execute(sa.text('DROP INDEX uq_workspace_memberships_one_active_owner'))
|
||||||
|
|
||||||
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
||||||
|
workspace_uuid = '00000000-0000-4000-8000-000000000001'
|
||||||
|
creator_uuid = '00000000-0000-4000-8000-000000000010'
|
||||||
|
promoted_uuid = '00000000-0000-4000-8000-000000000020'
|
||||||
|
async with session_factory() as session:
|
||||||
|
session.add_all(
|
||||||
|
[
|
||||||
|
User(
|
||||||
|
uuid=creator_uuid,
|
||||||
|
user='creator@example.test',
|
||||||
|
normalized_email='creator@example.test',
|
||||||
|
password='hash',
|
||||||
|
account_type='local',
|
||||||
|
),
|
||||||
|
User(
|
||||||
|
uuid=promoted_uuid,
|
||||||
|
user='promoted@example.test',
|
||||||
|
normalized_email='promoted@example.test',
|
||||||
|
password='hash',
|
||||||
|
account_type='local',
|
||||||
|
),
|
||||||
|
Workspace(
|
||||||
|
uuid=workspace_uuid,
|
||||||
|
instance_uuid='instance-test',
|
||||||
|
name='Workspace',
|
||||||
|
slug='workspace',
|
||||||
|
type='team',
|
||||||
|
status='active',
|
||||||
|
source='local',
|
||||||
|
created_by_account_uuid=creator_uuid,
|
||||||
|
),
|
||||||
|
WorkspaceMembership(
|
||||||
|
uuid='00000000-0000-4000-8000-000000000100',
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
account_uuid=creator_uuid,
|
||||||
|
role='owner',
|
||||||
|
status='active',
|
||||||
|
),
|
||||||
|
WorkspaceMembership(
|
||||||
|
uuid='00000000-0000-4000-8000-000000000200',
|
||||||
|
workspace_uuid=workspace_uuid,
|
||||||
|
account_uuid=promoted_uuid,
|
||||||
|
role='owner',
|
||||||
|
status='active',
|
||||||
|
),
|
||||||
|
]
|
||||||
|
)
|
||||||
|
await session.commit()
|
||||||
|
|
||||||
|
await run_alembic_stamp(engine, '0018_merge_launch_replay')
|
||||||
|
await run_alembic_upgrade(engine, 'head')
|
||||||
|
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
roles = dict(
|
||||||
|
(
|
||||||
|
await connection.execute(
|
||||||
|
sa.text(
|
||||||
|
'SELECT account_uuid, role FROM workspace_memberships '
|
||||||
|
'WHERE workspace_uuid = :workspace_uuid ORDER BY account_uuid'
|
||||||
|
),
|
||||||
|
{'workspace_uuid': workspace_uuid},
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
)
|
||||||
|
assert roles == {creator_uuid: 'owner', promoted_uuid: 'admin'}
|
||||||
|
indexes = await connection.run_sync(
|
||||||
|
lambda sync_connection: {
|
||||||
|
index['name'] for index in sa.inspect(sync_connection).get_indexes('workspace_memberships')
|
||||||
|
}
|
||||||
|
)
|
||||||
|
assert 'uq_workspace_memberships_one_active_owner' in indexes
|
||||||
|
|
||||||
|
with pytest.raises(sa.exc.IntegrityError):
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.execute(
|
||||||
|
sa.text("UPDATE workspace_memberships SET role = 'owner' WHERE account_uuid = :account_uuid"),
|
||||||
|
{'account_uuid': promoted_uuid},
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
@@ -22,6 +22,7 @@ from langbot.pkg.persistence.alembic_runner import (
|
|||||||
from langbot.pkg.utils import constants
|
from langbot.pkg.utils import constants
|
||||||
from langbot.pkg.utils import importutil
|
from langbot.pkg.utils import importutil
|
||||||
from langbot.pkg.workspace.collaboration import normalize_email
|
from langbot.pkg.workspace.collaboration import normalize_email
|
||||||
|
from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id
|
||||||
|
|
||||||
|
|
||||||
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||||
@@ -109,6 +110,7 @@ async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy
|
|||||||
.mappings()
|
.mappings()
|
||||||
.one()
|
.one()
|
||||||
)
|
)
|
||||||
|
assert workspace['uuid'] == workspace_uuid_from_instance_id('instance_migration_test')
|
||||||
assert workspace['instance_uuid'] == 'instance_migration_test'
|
assert workspace['instance_uuid'] == 'instance_migration_test'
|
||||||
assert workspace['slug'] == 'default'
|
assert workspace['slug'] == 'default'
|
||||||
assert workspace['status'] == 'active'
|
assert workspace['status'] == 'active'
|
||||||
@@ -149,6 +151,53 @@ async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_
|
|||||||
assert workspace_uuid_after == workspace_uuid_before
|
assert workspace_uuid_after == workspace_uuid_before
|
||||||
|
|
||||||
|
|
||||||
|
async def test_existing_oss_workspace_is_rekeyed_to_instance_identity(tmp_path):
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-rekey.db"}')
|
||||||
|
instance_id = 'instance_a711d9e4-0953-443f-a0e9-7dd50193a79f'
|
||||||
|
old_workspace_uuid = '11111111-1111-4111-8111-111111111111'
|
||||||
|
canonical_uuid = workspace_uuid_from_instance_id(instance_id)
|
||||||
|
schema = sa.MetaData()
|
||||||
|
sa.Table(
|
||||||
|
'metadata',
|
||||||
|
schema,
|
||||||
|
sa.Column('key', sa.String(255), primary_key=True),
|
||||||
|
sa.Column('value', sa.String(255)),
|
||||||
|
)
|
||||||
|
sa.Table(
|
||||||
|
'workspaces',
|
||||||
|
schema,
|
||||||
|
sa.Column('uuid', sa.String(36), primary_key=True),
|
||||||
|
sa.Column('instance_uuid', sa.String(255), nullable=False),
|
||||||
|
sa.Column('slug', sa.String(255), nullable=False),
|
||||||
|
sa.Column('source', sa.String(32), nullable=False),
|
||||||
|
)
|
||||||
|
sa.Table(
|
||||||
|
'tenant_rows',
|
||||||
|
schema,
|
||||||
|
sa.Column('id', sa.Integer, primary_key=True),
|
||||||
|
sa.Column('workspace_uuid', sa.String(36), sa.ForeignKey('workspaces.uuid'), nullable=False),
|
||||||
|
)
|
||||||
|
async with engine.begin() as conn:
|
||||||
|
await conn.run_sync(schema.create_all)
|
||||||
|
await conn.execute(sa.text("INSERT INTO metadata (key, value) VALUES ('instance_uuid', :value)"), {'value': instance_id})
|
||||||
|
await conn.execute(
|
||||||
|
sa.text("INSERT INTO workspaces (uuid, instance_uuid, slug, source) VALUES (:uuid, :instance, 'default', 'local')"),
|
||||||
|
{'uuid': old_workspace_uuid, 'instance': instance_id},
|
||||||
|
)
|
||||||
|
await conn.execute(
|
||||||
|
sa.text("INSERT INTO tenant_rows (id, workspace_uuid) VALUES (1, :uuid)"),
|
||||||
|
{'uuid': old_workspace_uuid},
|
||||||
|
)
|
||||||
|
await run_alembic_stamp(engine, '0016_support_admin_sessions')
|
||||||
|
|
||||||
|
await run_alembic_upgrade(engine, 'head')
|
||||||
|
|
||||||
|
async with engine.connect() as conn:
|
||||||
|
assert (await conn.execute(sa.text("SELECT uuid FROM workspaces"))).scalar_one() == canonical_uuid
|
||||||
|
assert (await conn.execute(sa.text("SELECT workspace_uuid FROM tenant_rows"))).scalar_one() == canonical_uuid
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
|
async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
|
||||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
|
||||||
try:
|
try:
|
||||||
@@ -362,6 +411,47 @@ async def test_persistence_startup_defers_workspace_tables_until_account_upgrade
|
|||||||
await engine.dispose()
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path):
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-rekey.db"}')
|
||||||
|
try:
|
||||||
|
await _create_legacy_schema(engine)
|
||||||
|
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||||
|
await run_alembic_upgrade(engine, '0016_support_admin_sessions')
|
||||||
|
|
||||||
|
async with engine.begin() as conn:
|
||||||
|
old_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||||
|
instance_uuid = await conn.scalar(sa.text("SELECT instance_uuid FROM workspaces WHERE source = 'local'"))
|
||||||
|
assert old_uuid
|
||||||
|
assert instance_uuid
|
||||||
|
await conn.execute(
|
||||||
|
sa.text(
|
||||||
|
"INSERT INTO workspace_metadata (workspace_uuid, key, value) "
|
||||||
|
"VALUES (:workspace_uuid, 'migration_probe', 'present')"
|
||||||
|
),
|
||||||
|
{'workspace_uuid': old_uuid},
|
||||||
|
)
|
||||||
|
await conn.execute(
|
||||||
|
sa.text(
|
||||||
|
"INSERT INTO metadata (key, value) VALUES ('oss_workspace_uuid', :workspace_uuid) "
|
||||||
|
"ON CONFLICT(key) DO UPDATE SET value = excluded.value"
|
||||||
|
),
|
||||||
|
{'workspace_uuid': old_uuid},
|
||||||
|
)
|
||||||
|
|
||||||
|
await run_alembic_upgrade(engine, 'head')
|
||||||
|
expected_uuid = workspace_uuid_from_instance_id(instance_uuid)
|
||||||
|
async with engine.connect() as conn:
|
||||||
|
assert await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'")) == expected_uuid
|
||||||
|
assert await conn.scalar(
|
||||||
|
sa.text("SELECT workspace_uuid FROM workspace_metadata WHERE key = 'migration_probe'")
|
||||||
|
) == expected_uuid
|
||||||
|
assert await conn.scalar(
|
||||||
|
sa.text("SELECT value FROM metadata WHERE key = 'oss_workspace_uuid'")
|
||||||
|
) == expected_uuid
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch):
|
async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch):
|
||||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}')
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}')
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -1,137 +0,0 @@
|
|||||||
"""PostgreSQL integration coverage for durable workspace quota locking.
|
|
||||||
|
|
||||||
Run with TEST_POSTGRES_URL=postgresql+asyncpg://... pytest ...
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import os
|
|
||||||
import uuid
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
import sqlalchemy as sa
|
|
||||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
||||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
||||||
|
|
||||||
from langbot.pkg.cloud import quotas as quota_module
|
|
||||||
from langbot.pkg.cloud.quotas import WorkspaceQuota, WorkspaceQuotaExceededError, require_resource_capacity
|
|
||||||
|
|
||||||
|
|
||||||
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
|
|
||||||
|
|
||||||
|
|
||||||
class _Base(DeclarativeBase):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class _Workspace(_Base):
|
|
||||||
__tablename__ = 'quota_integration_workspaces'
|
|
||||||
|
|
||||||
uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True)
|
|
||||||
|
|
||||||
|
|
||||||
class _Resource(_Base):
|
|
||||||
__tablename__ = 'quota_integration_resources'
|
|
||||||
|
|
||||||
uuid: Mapped[str] = mapped_column(sa.String(36), primary_key=True)
|
|
||||||
workspace_uuid: Mapped[str] = mapped_column(
|
|
||||||
sa.String(36),
|
|
||||||
sa.ForeignKey('quota_integration_workspaces.uuid', ondelete='CASCADE'),
|
|
||||||
nullable=False,
|
|
||||||
index=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
async def quota_postgres(monkeypatch):
|
|
||||||
url = os.environ.get('TEST_POSTGRES_URL')
|
|
||||||
if not url:
|
|
||||||
pytest.skip('TEST_POSTGRES_URL not set')
|
|
||||||
if url.startswith('postgresql://'):
|
|
||||||
url = url.replace('postgresql://', 'postgresql+asyncpg://', 1)
|
|
||||||
|
|
||||||
engine = create_async_engine(url, pool_size=5, max_overflow=0)
|
|
||||||
monkeypatch.setattr(quota_module.persistence_workspace, 'Workspace', _Workspace)
|
|
||||||
async with engine.begin() as connection:
|
|
||||||
await connection.run_sync(_Base.metadata.drop_all)
|
|
||||||
await connection.run_sync(_Base.metadata.create_all)
|
|
||||||
try:
|
|
||||||
yield url, engine
|
|
||||||
finally:
|
|
||||||
async with engine.begin() as connection:
|
|
||||||
await connection.run_sync(_Base.metadata.drop_all)
|
|
||||||
await engine.dispose()
|
|
||||||
|
|
||||||
|
|
||||||
async def test_workspace_row_lock_is_atomic_isolated_and_survives_pool_restart(quota_postgres) -> None:
|
|
||||||
url, engine = quota_postgres
|
|
||||||
workspace_a = str(uuid.uuid4())
|
|
||||||
workspace_b = str(uuid.uuid4())
|
|
||||||
quota = WorkspaceQuota(limit=1, requires_transaction_lock=True)
|
|
||||||
sessions = async_sessionmaker(engine, expire_on_commit=False)
|
|
||||||
|
|
||||||
async with engine.begin() as connection:
|
|
||||||
await connection.execute(sa.insert(_Workspace), [{'uuid': workspace_a}, {'uuid': workspace_b}])
|
|
||||||
|
|
||||||
lock_acquired = asyncio.Event()
|
|
||||||
release_first = asyncio.Event()
|
|
||||||
|
|
||||||
async def admit(workspace_uuid: str, *, hold: bool = False) -> None:
|
|
||||||
async with sessions() as session:
|
|
||||||
async with session.begin():
|
|
||||||
await require_resource_capacity(
|
|
||||||
session.execute,
|
|
||||||
workspace_uuid=workspace_uuid,
|
|
||||||
model=_Resource,
|
|
||||||
quota=quota,
|
|
||||||
resource_name='resources',
|
|
||||||
)
|
|
||||||
if hold:
|
|
||||||
lock_acquired.set()
|
|
||||||
await release_first.wait()
|
|
||||||
await session.execute(
|
|
||||||
sa.insert(_Resource).values(uuid=str(uuid.uuid4()), workspace_uuid=workspace_uuid)
|
|
||||||
)
|
|
||||||
|
|
||||||
first = asyncio.create_task(admit(workspace_a, hold=True))
|
|
||||||
await asyncio.wait_for(lock_acquired.wait(), timeout=2)
|
|
||||||
same_workspace = asyncio.create_task(admit(workspace_a))
|
|
||||||
other_workspace = asyncio.create_task(admit(workspace_b))
|
|
||||||
|
|
||||||
await asyncio.wait_for(other_workspace, timeout=2)
|
|
||||||
assert not same_workspace.done(), 'same-workspace transaction bypassed SELECT FOR UPDATE'
|
|
||||||
|
|
||||||
release_first.set()
|
|
||||||
await first
|
|
||||||
with pytest.raises(WorkspaceQuotaExceededError, match=r'Maximum number of resources \(1\) reached'):
|
|
||||||
await same_workspace
|
|
||||||
|
|
||||||
async with sessions() as session:
|
|
||||||
counts = dict(
|
|
||||||
(
|
|
||||||
await session.execute(
|
|
||||||
sa.select(_Resource.workspace_uuid, sa.func.count())
|
|
||||||
.group_by(_Resource.workspace_uuid)
|
|
||||||
.order_by(_Resource.workspace_uuid)
|
|
||||||
)
|
|
||||||
).all()
|
|
||||||
)
|
|
||||||
assert counts == {workspace_a: 1, workspace_b: 1}
|
|
||||||
|
|
||||||
await engine.dispose()
|
|
||||||
restarted_engine = create_async_engine(url, pool_size=2, max_overflow=0)
|
|
||||||
restarted_sessions = async_sessionmaker(restarted_engine, expire_on_commit=False)
|
|
||||||
try:
|
|
||||||
async with restarted_sessions() as session:
|
|
||||||
async with session.begin():
|
|
||||||
with pytest.raises(WorkspaceQuotaExceededError):
|
|
||||||
await require_resource_capacity(
|
|
||||||
session.execute,
|
|
||||||
workspace_uuid=workspace_a,
|
|
||||||
model=_Resource,
|
|
||||||
quota=quota,
|
|
||||||
resource_name='resources',
|
|
||||||
)
|
|
||||||
finally:
|
|
||||||
await restarted_engine.dispose()
|
|
||||||
@@ -27,10 +27,9 @@ def test_owner_has_every_fixed_permission():
|
|||||||
assert ctx.workspace.permissions == frozenset(permission.value for permission in authz.Permission)
|
assert ctx.workspace.permissions == frozenset(permission.value for permission in authz.Permission)
|
||||||
|
|
||||||
|
|
||||||
def test_admin_cannot_transfer_owner_delete_workspace_or_link_billing():
|
def test_admin_cannot_delete_workspace_or_link_billing():
|
||||||
ctx = _context(authz.WorkspaceRole.ADMIN)
|
ctx = _context(authz.WorkspaceRole.ADMIN)
|
||||||
|
|
||||||
assert not authz.has_permission(ctx, authz.Permission.OWNER_TRANSFER)
|
|
||||||
assert not authz.has_permission(ctx, authz.Permission.WORKSPACE_DELETE)
|
assert not authz.has_permission(ctx, authz.Permission.WORKSPACE_DELETE)
|
||||||
assert not authz.has_permission(ctx, authz.Permission.BILLING_LINK_MANAGE)
|
assert not authz.has_permission(ctx, authz.Permission.BILLING_LINK_MANAGE)
|
||||||
assert authz.has_permission(ctx, authz.Permission.MEMBER_INVITE)
|
assert authz.has_permission(ctx, authz.Permission.MEMBER_INVITE)
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ import quart
|
|||||||
|
|
||||||
from langbot.pkg.api.http.controller import group
|
from langbot.pkg.api.http.controller import group
|
||||||
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
|
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
|
||||||
from langbot.pkg.cloud.quotas import WorkspaceQuotaExceededError
|
|
||||||
from langbot.pkg.utils.bounded_executor import (
|
from langbot.pkg.utils.bounded_executor import (
|
||||||
BlockingWorkCapacityError,
|
BlockingWorkCapacityError,
|
||||||
current_blocking_work_scope,
|
current_blocking_work_scope,
|
||||||
@@ -49,16 +48,6 @@ class _BlockingCapacityRouterGroup(group.RouterGroup):
|
|||||||
raise BlockingWorkCapacityError('Workspace blocking executor capacity reached')
|
raise BlockingWorkCapacityError('Workspace blocking executor capacity reached')
|
||||||
|
|
||||||
|
|
||||||
class _QuotaRouterGroup(group.RouterGroup):
|
|
||||||
name = 'quota-test'
|
|
||||||
path = '/quota-test'
|
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
|
||||||
@self.route('', methods=['POST'], auth_type=group.AuthType.NONE)
|
|
||||||
async def _():
|
|
||||||
raise WorkspaceQuotaExceededError('bots', 2)
|
|
||||||
|
|
||||||
|
|
||||||
class _InvalidAccountRouterGroup(group.RouterGroup):
|
class _InvalidAccountRouterGroup(group.RouterGroup):
|
||||||
name = 'invalid-account-test'
|
name = 'invalid-account-test'
|
||||||
path = '/invalid-account-test'
|
path = '/invalid-account-test'
|
||||||
@@ -141,20 +130,6 @@ async def test_blocking_work_capacity_maps_to_retryable_http_response():
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
async def test_workspace_quota_maps_to_stable_conflict_response():
|
|
||||||
application = SimpleNamespace(logger=Mock())
|
|
||||||
quart_app = quart.Quart(__name__)
|
|
||||||
await _QuotaRouterGroup(application, quart_app).initialize()
|
|
||||||
|
|
||||||
response = await quart_app.test_client().post('/quota-test')
|
|
||||||
|
|
||||||
assert response.status_code == 409
|
|
||||||
assert await response.get_json() == {
|
|
||||||
'code': 'workspace_quota_exceeded',
|
|
||||||
'msg': 'Maximum number of bots (2) reached',
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
async def test_public_webhook_carries_scope_without_holding_database_session():
|
async def test_public_webhook_carries_scope_without_holding_database_session():
|
||||||
class ScopeOnlyPersistenceManager:
|
class ScopeOnlyPersistenceManager:
|
||||||
mode = SimpleNamespace(value='cloud_runtime')
|
mode = SimpleNamespace(value='cloud_runtime')
|
||||||
|
|||||||
@@ -311,9 +311,10 @@ class TestBotServiceCreateBot:
|
|||||||
ap.platform_mgr = SimpleNamespace()
|
ap.platform_mgr = SimpleNamespace()
|
||||||
ap.platform_mgr.load_bot = AsyncMock()
|
ap.platform_mgr.load_bot = AsyncMock()
|
||||||
|
|
||||||
# Mock the atomic count query to report 2 existing bots.
|
# Mock get_bots to return 2 bots already
|
||||||
mock_result = _create_mock_result()
|
bot1 = _create_mock_bot(bot_uuid='uuid-1')
|
||||||
mock_result.scalar_one = Mock(return_value=2)
|
bot2 = _create_mock_bot(bot_uuid='uuid-2')
|
||||||
|
mock_result = _create_mock_result([bot1, bot2])
|
||||||
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
|
ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
|
||||||
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'})
|
ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'uuid-1', 'name': 'Bot 1'})
|
||||||
|
|
||||||
|
|||||||
@@ -1,129 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from collections import defaultdict
|
|
||||||
from types import SimpleNamespace
|
|
||||||
from unittest.mock import AsyncMock
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
import sqlalchemy
|
|
||||||
|
|
||||||
from langbot.pkg.api.http.service.bot import BotService
|
|
||||||
from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot
|
|
||||||
|
|
||||||
|
|
||||||
INSTANCE_UUID = 'cloud-instance'
|
|
||||||
WORKSPACE_A = '11111111-1111-1111-1111-111111111111'
|
|
||||||
WORKSPACE_B = '22222222-2222-2222-2222-222222222222'
|
|
||||||
|
|
||||||
|
|
||||||
class _Provider:
|
|
||||||
async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
|
|
||||||
return EntitlementSnapshot(
|
|
||||||
instance_uuid=INSTANCE_UUID,
|
|
||||||
workspace_uuid=workspace_uuid,
|
|
||||||
entitlement_revision=1,
|
|
||||||
status='active',
|
|
||||||
not_before=0,
|
|
||||||
expires_at=4_102_444_800,
|
|
||||||
features={},
|
|
||||||
limits={'bots.max': 2},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class _Result:
|
|
||||||
def __init__(self, *, first=None, scalar=None) -> None:
|
|
||||||
self._first = first
|
|
||||||
self._scalar = scalar
|
|
||||||
|
|
||||||
def first(self):
|
|
||||||
return self._first
|
|
||||||
|
|
||||||
def scalar_one(self):
|
|
||||||
return self._scalar
|
|
||||||
|
|
||||||
|
|
||||||
class _TenantUow:
|
|
||||||
def __init__(self, manager: '_Persistence', workspace_uuid: str) -> None:
|
|
||||||
self.manager = manager
|
|
||||||
self.workspace_uuid = workspace_uuid
|
|
||||||
self.lock = manager.locks[workspace_uuid]
|
|
||||||
|
|
||||||
async def __aenter__(self):
|
|
||||||
await self.lock.acquire()
|
|
||||||
return self
|
|
||||||
|
|
||||||
async def __aexit__(self, exc_type, exc, tb):
|
|
||||||
self.lock.release()
|
|
||||||
|
|
||||||
async def execute(self, statement):
|
|
||||||
sql = str(statement)
|
|
||||||
if isinstance(statement, sqlalchemy.sql.dml.Insert):
|
|
||||||
assert statement.table.name == 'bots'
|
|
||||||
self.manager.bots[self.workspace_uuid].append(statement.compile().params)
|
|
||||||
return _Result()
|
|
||||||
if 'FROM workspaces' in sql:
|
|
||||||
assert statement._for_update_arg is not None
|
|
||||||
self.manager.workspace_locks_seen += 1
|
|
||||||
return _Result(first=(self.workspace_uuid,))
|
|
||||||
if 'count(' in sql.lower() and 'FROM bots' in sql:
|
|
||||||
return _Result(scalar=len(self.manager.bots[self.workspace_uuid]))
|
|
||||||
if 'FROM legacy_pipelines' in sql:
|
|
||||||
return _Result(first=None)
|
|
||||||
raise AssertionError(f'unexpected statement: {sql}')
|
|
||||||
|
|
||||||
|
|
||||||
class _Persistence:
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self.locks = defaultdict(asyncio.Lock)
|
|
||||||
self.bots = defaultdict(list)
|
|
||||||
self.workspace_locks_seen = 0
|
|
||||||
|
|
||||||
def tenant_uow(self, workspace_uuid: str) -> _TenantUow:
|
|
||||||
return _TenantUow(self, workspace_uuid)
|
|
||||||
|
|
||||||
async def execute_async(self, statement):
|
|
||||||
assert 'FROM legacy_pipelines' in str(statement)
|
|
||||||
return _Result(first=None)
|
|
||||||
|
|
||||||
|
|
||||||
async def _service(manager: _Persistence) -> BotService:
|
|
||||||
resolver = EntitlementResolver(INSTANCE_UUID, _Provider())
|
|
||||||
await resolver.reconcile_active_workspaces({WORKSPACE_A, WORKSPACE_B})
|
|
||||||
ap = SimpleNamespace(
|
|
||||||
entitlement_resolver=resolver,
|
|
||||||
persistence_mgr=manager,
|
|
||||||
instance_config=SimpleNamespace(data={'system': {'limitation': {'max_bots': 99}}}),
|
|
||||||
platform_mgr=SimpleNamespace(load_bot=AsyncMock()),
|
|
||||||
)
|
|
||||||
service = BotService(ap)
|
|
||||||
service.get_bot = AsyncMock(return_value={'uuid': 'created'})
|
|
||||||
return service
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_cloud_bot_quota_is_atomic_isolated_and_persists_across_service_restart() -> None:
|
|
||||||
manager = _Persistence()
|
|
||||||
service = await _service(manager)
|
|
||||||
|
|
||||||
async def create(workspace_uuid: str, index: int):
|
|
||||||
return await service.create_bot(workspace_uuid, {'name': f'bot-{index}'})
|
|
||||||
|
|
||||||
results = await asyncio.gather(
|
|
||||||
*(create(WORKSPACE_A, index) for index in range(8)),
|
|
||||||
*(create(WORKSPACE_B, index) for index in range(8)),
|
|
||||||
return_exceptions=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
successes = [result for result in results if isinstance(result, str)]
|
|
||||||
failures = [result for result in results if isinstance(result, ValueError)]
|
|
||||||
assert len(successes) == 4
|
|
||||||
assert len(failures) == 12
|
|
||||||
assert len(manager.bots[WORKSPACE_A]) == 2
|
|
||||||
assert len(manager.bots[WORKSPACE_B]) == 2
|
|
||||||
assert manager.workspace_locks_seen == 16
|
|
||||||
|
|
||||||
restarted_service = await _service(manager)
|
|
||||||
with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'):
|
|
||||||
await restarted_service.create_bot(WORKSPACE_A, {'name': 'after-restart'})
|
|
||||||
assert len(manager.bots[WORKSPACE_A]) == 2
|
|
||||||
@@ -0,0 +1,112 @@
|
|||||||
|
"""Cloud Runtime write protection for the managed LangBot Models catalog."""
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from langbot.pkg.api.http.service import model as model_service_module
|
||||||
|
from langbot.pkg.api.http.service.model import (
|
||||||
|
EmbeddingModelsService,
|
||||||
|
LLMModelsService,
|
||||||
|
RerankModelsService,
|
||||||
|
_assert_cloud_managed_provider_mutable,
|
||||||
|
)
|
||||||
|
from langbot.pkg.cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER
|
||||||
|
|
||||||
|
|
||||||
|
WORKSPACE = 'workspace-a'
|
||||||
|
PROVIDER = 'managed-provider'
|
||||||
|
MODEL = 'managed-model'
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_managed_provider_guard_is_cloud_only(monkeypatch) -> None:
|
||||||
|
async def managed_provider(_ap, _context, provider_uuid):
|
||||||
|
assert provider_uuid == PROVIDER
|
||||||
|
return {'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER}
|
||||||
|
|
||||||
|
monkeypatch.setattr(model_service_module, '_require_workspace_provider', managed_provider)
|
||||||
|
application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime')))
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match='managed by Cloud'):
|
||||||
|
await _assert_cloud_managed_provider_mutable(
|
||||||
|
application,
|
||||||
|
WORKSPACE,
|
||||||
|
PROVIDER,
|
||||||
|
)
|
||||||
|
|
||||||
|
application.persistence_mgr.mode.value = 'normal'
|
||||||
|
await _assert_cloud_managed_provider_mutable(
|
||||||
|
application,
|
||||||
|
WORKSPACE,
|
||||||
|
PROVIDER,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
('service_type', 'create_method', 'model_data'),
|
||||||
|
[
|
||||||
|
(LLMModelsService, 'create_llm_model', {'provider_uuid': PROVIDER, 'name': 'chat', 'abilities': []}),
|
||||||
|
(EmbeddingModelsService, 'create_embedding_model', {'provider_uuid': PROVIDER, 'name': 'embedding'}),
|
||||||
|
(RerankModelsService, 'create_rerank_model', {'provider_uuid': PROVIDER, 'name': 'rerank'}),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_all_model_types_reject_creation_under_managed_provider(
|
||||||
|
monkeypatch,
|
||||||
|
service_type,
|
||||||
|
create_method: str,
|
||||||
|
model_data: dict,
|
||||||
|
) -> None:
|
||||||
|
guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified'))
|
||||||
|
monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard)
|
||||||
|
application = SimpleNamespace(
|
||||||
|
persistence_mgr=SimpleNamespace(),
|
||||||
|
provider_service=SimpleNamespace(
|
||||||
|
get_provider=AsyncMock(return_value={'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER})
|
||||||
|
),
|
||||||
|
model_mgr=None,
|
||||||
|
)
|
||||||
|
service = service_type(application)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match='managed by Cloud'):
|
||||||
|
await getattr(service, create_method)(WORKSPACE, model_data)
|
||||||
|
|
||||||
|
guard.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
('service_type', 'get_method', 'write_method', 'payload'),
|
||||||
|
[
|
||||||
|
(LLMModelsService, 'get_llm_model', 'update_llm_model', {'name': 'changed'}),
|
||||||
|
(LLMModelsService, 'get_llm_model', 'delete_llm_model', None),
|
||||||
|
(EmbeddingModelsService, 'get_embedding_model', 'update_embedding_model', {'name': 'changed'}),
|
||||||
|
(EmbeddingModelsService, 'get_embedding_model', 'delete_embedding_model', None),
|
||||||
|
(RerankModelsService, 'get_rerank_model', 'update_rerank_model', {'name': 'changed'}),
|
||||||
|
(RerankModelsService, 'get_rerank_model', 'delete_rerank_model', None),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_all_model_types_reject_update_and_delete_for_managed_provider(
|
||||||
|
monkeypatch,
|
||||||
|
service_type,
|
||||||
|
get_method: str,
|
||||||
|
write_method: str,
|
||||||
|
payload: dict | None,
|
||||||
|
) -> None:
|
||||||
|
guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified'))
|
||||||
|
monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard)
|
||||||
|
application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime')))
|
||||||
|
service = service_type(application)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
service,
|
||||||
|
get_method,
|
||||||
|
AsyncMock(return_value={'uuid': MODEL, 'provider_uuid': PROVIDER, 'extra_args': {}}),
|
||||||
|
)
|
||||||
|
|
||||||
|
args = (WORKSPACE, MODEL) if payload is None else (WORKSPACE, MODEL, payload)
|
||||||
|
with pytest.raises(ValueError, match='managed by Cloud'):
|
||||||
|
await getattr(service, write_method)(*args)
|
||||||
|
|
||||||
|
guard.assert_awaited_once()
|
||||||
@@ -17,12 +17,14 @@ import pytest
|
|||||||
from unittest.mock import AsyncMock, Mock
|
from unittest.mock import AsyncMock, Mock
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from langbot.pkg.api.http.context import ExecutionContext
|
||||||
from langbot.pkg.api.http.service.model import (
|
from langbot.pkg.api.http.service.model import (
|
||||||
LLMModelsService,
|
LLMModelsService,
|
||||||
EmbeddingModelsService,
|
EmbeddingModelsService,
|
||||||
RerankModelsService,
|
RerankModelsService,
|
||||||
_parse_provider_api_keys,
|
_parse_provider_api_keys,
|
||||||
_runtime_model_data,
|
_runtime_model_data,
|
||||||
|
_serialize_llm_model,
|
||||||
_validate_provider_supports,
|
_validate_provider_supports,
|
||||||
)
|
)
|
||||||
from langbot.pkg.api.http.service import model as model_service_module
|
from langbot.pkg.api.http.service import model as model_service_module
|
||||||
@@ -64,15 +66,19 @@ def _create_mock_llm_model(
|
|||||||
abilities: list = None,
|
abilities: list = None,
|
||||||
context_length: int | None = None,
|
context_length: int | None = None,
|
||||||
extra_args: dict = None,
|
extra_args: dict = None,
|
||||||
|
reasoning_config: dict = None,
|
||||||
) -> Mock:
|
) -> Mock:
|
||||||
"""Helper to create mock LLMModel entity."""
|
"""Helper to create mock LLMModel entity."""
|
||||||
model = Mock(spec=LLMModel)
|
model = Mock(spec=LLMModel)
|
||||||
|
model.workspace_uuid = WORKSPACE_UUID
|
||||||
model.uuid = model_uuid
|
model.uuid = model_uuid
|
||||||
model.name = name
|
model.name = name
|
||||||
model.provider_uuid = provider_uuid
|
model.provider_uuid = provider_uuid
|
||||||
model.abilities = abilities or []
|
model.abilities = abilities or []
|
||||||
model.context_length = context_length
|
model.context_length = context_length
|
||||||
model.extra_args = extra_args or {}
|
model.extra_args = extra_args or {}
|
||||||
|
model.reasoning_config = reasoning_config or {'level': 'provider_default'}
|
||||||
|
model.prefered_ranking = 0
|
||||||
return model
|
return model
|
||||||
|
|
||||||
|
|
||||||
@@ -156,6 +162,26 @@ def _create_runtime_model_mgr() -> SimpleNamespace:
|
|||||||
return manager
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
def _create_reasoning_runtime_provider(capabilities: dict) -> SimpleNamespace:
|
||||||
|
execution_context = ExecutionContext(
|
||||||
|
instance_uuid='instance-test',
|
||||||
|
workspace_uuid=WORKSPACE_UUID,
|
||||||
|
placement_generation=1,
|
||||||
|
)
|
||||||
|
return SimpleNamespace(
|
||||||
|
execution_context=execution_context,
|
||||||
|
provider_entity=ModelProvider(
|
||||||
|
workspace_uuid=WORKSPACE_UUID,
|
||||||
|
uuid='provider-uuid',
|
||||||
|
name='Reasoning Provider',
|
||||||
|
requester='openai',
|
||||||
|
base_url='https://api.openai.com',
|
||||||
|
api_keys=[],
|
||||||
|
),
|
||||||
|
requester=SimpleNamespace(get_reasoning_capabilities=Mock(return_value=capabilities)),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class TestParseProviderApiKeys:
|
class TestParseProviderApiKeys:
|
||||||
"""Tests for _parse_provider_api_keys helper function."""
|
"""Tests for _parse_provider_api_keys helper function."""
|
||||||
|
|
||||||
@@ -209,6 +235,42 @@ class TestRuntimeModelData:
|
|||||||
assert result['extra_args'] == {'temp': 0.7}
|
assert result['extra_args'] == {'temp': 0.7}
|
||||||
|
|
||||||
|
|
||||||
|
class TestSerializeLLMModel:
|
||||||
|
def test_includes_runtime_reasoning_capabilities(self):
|
||||||
|
model = _create_mock_llm_model(
|
||||||
|
abilities=['reasoning'],
|
||||||
|
reasoning_config={'level': 'high'},
|
||||||
|
)
|
||||||
|
capabilities = {
|
||||||
|
'supported': True,
|
||||||
|
'levels': ['provider_default', 'low', 'high'],
|
||||||
|
'source': 'litellm',
|
||||||
|
}
|
||||||
|
runtime_model = SimpleNamespace(
|
||||||
|
model_entity=model,
|
||||||
|
provider=SimpleNamespace(
|
||||||
|
requester=SimpleNamespace(get_reasoning_capabilities=Mock(return_value=capabilities))
|
||||||
|
),
|
||||||
|
)
|
||||||
|
ap = SimpleNamespace(
|
||||||
|
persistence_mgr=SimpleNamespace(
|
||||||
|
serialize_model=Mock(
|
||||||
|
return_value={
|
||||||
|
'uuid': model.uuid,
|
||||||
|
'name': model.name,
|
||||||
|
'reasoning_config': {'level': 'high'},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
),
|
||||||
|
model_mgr=SimpleNamespace(llm_model_dict={('workspace', model.uuid): runtime_model}),
|
||||||
|
)
|
||||||
|
|
||||||
|
serialized = _serialize_llm_model(ap, model)
|
||||||
|
|
||||||
|
assert serialized['reasoning_config'] == {'level': 'high'}
|
||||||
|
assert serialized['reasoning_capabilities'] == capabilities
|
||||||
|
|
||||||
|
|
||||||
class TestLLMModelsServiceGetLLMModels:
|
class TestLLMModelsServiceGetLLMModels:
|
||||||
"""Tests for LLMModelsService.get_llm_models method."""
|
"""Tests for LLMModelsService.get_llm_models method."""
|
||||||
|
|
||||||
@@ -580,6 +642,66 @@ class TestLLMModelsServiceCreateLLMModel:
|
|||||||
ap.provider_service.find_or_create_provider.assert_called_once()
|
ap.provider_service.find_or_create_provider.assert_called_once()
|
||||||
assert result_uuid is not None
|
assert result_uuid is not None
|
||||||
|
|
||||||
|
async def test_create_llm_model_validates_explicit_reasoning_level(self):
|
||||||
|
ap = SimpleNamespace()
|
||||||
|
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock(return_value=_create_mock_result([])))
|
||||||
|
runtime_provider = _create_reasoning_runtime_provider(
|
||||||
|
{
|
||||||
|
'supported': True,
|
||||||
|
'levels': ['provider_default', 'low', 'high'],
|
||||||
|
'source': 'litellm',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
ap.model_mgr = _create_runtime_model_mgr()
|
||||||
|
ap.model_mgr.provider_dict = {'provider-uuid': runtime_provider}
|
||||||
|
|
||||||
|
service = LLMModelsService(ap)
|
||||||
|
await service.create_llm_model(
|
||||||
|
WORKSPACE_UUID,
|
||||||
|
{
|
||||||
|
'uuid': 'reasoning-model',
|
||||||
|
'name': 'Reasoning Model',
|
||||||
|
'provider_uuid': 'provider-uuid',
|
||||||
|
'abilities': ['reasoning'],
|
||||||
|
'reasoning_config': {'level': 'high'},
|
||||||
|
'extra_args': {},
|
||||||
|
},
|
||||||
|
preserve_uuid=True,
|
||||||
|
auto_set_to_default_pipeline=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
runtime_entity = ap.model_mgr.load_llm_model_with_provider.await_args.args[1]
|
||||||
|
assert runtime_entity.reasoning_config == {'level': 'high'}
|
||||||
|
|
||||||
|
async def test_create_llm_model_rejects_unsupported_reasoning_before_insert(self):
|
||||||
|
ap = SimpleNamespace()
|
||||||
|
ap.persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
|
||||||
|
runtime_provider = _create_reasoning_runtime_provider(
|
||||||
|
{
|
||||||
|
'supported': True,
|
||||||
|
'levels': ['provider_default'],
|
||||||
|
'source': 'manual',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
ap.model_mgr = _create_runtime_model_mgr()
|
||||||
|
ap.model_mgr.provider_dict = {'provider-uuid': runtime_provider}
|
||||||
|
|
||||||
|
service = LLMModelsService(ap)
|
||||||
|
with pytest.raises(ValueError, match='Available levels: provider_default'):
|
||||||
|
await service.create_llm_model(
|
||||||
|
WORKSPACE_UUID,
|
||||||
|
{
|
||||||
|
'name': 'Unknown Reasoning Model',
|
||||||
|
'provider_uuid': 'provider-uuid',
|
||||||
|
'abilities': ['reasoning'],
|
||||||
|
'reasoning_config': {'level': 'high'},
|
||||||
|
'extra_args': {},
|
||||||
|
},
|
||||||
|
auto_set_to_default_pipeline=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
class TestLLMModelsServiceUpdateLLMModel:
|
class TestLLMModelsServiceUpdateLLMModel:
|
||||||
"""Tests for LLMModelsService.update_llm_model method."""
|
"""Tests for LLMModelsService.update_llm_model method."""
|
||||||
@@ -595,7 +717,10 @@ class TestLLMModelsServiceUpdateLLMModel:
|
|||||||
ap.model_mgr.remove_llm_model = AsyncMock()
|
ap.model_mgr.remove_llm_model = AsyncMock()
|
||||||
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
|
ap.model_mgr.load_llm_model_with_provider = AsyncMock(return_value=Mock())
|
||||||
|
|
||||||
ap.persistence_mgr.execute_async = AsyncMock()
|
existing_model = _create_mock_llm_model()
|
||||||
|
ap.persistence_mgr.execute_async = AsyncMock(
|
||||||
|
side_effect=[_create_mock_result(first_item=existing_model), _create_mock_result()]
|
||||||
|
)
|
||||||
|
|
||||||
service = LLMModelsService(ap)
|
service = LLMModelsService(ap)
|
||||||
service.get_llm_model = AsyncMock(return_value=_existing_llm_data())
|
service.get_llm_model = AsyncMock(return_value=_existing_llm_data())
|
||||||
@@ -623,7 +748,8 @@ class TestLLMModelsServiceUpdateLLMModel:
|
|||||||
ap.model_mgr.provider_dict = {} # Empty
|
ap.model_mgr.provider_dict = {} # Empty
|
||||||
ap.model_mgr.remove_llm_model = AsyncMock()
|
ap.model_mgr.remove_llm_model = AsyncMock()
|
||||||
|
|
||||||
ap.persistence_mgr.execute_async = AsyncMock()
|
existing_model = _create_mock_llm_model()
|
||||||
|
ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=existing_model))
|
||||||
|
|
||||||
service = LLMModelsService(ap)
|
service = LLMModelsService(ap)
|
||||||
service.get_llm_model = AsyncMock(return_value=_existing_llm_data('nonexistent-provider'))
|
service.get_llm_model = AsyncMock(return_value=_existing_llm_data('nonexistent-provider'))
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ from langbot.pkg.workspace.errors import WorkspaceNotFoundError
|
|||||||
pytestmark = pytest.mark.asyncio
|
pytestmark = pytest.mark.asyncio
|
||||||
|
|
||||||
WORKSPACE_UUID = 'workspace-a'
|
WORKSPACE_UUID = 'workspace-a'
|
||||||
|
SYSTEM_REQUESTER = 'space-chat-completions'
|
||||||
|
|
||||||
|
|
||||||
def _create_mock_provider(
|
def _create_mock_provider(
|
||||||
@@ -1005,3 +1006,56 @@ class TestProviderSecretRoundtrip:
|
|||||||
)
|
)
|
||||||
|
|
||||||
ap.persistence_mgr.execute_async.assert_not_awaited()
|
ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
class TestCloudManagedProviderProtection:
|
||||||
|
@staticmethod
|
||||||
|
def _service() -> ModelProviderService:
|
||||||
|
ap = SimpleNamespace(
|
||||||
|
persistence_mgr=SimpleNamespace(
|
||||||
|
mode=SimpleNamespace(value='cloud_runtime'),
|
||||||
|
execute_async=AsyncMock(),
|
||||||
|
),
|
||||||
|
model_mgr=SimpleNamespace(),
|
||||||
|
)
|
||||||
|
return ModelProviderService(ap)
|
||||||
|
|
||||||
|
async def test_cloud_rejects_user_created_system_requester(self):
|
||||||
|
service = self._service()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match='reserved'):
|
||||||
|
await service.create_provider(
|
||||||
|
WORKSPACE_UUID,
|
||||||
|
{
|
||||||
|
'name': 'Fake LangBot Models',
|
||||||
|
'requester': SYSTEM_REQUESTER,
|
||||||
|
'base_url': 'https://example.invalid/v1',
|
||||||
|
'api_keys': ['fake'],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match='reserved'):
|
||||||
|
await service.find_or_create_provider(
|
||||||
|
WORKSPACE_UUID,
|
||||||
|
SYSTEM_REQUESTER,
|
||||||
|
'https://api.langbot.cloud/v1',
|
||||||
|
['fake'],
|
||||||
|
)
|
||||||
|
service.ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||||
|
|
||||||
|
async def test_cloud_rejects_update_and_delete_of_managed_provider(self):
|
||||||
|
service = self._service()
|
||||||
|
service.get_provider = AsyncMock(
|
||||||
|
return_value={'uuid': 'system-provider', 'requester': SYSTEM_REQUESTER}
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match='managed by Cloud'):
|
||||||
|
await service.update_provider(WORKSPACE_UUID, 'system-provider', {'name': 'Renamed'})
|
||||||
|
with pytest.raises(ValueError, match='managed by Cloud'):
|
||||||
|
await service.delete_provider(WORKSPACE_UUID, 'system-provider')
|
||||||
|
service.ap.persistence_mgr.execute_async.assert_not_awaited()
|
||||||
|
|
||||||
|
async def test_oss_does_not_reserve_space_requester(self):
|
||||||
|
ap = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss_compat')))
|
||||||
|
service = ModelProviderService(ap)
|
||||||
|
assert service._system_requester_is_reserved(SYSTEM_REQUESTER) is False
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ import time
|
|||||||
|
|
||||||
from langbot.pkg.api.http.service.space import SpaceService
|
from langbot.pkg.api.http.service.space import SpaceService
|
||||||
from langbot.pkg.entity.persistence.user import User
|
from langbot.pkg.entity.persistence.user import User
|
||||||
|
from langbot.pkg.utils import constants
|
||||||
|
|
||||||
|
|
||||||
pytestmark = pytest.mark.asyncio
|
pytestmark = pytest.mark.asyncio
|
||||||
@@ -573,10 +574,20 @@ class TestSpaceServiceExchangeOAuthCode:
|
|||||||
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
|
mock_session_obj.post.return_value.__aexit__ = AsyncMock(return_value=None)
|
||||||
|
|
||||||
# Execute
|
# Execute
|
||||||
result = await service.exchange_oauth_code('auth_code')
|
result = await service.exchange_oauth_code(
|
||||||
|
'auth_code',
|
||||||
|
['workspace-1'],
|
||||||
|
{'workspace-1': 1_700_000_000},
|
||||||
|
)
|
||||||
|
|
||||||
# Verify
|
# Verify
|
||||||
assert result['access_token'] == 'new_access_token'
|
assert result['access_token'] == 'new_access_token'
|
||||||
|
assert mock_session_obj.post.call_args.kwargs['json'] == {
|
||||||
|
'code': 'auth_code',
|
||||||
|
'instance_id': constants.instance_id,
|
||||||
|
'workspace_uuids': ['workspace-1'],
|
||||||
|
'workspace_created_ats': {'workspace-1': 1_700_000_000},
|
||||||
|
}
|
||||||
|
|
||||||
async def test_exchange_oauth_code_api_error(self):
|
async def test_exchange_oauth_code_api_error(self):
|
||||||
"""Raises ValueError on API error."""
|
"""Raises ValueError on API error."""
|
||||||
|
|||||||
@@ -377,7 +377,7 @@ class TestUserServiceAuthenticate:
|
|||||||
service = UserService(ap)
|
service = UserService(ap)
|
||||||
|
|
||||||
# Execute & Verify
|
# Execute & Verify
|
||||||
with pytest.raises(ValueError, match='请使用 Space 账户登录'):
|
with pytest.raises(ValueError, match='请使用 LangBot 账号登录'):
|
||||||
await service.authenticate('space@example.com', 'password')
|
await service.authenticate('space@example.com', 'password')
|
||||||
|
|
||||||
|
|
||||||
@@ -418,6 +418,7 @@ class TestUserServiceGenerateJwtToken:
|
|||||||
assert token is not None
|
assert token is not None
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
class TestUserServiceVerifyJwtToken:
|
class TestUserServiceVerifyJwtToken:
|
||||||
"""Tests for verify_jwt_token method."""
|
"""Tests for verify_jwt_token method."""
|
||||||
|
|
||||||
@@ -725,7 +726,7 @@ class TestUserServiceCreateOrUpdateSpaceUser:
|
|||||||
)
|
)
|
||||||
service = UserService(ap)
|
service = UserService(ap)
|
||||||
|
|
||||||
with pytest.raises(ControlPlaneDirectoryRequiredError, match='Space account'):
|
with pytest.raises(ControlPlaneDirectoryRequiredError, match='LangBot Account'):
|
||||||
await service.register_invited_account('invite-token', 'member@example.com', 'password')
|
await service.register_invited_account('invite-token', 'member@example.com', 'password')
|
||||||
|
|
||||||
async def test_create_or_update_new_space_user_first_init(self):
|
async def test_create_or_update_new_space_user_first_init(self):
|
||||||
|
|||||||
@@ -149,6 +149,36 @@ async def test_session_scope_matches_exact_tenant_placement_and_principal():
|
|||||||
assert sessions == {}
|
assert sessions == {}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_support_admin_sessions_are_scoped_to_the_persisted_grant():
|
||||||
|
def support_context(grant_jti_hash: str) -> RequestContext:
|
||||||
|
return RequestContext(
|
||||||
|
instance_uuid='instance-test',
|
||||||
|
placement_generation=1,
|
||||||
|
request_id='request-test',
|
||||||
|
auth_type='support-admin',
|
||||||
|
principal=PrincipalContext(
|
||||||
|
principal_type=PrincipalType.SUPPORT_ADMIN,
|
||||||
|
actor_account_uuid='support-actor',
|
||||||
|
support_session_id=grant_jti_hash,
|
||||||
|
),
|
||||||
|
workspace=WorkspaceContext(
|
||||||
|
workspace_uuid='workspace-a',
|
||||||
|
membership_uuid=None,
|
||||||
|
role='owner',
|
||||||
|
permissions=frozenset({'resource.manage'}),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
first_context = support_context('a' * 64)
|
||||||
|
second_context = support_context('b' * 64)
|
||||||
|
sessions: dict[str, dict] = {'session-test': {'status': 'waiting'}}
|
||||||
|
_bind_session_scope(sessions['session-test'], first_context)
|
||||||
|
|
||||||
|
assert _get_owned_session(sessions, 'session-test', second_context) is None
|
||||||
|
assert _pop_owned_session(sessions, 'session-test', second_context) is None
|
||||||
|
assert _get_owned_session(sessions, 'session-test', first_context) is sessions['session-test']
|
||||||
|
|
||||||
|
|
||||||
async def test_session_capacity_evicts_oldest_session_in_same_workspace():
|
async def test_session_capacity_evicts_oldest_session_in_same_workspace():
|
||||||
owner_context = _request_context()
|
owner_context = _request_context()
|
||||||
sessions: dict[str, dict] = {}
|
sessions: dict[str, dict] = {}
|
||||||
|
|||||||
@@ -124,24 +124,37 @@ async def test_background_plugin_operation_refences_captured_generation(plugin_r
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(plugin_router_cls):
|
async def test_background_plugin_operation_revalidates_and_runs_inside_tenant_uow(plugin_router_cls):
|
||||||
scopes = []
|
scopes = []
|
||||||
|
active_scope = None
|
||||||
|
|
||||||
|
transaction_active = False
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def tenant_uow(workspace_uuid):
|
async def tenant_scope(workspace_uuid):
|
||||||
|
nonlocal active_scope
|
||||||
scopes.append(workspace_uuid)
|
scopes.append(workspace_uuid)
|
||||||
yield
|
active_scope = workspace_uuid
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
active_scope = None
|
||||||
|
|
||||||
connector = SimpleNamespace(
|
connector = SimpleNamespace(
|
||||||
require_workspace_context=AsyncMock(side_effect=lambda context: context),
|
require_workspace_context=AsyncMock(side_effect=lambda context: context),
|
||||||
)
|
)
|
||||||
operation = AsyncMock(return_value='done')
|
|
||||||
|
async def operation():
|
||||||
|
assert active_scope == CONTEXT.workspace_uuid
|
||||||
|
assert transaction_active is False
|
||||||
|
return 'done'
|
||||||
|
|
||||||
router = object.__new__(plugin_router_cls)
|
router = object.__new__(plugin_router_cls)
|
||||||
router.ap = SimpleNamespace(
|
router.ap = SimpleNamespace(
|
||||||
plugin_connector=connector,
|
plugin_connector=connector,
|
||||||
persistence_mgr=SimpleNamespace(
|
persistence_mgr=SimpleNamespace(
|
||||||
mode=SimpleNamespace(value='cloud_runtime'),
|
mode=SimpleNamespace(value='cloud_runtime'),
|
||||||
tenant_uow=tenant_uow,
|
tenant_scope=tenant_scope,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -150,4 +163,3 @@ async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(p
|
|||||||
assert result == 'done'
|
assert result == 'done'
|
||||||
assert scopes == [CONTEXT.workspace_uuid]
|
assert scopes == [CONTEXT.workspace_uuid]
|
||||||
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
|
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
|
||||||
operation.assert_awaited_once()
|
|
||||||
|
|||||||
@@ -24,7 +24,7 @@ from langbot.pkg.box.connector import BoxRuntimeConnector
|
|||||||
_CONTROL_TOKEN = 'box-control-token-that-is-longer-than-32-bytes'
|
_CONTROL_TOKEN = 'box-control-token-that-is-longer-than-32-bytes'
|
||||||
|
|
||||||
|
|
||||||
def make_app(logger: Mock, runtime_endpoint: str = ''):
|
def make_app(logger: Mock, runtime_endpoint: str = '', *, cloud: bool = False):
|
||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
logger=logger,
|
logger=logger,
|
||||||
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
|
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
|
||||||
@@ -42,6 +42,7 @@ def make_app(logger: Mock, runtime_endpoint: str = ''):
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
),
|
),
|
||||||
|
deployment=SimpleNamespace(mode='cloud' if cloud else 'oss'),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -306,10 +307,27 @@ def test_box_runtime_connector_rejects_relay_context_from_other_instance(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_external_box_runtime_fails_closed_without_control_token(monkeypatch: pytest.MonkeyPatch):
|
def test_external_box_runtime_control_headers_are_tokenless_when_secret_is_unset(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
):
|
||||||
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
|
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
|
||||||
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
|
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
|
||||||
|
|
||||||
|
assert connector.get_control_headers() == {BOX_INSTANCE_HEADER: 'instance-a'}
|
||||||
|
|
||||||
|
|
||||||
|
def test_cloud_box_runtime_rejects_missing_control_secret(monkeypatch: pytest.MonkeyPatch):
|
||||||
|
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
|
||||||
|
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410', cloud=True))
|
||||||
|
|
||||||
|
with pytest.raises(BoxRuntimeUnavailableError, match=BOX_CONTROL_TOKEN_ENV):
|
||||||
|
connector.get_control_headers()
|
||||||
|
|
||||||
|
|
||||||
|
def test_external_box_runtime_rejects_invalid_configured_control_token(monkeypatch: pytest.MonkeyPatch):
|
||||||
|
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, 'too-short')
|
||||||
|
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
|
||||||
|
|
||||||
with pytest.raises(BoxRuntimeUnavailableError, match=BOX_CONTROL_TOKEN_ENV):
|
with pytest.raises(BoxRuntimeUnavailableError, match=BOX_CONTROL_TOKEN_ENV):
|
||||||
connector.get_control_headers()
|
connector.get_control_headers()
|
||||||
|
|
||||||
|
|||||||
@@ -66,6 +66,10 @@ class _Provider:
|
|||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.manifest_provider = _Manifest()
|
self.manifest_provider = _Manifest()
|
||||||
|
|
||||||
|
async def fetch_model_catalog(self, instance_uuid: str):
|
||||||
|
del instance_uuid
|
||||||
|
raise AssertionError('not used by bootstrap contract tests')
|
||||||
|
|
||||||
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
|
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
|
||||||
del instance_config
|
del instance_config
|
||||||
return VerifiedCloudDeployment(
|
return VerifiedCloudDeployment(
|
||||||
@@ -79,6 +83,7 @@ class _Provider:
|
|||||||
entitlement_provider=_Entitlements(),
|
entitlement_provider=_Entitlements(),
|
||||||
directory_provider=_Directory(),
|
directory_provider=_Directory(),
|
||||||
manifest_provider=self.manifest_provider,
|
manifest_provider=self.manifest_provider,
|
||||||
|
model_catalog_provider=self,
|
||||||
verification_key_id='root-2026',
|
verification_key_id='root-2026',
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -103,7 +108,7 @@ def _cloud_config() -> dict:
|
|||||||
'use': 'pgvector',
|
'use': 'pgvector',
|
||||||
'pgvector': {
|
'pgvector': {
|
||||||
'use_business_database': True,
|
'use_business_database': True,
|
||||||
'allowed_dimensions': [384, 768, 1536],
|
'allowed_dimensions': [384, 768, 1536, 3072],
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
'mcp': {'stdio': {'enabled': False}},
|
'mcp': {'stdio': {'enabled': False}},
|
||||||
@@ -211,7 +216,6 @@ async def test_cloud_directory_capacity_contract_is_fail_closed(directory_config
|
|||||||
[
|
[
|
||||||
({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
|
({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
|
||||||
({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
|
({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
|
||||||
({'use_business_database': True, 'allowed_dimensions': [3072]}, 'allowed_dimensions'),
|
|
||||||
({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
|
({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -181,6 +181,39 @@ def _delta(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_directory_delta_requests_model_catalog_sync_after_commit(projection_context):
|
||||||
|
application, _session_factory = projection_context
|
||||||
|
request_sync = Mock()
|
||||||
|
application.cloud_model_catalog_service = SimpleNamespace(request_sync=request_sync)
|
||||||
|
event = DirectoryEvent(
|
||||||
|
cursor=2,
|
||||||
|
uuid='20000000-0000-4000-8000-000000000002',
|
||||||
|
aggregate_uuid=WORKSPACE_UUID,
|
||||||
|
event_type='directory.changed',
|
||||||
|
revision=2,
|
||||||
|
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
|
||||||
|
created_at=datetime.datetime(2026, 7, 24, 12, 30, tzinfo=datetime.UTC),
|
||||||
|
)
|
||||||
|
batch = DirectoryEventBatch(
|
||||||
|
instance_uuid=INSTANCE_UUID,
|
||||||
|
after_cursor=1,
|
||||||
|
cursor=2,
|
||||||
|
high_water_cursor=2,
|
||||||
|
events=[event],
|
||||||
|
)
|
||||||
|
service = DirectoryProjectionService(
|
||||||
|
application,
|
||||||
|
_Provider([_snapshot(1)], [batch], [_delta(workspaces=[_workspace(revision=2)])]),
|
||||||
|
INSTANCE_UUID,
|
||||||
|
)
|
||||||
|
await service.initialize()
|
||||||
|
request_sync.reset_mock()
|
||||||
|
|
||||||
|
await service.sync_once()
|
||||||
|
|
||||||
|
request_sync.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
async def test_initial_snapshot_projects_core_owned_rows(projection_context):
|
async def test_initial_snapshot_projects_core_owned_rows(projection_context):
|
||||||
application, session_factory = projection_context
|
application, session_factory = projection_context
|
||||||
reconcile_execution_projection = Mock()
|
reconcile_execution_projection = Mock()
|
||||||
@@ -1023,7 +1056,7 @@ async def test_snapshot_for_another_instance_is_rejected(projection_context):
|
|||||||
await service.initialize()
|
await service.initialize()
|
||||||
|
|
||||||
|
|
||||||
async def test_core_owned_membership_survives_directory_updates_and_omission(projection_context):
|
async def test_directory_revision_zero_membership_is_adopted(projection_context):
|
||||||
application, session_factory = projection_context
|
application, session_factory = projection_context
|
||||||
service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
|
service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
|
||||||
await service.initialize()
|
await service.initialize()
|
||||||
@@ -1034,29 +1067,125 @@ async def test_core_owned_membership_survives_directory_updates_and_omission(pro
|
|||||||
membership.role = 'viewer'
|
membership.role = 'viewer'
|
||||||
membership.status = 'active'
|
membership.status = 'active'
|
||||||
membership.projection_revision = 0
|
membership.projection_revision = 0
|
||||||
session.add(
|
|
||||||
WorkspaceMembership(
|
|
||||||
uuid=SECOND_MEMBERSHIP_UUID,
|
|
||||||
workspace_uuid=WORKSPACE_UUID,
|
|
||||||
account_uuid='20000000-0000-0000-0000-000000000099',
|
|
||||||
role='viewer',
|
|
||||||
status='active',
|
|
||||||
joined_at=membership.joined_at,
|
|
||||||
projection_revision=0,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
projected_member = _member(revision=2).model_copy(update={'role': 'owner', 'membership_status': 'removed'})
|
projected_member = _member(revision=2).model_copy(update={'role': 'owner', 'membership_status': 'removed'})
|
||||||
projected_workspace = _workspace(revision=2).model_copy(update={'members': (projected_member,)})
|
projected_workspace = _workspace(revision=2).model_copy(update={'members': (projected_member,)})
|
||||||
await service.apply_snapshot(_snapshot(2, workspaces=[projected_workspace]))
|
await service.apply_snapshot(_snapshot(2, workspaces=[projected_workspace]))
|
||||||
|
|
||||||
async with session_factory() as session:
|
async with session_factory() as session:
|
||||||
memberships = {
|
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
|
||||||
membership.uuid: membership
|
assert membership.source == 'cloud_projection'
|
||||||
for membership in (await session.scalars(sqlalchemy.select(WorkspaceMembership))).all()
|
assert membership.role == 'owner'
|
||||||
}
|
assert membership.status == 'removed'
|
||||||
assert memberships[MEMBERSHIP_UUID].role == 'viewer'
|
assert membership.projection_revision == 2
|
||||||
assert memberships[MEMBERSHIP_UUID].status == 'active'
|
|
||||||
assert memberships[MEMBERSHIP_UUID].projection_revision == 0
|
|
||||||
assert memberships[SECOND_MEMBERSHIP_UUID].status == 'active'
|
async def test_directory_revision_zero_membership_omitted_from_snapshot_is_removed(projection_context):
|
||||||
assert memberships[SECOND_MEMBERSHIP_UUID].projection_revision == 0
|
application, session_factory = projection_context
|
||||||
|
service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
|
||||||
|
await service.initialize()
|
||||||
|
|
||||||
|
historical_account_uuid = '20000000-0000-0000-0000-000000000099'
|
||||||
|
async with session_factory() as session:
|
||||||
|
async with session.begin():
|
||||||
|
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
|
||||||
|
session.add(
|
||||||
|
User(
|
||||||
|
uuid=historical_account_uuid,
|
||||||
|
user='Historical Space Member',
|
||||||
|
normalized_email='historical@example.com',
|
||||||
|
password='',
|
||||||
|
status='active',
|
||||||
|
source='cloud_projection',
|
||||||
|
projection_revision=1,
|
||||||
|
account_type='space',
|
||||||
|
space_account_uuid=historical_account_uuid,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
session.add(
|
||||||
|
WorkspaceMembership(
|
||||||
|
uuid=SECOND_MEMBERSHIP_UUID,
|
||||||
|
workspace_uuid=WORKSPACE_UUID,
|
||||||
|
account_uuid=historical_account_uuid,
|
||||||
|
role='viewer',
|
||||||
|
status='active',
|
||||||
|
source='cloud_projection',
|
||||||
|
joined_at=membership.joined_at,
|
||||||
|
projection_revision=0,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.apply_snapshot(_snapshot(2))
|
||||||
|
|
||||||
|
async with session_factory() as session:
|
||||||
|
historical = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
|
||||||
|
assert historical.status == 'removed'
|
||||||
|
assert historical.projection_revision == 2
|
||||||
|
|
||||||
|
|
||||||
|
async def test_cloud_account_core_invitation_membership_survives_directory_omission(projection_context):
|
||||||
|
application, session_factory = projection_context
|
||||||
|
service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
|
||||||
|
await service.initialize()
|
||||||
|
|
||||||
|
invited_account_uuid = '20000000-0000-0000-0000-000000000098'
|
||||||
|
async with session_factory() as session:
|
||||||
|
async with session.begin():
|
||||||
|
projected_membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
|
||||||
|
session.add(
|
||||||
|
User(
|
||||||
|
uuid=invited_account_uuid,
|
||||||
|
user='Invited Cloud Account',
|
||||||
|
normalized_email='invited-cloud@example.com',
|
||||||
|
password='',
|
||||||
|
status='active',
|
||||||
|
source='cloud_projection',
|
||||||
|
projection_revision=1,
|
||||||
|
account_type='space',
|
||||||
|
space_account_uuid=invited_account_uuid,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
session.add(
|
||||||
|
WorkspaceMembership(
|
||||||
|
uuid=SECOND_MEMBERSHIP_UUID,
|
||||||
|
workspace_uuid=WORKSPACE_UUID,
|
||||||
|
account_uuid=invited_account_uuid,
|
||||||
|
role='viewer',
|
||||||
|
status='active',
|
||||||
|
source='local',
|
||||||
|
joined_at=projected_membership.joined_at,
|
||||||
|
projection_revision=0,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.apply_snapshot(_snapshot(2))
|
||||||
|
|
||||||
|
async with session_factory() as session:
|
||||||
|
membership = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
|
||||||
|
assert membership.source == 'local'
|
||||||
|
assert membership.status == 'active'
|
||||||
|
assert membership.projection_revision == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_directory_does_not_adopt_local_membership_with_different_uuid_for_same_cloud_account(projection_context):
|
||||||
|
application, session_factory = projection_context
|
||||||
|
service = DirectoryProjectionService(application, _Provider([_snapshot(1)]), INSTANCE_UUID)
|
||||||
|
await service.initialize()
|
||||||
|
|
||||||
|
async with session_factory() as session:
|
||||||
|
async with session.begin():
|
||||||
|
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
|
||||||
|
membership.uuid = SECOND_MEMBERSHIP_UUID
|
||||||
|
membership.source = 'local'
|
||||||
|
membership.projection_revision = 0
|
||||||
|
|
||||||
|
projected_member = _member(revision=2).model_copy(update={'role': 'owner', 'membership_status': 'removed'})
|
||||||
|
projected_workspace = _workspace(revision=2).model_copy(update={'members': (projected_member,)})
|
||||||
|
await service.apply_snapshot(_snapshot(2, workspaces=[projected_workspace]))
|
||||||
|
|
||||||
|
async with session_factory() as session:
|
||||||
|
membership = await session.get(WorkspaceMembership, SECOND_MEMBERSHIP_UUID)
|
||||||
|
assert membership.source == 'local'
|
||||||
|
assert membership.role == 'developer'
|
||||||
|
assert membership.status == 'active'
|
||||||
|
assert membership.projection_revision == 0
|
||||||
|
|||||||
@@ -0,0 +1,466 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from datetime import UTC, datetime
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import sqlalchemy
|
||||||
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
|
from langbot.pkg.cloud.model_catalog import (
|
||||||
|
CloudModelCatalogSnapshot,
|
||||||
|
CloudModelCatalogSyncService,
|
||||||
|
system_model_uuid,
|
||||||
|
system_provider_uuid,
|
||||||
|
)
|
||||||
|
from langbot.pkg.entity.persistence.base import Base
|
||||||
|
from langbot.pkg.entity.persistence.model import EmbeddingModel, LLMModel, ModelProvider
|
||||||
|
from langbot.pkg.entity.persistence.workspace import Workspace
|
||||||
|
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
|
||||||
|
|
||||||
|
|
||||||
|
pytestmark = pytest.mark.asyncio
|
||||||
|
INSTANCE_UUID = 'instance-model-catalog'
|
||||||
|
WORKSPACE_A = '00000000-0000-4000-8000-000000000001'
|
||||||
|
WORKSPACE_B = '00000000-0000-4000-8000-000000000002'
|
||||||
|
OWNER_A = '10000000-0000-4000-8000-000000000001'
|
||||||
|
OWNER_B = '10000000-0000-4000-8000-000000000002'
|
||||||
|
|
||||||
|
|
||||||
|
class _CatalogProvider:
|
||||||
|
def __init__(self, snapshot: CloudModelCatalogSnapshot) -> None:
|
||||||
|
self.snapshot = snapshot
|
||||||
|
|
||||||
|
async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot:
|
||||||
|
assert instance_uuid == INSTANCE_UUID
|
||||||
|
return self.snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def _snapshot(
|
||||||
|
*,
|
||||||
|
key_a: str | None = 'owner-a-key',
|
||||||
|
model_id: str = 'gpt-test',
|
||||||
|
include_embedding: bool = True,
|
||||||
|
) -> CloudModelCatalogSnapshot:
|
||||||
|
models = [
|
||||||
|
{
|
||||||
|
'uuid': 'upstream-chat',
|
||||||
|
'model_id': model_id,
|
||||||
|
'category': 'chat',
|
||||||
|
'llm_abilities': ['chat', 'vision'],
|
||||||
|
'is_featured': True,
|
||||||
|
'featured_order': 7,
|
||||||
|
}
|
||||||
|
]
|
||||||
|
if include_embedding:
|
||||||
|
models.append(
|
||||||
|
{
|
||||||
|
'uuid': 'upstream-embedding',
|
||||||
|
'model_id': 'embedding-test',
|
||||||
|
'category': 'embedding',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return CloudModelCatalogSnapshot.model_validate(
|
||||||
|
{
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'generated_at': datetime.now(UTC),
|
||||||
|
'base_url': 'https://api.langbot.cloud/v1/',
|
||||||
|
'models': models,
|
||||||
|
'workspaces': [
|
||||||
|
{
|
||||||
|
'workspace_uuid': WORKSPACE_A,
|
||||||
|
'owner_account_uuid': OWNER_A,
|
||||||
|
'api_key': key_a,
|
||||||
|
'credits': 25000,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
'workspace_uuid': WORKSPACE_B,
|
||||||
|
'owner_account_uuid': OWNER_B,
|
||||||
|
'api_key': 'owner-b-key',
|
||||||
|
'credits': 5000,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_catalog_snapshot_treats_null_model_abilities_as_empty() -> None:
|
||||||
|
payload = _snapshot().model_dump(mode='json')
|
||||||
|
payload['models'][0]['llm_abilities'] = None
|
||||||
|
|
||||||
|
snapshot = CloudModelCatalogSnapshot.model_validate(payload)
|
||||||
|
|
||||||
|
assert snapshot.models[0].llm_abilities == ()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_catalog_reconciles_every_workspace_idempotently_and_tracks_owner_and_downlisting(tmp_path) -> None:
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog.db"}')
|
||||||
|
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||||
|
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||||
|
bindings = [
|
||||||
|
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1),
|
||||||
|
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1),
|
||||||
|
]
|
||||||
|
workspace_service = SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings))
|
||||||
|
reload_counter = _AsyncCounter()
|
||||||
|
runtime_reload = SimpleNamespace(load_models_from_db=reload_counter)
|
||||||
|
app = SimpleNamespace(
|
||||||
|
persistence_mgr=manager,
|
||||||
|
workspace_service=workspace_service,
|
||||||
|
model_mgr=runtime_reload,
|
||||||
|
logger=logging.getLogger(__name__),
|
||||||
|
)
|
||||||
|
provider = _CatalogProvider(_snapshot())
|
||||||
|
service = CloudModelCatalogSyncService(app, provider, INSTANCE_UUID)
|
||||||
|
|
||||||
|
try:
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.run_sync(Base.metadata.create_all)
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(Workspace),
|
||||||
|
[
|
||||||
|
{
|
||||||
|
'uuid': WORKSPACE_A,
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'name': 'A',
|
||||||
|
'slug': 'a',
|
||||||
|
'source': 'cloud_projection',
|
||||||
|
},
|
||||||
|
{
|
||||||
|
'uuid': WORKSPACE_B,
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'name': 'B',
|
||||||
|
'slug': 'b',
|
||||||
|
'source': 'cloud_projection',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(ModelProvider).values(
|
||||||
|
uuid='custom-provider',
|
||||||
|
workspace_uuid=WORKSPACE_A,
|
||||||
|
name='Custom',
|
||||||
|
requester='openai-chat-completions',
|
||||||
|
base_url='https://custom.example/v1',
|
||||||
|
api_keys=['custom-key'],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(LLMModel).values(
|
||||||
|
uuid='custom-model',
|
||||||
|
workspace_uuid=WORKSPACE_A,
|
||||||
|
name='custom-model',
|
||||||
|
provider_uuid='custom-provider',
|
||||||
|
abilities=['chat'],
|
||||||
|
extra_args={},
|
||||||
|
prefered_ranking=0,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
first = await service.sync_once()
|
||||||
|
assert first == {'workspaces': 2, 'created': 6, 'updated': 0, 'deleted': 0}
|
||||||
|
assert reload_counter.calls == 1
|
||||||
|
assert service.get_workspace_credits(WORKSPACE_A) == 25000
|
||||||
|
assert service.get_workspace_credits(WORKSPACE_B) == 5000
|
||||||
|
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
providers = (
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.select(
|
||||||
|
ModelProvider.uuid,
|
||||||
|
ModelProvider.workspace_uuid,
|
||||||
|
ModelProvider.api_keys,
|
||||||
|
).where(ModelProvider.requester == 'space-chat-completions')
|
||||||
|
)
|
||||||
|
).all()
|
||||||
|
assert {item.workspace_uuid for item in providers} == {WORKSPACE_A, WORKSPACE_B}
|
||||||
|
assert {item.uuid for item in providers} == {
|
||||||
|
system_provider_uuid(WORKSPACE_A),
|
||||||
|
system_provider_uuid(WORKSPACE_B),
|
||||||
|
}
|
||||||
|
assert {item.workspace_uuid: item.api_keys for item in providers} == {
|
||||||
|
WORKSPACE_A: ['owner-a-key'],
|
||||||
|
WORKSPACE_B: ['owner-b-key'],
|
||||||
|
}
|
||||||
|
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(LLMModel)) == 3
|
||||||
|
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 2
|
||||||
|
|
||||||
|
second = await service.sync_once()
|
||||||
|
assert second == {'workspaces': 2, 'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
assert reload_counter.calls == 1
|
||||||
|
|
||||||
|
provider.snapshot = _snapshot(
|
||||||
|
key_a='new-owner-key',
|
||||||
|
model_id='gpt-renamed',
|
||||||
|
include_embedding=False,
|
||||||
|
)
|
||||||
|
third = await service.sync_once()
|
||||||
|
assert third == {'workspaces': 2, 'created': 0, 'updated': 3, 'deleted': 2}
|
||||||
|
assert reload_counter.calls == 2
|
||||||
|
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
provider_a_keys = await connection.scalar(
|
||||||
|
sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A))
|
||||||
|
)
|
||||||
|
assert provider_a_keys == ['new-owner-key']
|
||||||
|
system_model_names = (
|
||||||
|
(
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.select(LLMModel.name).where(
|
||||||
|
LLMModel.provider_uuid.in_(
|
||||||
|
[system_provider_uuid(WORKSPACE_A), system_provider_uuid(WORKSPACE_B)]
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
.scalars()
|
||||||
|
.all()
|
||||||
|
)
|
||||||
|
assert set(system_model_names) == {'gpt-renamed'}
|
||||||
|
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 0
|
||||||
|
assert (
|
||||||
|
await connection.scalar(
|
||||||
|
sqlalchemy.select(sqlalchemy.func.count())
|
||||||
|
.select_from(ModelProvider)
|
||||||
|
.where(ModelProvider.uuid == 'custom-provider')
|
||||||
|
)
|
||||||
|
== 1
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
await connection.scalar(
|
||||||
|
sqlalchemy.select(sqlalchemy.func.count())
|
||||||
|
.select_from(LLMModel)
|
||||||
|
.where(LLMModel.uuid == 'custom-model')
|
||||||
|
)
|
||||||
|
== 1
|
||||||
|
)
|
||||||
|
|
||||||
|
provider.snapshot = _snapshot(key_a=None, model_id='gpt-renamed', include_embedding=False)
|
||||||
|
fourth = await service.sync_once()
|
||||||
|
assert fourth == {'workspaces': 2, 'created': 0, 'updated': 1, 'deleted': 0}
|
||||||
|
assert reload_counter.calls == 3
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
provider_a_keys = await connection.scalar(
|
||||||
|
sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A))
|
||||||
|
)
|
||||||
|
assert provider_a_keys == []
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
def test_workspace_scoped_ids_are_stable_and_secrets_are_redacted() -> None:
|
||||||
|
assert system_provider_uuid(WORKSPACE_A) == system_provider_uuid(WORKSPACE_A)
|
||||||
|
assert system_provider_uuid(WORKSPACE_A) != system_provider_uuid(WORKSPACE_B)
|
||||||
|
assert system_model_uuid(WORKSPACE_A, 'chat', 'upstream') != system_model_uuid(WORKSPACE_B, 'chat', 'upstream')
|
||||||
|
snapshot = _snapshot()
|
||||||
|
assert 'owner-a-key' not in repr(snapshot)
|
||||||
|
|
||||||
|
|
||||||
|
async def test_snapshot_must_cover_every_active_workspace() -> None:
|
||||||
|
snapshot = _snapshot().model_copy(update={'workspaces': _snapshot().workspaces[:1]})
|
||||||
|
app = SimpleNamespace(
|
||||||
|
workspace_service=SimpleNamespace(
|
||||||
|
list_active_execution_bindings=lambda: _async_value(
|
||||||
|
[SimpleNamespace(workspace_uuid=WORKSPACE_A), SimpleNamespace(workspace_uuid=WORKSPACE_B)]
|
||||||
|
)
|
||||||
|
),
|
||||||
|
logger=logging.getLogger(__name__),
|
||||||
|
)
|
||||||
|
service = CloudModelCatalogSyncService(app, _CatalogProvider(snapshot), INSTANCE_UUID)
|
||||||
|
with pytest.raises(ValueError, match='missing billing projections for 1 active Workspaces'):
|
||||||
|
await service.sync_once()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_periodic_sync_discovers_workspace_created_after_startup_cache_release(tmp_path) -> None:
|
||||||
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog-new-workspace.db"}')
|
||||||
|
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
|
||||||
|
manager.db = SimpleNamespace(get_engine=lambda: engine)
|
||||||
|
startup_bindings = [
|
||||||
|
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1)
|
||||||
|
]
|
||||||
|
live_bindings = [
|
||||||
|
*startup_bindings,
|
||||||
|
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1),
|
||||||
|
]
|
||||||
|
|
||||||
|
class _WorkspaceService:
|
||||||
|
startup_released = False
|
||||||
|
|
||||||
|
async def list_active_execution_bindings(self):
|
||||||
|
return list(live_bindings if self.startup_released else startup_bindings)
|
||||||
|
|
||||||
|
def release_startup_execution_bindings(self):
|
||||||
|
self.startup_released = True
|
||||||
|
|
||||||
|
workspace_service = _WorkspaceService()
|
||||||
|
app = SimpleNamespace(
|
||||||
|
persistence_mgr=manager,
|
||||||
|
workspace_service=workspace_service,
|
||||||
|
model_mgr=SimpleNamespace(load_models_from_db=_AsyncCounter()),
|
||||||
|
logger=logging.getLogger(__name__),
|
||||||
|
)
|
||||||
|
service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID)
|
||||||
|
|
||||||
|
try:
|
||||||
|
async with engine.begin() as connection:
|
||||||
|
await connection.run_sync(Base.metadata.create_all)
|
||||||
|
await connection.execute(
|
||||||
|
sqlalchemy.insert(Workspace),
|
||||||
|
[
|
||||||
|
{
|
||||||
|
'uuid': WORKSPACE_A,
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'name': 'A',
|
||||||
|
'slug': 'a',
|
||||||
|
'source': 'cloud_projection',
|
||||||
|
},
|
||||||
|
{
|
||||||
|
'uuid': WORKSPACE_B,
|
||||||
|
'instance_uuid': INSTANCE_UUID,
|
||||||
|
'name': 'B',
|
||||||
|
'slug': 'b',
|
||||||
|
'source': 'cloud_projection',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
await service.initialize()
|
||||||
|
workspace_service.release_startup_execution_bindings()
|
||||||
|
await service.sync_once()
|
||||||
|
|
||||||
|
async with engine.connect() as connection:
|
||||||
|
provider_b = await connection.scalar(
|
||||||
|
sqlalchemy.select(ModelProvider).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_B))
|
||||||
|
)
|
||||||
|
assert provider_b is not None
|
||||||
|
finally:
|
||||||
|
await engine.dispose()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_catalog_run_wakes_immediately_when_directory_changes() -> None:
|
||||||
|
sync_started = asyncio.Event()
|
||||||
|
|
||||||
|
class _WakeService(CloudModelCatalogSyncService):
|
||||||
|
async def sync_once(self, *, reload_runtime: bool = True):
|
||||||
|
del reload_runtime
|
||||||
|
sync_started.set()
|
||||||
|
return {'workspaces': 0, 'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
|
||||||
|
app = SimpleNamespace(logger=logging.getLogger(__name__))
|
||||||
|
service = _WakeService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID, sync_interval_seconds=3600)
|
||||||
|
task = asyncio.create_task(service.run())
|
||||||
|
try:
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
service.request_sync()
|
||||||
|
await asyncio.wait_for(sync_started.wait(), timeout=0.2)
|
||||||
|
finally:
|
||||||
|
task.cancel()
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
|
||||||
|
|
||||||
|
async def _async_value(value):
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
class _AsyncCounter:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.calls = 0
|
||||||
|
|
||||||
|
async def __call__(self) -> None:
|
||||||
|
self.calls += 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_partial_workspace_failure_reloads_already_committed_changes() -> None:
|
||||||
|
bindings = [
|
||||||
|
SimpleNamespace(workspace_uuid=WORKSPACE_A),
|
||||||
|
SimpleNamespace(workspace_uuid=WORKSPACE_B),
|
||||||
|
]
|
||||||
|
reload_counter = _AsyncCounter()
|
||||||
|
app = SimpleNamespace(
|
||||||
|
workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)),
|
||||||
|
model_mgr=SimpleNamespace(load_models_from_db=reload_counter),
|
||||||
|
logger=logging.getLogger(__name__),
|
||||||
|
)
|
||||||
|
service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID)
|
||||||
|
calls = 0
|
||||||
|
|
||||||
|
async def sync_workspace(*_args):
|
||||||
|
nonlocal calls
|
||||||
|
calls += 1
|
||||||
|
if calls == 1:
|
||||||
|
return {'created': 1, 'updated': 0, 'deleted': 0}
|
||||||
|
raise RuntimeError('second Workspace failed')
|
||||||
|
|
||||||
|
service._sync_workspace = sync_workspace # type: ignore[method-assign]
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match='second Workspace failed'):
|
||||||
|
await service.sync_once()
|
||||||
|
assert service.get_workspace_credits(WORKSPACE_A) == 25000
|
||||||
|
assert service.get_workspace_credits(WORKSPACE_B) is None
|
||||||
|
assert reload_counter.calls == 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_failed_runtime_reload_is_retried_after_noop_sync() -> None:
|
||||||
|
bindings = [SimpleNamespace(workspace_uuid=WORKSPACE_A)]
|
||||||
|
|
||||||
|
class _FlakyReload:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.calls = 0
|
||||||
|
|
||||||
|
async def __call__(self) -> None:
|
||||||
|
self.calls += 1
|
||||||
|
if self.calls == 1:
|
||||||
|
raise RuntimeError('reload failed')
|
||||||
|
|
||||||
|
runtime_reload = _FlakyReload()
|
||||||
|
app = SimpleNamespace(
|
||||||
|
workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)),
|
||||||
|
model_mgr=SimpleNamespace(load_models_from_db=runtime_reload),
|
||||||
|
logger=logging.getLogger(__name__),
|
||||||
|
)
|
||||||
|
service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID)
|
||||||
|
calls = 0
|
||||||
|
|
||||||
|
async def sync_workspace(*_args):
|
||||||
|
nonlocal calls
|
||||||
|
calls += 1
|
||||||
|
if calls == 1:
|
||||||
|
return {'created': 1, 'updated': 0, 'deleted': 0}
|
||||||
|
return {'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
|
||||||
|
service._sync_workspace = sync_workspace # type: ignore[method-assign]
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match='reload failed'):
|
||||||
|
await service.sync_once()
|
||||||
|
summary = await service.sync_once()
|
||||||
|
assert summary == {'workspaces': 1, 'created': 0, 'updated': 0, 'deleted': 0}
|
||||||
|
assert runtime_reload.calls == 2
|
||||||
|
|
||||||
|
|
||||||
|
async def test_background_sync_log_redacts_exception_message(caplog) -> None:
|
||||||
|
secret = 'owner-secret-api-key'
|
||||||
|
attempted = asyncio.Event()
|
||||||
|
|
||||||
|
class _FailingProvider:
|
||||||
|
async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot:
|
||||||
|
del instance_uuid
|
||||||
|
attempted.set()
|
||||||
|
raise RuntimeError(f'database parameters include {secret}')
|
||||||
|
|
||||||
|
app = SimpleNamespace(logger=logging.getLogger(__name__))
|
||||||
|
service = CloudModelCatalogSyncService(app, _FailingProvider(), INSTANCE_UUID)
|
||||||
|
service.sync_interval_seconds = 0.001
|
||||||
|
task = asyncio.create_task(service.run())
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(attempted.wait(), timeout=1)
|
||||||
|
await asyncio.sleep(0.01)
|
||||||
|
finally:
|
||||||
|
task.cancel()
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await task
|
||||||
|
|
||||||
|
assert secret not in caplog.text
|
||||||
|
assert 'Cloud model catalog synchronization failed (RuntimeError)' in caplog.text
|
||||||
@@ -11,6 +11,7 @@ from cryptography.hazmat.primitives import serialization
|
|||||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
|
||||||
|
|
||||||
from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService
|
from langbot.pkg.cloud.launch import SpaceLaunchError, SpaceLaunchService
|
||||||
|
from langbot.pkg.cloud.support_admin import SupportAdminReplayError
|
||||||
|
|
||||||
|
|
||||||
pytestmark = pytest.mark.asyncio
|
pytestmark = pytest.mark.asyncio
|
||||||
@@ -48,7 +49,6 @@ def _claims(*, now: int, jti: str | None = None, workspace_uuid: str = WORKSPACE
|
|||||||
'payload': {
|
'payload': {
|
||||||
'account_uuid': ACCOUNT_UUID,
|
'account_uuid': ACCOUNT_UUID,
|
||||||
'workspace_uuid': workspace_uuid,
|
'workspace_uuid': workspace_uuid,
|
||||||
'return_path': '/',
|
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -58,9 +58,21 @@ def _service(private_key: Ed25519PrivateKey, *, now: int) -> SpaceLaunchService:
|
|||||||
encoding=serialization.Encoding.Raw,
|
encoding=serialization.Encoding.Raw,
|
||||||
format=serialization.PublicFormat.Raw,
|
format=serialization.PublicFormat.Raw,
|
||||||
)
|
)
|
||||||
|
consumed: set[str] = set()
|
||||||
|
|
||||||
|
class DurableSupportAdminService:
|
||||||
|
async def consume_launch_grant(self, **kwargs):
|
||||||
|
grant_hash = kwargs['grant_jti_hash']
|
||||||
|
if grant_hash in consumed:
|
||||||
|
raise SupportAdminReplayError('already consumed')
|
||||||
|
consumed.add(grant_hash)
|
||||||
|
return SimpleNamespace(token='support-admin-token')
|
||||||
|
|
||||||
app = SimpleNamespace(
|
app = SimpleNamespace(
|
||||||
deployment=SimpleNamespace(multi_workspace_enabled=True, verification_key_id=KEY_ID),
|
deployment=SimpleNamespace(multi_workspace_enabled=True, verification_key_id=KEY_ID),
|
||||||
workspace_service=SimpleNamespace(instance_uuid=INSTANCE_UUID),
|
workspace_service=SimpleNamespace(instance_uuid=INSTANCE_UUID),
|
||||||
|
logger=SimpleNamespace(info=lambda *args, **kwargs: None),
|
||||||
|
support_admin_session_service=DurableSupportAdminService(),
|
||||||
instance_config=SimpleNamespace(
|
instance_config=SimpleNamespace(
|
||||||
data={
|
data={
|
||||||
'space': {
|
'space': {
|
||||||
@@ -82,30 +94,85 @@ async def test_consumes_valid_workspace_launch_assertion_once():
|
|||||||
|
|
||||||
launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
assert launch == {
|
assert launch == {'account_uuid': ACCOUNT_UUID, 'workspace_uuid': WORKSPACE_UUID}
|
||||||
'account_uuid': ACCOUNT_UUID,
|
|
||||||
'workspace_uuid': WORKSPACE_UUID,
|
|
||||||
'return_path': '/',
|
|
||||||
}
|
|
||||||
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
||||||
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
|
||||||
async def test_consumed_assertion_remains_blocked_through_clock_skew_window():
|
|
||||||
|
|
||||||
|
async def test_consumes_admin_owner_launch_once_and_validates_claims():
|
||||||
private_key = Ed25519PrivateKey.generate()
|
private_key = Ed25519PrivateKey.generate()
|
||||||
now = int(time.time())
|
now = int(time.time())
|
||||||
service = _service(private_key, now=now)
|
service = _service(private_key, now=now)
|
||||||
claims = _claims(now=now)
|
claims = _claims(now=now)
|
||||||
claims['iat'] = now - 10
|
claims['kind'] = 'workspace.support_admin_launch'
|
||||||
claims['nbf'] = now - 10
|
claims['payload'].update(
|
||||||
claims['exp'] = now - 1
|
{
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': '33333333-3333-4333-8333-333333333333',
|
||||||
|
'effective_role': 'owner',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
claims['payload'].pop('account_uuid')
|
||||||
token = _sign(private_key, claims)
|
token = _sign(private_key, claims)
|
||||||
|
|
||||||
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
launch = await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
assert launch == {
|
||||||
|
'workspace_uuid': WORKSPACE_UUID,
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'actor_account_uuid': '33333333-3333-4333-8333-333333333333',
|
||||||
|
'effective_role': 'owner',
|
||||||
|
'grant_jti_hash': launch['grant_jti_hash'],
|
||||||
|
'support_admin_token': 'support-admin-token',
|
||||||
|
}
|
||||||
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
with pytest.raises(SpaceLaunchError, match='already been consumed'):
|
||||||
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
invalid = _claims(now=now)
|
||||||
|
invalid['kind'] = 'workspace.support_admin_launch'
|
||||||
|
invalid['payload'].update(
|
||||||
|
{
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': '33333333-3333-4333-8333-333333333333',
|
||||||
|
'effective_role': 'member',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
invalid['payload'].pop('account_uuid')
|
||||||
|
with pytest.raises(SpaceLaunchError, match='effective role'):
|
||||||
|
await service.consume_assertion(_sign(private_key, invalid), expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
too_long = _claims(now=now)
|
||||||
|
too_long['kind'] = 'workspace.support_admin_launch'
|
||||||
|
too_long['exp'] = now + 91
|
||||||
|
too_long['payload'].update(
|
||||||
|
{
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': '33333333-3333-4333-8333-333333333333',
|
||||||
|
'effective_role': 'owner',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
too_long['payload'].pop('account_uuid')
|
||||||
|
with pytest.raises(SpaceLaunchError, match='lifetime exceeds 90 seconds'):
|
||||||
|
await service.consume_assertion(_sign(private_key, too_long), expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
impersonating = _claims(now=now)
|
||||||
|
impersonating['kind'] = 'workspace.support_admin_launch'
|
||||||
|
impersonating['payload'].update(
|
||||||
|
{
|
||||||
|
'launch_mode': 'support_admin',
|
||||||
|
'principal_type': 'support_admin',
|
||||||
|
'actor_account_uuid': '33333333-3333-4333-8333-333333333333',
|
||||||
|
'effective_role': 'owner',
|
||||||
|
}
|
||||||
|
)
|
||||||
|
with pytest.raises(SpaceLaunchError, match='customer Account'):
|
||||||
|
await service.consume_assertion(_sign(private_key, impersonating), expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
|
||||||
async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch):
|
async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch):
|
||||||
private_key = Ed25519PrivateKey.generate()
|
private_key = Ed25519PrivateKey.generate()
|
||||||
@@ -183,15 +250,3 @@ async def test_rejects_invalid_signature_and_non_cloud_mode():
|
|||||||
oss_service.ap.deployment.multi_workspace_enabled = False
|
oss_service.ap.deployment.multi_workspace_enabled = False
|
||||||
with pytest.raises(SpaceLaunchError, match='verified Cloud mode'):
|
with pytest.raises(SpaceLaunchError, match='verified Cloud mode'):
|
||||||
await oss_service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
await oss_service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_rejects_unsafe_signed_return_path() -> None:
|
|
||||||
private_key = Ed25519PrivateKey.generate()
|
|
||||||
now = int(time.time())
|
|
||||||
service = _service(private_key, now=now)
|
|
||||||
claims = _claims(now=now)
|
|
||||||
claims['payload']['return_path'] = '//evil.example'
|
|
||||||
|
|
||||||
with pytest.raises(SpaceLaunchError, match='return path'):
|
|
||||||
await service.consume_assertion(_sign(private_key, claims))
|
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user