mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-13 07:22:24 -06:00
Compare commits
46 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3dc10541c6 | |||
| 03e9b565b8 | |||
| b30e1394e0 | |||
| af0bf5270c | |||
| d100ac92d9 | |||
| 57df445224 | |||
| f978e7facd | |||
| 0872f5f5ba | |||
| 205e7818f8 | |||
| caf449e048 | |||
| 2bfc0f2c5d | |||
| c67aba0127 | |||
| db0baefeb2 | |||
| 38fc933c1d | |||
| 39b39fb79d | |||
| 0923add7db | |||
| 01cec062d9 | |||
| 830eb8ba00 | |||
| 5bbf2e65eb | |||
| 4d402fea6b | |||
| 46d14ddd86 | |||
| 7c4157f78d | |||
| 3856d80709 | |||
| 8142d2f1ad | |||
| c234d66ebf | |||
| b180770eff | |||
| 9d2e11f2be | |||
| 57080f4615 | |||
| 45f27fb2a7 | |||
| ebc8e75285 | |||
| 485af92f7f | |||
| 664d44c109 | |||
| 3cf9485169 | |||
| d43b9d1647 | |||
| ea8d9d1798 | |||
| d9aa50dca9 | |||
| 6f89d0cc13 | |||
| 62d2a0fe6a | |||
| 5df37f83a7 | |||
| 651c4d98cd | |||
| e901e859c7 | |||
| 200dcfeac5 | |||
| 8c414feba2 | |||
| d7cea053b6 | |||
| c45e98462b | |||
| e17cbe35a5 |
+34
-14
@@ -1,29 +1,49 @@
|
||||
# =============================================================================
|
||||
# Turnstone Environment Variables
|
||||
# Copy to .env and adjust values for your deployment
|
||||
# Copy to .env and adjust values for your deployment.
|
||||
#
|
||||
# Usage:
|
||||
# Single node: docker compose --profile production up
|
||||
# 10-node cluster: docker compose --profile cluster up
|
||||
# =============================================================================
|
||||
|
||||
# -- LLM Backend --------------------------------------------------------------
|
||||
LLM_BASE_URL=http://host.docker.internal:8000/v1
|
||||
OPENAI_API_KEY=sk-...
|
||||
# ANTHROPIC_API_KEY=sk-ant-... # Set instead for Anthropic provider
|
||||
# TAVILY_API_KEY=tvly-... # For web search fallback (local models only)
|
||||
OPENAI_API_KEY=dummy
|
||||
# ANTHROPIC_API_KEY=sk-ant-...# Set instead of OPENAI_API_KEY for Anthropic
|
||||
# TAVILY_API_KEY=tvly-... # Web search fallback (local models only)
|
||||
# MODEL=# Override default model alias
|
||||
|
||||
# -- Database (production profile) --------------------------------------------
|
||||
# -- Authentication (required) ------------------------------------------------
|
||||
# Generate with: python -c "import secrets; print(secrets.token_hex(32))"
|
||||
TURNSTONE_JWT_SECRET=changeme-to-32-bytes-of-hex
|
||||
|
||||
# -- Database ------------------------------------------------------------------
|
||||
# Single-node default is SQLite (zero config). Set these for PostgreSQL:
|
||||
# DB_BACKEND=postgresql
|
||||
# POSTGRES_USER=turnstone
|
||||
# POSTGRES_PASSWORD=changeme
|
||||
# DATABASE_URL=postgresql+psycopg://turnstone:changeme@postgres:5432/turnstone
|
||||
|
||||
# -- Redis ---------------------------------------------------------------------
|
||||
# REDIS_PASSWORD=
|
||||
# REDIS_PORT=6379
|
||||
|
||||
# -- Authentication ------------------------------------------------------------
|
||||
# TURNSTONE_AUTH_ENABLED=true
|
||||
# TURNSTONE_AUTH_TOKEN=your-secret-token
|
||||
# TURNSTONE_JWT_SECRET=python -c "import secrets; print(secrets.token_hex(32))"
|
||||
|
||||
# -- Ports ---------------------------------------------------------------------
|
||||
# SERVER_PORT=8080
|
||||
# CONSOLE_PORT=8090
|
||||
|
||||
# -- Workspace -----------------------------------------------------------------
|
||||
# Bind-mount a host directory into the container at /workspace.
|
||||
# The model can read/write files here. Default: empty Docker volume.
|
||||
# WORKSPACE_MOUNT=/path/to/your/project
|
||||
|
||||
# -- Agent behavior ------------------------------------------------------------
|
||||
# SKIP_PERMISSIONS=true # Auto-approve all tool calls (dev only)
|
||||
# MCP_CONFIG=/workspace/mcp.json# MCP server configuration file
|
||||
|
||||
# -- Discord channel gateway ---------------------------------------------------
|
||||
# TURNSTONE_DISCORD_TOKEN=
|
||||
# TURNSTONE_DISCORD_GUILD=0
|
||||
|
||||
# -- Cluster (profile: cluster) -----------------------------------------------
|
||||
# These are set per-node in compose.yaml; only override for custom topologies.
|
||||
# TURNSTONE_NODE_ID=node-1
|
||||
# TURNSTONE_ADVERTISE_URL=http://server-1:8080
|
||||
|
||||
|
||||
+10
-3
@@ -41,6 +41,14 @@
|
||||
"matchStrings": ["mermaid-(?<currentValue>[\\d.]+)/"],
|
||||
"depNameTemplate": "mermaid",
|
||||
"datasourceTemplate": "npm"
|
||||
},
|
||||
{
|
||||
"customType": "regex",
|
||||
"description": "Track vendored hls.js version",
|
||||
"managerFilePatterns": ["/pyproject\\.toml$/"],
|
||||
"matchStrings": ["hls-(?<currentValue>[\\d.]+)/"],
|
||||
"depNameTemplate": "hls.js",
|
||||
"datasourceTemplate": "npm"
|
||||
}
|
||||
],
|
||||
"packageRules": [
|
||||
@@ -83,7 +91,7 @@
|
||||
{
|
||||
"description": "Infrastructure dependencies",
|
||||
"groupName": "Infrastructure",
|
||||
"matchPackageNames": ["structlog", "redis", "croniter", "discord.py"],
|
||||
"matchPackageNames": ["structlog", "croniter", "discord.py"],
|
||||
"schedule": ["before 9am on the first day of the month"],
|
||||
"automerge": true,
|
||||
"matchUpdateTypes": ["patch"]
|
||||
@@ -91,7 +99,7 @@
|
||||
{
|
||||
"description": "Vendored JS — CI workflow downloads files automatically",
|
||||
"groupName": "Vendored JS",
|
||||
"matchPackageNames": ["katex", "highlight.js", "mermaid"],
|
||||
"matchPackageNames": ["katex", "highlight.js", "mermaid", "hls.js"],
|
||||
"schedule": ["before 9am on the first day of the month"],
|
||||
"automerge": false
|
||||
},
|
||||
@@ -101,7 +109,6 @@
|
||||
"matchPackageNames": [
|
||||
"ruff",
|
||||
"mypy",
|
||||
"types-redis",
|
||||
"pytest",
|
||||
"pytest-cov",
|
||||
"pre-commit"
|
||||
|
||||
@@ -2,9 +2,13 @@ name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
branches: [main, "stable/*"]
|
||||
tags: ["v*"]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
branches: [main, "stable/*"]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
@@ -25,8 +29,8 @@ jobs:
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
with:
|
||||
python-version: "3.14"
|
||||
- run: pip install mypy types-redis
|
||||
- run: pip install -e ".[mq]"
|
||||
- run: pip install mypy
|
||||
- run: pip install -e ".[all]"
|
||||
- run: mypy turnstone/
|
||||
|
||||
test:
|
||||
@@ -39,7 +43,7 @@ jobs:
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- run: pip install -e ".[test,mq]"
|
||||
- run: pip install -e ".[test]"
|
||||
- run: pytest tests/ -m "not live" --cov=turnstone --cov-report=term-missing --cov-report=xml -q
|
||||
- uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7
|
||||
if: always()
|
||||
@@ -68,11 +72,58 @@ jobs:
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
with:
|
||||
python-version: "3.14"
|
||||
- run: pip install -e ".[test,mq,postgres]"
|
||||
- run: pip install -e ".[test,postgres]"
|
||||
- run: pytest tests/ -m "not live" --storage-backend=postgresql -q
|
||||
env:
|
||||
TURNSTONE_TEST_PG_URL: postgresql+psycopg://postgres:postgres@localhost:5432/turnstone_test
|
||||
|
||||
wheel-completeness:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
with:
|
||||
python-version: "3.14"
|
||||
- run: pip install build
|
||||
- run: python -m build --wheel
|
||||
- name: Check all data files are in wheel
|
||||
run: |
|
||||
SOURCE=$(find turnstone -type f \
|
||||
! -name '*.py' ! -name '*.pyc' ! -path '*__pycache__*' \
|
||||
| sort)
|
||||
WHEEL=$(python -m zipfile -l dist/*.whl \
|
||||
| awk '{print $1}' \
|
||||
| grep -v '\.py$' | grep -v '\.dist-info' | grep -v '\.pyc' | grep -v '^File$' \
|
||||
| sort)
|
||||
|
||||
# Files intentionally excluded from the wheel (one per line)
|
||||
ALLOW="
|
||||
turnstone/core/storage/migrations/script.py.mako
|
||||
"
|
||||
|
||||
MISSING=$(comm -23 <(echo "$SOURCE") <(echo "$WHEEL") \
|
||||
| grep -vFxf <(echo "$ALLOW" | sed '/^[[:space:]]*$/d; s/^[[:space:]]*//' ) || true)
|
||||
|
||||
if [ -n "$MISSING" ]; then
|
||||
echo "::error::Data files in source tree but missing from wheel:"
|
||||
echo "$MISSING"
|
||||
echo ""
|
||||
echo "Add them to [tool.hatch.build.targets.wheel] in pyproject.toml"
|
||||
echo "or to the ALLOW list in this job if intentionally excluded."
|
||||
exit 1
|
||||
fi
|
||||
echo "All source data files present in wheel"
|
||||
- name: Smoke-test entry points from installed wheel
|
||||
run: |
|
||||
python -m venv /tmp/smoke
|
||||
/tmp/smoke/bin/pip install dist/*.whl
|
||||
/tmp/smoke/bin/turnstone --help
|
||||
/tmp/smoke/bin/turnstone-server --help
|
||||
/tmp/smoke/bin/turnstone-console --help
|
||||
/tmp/smoke/bin/turnstone-admin --help
|
||||
/tmp/smoke/bin/turnstone-channel --help
|
||||
/tmp/smoke/bin/turnstone-bootstrap --help
|
||||
|
||||
lock-check:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
name: Publish Docker Image
|
||||
|
||||
on:
|
||||
workflow_run:
|
||||
workflows: ["CI"]
|
||||
types: [completed]
|
||||
|
||||
concurrency:
|
||||
group: docker-${{ github.event.workflow_run.head_sha }}
|
||||
cancel-in-progress: true
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
|
||||
env:
|
||||
REGISTRY: ghcr.io
|
||||
IMAGE_NAME: ${{ github.repository }}
|
||||
|
||||
jobs:
|
||||
docker:
|
||||
if: >-
|
||||
github.event.workflow_run.conclusion == 'success' &&
|
||||
github.event.workflow_run.head_repository.full_name == github.repository
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
|
||||
with:
|
||||
ref: ${{ github.event.workflow_run.head_sha }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Resolve release tag
|
||||
id: tag
|
||||
run: |
|
||||
TAG=$(git tag --points-at HEAD | grep '^v' | head -1)
|
||||
if [ -z "$TAG" ]; then
|
||||
echo "No v* tag at HEAD — skipping publish"
|
||||
echo "skip=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "tag=${TAG}" >> "$GITHUB_OUTPUT"
|
||||
echo "skip=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Log in to GHCR
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
uses: docker/login-action@4907a6ddec9925e35a0a9e82d7399ccc52663121 # v4
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Compute Docker tags
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
id: tags
|
||||
env:
|
||||
REF: ${{ steps.tag.outputs.tag }}
|
||||
run: |
|
||||
VERSION="${REF#v}"
|
||||
FULL="${REGISTRY}/${IMAGE_NAME}"
|
||||
FULL="${FULL,,}"
|
||||
|
||||
if echo "$VERSION" | grep -qE '(a|b|rc)[0-9]+$'; then
|
||||
TAGS="${FULL}:${VERSION},${FULL}:experimental"
|
||||
else
|
||||
MINOR="${VERSION%.*}"
|
||||
TAGS="${FULL}:${VERSION},${FULL}:${MINOR},${FULL}:stable,${FULL}:latest"
|
||||
fi
|
||||
echo "tags=${TAGS}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd # v4
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
|
||||
- name: Build and push
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
uses: docker/build-push-action@d08e5c354a6adb9ed34480a06d141179aa583294 # v7
|
||||
with:
|
||||
context: .
|
||||
push: true
|
||||
tags: ${{ steps.tags.outputs.tags }}
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
@@ -2,7 +2,7 @@ name: Docker Security Scan
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
branches: [main, "stable/*"]
|
||||
schedule:
|
||||
- cron: "0 6 * * 1" # Weekly Monday 06:00 UTC
|
||||
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
name: Publish to PyPI
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
workflow_run:
|
||||
workflows: ["CI"]
|
||||
types: [completed]
|
||||
|
||||
concurrency:
|
||||
group: publish-${{ github.event.workflow_run.head_sha }}
|
||||
cancel-in-progress: true
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
@@ -10,20 +15,43 @@ permissions:
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
if: github.event.workflow_run.conclusion == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
environment: pypi
|
||||
steps:
|
||||
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
|
||||
with:
|
||||
ref: ${{ github.event.workflow_run.head_sha }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Resolve release tag
|
||||
id: tag
|
||||
run: |
|
||||
TAG=$(git tag --points-at HEAD | grep '^v' | head -1)
|
||||
if [ -z "$TAG" ]; then
|
||||
echo "No v* tag at HEAD — skipping publish"
|
||||
echo "skip=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "tag=${TAG}" >> "$GITHUB_OUTPUT"
|
||||
echo "skip=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
with:
|
||||
python-version: "3.14"
|
||||
- run: pip install build
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
- run: python -m build
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
- uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e # release/v1
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
|
||||
- name: Create GitHub Release
|
||||
if: steps.tag.outputs.skip == 'false'
|
||||
uses: softprops/action-gh-release@153bb8e04406b158c6c84fc1615b65b24149a1fe # v2
|
||||
with:
|
||||
tag_name: ${{ steps.tag.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
draft: false
|
||||
prerelease: ${{ contains(github.ref, '-') }}
|
||||
prerelease: ${{ contains(steps.tag.outputs.tag, 'a') || contains(steps.tag.outputs.tag, 'b') || contains(steps.tag.outputs.tag, 'rc') }}
|
||||
|
||||
@@ -17,3 +17,24 @@ CVE-2026-27135
|
||||
# Affects libsystemd0, libudev1
|
||||
# https://avd.aquasec.com/nvd/cve-2026-29111
|
||||
CVE-2026-29111
|
||||
|
||||
# glibc iconv() DoS — fix_deferred, no patched libc in Debian 13 yet
|
||||
# Affects libc-bin, libc6
|
||||
# https://avd.aquasec.com/nvd/cve-2026-4046
|
||||
CVE-2026-4046
|
||||
|
||||
# minimatch ReDoS — transitive npm dep (MCP server), no direct exposure
|
||||
# https://avd.aquasec.com/nvd/cve-2026-27903
|
||||
CVE-2026-27903
|
||||
# https://avd.aquasec.com/nvd/cve-2026-27904
|
||||
CVE-2026-27904
|
||||
|
||||
# picomatch ReDoS — transitive npm dep, no direct exposure
|
||||
# https://avd.aquasec.com/nvd/cve-2026-33671
|
||||
CVE-2026-33671
|
||||
|
||||
# node-tar path traversal — transitive npm dep, not used to extract untrusted archives
|
||||
# https://avd.aquasec.com/nvd/cve-2026-29786
|
||||
CVE-2026-29786
|
||||
# https://avd.aquasec.com/nvd/cve-2026-31802
|
||||
CVE-2026-31802
|
||||
|
||||
+10
-1
@@ -8,7 +8,7 @@ FROM python:3.14-slim
|
||||
LABEL org.opencontainers.image.title="turnstone" \
|
||||
org.opencontainers.image.description="Multi-node AI orchestration platform"
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.11.2 /uv /usr/local/bin/uv
|
||||
COPY --from=ghcr.io/astral-sh/uv:0.11.3 /uv /usr/local/bin/uv
|
||||
|
||||
# Remove the slim image's man page exclusion so man-db has actual content
|
||||
RUN rm -f /etc/dpkg/dpkg.cfg.d/docker
|
||||
@@ -18,6 +18,12 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco
|
||||
libpq5 git curl jq man-db manpages procps file \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Node.js LTS (for npx-based MCP servers like @modelcontextprotocol/server-github)
|
||||
COPY --from=node:24-slim /usr/local/bin/node /usr/local/bin/node
|
||||
COPY --from=node:24-slim /usr/local/lib/node_modules /usr/local/lib/node_modules
|
||||
RUN ln -s ../lib/node_modules/npm/bin/npm-cli.js /usr/local/bin/npm \
|
||||
&& ln -s ../lib/node_modules/npm/bin/npx-cli.js /usr/local/bin/npx
|
||||
|
||||
# Non-root user
|
||||
RUN useradd --create-home --shell /bin/bash turnstone
|
||||
|
||||
@@ -49,6 +55,9 @@ COPY docker/entrypoint.sh /usr/local/bin/entrypoint.sh
|
||||
WORKDIR /data
|
||||
RUN chown turnstone:turnstone /data
|
||||
|
||||
# Workspace mount point — bind-mount a host directory here
|
||||
RUN mkdir -p /workspace && chown turnstone:turnstone /workspace
|
||||
|
||||
USER turnstone
|
||||
|
||||
ENTRYPOINT ["entrypoint.sh"]
|
||||
|
||||
@@ -7,10 +7,21 @@
|
||||
|
||||
Multi-node AI orchestration platform. Deploy tool-using AI agents across a cluster of servers with direct HTTP routing, interactive interfaces, and enterprise governance.
|
||||
|
||||
> **Beta — Use at your own risk.** APIs, configuration formats, and database schemas may change between versions without migration paths.
|
||||
<p align="center">
|
||||
<img src="docs/assets/hero.png" alt="Turnstone console — multi-workstream AI orchestration with mermaid diagrams" width="960"/>
|
||||
</p>
|
||||
|
||||
Named after the [Ruddy Turnstone](https://en.wikipedia.org/wiki/Ruddy_turnstone) (*Arenaria interpres*) — a shorebird that flips stones to discover what's hiding underneath.
|
||||
|
||||
### Release Tracks
|
||||
|
||||
| Track | Install | Docker | Description |
|
||||
|-------|---------|--------|-------------|
|
||||
| **Stable** | `pip install turnstone` | `ghcr.io/turnstonelabs/turnstone:stable` | Production-grade. Bugfixes only. |
|
||||
| **Experimental** | `pip install turnstone --pre` | `ghcr.io/turnstonelabs/turnstone:experimental` | New features. May have rough edges. |
|
||||
|
||||
See [docs/releasing.md](docs/releasing.md) for the full release process.
|
||||
|
||||
## What it does
|
||||
|
||||
Turnstone gives LLMs tools — shell, files, search, web, planning — and orchestrates multi-turn conversations where the model investigates, acts, and reports.
|
||||
|
||||
@@ -95,3 +95,22 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
================================================================================
|
||||
|
||||
hls.js 1.6.15
|
||||
https://github.com/video-dev/hls.js
|
||||
|
||||
Copyright 2017 Dailymotion
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
|
||||
+23
-63
@@ -1,12 +1,16 @@
|
||||
# =============================================================================
|
||||
# Turnstone Docker Compose Stack
|
||||
# Turnstone Docker Compose Stack — Development
|
||||
#
|
||||
# This file is for local development from a git clone. It builds images
|
||||
# locally from the Dockerfile. If you installed via pip/pipx, run
|
||||
# `turnstone-bootstrap` instead — it writes a production compose.yaml
|
||||
# that pulls pre-built images from ghcr.io.
|
||||
#
|
||||
# Usage:
|
||||
# Infra only: docker compose up
|
||||
# Single node: docker compose --profile production up
|
||||
# Production (PG): DB_BACKEND=postgresql docker compose --profile production up
|
||||
# 10-node cluster: docker compose --profile cluster up
|
||||
# Cluster + DDG: docker compose --profile ddgCluster up
|
||||
# =============================================================================
|
||||
|
||||
name: turnstone
|
||||
@@ -17,6 +21,7 @@ networks:
|
||||
|
||||
volumes:
|
||||
turnstone-data:
|
||||
workspace:
|
||||
postgres-data:
|
||||
|
||||
services:
|
||||
@@ -28,7 +33,6 @@ services:
|
||||
profiles:
|
||||
- production
|
||||
- cluster
|
||||
- ddgCluster
|
||||
command:
|
||||
- postgres
|
||||
- -c
|
||||
@@ -61,9 +65,7 @@ services:
|
||||
# turnstone-server — Web UI + chat workstreams + LLM interaction
|
||||
# -------------------------------------------------------------------
|
||||
server:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
image: turnstone:local
|
||||
profiles:
|
||||
- production
|
||||
command:
|
||||
@@ -82,15 +84,14 @@ services:
|
||||
- "${SERVER_PORT:-8080}:8080"
|
||||
volumes:
|
||||
- turnstone-data:/data
|
||||
- ./docker/mcp-ddg.json:/etc/turnstone/mcp-ddg.json:ro
|
||||
- ${WORKSPACE_MOUNT:-workspace}:/workspace
|
||||
environment:
|
||||
- LLM_BASE_URL=${LLM_BASE_URL:-http://host.docker.internal:8000/v1}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-dummy}
|
||||
- TAVILY_API_KEY=${TAVILY_API_KEY:-}
|
||||
- SKIP_PERMISSIONS=${SKIP_PERMISSIONS:-}
|
||||
- TURNSTONE_AUTH_ENABLED=${TURNSTONE_AUTH_ENABLED:-}
|
||||
- TURNSTONE_AUTH_TOKEN=${TURNSTONE_AUTH_TOKEN:-}
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:-}
|
||||
# Generate with: python -c "import secrets; print(secrets.token_hex(32))"
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:?Set TURNSTONE_JWT_SECRET in .env}
|
||||
- MODEL=${MODEL:-}
|
||||
- MCP_CONFIG=${MCP_CONFIG:-}
|
||||
- TURNSTONE_DB_BACKEND=${DB_BACKEND:-sqlite}
|
||||
@@ -105,9 +106,6 @@ services:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
ddg-search:
|
||||
condition: service_healthy
|
||||
required: false
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "/usr/local/bin/healthcheck.py", "http://127.0.0.1:8080/health"]
|
||||
interval: 10s
|
||||
@@ -120,6 +118,7 @@ services:
|
||||
# turnstone-console — Cluster dashboard
|
||||
# -------------------------------------------------------------------
|
||||
console:
|
||||
image: turnstone:local
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
@@ -130,9 +129,8 @@ services:
|
||||
ports:
|
||||
- "${CONSOLE_PORT:-8090}:8090"
|
||||
environment:
|
||||
- TURNSTONE_AUTH_ENABLED=${TURNSTONE_AUTH_ENABLED:-}
|
||||
- TURNSTONE_AUTH_TOKEN=${TURNSTONE_AUTH_TOKEN:-}
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:-}
|
||||
# Generate with: python -c "import secrets; print(secrets.token_hex(32))"
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:?Set TURNSTONE_JWT_SECRET in .env}
|
||||
- TURNSTONE_DB_BACKEND=${DB_BACKEND:-sqlite}
|
||||
- TURNSTONE_DB_URL=${DATABASE_URL:-}
|
||||
- TURNSTONE_CONSOLE_URL=http://console:8090
|
||||
@@ -151,13 +149,10 @@ services:
|
||||
# Requires TURNSTONE_DISCORD_TOKEN to enable Discord adapter
|
||||
# -------------------------------------------------------------------
|
||||
channel:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
image: turnstone:local
|
||||
profiles:
|
||||
- production
|
||||
- cluster
|
||||
- ddgCluster
|
||||
command:
|
||||
- sh
|
||||
- -c
|
||||
@@ -168,10 +163,10 @@ services:
|
||||
environment:
|
||||
- TURNSTONE_DISCORD_TOKEN=${TURNSTONE_DISCORD_TOKEN:-}
|
||||
- TURNSTONE_DISCORD_GUILD=${TURNSTONE_DISCORD_GUILD:-0}
|
||||
- TURNSTONE_AUTH_TOKEN=${TURNSTONE_AUTH_TOKEN:-}
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:-}
|
||||
# Generate with: python -c "import secrets; print(secrets.token_hex(32))"
|
||||
- TURNSTONE_JWT_SECRET=${TURNSTONE_JWT_SECRET:?Set TURNSTONE_JWT_SECRET in .env}
|
||||
- TURNSTONE_DB_BACKEND=${DB_BACKEND:-postgresql}
|
||||
- TURNSTONE_DB_URL=${DATABASE_URL:-postgresql://${POSTGRES_USER:-turnstone}:${POSTGRES_PASSWORD:-turnstone}@postgres:5432/turnstone}
|
||||
- TURNSTONE_DB_URL=${DATABASE_URL:-postgresql+psycopg://${POSTGRES_USER:-turnstone}:${POSTGRES_PASSWORD:-turnstone}@postgres:5432/turnstone}
|
||||
- TURNSTONE_CHANNEL_ADVERTISE_URL=http://channel:8091
|
||||
networks:
|
||||
- turnstone-net
|
||||
@@ -181,39 +176,6 @@ services:
|
||||
required: false
|
||||
restart: unless-stopped
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# ddg-search — DuckDuckGo Search MCP server (HTTP transport)
|
||||
# Provides web search + content fetch tools to turnstone via MCP.
|
||||
# No API key required.
|
||||
#
|
||||
# Start with: MCP_CONFIG=/etc/turnstone/mcp-ddg.json \
|
||||
# docker compose --profile ddgCluster up
|
||||
# -------------------------------------------------------------------
|
||||
ddg-search:
|
||||
image: python:3.14-slim
|
||||
profiles:
|
||||
- ddgCluster
|
||||
command:
|
||||
- sh
|
||||
- -c
|
||||
- >-
|
||||
pip install --no-cache-dir duckduckgo-mcp-server &&
|
||||
python -c "from mcp.server.transport_security import TransportSecuritySettings; import duckduckgo_mcp_server.server as s; s.safe_search=s.SafeSearchMode.OFF; s.mcp.settings.host='0.0.0.0'; s.mcp.settings.port=3000; s.mcp.settings.transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False); s.mcp.run(transport='streamable-http')"
|
||||
networks:
|
||||
- turnstone-net
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "python -c \"import socket; s=socket.create_connection(('0.0.0.0',3000),2); s.close()\""]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
memory: 256M
|
||||
cpus: '0.25'
|
||||
restart: unless-stopped
|
||||
|
||||
# ===================================================================
|
||||
# 10-node cluster (profile: cluster)
|
||||
#
|
||||
@@ -228,7 +190,7 @@ services:
|
||||
server-1: &cluster-server
|
||||
image: turnstone:local
|
||||
build: { context: ., dockerfile: Dockerfile }
|
||||
profiles: [cluster, ddgCluster]
|
||||
profiles: [cluster]
|
||||
command:
|
||||
- sh
|
||||
- -c
|
||||
@@ -243,26 +205,24 @@ services:
|
||||
$${MCP_CONFIG:+--mcp-config $$MCP_CONFIG}
|
||||
volumes:
|
||||
- turnstone-data:/data
|
||||
- ./docker/mcp-ddg.json:/etc/turnstone/mcp-ddg.json:ro
|
||||
- ${WORKSPACE_MOUNT:-workspace}:/workspace
|
||||
environment: &cluster-server-env
|
||||
LLM_BASE_URL: ${LLM_BASE_URL:-http://host.docker.internal:8000/v1}
|
||||
OPENAI_API_KEY: ${OPENAI_API_KEY:-dummy}
|
||||
TAVILY_API_KEY: ${TAVILY_API_KEY:-}
|
||||
SKIP_PERMISSIONS: ${SKIP_PERMISSIONS:-}
|
||||
TURNSTONE_AUTH_ENABLED: ${TURNSTONE_AUTH_ENABLED:-}
|
||||
TURNSTONE_AUTH_TOKEN: ${TURNSTONE_AUTH_TOKEN:-}
|
||||
TURNSTONE_JWT_SECRET: ${TURNSTONE_JWT_SECRET:-}
|
||||
# Generate with: python -c "import secrets; print(secrets.token_hex(32))"
|
||||
TURNSTONE_JWT_SECRET: ${TURNSTONE_JWT_SECRET:?Set TURNSTONE_JWT_SECRET in .env}
|
||||
MODEL: ${MODEL:-}
|
||||
MCP_CONFIG: ${MCP_CONFIG:-}
|
||||
TURNSTONE_DB_BACKEND: ${DB_BACKEND:-postgresql}
|
||||
TURNSTONE_DB_URL: ${DATABASE_URL:-postgresql://${POSTGRES_USER:-turnstone}:${POSTGRES_PASSWORD:?}@postgres:5432/turnstone}
|
||||
TURNSTONE_DB_URL: ${DATABASE_URL:-postgresql+psycopg://${POSTGRES_USER:-turnstone}:${POSTGRES_PASSWORD:?}@postgres:5432/turnstone}
|
||||
TURNSTONE_NODE_ID: node-1
|
||||
TURNSTONE_ADVERTISE_URL: http://server-1:8080
|
||||
extra_hosts: ["host.docker.internal:host-gateway"]
|
||||
networks: [turnstone-net]
|
||||
depends_on:
|
||||
postgres: { condition: service_healthy }
|
||||
ddg-search: { condition: service_healthy, required: false }
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "/usr/local/bin/healthcheck.py", "http://127.0.0.1:8080/health"]
|
||||
interval: 10s
|
||||
|
||||
@@ -36,13 +36,13 @@ spec:
|
||||
- secretRef:
|
||||
name: {{ include "turnstone.llm.secretName" . }}
|
||||
optional: true
|
||||
{{- if and .Values.auth.enabled .Values.auth.existingSecret }}
|
||||
{{- if or .Values.auth.existingSecret .Values.auth.jwtSecret }}
|
||||
env:
|
||||
- name: TURNSTONE_AUTH_TOKEN
|
||||
- name: TURNSTONE_JWT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.auth.existingSecret }}
|
||||
key: TURNSTONE_AUTH_TOKEN
|
||||
name: {{ include "turnstone.auth.secretName" . }}
|
||||
key: TURNSTONE_JWT_SECRET
|
||||
{{- end }}
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
|
||||
@@ -41,12 +41,12 @@ spec:
|
||||
env:
|
||||
- name: TURNSTONE_DB_URL
|
||||
value: "postgresql+psycopg://$(TURNSTONE_DB_USER):$(POSTGRES_PASSWORD)@$(TURNSTONE_DB_HOST):$(TURNSTONE_DB_PORT)/$(TURNSTONE_DB_NAME)"
|
||||
{{- if and .Values.auth.enabled .Values.auth.existingSecret }}
|
||||
- name: TURNSTONE_AUTH_TOKEN
|
||||
{{- if or .Values.auth.existingSecret .Values.auth.jwtSecret }}
|
||||
- name: TURNSTONE_JWT_SECRET
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.auth.existingSecret }}
|
||||
key: TURNSTONE_AUTH_TOKEN
|
||||
name: {{ include "turnstone.auth.secretName" . }}
|
||||
key: TURNSTONE_JWT_SECRET
|
||||
{{- end }}
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
|
||||
@@ -15,7 +15,7 @@ data:
|
||||
{{- else if and (not .Values.postgresql.enabled) .Values.database.external.password }}
|
||||
POSTGRES_PASSWORD: {{ .Values.database.external.password | b64enc | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.auth.enabled .Values.auth.token (not .Values.auth.existingSecret) }}
|
||||
TURNSTONE_AUTH_TOKEN: {{ .Values.auth.token | b64enc | quote }}
|
||||
{{- if and .Values.auth.jwtSecret (not .Values.auth.existingSecret) }}
|
||||
TURNSTONE_JWT_SECRET: {{ .Values.auth.jwtSecret | b64enc | quote }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
@@ -59,10 +59,9 @@ llm:
|
||||
apiKey: ""
|
||||
existingSecret: ""
|
||||
|
||||
# -- Authentication
|
||||
# -- Authentication (always enabled, JWT secret required)
|
||||
auth:
|
||||
enabled: false
|
||||
token: ""
|
||||
jwtSecret: ""
|
||||
existingSecret: ""
|
||||
|
||||
# -- Ingress configuration
|
||||
|
||||
@@ -40,8 +40,8 @@ resource "aws_iam_role_policy" "ecs_execution_secrets" {
|
||||
[
|
||||
aws_secretsmanager_secret.openai_api_key.arn,
|
||||
aws_secretsmanager_secret.db_password.arn,
|
||||
aws_secretsmanager_secret.jwt_secret.arn,
|
||||
],
|
||||
var.auth_token != "" ? [aws_secretsmanager_secret.auth_token[0].arn] : [],
|
||||
)
|
||||
},
|
||||
]
|
||||
|
||||
@@ -41,20 +41,26 @@ locals {
|
||||
},
|
||||
]
|
||||
|
||||
auth_env = var.auth_token != "" ? [
|
||||
{ name = "TURNSTONE_AUTH_ENABLED", value = "true" },
|
||||
] : []
|
||||
|
||||
auth_secrets = var.auth_token != "" ? [
|
||||
auth_secrets = [
|
||||
{
|
||||
name = "TURNSTONE_AUTH_TOKEN"
|
||||
valueFrom = aws_secretsmanager_secret_version.auth_token[0].arn
|
||||
name = "TURNSTONE_JWT_SECRET"
|
||||
valueFrom = aws_secretsmanager_secret_version.jwt_secret.arn
|
||||
},
|
||||
] : []
|
||||
]
|
||||
}
|
||||
|
||||
# ---------- Secrets Manager ----------
|
||||
|
||||
resource "aws_secretsmanager_secret" "jwt_secret" {
|
||||
name = "${var.name_prefix}-${var.environment}-jwt-secret"
|
||||
tags = local.common_tags
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret_version" "jwt_secret" {
|
||||
secret_id = aws_secretsmanager_secret.jwt_secret.id
|
||||
secret_string = var.jwt_secret
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "openai_api_key" {
|
||||
name = "${var.name_prefix}-${var.environment}-openai-api-key"
|
||||
tags = local.common_tags
|
||||
@@ -65,17 +71,7 @@ resource "aws_secretsmanager_secret_version" "openai_api_key" {
|
||||
secret_string = var.openai_api_key
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "auth_token" {
|
||||
count = var.auth_token != "" ? 1 : 0
|
||||
name = "${var.name_prefix}-${var.environment}-auth-token"
|
||||
tags = local.common_tags
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret_version" "auth_token" {
|
||||
count = var.auth_token != "" ? 1 : 0
|
||||
secret_id = aws_secretsmanager_secret.auth_token[0].id
|
||||
secret_string = var.auth_token
|
||||
}
|
||||
|
||||
resource "aws_secretsmanager_secret" "db_password" {
|
||||
name = "${var.name_prefix}-${var.environment}-db-password"
|
||||
@@ -140,7 +136,7 @@ resource "aws_ecs_task_definition" "server" {
|
||||
{ containerPort = 8080, protocol = "tcp" },
|
||||
]
|
||||
|
||||
environment = concat(local.common_env, local.auth_env)
|
||||
environment = local.common_env
|
||||
secrets = concat(local.common_secrets, local.auth_secrets)
|
||||
|
||||
logConfiguration = {
|
||||
@@ -209,7 +205,7 @@ resource "aws_ecs_task_definition" "console" {
|
||||
{ containerPort = 8090, protocol = "tcp" },
|
||||
]
|
||||
|
||||
environment = concat(local.common_env, local.auth_env)
|
||||
environment = local.common_env
|
||||
secrets = concat(local.common_secrets, local.auth_secrets)
|
||||
|
||||
logConfiguration = {
|
||||
|
||||
@@ -90,11 +90,10 @@ variable "name_prefix" {
|
||||
default = "turnstone"
|
||||
}
|
||||
|
||||
variable "auth_token" {
|
||||
description = "Optional authentication token for the Turnstone API. Empty string disables auth."
|
||||
variable "jwt_secret" {
|
||||
description = "JWT signing secret for Turnstone auth (required, min 32 characters)."
|
||||
type = string
|
||||
sensitive = true
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "certificate_arn" {
|
||||
|
||||
@@ -1,7 +0,0 @@
|
||||
{
|
||||
"mcpServers": {
|
||||
"ddg": {
|
||||
"url": "http://ddg-search:3000/mcp"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -56,7 +56,7 @@ console.log(result.content);
|
||||
|
||||
## Authentication
|
||||
|
||||
When auth is enabled (`[auth].enabled = true` or `TURNSTONE_AUTH_ENABLED=1`), all API endpoints except public paths require a valid token.
|
||||
Auth is always enabled. All API endpoints except public paths require a valid token.
|
||||
|
||||
### Sending Credentials
|
||||
|
||||
@@ -65,15 +65,14 @@ Include a token in one of two ways:
|
||||
- **Bearer header**: `Authorization: Bearer <token>`
|
||||
- **Cookie**: `turnstone_auth=<token>` (set automatically by the login endpoint)
|
||||
|
||||
The server accepts three token types:
|
||||
The server accepts two token types:
|
||||
|
||||
| Type | Format | Example |
|
||||
|------|--------|---------|
|
||||
| JWT | Base64 segments separated by dots | `eyJhbG...` |
|
||||
| API token | `ts_` prefix + 64 hex chars | `ts_a1b2c3d4...` |
|
||||
| Config token | Arbitrary string from `config.toml` | `my-secret-token` |
|
||||
|
||||
JWTs are the recommended credential for browser sessions. API tokens are suitable for programmatic access and CI/CD. Config tokens are a simple option for single-node deployments.
|
||||
JWTs are the recommended credential for browser sessions. API tokens are suitable for programmatic access and CI/CD.
|
||||
|
||||
### `POST /v1/api/auth/login`
|
||||
|
||||
|
||||
+19
-8
@@ -547,6 +547,21 @@ expanded tools).
|
||||
**Tool naming:** `mcp__{server}__{tool}` — double underscore delimiter, validated
|
||||
at connection time (server names with `__` are rejected).
|
||||
|
||||
**Resilience:** Each MCP server has an independent circuit breaker that opens
|
||||
after 3 consecutive transport failures (timeouts, broken pipes, connection
|
||||
resets). Cooldown uses capped exponential backoff (30 s base, 5 min max) with
|
||||
per-server jitter to avoid thundering herd. Protocol-level errors (`McpError`)
|
||||
from a healthy connection do not trip the breaker. When the cooldown expires
|
||||
(half-open), the next operation attempt triggers automatic reconnection. Manual
|
||||
`/mcp refresh` also clears the circuit on success. All sync bridge methods
|
||||
(`call_tool_sync`, `read_resource_sync`, `get_prompt_sync`, `refresh_sync`)
|
||||
cancel orphaned futures on timeout to prevent coroutine accumulation on the
|
||||
background event loop. Push notification refreshes are debounced (5 s per
|
||||
server) to protect against notification storms. The periodic refresh loop
|
||||
attempts reconnection for disconnected servers with exponential backoff
|
||||
(60 s–1 h). Transport stream references are pre-closed before stack teardown to
|
||||
work around the MCP SDK's anyio cancel-scope CPU busy-loop (SDK #2147).
|
||||
|
||||
**Error isolation:** Per-server connection/refresh failures are caught and logged; other
|
||||
servers are unaffected. Tool execution errors return error strings to the LLM
|
||||
rather than crashing the session.
|
||||
@@ -1016,13 +1031,10 @@ limits using a token-bucket algorithm. Each IP gets a `TokenBucket` with
|
||||
Turnstone supports three authentication mechanisms, unified behind an
|
||||
`AuthResult` dataclass that carries `user_id`, `scopes`, and `token_source`:
|
||||
|
||||
1. **Config-file tokens** — static secrets in `config.toml` `[[auth.tokens]]`
|
||||
or the `TURNSTONE_AUTH_TOKEN` env var. Validated in-memory via
|
||||
`hmac.compare_digest`. Map to scopes through their role (`read` or `full`).
|
||||
2. **API tokens** — database-backed, prefixed `ts_`, stored as SHA-256 hashes
|
||||
1. **API tokens** — database-backed, prefixed `ts_`, stored as SHA-256 hashes
|
||||
in the `api_tokens` table. Can be exchanged for JWTs via
|
||||
`POST /v1/api/auth/login`.
|
||||
3. **JWTs** — short-lived HMAC-SHA256 session tokens (default 24h) issued after
|
||||
2. **JWTs** — short-lived HMAC-SHA256 session tokens (default 24h) issued after
|
||||
successful credential validation. Contain `sub` (user_id), `scopes`, and
|
||||
`src` (origin) in claims.
|
||||
|
||||
@@ -1046,9 +1058,8 @@ Three hierarchical scopes control endpoint access:
|
||||
2. **Token extraction** — `Authorization: Bearer <token>` header first, then
|
||||
`turnstone_auth` cookie as fallback.
|
||||
3. **Token type detection** — dots in the token indicate JWT; `ts_` prefix
|
||||
indicates API token; otherwise config-file token.
|
||||
4. **Validation** — JWT signature check, API token hash lookup in storage, or
|
||||
config-token hmac comparison.
|
||||
indicates API token.
|
||||
4. **Validation** — JWT signature check or API token hash lookup in storage.
|
||||
5. **Scope check** — `required_scope(method, path)` determines the minimum
|
||||
scope; the request is rejected with 403 if the token lacks it.
|
||||
6. **Context propagation** — on success, `ctx_user_id` is set so structured
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:75c1832b6079e8628f4bbf4ce98d37880c4de133636b7555e3869990b046ddc6
|
||||
size 567704
|
||||
+5
-6
@@ -193,7 +193,6 @@ Plan review requests are displayed as a blue embed with:
|
||||
| `--auto-approve` | — | `false` | Auto-approve ALL tool calls (skips approval buttons entirely) |
|
||||
| `--http-host` | — | `127.0.0.1` | HTTP server bind address for notify endpoint |
|
||||
| `--http-port` | `TURNSTONE_CHANNEL_PORT` | `8091` | HTTP server port |
|
||||
| `--auth-token` | `TURNSTONE_CHANNEL_AUTH_TOKEN` | — | Static auth token for `/v1/api/notify` (alternative to JWT) |
|
||||
| `--log-level` | `TURNSTONE_LOG_LEVEL` | `INFO` | Log level |
|
||||
| `--log-format` | `TURNSTONE_LOG_FORMAT` | `auto` | Log format (`auto`/`json`/`text`) |
|
||||
|
||||
@@ -321,11 +320,11 @@ The `services` table schema:
|
||||
### Security
|
||||
|
||||
- **Authentication** — the gateway's `POST /v1/api/notify` endpoint
|
||||
requires authentication. Configure either `TURNSTONE_JWT_SECRET`
|
||||
(the server mints JWTs with `aud: turnstone-channel` automatically)
|
||||
or a static token via `--auth-token`. If neither is set, the
|
||||
gateway fails closed and rejects all requests with 401. Server JWTs
|
||||
(`aud: turnstone-server`) are rejected.
|
||||
requires authentication. Configure `TURNSTONE_JWT_SECRET` so the
|
||||
server can mint JWTs with `aud: turnstone-channel` automatically.
|
||||
If the secret is not set, the gateway fails closed and rejects all
|
||||
requests with 401. Server JWTs (`aud: turnstone-server`) are
|
||||
rejected.
|
||||
- **Rate limit** — maximum 5 notifications per turn. The counter only
|
||||
increments on successful delivery, so failures don't consume the
|
||||
budget.
|
||||
|
||||
+1
-2
@@ -628,7 +628,6 @@ CLI flags for `turnstone-console`:
|
||||
|------|---------|-------------|
|
||||
| `--host` | `0.0.0.0` | Bind host |
|
||||
| `--port` | `8090` | HTTP port |
|
||||
| `--auth-token` | `$TURNSTONE_AUTH_TOKEN` | Bearer token for server node communication and proxy |
|
||||
| `--log-level` | `INFO` | Log level |
|
||||
|
||||
Config file (`~/.config/turnstone/config.toml`):
|
||||
@@ -649,7 +648,7 @@ url = "http://localhost:8090" # used by CLI /cluster commands
|
||||
turnstone-server --port 8080
|
||||
|
||||
# Start cluster console (one instance)
|
||||
turnstone-console --port 8090 --auth-token "$TURNSTONE_AUTH_TOKEN"
|
||||
turnstone-console --port 8090
|
||||
```
|
||||
|
||||
Open `http://localhost:8090` for the cluster dashboard. Create workstreams via the "+ new" button. Click any workstream to open the proxied server UI — no direct access to server ports required.
|
||||
|
||||
@@ -152,10 +152,34 @@ MCPMgr -> MCPSrv : prompts/get
|
||||
MCPSrv --> MCPMgr : GetPromptResult
|
||||
MCPMgr --> Session : messages [{role, content}]
|
||||
|
||||
== Resilience: Circuit Breaker & Stream Safety ==
|
||||
|
||||
note over MCPMgr
|
||||
**Per-server circuit breaker**
|
||||
CLOSED --(3 failures)--> OPEN
|
||||
OPEN --(cooldown expires)--> half-open probe
|
||||
Probe success --> CLOSED (trip_count decays by 1)
|
||||
Probe failure --> OPEN (cooldown doubles, max 5 min)
|
||||
|
||||
McpError (protocol) does NOT trip breaker.
|
||||
BrokenPipeError / EOFError evicts dead session.
|
||||
All sync methods cancel orphaned futures on timeout.
|
||||
Transport streams pre-closed before stack teardown
|
||||
to avoid anyio cancel-scope CPU busy-loop (SDK #2147).
|
||||
end note
|
||||
|
||||
Session -> MCPMgr : call_tool_sync()
|
||||
MCPMgr -> MCPMgr : _cb_gate(server)\n[reject if circuit open]
|
||||
MCPMgr -> MCPMgr : _cb_auto_reconnect()\n[if session gone + cooldown expired]
|
||||
MCPMgr -> MCPSrv : tools/call
|
||||
MCPSrv --> MCPMgr : result or error
|
||||
MCPMgr -> MCPMgr : _cb_record_success()\nor _cb_record_failure()
|
||||
|
||||
== Three-Tier Refresh ==
|
||||
|
||||
group Push Notifications
|
||||
group Push Notifications (debounced 5s per server)
|
||||
MCPSrv -> MCPMgr : ToolListChangedNotification
|
||||
MCPMgr -> MCPMgr : debounce check\n(skip if < 5s since last)
|
||||
MCPMgr -> MCPMgr : _refresh_server_tools()
|
||||
|
||||
MCPSrv -> MCPMgr : ResourceListChangedNotification
|
||||
@@ -172,6 +196,9 @@ group Periodic Polling (default 4h)
|
||||
Only polls capabilities
|
||||
without push support.
|
||||
Staggered per-server.
|
||||
Disconnected servers get
|
||||
reconnect attempts with
|
||||
exponential backoff (60s-1h).
|
||||
end note
|
||||
end
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a6b7769aa7e732ffbeb1eb7f5b65273a135fb3a78d9802ec36d3b92801c34f6b
|
||||
size 427745
|
||||
oid sha256:7623df33be9baf7647ca1c2450640df57e1cd73e8be1f8168aae16e546ad683c
|
||||
size 459941
|
||||
|
||||
+3
-3
@@ -72,11 +72,11 @@ All configuration is via environment variables in `.env` (copy from `.env.exampl
|
||||
|
||||
### Auth
|
||||
|
||||
Auth is always enabled. `TURNSTONE_JWT_SECRET` is required.
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `TURNSTONE_AUTH_ENABLED` | — | Set to `1` to require authentication |
|
||||
| `TURNSTONE_AUTH_TOKEN` | — | Config-file token for server/console (backward compat, works alongside JWT) |
|
||||
| `TURNSTONE_JWT_SECRET` | — | Secret key for signing JWTs (required when using user identity / JWT auth) |
|
||||
| `TURNSTONE_JWT_SECRET` | — | Secret key for signing JWTs (required) |
|
||||
|
||||
### Database
|
||||
|
||||
|
||||
+5
-5
@@ -38,7 +38,7 @@ are set.
|
||||
| `TURNSTONE_OIDC_PROVIDER_NAME` | No | `SSO` | Display name for the login button (e.g. "Google", "Okta") |
|
||||
| `TURNSTONE_OIDC_ROLE_CLAIM` | No | — | ID token claim containing role/group values (see [Role Mapping](#role-mapping)) |
|
||||
| `TURNSTONE_OIDC_ROLE_MAP` | No | — | Mapping from claim values to Turnstone role IDs (see [Role Mapping](#role-mapping)) |
|
||||
| `TURNSTONE_OIDC_PASSWORD_ENABLED` | No | `true` | Set to `false` to hide the password form and block all username/password logins (including admin). API tokens and config-file tokens still work. |
|
||||
| `TURNSTONE_OIDC_PASSWORD_ENABLED` | No | `true` | Set to `false` to hide the password form and block all username/password logins (including admin). API tokens continue to work. |
|
||||
| `TURNSTONE_OIDC_REDIRECT_BASE` | No | — | Externally-reachable origin for the OIDC redirect URI (e.g. `https://app.example.com`). Recommended when running behind a reverse proxy. When unset, derived from the request Host header. |
|
||||
|
||||
OIDC is enabled when all three required fields (issuer, client ID, client
|
||||
@@ -246,10 +246,10 @@ password) before OIDC is enabled. The setup wizard always works
|
||||
regardless of this setting because it is only available when zero users
|
||||
exist in the database.
|
||||
|
||||
API token login (`POST /v1/api/auth/login` with a `ts_` token) and
|
||||
config-file tokens (`Authorization: Bearer tok_xxx`) continue to work
|
||||
regardless of this setting. OIDC-only mode affects password-based
|
||||
authentication only.
|
||||
API token login (`POST /v1/api/auth/login` with a `ts_` token)
|
||||
continues to work regardless of this setting. JWTs and API tokens are
|
||||
the supported authentication methods. OIDC-only mode affects
|
||||
password-based authentication only.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# Release Process
|
||||
|
||||
Turnstone uses two parallel release tracks published from a single PyPI package.
|
||||
|
||||
## Release Tracks
|
||||
|
||||
| Track | Versions | Branch | Docker tags | PyPI install |
|
||||
|-------|----------|--------|-------------|--------------|
|
||||
| **Stable** | `1.0.0`, `1.0.1` | `stable/1.0` | `:1.0.1`, `:1.0`, `:stable`, `:latest` | `pip install turnstone` |
|
||||
| **Experimental** | `1.1.0a1`, `1.1.0a2` | `main` | `:1.1.0a1`, `:experimental` | `pip install turnstone --pre` |
|
||||
|
||||
- **Stable** receives bugfixes only. Production-grade.
|
||||
- **Experimental** receives new features. May be rough around the edges.
|
||||
- When experimental matures, it is promoted to stable. The previous stable branch stops receiving patches.
|
||||
|
||||
## Version Scheme
|
||||
|
||||
[PEP 440](https://peps.python.org/pep-0440/) pre-release suffixes on a single package:
|
||||
|
||||
- `1.0.0` — stable release
|
||||
- `1.1.0a1` — alpha (experimental)
|
||||
- `1.1.0b1` — beta (experimental, more stable)
|
||||
- `1.1.0rc1` — release candidate (experimental, nearly stable)
|
||||
- `1.1.0` — promoted to stable
|
||||
|
||||
## Releasing an Experimental Version (from main)
|
||||
|
||||
```bash
|
||||
scripts/release.sh 1.1.0a2 --push
|
||||
```
|
||||
|
||||
This bumps `pyproject.toml` + `turnstone/__init__.py`, regenerates `uv.lock`, commits, tags `v1.1.0a2`, and pushes. CI runs, then publish + Docker workflows fire automatically.
|
||||
|
||||
## Releasing a Stable Patch (from stable/X.Y)
|
||||
|
||||
```bash
|
||||
git checkout stable/1.0
|
||||
git cherry-pick <commit-hash> # bugfix from main
|
||||
scripts/release.sh 1.0.2 --push
|
||||
```
|
||||
|
||||
## Promoting Experimental to Stable
|
||||
|
||||
When `main` is ready for a stable release:
|
||||
|
||||
```bash
|
||||
# 1. Tag the stable release on main
|
||||
scripts/release.sh 1.1.0 --push
|
||||
|
||||
# 2. Create the stable maintenance branch from that tag
|
||||
git branch stable/1.1 v1.1.0
|
||||
git push origin stable/1.1
|
||||
|
||||
# 3. Start the next experimental cycle on main
|
||||
scripts/release.sh 1.2.0a1 --push
|
||||
```
|
||||
|
||||
The previous `stable/1.0` branch stops receiving patches at this point.
|
||||
|
||||
## CI/CD Pipeline
|
||||
|
||||
All releases are gated on CI success:
|
||||
|
||||
1. `git push` with `v*` tag triggers **CI** (lint, typecheck, test, test-postgres, lock-check, security audit)
|
||||
2. On CI success, **Publish to PyPI** fires via `workflow_run`
|
||||
3. On CI success, **Publish Docker Image** fires via `workflow_run`
|
||||
|
||||
Pre-release tags (`a`, `b`, `rc` suffixes) produce:
|
||||
- PyPI: pre-release version (not installed by default)
|
||||
- GitHub Release: marked as pre-release
|
||||
- Docker: `:experimental` alias + exact version tag
|
||||
|
||||
Stable tags produce:
|
||||
- PyPI: stable version (default `pip install`)
|
||||
- GitHub Release: full release
|
||||
- Docker: `:stable`, `:latest`, `:X.Y`, `:X.Y.Z` tags
|
||||
|
||||
## Dependency Updates
|
||||
|
||||
Renovate targets `main` (experimental) only. Stable branches receive manual dependency updates via cherry-pick when security-relevant.
|
||||
+2
-2
@@ -332,6 +332,6 @@ client.login(token="ts_abc123...")
|
||||
- `client.logout()` clears the stored JWT from the client.
|
||||
- If a request returns 401, the SDK raises `TurnstoneAPIError` — the caller is responsible for re-authenticating.
|
||||
|
||||
### Backward Compatibility
|
||||
### Token Types
|
||||
|
||||
The config-file token (`TURNSTONE_AUTH_TOKEN`) still works as a simple Bearer token for environments that do not use the user/JWT system. When the server receives a non-JWT Bearer token, it falls back to the legacy token check.
|
||||
The SDK accepts any Bearer token — JWTs (from `ServiceTokenManager` or login) and API tokens (`ts_` prefix) are both supported. Use `token_factory` for auto-rotating JWTs or a static `token` for API tokens.
|
||||
|
||||
+11
-54
@@ -8,23 +8,6 @@ credentials while individual server nodes validate JWTs locally.
|
||||
|
||||
## Token Types
|
||||
|
||||
### Config-file tokens
|
||||
|
||||
Static tokens defined in `config.toml` or the `TURNSTONE_AUTH_TOKEN`
|
||||
environment variable. Validated in-memory using `hmac.compare_digest`
|
||||
(timing-safe). Each token maps to a role that determines its scopes.
|
||||
|
||||
```toml
|
||||
[[auth.tokens]]
|
||||
value = "tok_legacy"
|
||||
role = "full" # full → {read, write, approve}
|
||||
```
|
||||
|
||||
Role mappings: `"read"` → `{read}`, `"full"` → `{read, write, approve}`.
|
||||
|
||||
Config tokens are sent directly as `Authorization: Bearer tok_legacy`
|
||||
on every request. No JWT exchange is needed.
|
||||
|
||||
### API tokens
|
||||
|
||||
Database-backed tokens prefixed with `ts_`. Created via the admin CLI
|
||||
@@ -149,15 +132,6 @@ The API token is hashed, looked up in the database, and exchanged for a
|
||||
JWT with the token's scopes. This is the recommended flow for SDKs and
|
||||
automated clients that need cookie-based sessions.
|
||||
|
||||
### Config-file tokens (direct)
|
||||
|
||||
Config tokens are validated per-request via `hmac.compare_digest`. No
|
||||
login exchange is needed — include the token as a `Bearer` header:
|
||||
|
||||
```
|
||||
Authorization: Bearer tok_legacy
|
||||
```
|
||||
|
||||
### First-time setup
|
||||
|
||||
When no users exist in the database:
|
||||
@@ -276,7 +250,7 @@ Setting `TURNSTONE_OIDC_PASSWORD_ENABLED=false` hides the password
|
||||
form on the login page and blocks password-based login at the API
|
||||
level. The setup wizard always works regardless of this setting — the
|
||||
first admin user is created with a password before OIDC is relevant.
|
||||
API tokens and config-file tokens are unaffected by this setting.
|
||||
API tokens are unaffected by this setting.
|
||||
|
||||
#### Known limitations
|
||||
|
||||
@@ -297,8 +271,6 @@ and classifies the token:
|
||||
|
||||
1. **Contains `.`** → JWT → validate HS256 signature and expiry
|
||||
2. **Starts with `ts_`** → API token → SHA-256 hash, database lookup
|
||||
3. **Otherwise** → config-file token → `hmac.compare_digest` against
|
||||
each configured token
|
||||
|
||||
If a session cookie is present and no `Authorization` header is sent,
|
||||
the cookie value is treated as a JWT (step 1).
|
||||
@@ -332,16 +304,10 @@ deployments.
|
||||
| Signing secret | `[auth] jwt_secret` | `TURNSTONE_JWT_SECRET` | Auto-generated ephemeral (warning logged) |
|
||||
| Expiry | `[auth] jwt_expiry_hours` | — | 24 hours |
|
||||
| Algorithm | — | — | HS256 (not configurable) |
|
||||
| Minimum secret length | — | — | 32 characters (warning if shorter) |
|
||||
| Minimum secret length | — | — | 32 characters (exits if shorter) |
|
||||
|
||||
All service nodes that need to validate JWTs must share the same signing
|
||||
secret. If no secret is configured, an ephemeral key is generated at
|
||||
startup and a warning is logged — JWTs will not survive restarts or work
|
||||
across nodes.
|
||||
|
||||
The console **requires** `TURNSTONE_JWT_SECRET` when no `--auth-token`
|
||||
is provided. It exits with an error if the secret is missing, since
|
||||
ephemeral secrets would silently break inter-service communication.
|
||||
All services require `TURNSTONE_JWT_SECRET` and exit at startup if it is
|
||||
missing or shorter than 32 characters.
|
||||
|
||||
---
|
||||
|
||||
@@ -442,16 +408,15 @@ Console (cluster-wide) Server (per-node)
|
||||
┌──────────────────────┐ ┌──────────────────────┐
|
||||
│ User/Token CRUD (DB) │ │ JWT validation only │
|
||||
│ Login: creds → JWT │ │ (shared signing key) │
|
||||
│ Admin API endpoints │ │ Config tokens: hmac │
|
||||
│ Storage: users, │ │ No auth DB needed │
|
||||
│ Admin API endpoints │ │ No auth DB needed │
|
||||
│ Storage: users, │ │ │
|
||||
│ api_tokens tables │ │ │
|
||||
└──────────────────────┘ └──────────────────────┘
|
||||
```
|
||||
|
||||
The console owns the credential database and handles all user/token
|
||||
CRUD. Individual server nodes only need the JWT signing secret to
|
||||
validate session tokens. Config-file tokens are validated locally
|
||||
without any database.
|
||||
validate session tokens.
|
||||
|
||||
### Proxy auth forwarding
|
||||
|
||||
@@ -478,8 +443,7 @@ distinguish proxied requests from direct logins in audit logs.
|
||||
|
||||
When no user context is available (auth disabled, or internal requests),
|
||||
the proxy falls back to a `ServiceTokenManager` with service identity
|
||||
`console-proxy` and full scopes. If `--auth-token` is provided, that
|
||||
static token is used as a final fallback.
|
||||
`console-proxy` and full scopes.
|
||||
|
||||
### Service-to-service authentication
|
||||
|
||||
@@ -518,22 +482,17 @@ channel gateway endpoint, and vice versa.
|
||||
|
||||
```toml
|
||||
[auth]
|
||||
enabled = true
|
||||
jwt_secret = "your-secret-key-here"
|
||||
jwt_expiry_hours = 24
|
||||
|
||||
[[auth.tokens]]
|
||||
value = "tok_legacy"
|
||||
role = "full"
|
||||
```
|
||||
|
||||
### Environment variables
|
||||
|
||||
Auth is always enabled. `TURNSTONE_JWT_SECRET` is required.
|
||||
|
||||
| Variable | Description |
|
||||
|----------|-------------|
|
||||
| `TURNSTONE_AUTH_ENABLED=1` | Enable authentication |
|
||||
| `TURNSTONE_AUTH_TOKEN=tok_xxx` | Register a config-file token with `full` access |
|
||||
| `TURNSTONE_JWT_SECRET=xxx` | JWT signing secret (must match across nodes) |
|
||||
| `TURNSTONE_JWT_SECRET=xxx` | JWT signing secret (required, must match across nodes) |
|
||||
| `TURNSTONE_CORS_ORIGINS=` | CORS allowed origins (comma-separated; empty = same-origin only) |
|
||||
|
||||
---
|
||||
@@ -571,8 +530,6 @@ and browsers enforce same-origin policy.
|
||||
|
||||
## Security Properties
|
||||
|
||||
- **Timing-safe comparison** for config-file tokens via
|
||||
`hmac.compare_digest` — no timing side-channel.
|
||||
- **Hash-based lookup** for API tokens — the database stores only
|
||||
SHA-256 hashes, eliminating timing attacks on token comparison.
|
||||
- **Local JWT validation** — no network call or database query needed
|
||||
|
||||
+1
-1
@@ -105,7 +105,7 @@ turnstone-admin tls-ca-cert --out ca.pem --console-url http://console:8080
|
||||
turnstone-admin tls-issue worker-1.internal --out /certs --console-url http://console:8080
|
||||
|
||||
# List issued certs
|
||||
turnstone-admin tls-list --console-url http://console:8080 --auth-token $TOKEN
|
||||
turnstone-admin tls-list --console-url http://console:8080
|
||||
```
|
||||
|
||||
### Console URL Discovery
|
||||
|
||||
+1
-1
@@ -593,7 +593,7 @@ current turn and letting it search for them on demand.
|
||||
Tool search uses the best available mechanism for each provider:
|
||||
|
||||
1. **Anthropic (native)** -- Models that support it receive `defer_loading: true`
|
||||
on deferred tool definitions plus the `tool_search_tool_bm25_20251119` server-side
|
||||
on deferred tool definitions plus the `tool_search_tool_bm25` server-side
|
||||
search tool. Anthropic's API handles search and expansion transparently.
|
||||
|
||||
2. **OpenAI GPT-5.4+ (native)** -- Models with hosted tool search receive
|
||||
|
||||
@@ -1,10 +1,14 @@
|
||||
# MCP Cluster Ops
|
||||
|
||||
An MCP server that exposes tools for executing commands across a [Turnstone](https://github.com/turnstonelabs/turnstone) cluster. Serves as a reference implementation for both MCP server patterns and Turnstone SDK usage.
|
||||
An MCP server that exposes tools for executing commands across a Turnstone cluster. Serves as a reference implementation for both MCP server patterns and Turnstone SDK usage.
|
||||
|
||||
## How it works
|
||||
|
||||
This server uses Turnstone's SDK client (`TurnstoneServer`) to dispatch shell commands to specific nodes via HTTP. Remote agents execute the command and the raw bash output is captured directly from the `ToolResultEvent` stream — bypassing the costly "agent reads output → re-generates output as completion tokens" round-trip.
|
||||
This server uses the Turnstone console SDK (`TurnstoneConsole`) for node discovery and routing, and `TurnstoneServer` for per-node SSE streaming. The dispatch flow for each command is:
|
||||
|
||||
1. **Route** — `TurnstoneConsole.route_create_workstream(target_node=..., auto_approve=True)` creates a workstream pinned to the target node via the console's hash-ring routing proxy, returning `ws_id` and `node_url`.
|
||||
2. **Execute** — `TurnstoneServer(node_url, token=...)` connects directly to the node's SSE stream using the same `TURNSTONE_API_TOKEN`. `send_and_wait(prompt, ws_id)` runs the command and the raw bash output is captured from the `ToolResultEvent` — bypassing the costly "agent reads output then re-generates output as completion tokens" round-trip.
|
||||
3. **Cleanup** — `TurnstoneConsole.route_close(ws_id)` closes the workstream.
|
||||
|
||||
Multi-node dispatches run in parallel via `asyncio.gather`, so total wall time is bounded by the slowest node rather than the sum.
|
||||
|
||||
@@ -19,7 +23,7 @@ Multi-node dispatches run in parallel via `asyncio.gather`, so total wall time i
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- A running Turnstone cluster (at least one `turnstone-server`)
|
||||
- A running Turnstone cluster with at least one `turnstone-server` and a `turnstone-console`
|
||||
- Python 3.11+
|
||||
|
||||
## Installation
|
||||
@@ -35,8 +39,8 @@ pip install -e ./examples/mcp-cluster-ops
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `TURNSTONE_SERVER_URL` | `http://localhost:8080` | Server URL |
|
||||
| `TURNSTONE_API_TOKEN` | _(none)_ | API token for authentication |
|
||||
| `TURNSTONE_CONSOLE_URL` | `http://localhost:8090` | Console URL for node discovery and routing |
|
||||
| `TURNSTONE_API_TOKEN` | _(none)_ | API token / JWT for authentication |
|
||||
| `MCP_CLUSTER_OPS_TIMEOUT` | `120` | Default command timeout (seconds, clamped 5-3600) |
|
||||
| `MCP_CLUSTER_OPS_MAX_OUTPUT` | `8192` | Max output bytes per node (0 = unlimited) |
|
||||
| `MCP_CLUSTER_OPS_MAX_NODES` | `32` | Max concurrent node dispatches |
|
||||
@@ -51,7 +55,7 @@ pip install -e ./examples/mcp-cluster-ops
|
||||
command = "mcp-cluster-ops"
|
||||
|
||||
[mcp.servers.cluster-ops.env]
|
||||
TURNSTONE_SERVER_URL = "http://turnstone.example.com:8080"
|
||||
TURNSTONE_CONSOLE_URL = "http://console.example.com:8090"
|
||||
```
|
||||
|
||||
**JSON** (via `--mcp-config`):
|
||||
@@ -62,7 +66,7 @@ TURNSTONE_SERVER_URL = "http://turnstone.example.com:8080"
|
||||
"cluster-ops": {
|
||||
"command": "mcp-cluster-ops",
|
||||
"env": {
|
||||
"TURNSTONE_SERVER_URL": "http://turnstone.example.com:8080"
|
||||
"TURNSTONE_CONSOLE_URL": "http://console.example.com:8090"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
"""MCP server for Turnstone cluster operations.
|
||||
|
||||
Exposes tools to execute commands on specific nodes in a Turnstone cluster.
|
||||
Uses the SDK client (``TurnstoneServer``) for direct node targeting via HTTP.
|
||||
Uses the SDK console client (``TurnstoneConsole``) for node discovery and
|
||||
routing, and ``TurnstoneServer`` for per-node SSE streaming.
|
||||
|
||||
Usage::
|
||||
|
||||
@@ -14,12 +15,12 @@ Configure in ``~/.config/turnstone/config.toml``::
|
||||
command = "mcp-cluster-ops"
|
||||
|
||||
[mcp.servers.cluster-ops.env]
|
||||
TURNSTONE_SERVER_URL = "http://localhost:8080"
|
||||
TURNSTONE_CONSOLE_URL = "http://localhost:8090"
|
||||
|
||||
Environment variables
|
||||
---------------------
|
||||
TURNSTONE_SERVER_URL Server URL (default: http://localhost:8080)
|
||||
TURNSTONE_API_TOKEN API token for authentication (default: none)
|
||||
TURNSTONE_CONSOLE_URL Console URL (default: http://localhost:8090)
|
||||
TURNSTONE_API_TOKEN API token / JWT for authentication (default: none)
|
||||
MCP_CLUSTER_OPS_TIMEOUT Default command timeout in seconds (default: 120)
|
||||
MCP_CLUSTER_OPS_MAX_OUTPUT Max output bytes per node (default: 8192, 0=unlimited)
|
||||
|
||||
@@ -43,7 +44,7 @@ from contextlib import asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from turnstone.sdk import TurnResult, TurnstoneServer
|
||||
from turnstone.sdk import TurnResult, TurnstoneConsole, TurnstoneServer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
@@ -66,15 +67,12 @@ _MAX_TIMEOUT = 3600
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _server_kwargs() -> dict[str, Any]:
|
||||
"""Build TurnstoneServer connection kwargs from environment variables."""
|
||||
kwargs: dict[str, Any] = {
|
||||
"base_url": os.environ.get("TURNSTONE_SERVER_URL", "http://localhost:8080"),
|
||||
def _console_kwargs() -> dict[str, Any]:
|
||||
"""Build TurnstoneConsole connection kwargs from environment variables."""
|
||||
return {
|
||||
"base_url": os.environ.get("TURNSTONE_CONSOLE_URL", "http://localhost:8090"),
|
||||
"token": os.environ.get("TURNSTONE_API_TOKEN", ""),
|
||||
}
|
||||
token = os.environ.get("TURNSTONE_API_TOKEN")
|
||||
if token:
|
||||
kwargs["token"] = token
|
||||
return kwargs
|
||||
|
||||
|
||||
def _exec_prompt(command: str) -> str:
|
||||
@@ -141,6 +139,11 @@ def _validate_command(command: str) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _extract_node_ids(nodes: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract unique, non-empty node IDs from a list of node dicts."""
|
||||
return list(dict.fromkeys(n["node_id"].strip() for n in nodes if n.get("node_id", "").strip()))
|
||||
|
||||
|
||||
def _format_node_result(
|
||||
node_id: str,
|
||||
result: TurnResult,
|
||||
@@ -165,12 +168,12 @@ def _format_node_result(
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core dispatch functions (testable with mocked TurnstoneServer)
|
||||
# Core dispatch functions (testable with mocked SDK clients)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _exec_on_node_sync(
|
||||
server_kw: dict[str, Any],
|
||||
console_kw: dict[str, Any],
|
||||
node_id: str,
|
||||
command: str,
|
||||
timeout: float,
|
||||
@@ -178,22 +181,35 @@ def _exec_on_node_sync(
|
||||
"""Dispatch *command* to *node_id* and block until complete.
|
||||
|
||||
Runs inside ``asyncio.to_thread`` so it does not block the event loop.
|
||||
Each call creates its own ``TurnstoneServer`` client to avoid state
|
||||
conflicts between concurrent dispatches.
|
||||
|
||||
Flow:
|
||||
1. Create a workstream on the target node via the console routing proxy
|
||||
2. Connect directly to the node's SSE stream to send + collect output
|
||||
3. Close the workstream via the routing proxy
|
||||
"""
|
||||
prompt = _exec_prompt(command)
|
||||
with TurnstoneServer(**server_kw) as client:
|
||||
result = client.send_and_wait(
|
||||
message=prompt,
|
||||
target_node=node_id,
|
||||
auto_approve=True,
|
||||
timeout=timeout,
|
||||
)
|
||||
ws_id = ""
|
||||
with TurnstoneConsole(**console_kw) as console:
|
||||
try:
|
||||
route_resp = console.route_create_workstream(
|
||||
target_node=node_id,
|
||||
auto_approve=True,
|
||||
)
|
||||
ws_id = route_resp["ws_id"]
|
||||
node_url: str = route_resp["node_url"]
|
||||
with TurnstoneServer(
|
||||
base_url=node_url,
|
||||
token=console_kw["token"],
|
||||
) as server:
|
||||
result = server.send_and_wait(prompt, ws_id, timeout=timeout)
|
||||
finally:
|
||||
if ws_id:
|
||||
console.route_close(ws_id)
|
||||
return node_id, result
|
||||
|
||||
|
||||
async def _dispatch_parallel(
|
||||
server_kw: dict[str, Any],
|
||||
console_kw: dict[str, Any],
|
||||
node_ids: list[str],
|
||||
command: str,
|
||||
timeout: float,
|
||||
@@ -204,7 +220,7 @@ async def _dispatch_parallel(
|
||||
Total wall time is bounded by the slowest node.
|
||||
"""
|
||||
tasks = [
|
||||
asyncio.to_thread(_exec_on_node_sync, server_kw, nid, command, timeout) for nid in node_ids
|
||||
asyncio.to_thread(_exec_on_node_sync, console_kw, nid, command, timeout) for nid in node_ids
|
||||
]
|
||||
outcomes = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
@@ -220,16 +236,22 @@ async def _dispatch_parallel(
|
||||
return results
|
||||
|
||||
|
||||
def _list_nodes_sync(server_kw: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""List active cluster nodes (blocking)."""
|
||||
with TurnstoneServer(**server_kw) as client:
|
||||
nodes: list[dict[str, Any]] = client.list_nodes()
|
||||
return nodes
|
||||
def _list_nodes_sync(console_kw: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""List active cluster nodes (blocking), paginating if needed."""
|
||||
page_size = 100
|
||||
nodes: list[dict[str, Any]] = []
|
||||
with TurnstoneConsole(**console_kw) as console:
|
||||
while True:
|
||||
resp = console.nodes(limit=page_size, offset=len(nodes))
|
||||
nodes.extend(n.model_dump() for n in resp.nodes)
|
||||
if len(nodes) >= resp.total or not resp.nodes:
|
||||
break
|
||||
return nodes
|
||||
|
||||
|
||||
async def _list_nodes_impl(server_kw: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
async def _list_nodes_impl(console_kw: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""List active cluster nodes."""
|
||||
return await asyncio.to_thread(_list_nodes_sync, server_kw)
|
||||
return await asyncio.to_thread(_list_nodes_sync, console_kw)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -239,9 +261,9 @@ async def _list_nodes_impl(server_kw: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
|
||||
@asynccontextmanager
|
||||
async def _lifespan(server: FastMCP[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]:
|
||||
"""Lifespan context — stores server connection kwargs for tool handlers."""
|
||||
kw = _server_kwargs()
|
||||
yield {"server_kwargs": kw}
|
||||
"""Lifespan context — stores console connection kwargs for tool handlers."""
|
||||
kw = _console_kwargs()
|
||||
yield {"console_kwargs": kw}
|
||||
|
||||
|
||||
mcp = FastMCP(
|
||||
@@ -263,8 +285,8 @@ async def list_nodes(ctx: Context[Any, Any, Any]) -> str:
|
||||
Call this before dispatching work to discover available node IDs.
|
||||
Returns a JSON array of node metadata objects.
|
||||
"""
|
||||
server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
|
||||
nodes = await _list_nodes_impl(server_kw)
|
||||
console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
|
||||
nodes = await _list_nodes_impl(console_kw)
|
||||
return json.dumps(nodes, indent=2)
|
||||
|
||||
|
||||
@@ -291,13 +313,16 @@ async def run_on_node(
|
||||
if cmd_err:
|
||||
return json.dumps({"error": cmd_err})
|
||||
|
||||
server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
|
||||
console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
|
||||
max_output = _DEFAULT_MAX_OUTPUT
|
||||
|
||||
log.info("run_on_node node=%s cmd=%r", node_id, command)
|
||||
_, result = await asyncio.to_thread(
|
||||
_exec_on_node_sync, server_kw, node_id, command, _clamp_timeout(timeout)
|
||||
)
|
||||
try:
|
||||
_, result = await asyncio.to_thread(
|
||||
_exec_on_node_sync, console_kw, node_id, command, _clamp_timeout(timeout)
|
||||
)
|
||||
except Exception as exc:
|
||||
return json.dumps({"node": node_id, "ok": False, "error": str(exc)}, indent=2)
|
||||
formatted = _format_node_result(node_id, result, max_output)
|
||||
return json.dumps(formatted, indent=2)
|
||||
|
||||
@@ -323,7 +348,7 @@ async def run_on_nodes(
|
||||
if cmd_err:
|
||||
return json.dumps({"error": cmd_err})
|
||||
|
||||
server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
|
||||
console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
|
||||
max_output = _DEFAULT_MAX_OUTPUT
|
||||
|
||||
clean_ids = list(dict.fromkeys(nid.strip() for nid in node_ids if nid.strip()))
|
||||
@@ -336,7 +361,7 @@ async def run_on_nodes(
|
||||
|
||||
log.info("run_on_nodes nodes=%s cmd=%r", clean_ids, command)
|
||||
results = await _dispatch_parallel(
|
||||
server_kw, clean_ids, command, _clamp_timeout(timeout), max_output
|
||||
console_kw, clean_ids, command, _clamp_timeout(timeout), max_output
|
||||
)
|
||||
return json.dumps(results, indent=2)
|
||||
|
||||
@@ -361,18 +386,14 @@ async def run_on_all_nodes(
|
||||
if cmd_err:
|
||||
return json.dumps({"error": cmd_err})
|
||||
|
||||
server_kw: dict[str, Any] = ctx.request_context.lifespan_context["server_kwargs"]
|
||||
console_kw: dict[str, Any] = ctx.request_context.lifespan_context["console_kwargs"]
|
||||
max_output = _DEFAULT_MAX_OUTPUT
|
||||
|
||||
nodes = await _list_nodes_impl(server_kw)
|
||||
nodes = await _list_nodes_impl(console_kw)
|
||||
if not nodes:
|
||||
return json.dumps({"error": "No active nodes found in cluster"})
|
||||
|
||||
node_ids = list(
|
||||
dict.fromkeys(
|
||||
nid.strip() for n in nodes if (nid := n.get("node_id") or n.get("id")) and nid.strip()
|
||||
)
|
||||
)
|
||||
node_ids = _extract_node_ids(nodes)
|
||||
if not node_ids:
|
||||
return json.dumps({"error": "No nodes with identifiable IDs found"})
|
||||
if len(node_ids) > _MAX_CONCURRENT_NODES:
|
||||
@@ -381,7 +402,7 @@ async def run_on_all_nodes(
|
||||
)
|
||||
log.info("run_on_all_nodes nodes=%s cmd=%r", node_ids, command)
|
||||
results = await _dispatch_parallel(
|
||||
server_kw, node_ids, command, _clamp_timeout(timeout), max_output
|
||||
console_kw, node_ids, command, _clamp_timeout(timeout), max_output
|
||||
)
|
||||
return json.dumps(results, indent=2)
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from turnstone.sdk import TurnResult
|
||||
from mcp_cluster_ops.server import (
|
||||
_clamp_timeout,
|
||||
_exec_prompt,
|
||||
_extract_node_ids,
|
||||
_extract_output,
|
||||
_format_node_result,
|
||||
_truncate,
|
||||
@@ -190,3 +191,53 @@ class TestClampTimeout:
|
||||
|
||||
def test_negative(self):
|
||||
assert _clamp_timeout(-1) == 5.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _extract_node_ids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractNodeIds:
|
||||
def test_normal(self):
|
||||
nodes = [
|
||||
{"node_id": "a", "server_url": "http://a:8080"},
|
||||
{"node_id": "b", "server_url": "http://b:8080"},
|
||||
]
|
||||
assert _extract_node_ids(nodes) == ["a", "b"]
|
||||
|
||||
def test_deduplicates(self):
|
||||
nodes = [
|
||||
{"node_id": "a"},
|
||||
{"node_id": "a"},
|
||||
{"node_id": "b"},
|
||||
]
|
||||
assert _extract_node_ids(nodes) == ["a", "b"]
|
||||
|
||||
def test_strips_whitespace(self):
|
||||
nodes = [{"node_id": " a "}, {"node_id": "b "}]
|
||||
assert _extract_node_ids(nodes) == ["a", "b"]
|
||||
|
||||
def test_skips_empty(self):
|
||||
nodes = [
|
||||
{"node_id": "a"},
|
||||
{"node_id": ""},
|
||||
{"node_id": " "},
|
||||
{"node_id": "b"},
|
||||
]
|
||||
assert _extract_node_ids(nodes) == ["a", "b"]
|
||||
|
||||
def test_skips_missing_key(self):
|
||||
nodes = [
|
||||
{"node_id": "a"},
|
||||
{"server_url": "http://orphan:8080"},
|
||||
{"node_id": "b"},
|
||||
]
|
||||
assert _extract_node_ids(nodes) == ["a", "b"]
|
||||
|
||||
def test_empty_list(self):
|
||||
assert _extract_node_ids([]) == []
|
||||
|
||||
def test_all_empty_ids(self):
|
||||
nodes = [{"node_id": ""}, {"node_id": " "}]
|
||||
assert _extract_node_ids(nodes) == []
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
"""Tests for MCP tool handlers with mocked TurnstoneServer."""
|
||||
"""Tests for MCP tool handlers with mocked SDK clients."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from turnstone.sdk import TurnResult
|
||||
|
||||
from mcp_cluster_ops.server import (
|
||||
@@ -14,6 +16,22 @@ from mcp_cluster_ops.server import (
|
||||
_list_nodes_impl,
|
||||
)
|
||||
|
||||
_CONSOLE_KW: dict[str, Any] = {"base_url": "http://localhost:8090", "token": ""}
|
||||
_CONSOLE_KW_AUTH: dict[str, Any] = {"base_url": "http://localhost:8090", "token": "tok_test"}
|
||||
|
||||
|
||||
def _mock_console_ctx(mock_cls: MagicMock, mock_client: MagicMock) -> None:
|
||||
"""Wire up a TurnstoneConsole mock as a context manager."""
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
|
||||
def _mock_server_ctx(mock_cls: MagicMock, mock_server: MagicMock) -> None:
|
||||
"""Wire up a TurnstoneServer mock as a context manager."""
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_server)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _list_nodes_impl
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -21,26 +39,71 @@ from mcp_cluster_ops.server import (
|
||||
|
||||
class TestListNodesImpl:
|
||||
def test_returns_nodes(self):
|
||||
nodes = [{"node_id": "a", "model": "gpt-5"}, {"node_id": "b", "model": "gpt-5"}]
|
||||
with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.list_nodes.return_value = nodes
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
mock_node_a = MagicMock()
|
||||
mock_node_a.model_dump.return_value = {"node_id": "a", "server_url": "http://a:8080"}
|
||||
mock_node_b = MagicMock()
|
||||
mock_node_b.model_dump.return_value = {"node_id": "b", "server_url": "http://b:8080"}
|
||||
|
||||
result = asyncio.run(_list_nodes_impl({"host": "localhost"}))
|
||||
assert result == nodes
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.nodes = [mock_node_a, mock_node_b]
|
||||
mock_resp.total = 2
|
||||
|
||||
with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.nodes.return_value = mock_resp
|
||||
_mock_console_ctx(mock_cls, mock_client)
|
||||
|
||||
result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
|
||||
assert len(result) == 2
|
||||
assert result[0]["node_id"] == "a"
|
||||
assert result[1]["node_id"] == "b"
|
||||
|
||||
def test_empty_cluster(self):
|
||||
with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.list_nodes.return_value = []
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.nodes = []
|
||||
mock_resp.total = 0
|
||||
|
||||
result = asyncio.run(_list_nodes_impl({"host": "localhost"}))
|
||||
with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.nodes.return_value = mock_resp
|
||||
_mock_console_ctx(mock_cls, mock_client)
|
||||
|
||||
result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
|
||||
assert result == []
|
||||
|
||||
def test_paginates_large_clusters(self):
|
||||
"""Clusters with >100 nodes are fetched across multiple pages."""
|
||||
|
||||
def _make_node(nid: str) -> MagicMock:
|
||||
m = MagicMock()
|
||||
m.model_dump.return_value = {"node_id": nid}
|
||||
return m
|
||||
|
||||
page1_nodes = [_make_node(f"n-{i}") for i in range(100)]
|
||||
page2_nodes = [_make_node(f"n-{i}") for i in range(100, 150)]
|
||||
|
||||
page1_resp = MagicMock()
|
||||
page1_resp.nodes = page1_nodes
|
||||
page1_resp.total = 150
|
||||
|
||||
page2_resp = MagicMock()
|
||||
page2_resp.nodes = page2_nodes
|
||||
page2_resp.total = 150
|
||||
|
||||
with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.nodes.side_effect = [page1_resp, page2_resp]
|
||||
_mock_console_ctx(mock_cls, mock_client)
|
||||
|
||||
result = asyncio.run(_list_nodes_impl(_CONSOLE_KW))
|
||||
assert len(result) == 150
|
||||
assert result[0]["node_id"] == "n-0"
|
||||
assert result[149]["node_id"] == "n-149"
|
||||
assert mock_client.nodes.call_count == 2
|
||||
# Verify offset was passed correctly
|
||||
mock_client.nodes.assert_any_call(limit=100, offset=0)
|
||||
mock_client.nodes.assert_any_call(limit=100, offset=100)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _exec_on_node_sync
|
||||
@@ -50,35 +113,111 @@ class TestListNodesImpl:
|
||||
class TestExecOnNodeSync:
|
||||
def test_success(self):
|
||||
turn_result = TurnResult(
|
||||
ws_id="ws-123",
|
||||
tool_results=[("bash", "hello world")],
|
||||
)
|
||||
with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.send_and_wait.return_value = turn_result
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
with (
|
||||
patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
|
||||
patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
|
||||
):
|
||||
mock_console = MagicMock()
|
||||
mock_console.route_create_workstream.return_value = {
|
||||
"ws_id": "ws-123",
|
||||
"node_url": "http://node-1:8080",
|
||||
"node_id": "node-1",
|
||||
"name": "ws-123",
|
||||
}
|
||||
_mock_console_ctx(mock_console_cls, mock_console)
|
||||
|
||||
node_id, result = _exec_on_node_sync(
|
||||
{"host": "localhost"}, "node-1", "echo hello", 60.0
|
||||
)
|
||||
mock_server = MagicMock()
|
||||
mock_server.send_and_wait.return_value = turn_result
|
||||
_mock_server_ctx(mock_server_cls, mock_server)
|
||||
|
||||
node_id, result = _exec_on_node_sync(_CONSOLE_KW_AUTH, "node-1", "echo hello", 60.0)
|
||||
assert node_id == "node-1"
|
||||
assert result.ok
|
||||
mock_client.send_and_wait.assert_called_once()
|
||||
call_kwargs = mock_client.send_and_wait.call_args
|
||||
assert call_kwargs.kwargs["target_node"] == "node-1"
|
||||
assert call_kwargs.kwargs["auto_approve"] is True
|
||||
|
||||
# Verify console created ws on the right node
|
||||
mock_console.route_create_workstream.assert_called_once_with(
|
||||
target_node="node-1",
|
||||
auto_approve=True,
|
||||
)
|
||||
|
||||
# Verify server connected to the node URL with the token
|
||||
mock_server_cls.assert_called_once_with(
|
||||
base_url="http://node-1:8080",
|
||||
token="tok_test",
|
||||
)
|
||||
|
||||
# Verify send_and_wait got the right ws_id
|
||||
call_kwargs = mock_server.send_and_wait.call_args
|
||||
assert call_kwargs.args[1] == "ws-123"
|
||||
|
||||
# Verify workstream was closed
|
||||
mock_console.route_close.assert_called_once_with("ws-123")
|
||||
|
||||
def test_timeout(self):
|
||||
turn_result = TurnResult(timed_out=True)
|
||||
with patch("mcp_cluster_ops.server.TurnstoneServer") as mock_cls:
|
||||
mock_client = MagicMock()
|
||||
mock_client.send_and_wait.return_value = turn_result
|
||||
mock_cls.return_value.__enter__ = MagicMock(return_value=mock_client)
|
||||
mock_cls.return_value.__exit__ = MagicMock(return_value=False)
|
||||
turn_result = TurnResult(ws_id="ws-456", timed_out=True)
|
||||
with (
|
||||
patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
|
||||
patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
|
||||
):
|
||||
mock_console = MagicMock()
|
||||
mock_console.route_create_workstream.return_value = {
|
||||
"ws_id": "ws-456",
|
||||
"node_url": "http://node-1:8080",
|
||||
"node_id": "node-1",
|
||||
}
|
||||
_mock_console_ctx(mock_console_cls, mock_console)
|
||||
|
||||
_, result = _exec_on_node_sync({"host": "localhost"}, "node-1", "sleep 9999", 1.0)
|
||||
mock_server = MagicMock()
|
||||
mock_server.send_and_wait.return_value = turn_result
|
||||
_mock_server_ctx(mock_server_cls, mock_server)
|
||||
|
||||
_, result = _exec_on_node_sync(_CONSOLE_KW, "node-1", "sleep 9999", 1.0)
|
||||
assert result.timed_out
|
||||
assert not result.ok
|
||||
# Workstream still closed even on timeout
|
||||
mock_console.route_close.assert_called_once_with("ws-456")
|
||||
|
||||
def test_send_failure_still_closes_workstream(self):
|
||||
"""Workstream must be closed even if send_and_wait raises."""
|
||||
with (
|
||||
patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls,
|
||||
patch("mcp_cluster_ops.server.TurnstoneServer") as mock_server_cls,
|
||||
):
|
||||
mock_console = MagicMock()
|
||||
mock_console.route_create_workstream.return_value = {
|
||||
"ws_id": "ws-789",
|
||||
"node_url": "http://node-1:8080",
|
||||
"node_id": "node-1",
|
||||
}
|
||||
_mock_console_ctx(mock_console_cls, mock_console)
|
||||
|
||||
mock_server = MagicMock()
|
||||
mock_server.send_and_wait.side_effect = ConnectionError("lost connection")
|
||||
_mock_server_ctx(mock_server_cls, mock_server)
|
||||
|
||||
with contextlib.suppress(ConnectionError):
|
||||
_exec_on_node_sync(_CONSOLE_KW, "node-1", "echo hi", 60.0)
|
||||
|
||||
mock_console.route_close.assert_called_once_with("ws-789")
|
||||
|
||||
def test_malformed_route_response_no_leak(self):
|
||||
"""If route response is missing ws_id, no route_close is attempted."""
|
||||
with patch("mcp_cluster_ops.server.TurnstoneConsole") as mock_console_cls:
|
||||
mock_console = MagicMock()
|
||||
mock_console.route_create_workstream.return_value = {
|
||||
# Missing "ws_id" and "node_url"
|
||||
"node_id": "node-1",
|
||||
}
|
||||
_mock_console_ctx(mock_console_cls, mock_console)
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
_exec_on_node_sync(_CONSOLE_KW, "node-1", "echo hi", 60.0)
|
||||
|
||||
# route_close must NOT be called — ws_id was never assigned
|
||||
mock_console.route_close.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -88,40 +227,32 @@ class TestExecOnNodeSync:
|
||||
|
||||
class TestDispatchParallel:
|
||||
def test_parallel_success(self):
|
||||
def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
return (node_id, TurnResult(tool_results=[("bash", f"output-{node_id}")]))
|
||||
def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
return (
|
||||
node_id,
|
||||
TurnResult(tool_results=[("bash", f"output-{node_id}")]),
|
||||
)
|
||||
|
||||
with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
|
||||
results = asyncio.run(
|
||||
_dispatch_parallel(
|
||||
{"host": "localhost"},
|
||||
["a", "b", "c"],
|
||||
"echo hi",
|
||||
60.0,
|
||||
8192,
|
||||
)
|
||||
_dispatch_parallel(_CONSOLE_KW, ["a", "b", "c"], "echo hi", 60.0, 8192)
|
||||
)
|
||||
assert len(results) == 3
|
||||
assert all(r["ok"] for r in results)
|
||||
outputs = {r["node"]: r["output"] for r in results}
|
||||
assert outputs["a"] == "output-a"
|
||||
assert outputs["b"] == "output-b"
|
||||
assert outputs["c"] == "output-c"
|
||||
|
||||
def test_partial_failure(self):
|
||||
def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
if node_id == "bad":
|
||||
raise ConnectionError("connection refused")
|
||||
return (node_id, TurnResult(tool_results=[("bash", "ok")]))
|
||||
|
||||
with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
|
||||
results = asyncio.run(
|
||||
_dispatch_parallel(
|
||||
{"host": "localhost"},
|
||||
["good", "bad"],
|
||||
"echo hi",
|
||||
60.0,
|
||||
8192,
|
||||
)
|
||||
_dispatch_parallel(_CONSOLE_KW, ["good", "bad"], "echo hi", 60.0, 8192)
|
||||
)
|
||||
assert len(results) == 2
|
||||
good = next(r for r in results if r["node"] == "good")
|
||||
@@ -131,18 +262,12 @@ class TestDispatchParallel:
|
||||
assert "connection refused" in bad["error"]
|
||||
|
||||
def test_all_fail(self):
|
||||
def fake_exec(server_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
def fake_exec(console_kw: Any, node_id: str, command: str, timeout: float) -> Any:
|
||||
raise RuntimeError(f"fail-{node_id}")
|
||||
|
||||
with patch("mcp_cluster_ops.server._exec_on_node_sync", side_effect=fake_exec):
|
||||
results = asyncio.run(
|
||||
_dispatch_parallel(
|
||||
{"host": "localhost"},
|
||||
["a", "b"],
|
||||
"echo hi",
|
||||
60.0,
|
||||
8192,
|
||||
)
|
||||
_dispatch_parallel(_CONSOLE_KW, ["a", "b"], "echo hi", 60.0, 8192)
|
||||
)
|
||||
assert all(not r["ok"] for r in results)
|
||||
assert "fail-a" in results[0]["error"]
|
||||
|
||||
+6
-3
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "turnstone"
|
||||
version = "0.9.8"
|
||||
version = "1.1.1"
|
||||
description = "Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."
|
||||
readme = "README.md"
|
||||
license = "BUSL-1.1"
|
||||
@@ -12,7 +12,7 @@ requires-python = ">=3.11"
|
||||
authors = [{name = "Patrick Buckley", email = "buckleypm@gmail.com"}]
|
||||
keywords = ["ai", "chat", "llm", "agent", "tools", "openai"]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Development Status :: 5 - Production/Stable",
|
||||
"Environment :: Console",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
@@ -67,6 +67,7 @@ turnstone-bootstrap = "turnstone.bootstrap:main"
|
||||
[tool.hatch.build.targets.wheel]
|
||||
include = [
|
||||
"turnstone/**/*.py",
|
||||
"turnstone/prompts/**/*.md",
|
||||
"turnstone/tools/*.json",
|
||||
"turnstone/ui/static/*.html",
|
||||
"turnstone/ui/static/*.css",
|
||||
@@ -78,8 +79,10 @@ include = [
|
||||
"turnstone/shared_static/*.js",
|
||||
"turnstone/shared_static/katex-0.16.44/**/*",
|
||||
"turnstone/shared_static/hljs-11.11.1/**/*",
|
||||
"turnstone/shared_static/mermaid-11.13.0/**/*",
|
||||
"turnstone/shared_static/mermaid-11.14.0/**/*",
|
||||
"turnstone/shared_static/hls-1.6.15/**/*",
|
||||
"turnstone/sdk/py.typed",
|
||||
"turnstone/deploy/*.yaml",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
|
||||
Executable
+62
@@ -0,0 +1,62 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Bump version, regenerate lockfile, commit, and tag.
|
||||
#
|
||||
# Usage:
|
||||
# scripts/release.sh 1.0.0 # stable release
|
||||
# scripts/release.sh 1.1.0a1 # experimental pre-release
|
||||
# scripts/release.sh 1.0.1 --push # bump + push tag to origin
|
||||
#
|
||||
set -euo pipefail
|
||||
|
||||
VERSION="${1:?Usage: scripts/release.sh VERSION [--push]}"
|
||||
PUSH="${2:-}"
|
||||
|
||||
# Validate PEP 440 version
|
||||
if ! echo "$VERSION" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+(a[0-9]+|b[0-9]+|rc[0-9]+)?$'; then
|
||||
echo "error: invalid PEP 440 version: $VERSION" >&2
|
||||
echo " examples: 1.0.0, 1.1.0a1, 1.0.1rc2" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
TAG="v${VERSION}"
|
||||
|
||||
# Check for clean working tree
|
||||
if ! git diff --quiet || ! git diff --cached --quiet; then
|
||||
echo "error: working tree is dirty — commit or stash first" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Check tag doesn't already exist
|
||||
if git rev-parse "$TAG" >/dev/null 2>&1; then
|
||||
echo "error: tag $TAG already exists" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Detect current version
|
||||
CURRENT=$(grep -oP '(?<=^version = ")[^"]+' pyproject.toml)
|
||||
echo "Bumping $CURRENT → $VERSION"
|
||||
|
||||
# Update version in both files
|
||||
sed -i "s/^version = \".*\"/version = \"$VERSION\"/" pyproject.toml
|
||||
sed -i "s/^__version__ = \".*\"/__version__ = \"$VERSION\"/" turnstone/__init__.py
|
||||
|
||||
# Regenerate lockfile
|
||||
echo "Regenerating uv.lock..."
|
||||
uv lock
|
||||
|
||||
# Commit and tag
|
||||
git add pyproject.toml turnstone/__init__.py uv.lock
|
||||
git commit -m "chore: bump version to $VERSION"
|
||||
git tag "$TAG"
|
||||
|
||||
echo ""
|
||||
echo "Created commit and tag $TAG"
|
||||
|
||||
if [ "$PUSH" = "--push" ]; then
|
||||
BRANCH=$(git rev-parse --abbrev-ref HEAD)
|
||||
echo "Pushing $BRANCH + $TAG to origin..."
|
||||
git push origin "$BRANCH" "$TAG"
|
||||
else
|
||||
echo "Run 'git push origin <branch> $TAG' to publish"
|
||||
fi
|
||||
@@ -5,6 +5,7 @@
|
||||
# scripts/update-vendored-js.sh katex 0.16.39
|
||||
# scripts/update-vendored-js.sh hljs 11.12.0
|
||||
# scripts/update-vendored-js.sh mermaid 11.14.0
|
||||
# scripts/update-vendored-js.sh hls 1.6.15
|
||||
#
|
||||
# This script:
|
||||
# 1. Downloads the new version from CDN
|
||||
@@ -18,7 +19,7 @@ STATIC_DIR="turnstone/shared_static"
|
||||
CDN="https://cdn.jsdelivr.net/npm"
|
||||
|
||||
usage() {
|
||||
echo "Usage: $0 <katex|hljs|mermaid> <version>"
|
||||
echo "Usage: $0 <katex|hljs|mermaid|hls> <version>"
|
||||
echo "Example: $0 katex 0.16.39"
|
||||
exit 1
|
||||
}
|
||||
@@ -147,6 +148,32 @@ case "$LIB" in
|
||||
echo "Done. Old directory removed: ${OLD_DIR}"
|
||||
;;
|
||||
|
||||
hls)
|
||||
OLD_VERSION=$(detect_old_version "hls")
|
||||
check_same_version "$OLD_VERSION" "$VERSION" "hls"
|
||||
OLD_DIR="${STATIC_DIR}/hls-${OLD_VERSION}"
|
||||
NEW_DIR="${STATIC_DIR}/hls-${VERSION}"
|
||||
|
||||
echo "Updating hls.js ${OLD_VERSION} -> ${VERSION}"
|
||||
mkdir -p "${NEW_DIR}"
|
||||
|
||||
echo " Downloading hls.min.js..."
|
||||
curl -sSfL "${CDN}/hls.js@${VERSION}/dist/hls.min.js" -o "${NEW_DIR}/hls.min.js"
|
||||
|
||||
echo " Downloading LICENSE..."
|
||||
if ! curl -sSfL "${CDN}/hls.js@${VERSION}/LICENSE" -o "${NEW_DIR}/LICENSE" 2>/dev/null; then
|
||||
if [[ -f "${OLD_DIR}/LICENSE" ]]; then
|
||||
cp "${OLD_DIR}/LICENSE" "${NEW_DIR}/LICENSE"
|
||||
else
|
||||
echo " WARNING: Could not obtain LICENSE for hls.js ${VERSION}"
|
||||
fi
|
||||
fi
|
||||
|
||||
update_refs "hls-${OLD_VERSION}" "hls-${VERSION}"
|
||||
rm -rf "${OLD_DIR}"
|
||||
echo "Done. Old directory removed: ${OLD_DIR}"
|
||||
;;
|
||||
|
||||
*)
|
||||
echo "Unknown library: ${LIB}"
|
||||
usage
|
||||
|
||||
Generated
+10
-10
@@ -14,22 +14,22 @@
|
||||
}
|
||||
},
|
||||
"node_modules/@emnapi/core": {
|
||||
"version": "1.9.1",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.9.1.tgz",
|
||||
"integrity": "sha512-mukuNALVsoix/w1BJwFzwXBN/dHeejQtuVzcDsfOEsdpCumXb/E9j8w11h5S54tT1xhifGfbbSm/ICrObRb3KA==",
|
||||
"version": "1.9.2",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/core/-/core-1.9.2.tgz",
|
||||
"integrity": "sha512-UC+ZhH3XtczQYfOlu3lNEkdW/p4dsJ1r/bP7H8+rhao3TTTMO1ATq/4DdIi23XuGoFY+Cz0JmCbdVl0hz9jZcA==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"@emnapi/wasi-threads": "1.2.0",
|
||||
"@emnapi/wasi-threads": "1.2.1",
|
||||
"tslib": "^2.4.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@emnapi/runtime": {
|
||||
"version": "1.9.1",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.9.1.tgz",
|
||||
"integrity": "sha512-VYi5+ZVLhpgK4hQ0TAjiQiZ6ol0oe4mBx7mVv7IflsiEp0OWoVsp/+f9Vc1hOhE0TtkORVrI1GvzyreqpgWtkA==",
|
||||
"version": "1.9.2",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/runtime/-/runtime-1.9.2.tgz",
|
||||
"integrity": "sha512-3U4+MIWHImeyu1wnmVygh5WlgfYDtyf0k8AbLhMFxOipihf6nrWC4syIm/SwEeec0mNSafiiNnMJwbza/Is6Lw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
@@ -39,9 +39,9 @@
|
||||
}
|
||||
},
|
||||
"node_modules/@emnapi/wasi-threads": {
|
||||
"version": "1.2.0",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.0.tgz",
|
||||
"integrity": "sha512-N10dEJNSsUx41Z6pZsXU8FjPjpBEplgH24sfkmITrBED1/U2Esum9F3lfLrMjKHHjmi557zQn7kR9R+XWXu5Rg==",
|
||||
"version": "1.2.1",
|
||||
"resolved": "https://registry.npmjs.org/@emnapi/wasi-threads/-/wasi-threads-1.2.1.tgz",
|
||||
"integrity": "sha512-uTII7OYF+/Mes/MrcIOYp5yOtSMLBWSIoLPpcgwipoiKbli6k322tcoFsxoIIxPDqW01SQGAgko4EzZi2BNv2w==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*
|
||||
* const client = new TurnstoneServer({
|
||||
* baseUrl: "http://localhost:8080",
|
||||
* token: "tok_xxx",
|
||||
* token: "ts_your_api_token",
|
||||
* });
|
||||
*
|
||||
* const ws = await client.createWorkstream({ name: "demo" });
|
||||
|
||||
@@ -6,6 +6,37 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Shared test auth — JWT-based
|
||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
|
||||
def _server_jwt() -> str:
|
||||
from turnstone.core.auth import JWT_AUD_SERVER, create_jwt
|
||||
|
||||
return create_jwt(
|
||||
user_id="test-versioning",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="test",
|
||||
secret=_TEST_JWT_SECRET,
|
||||
audience=JWT_AUD_SERVER,
|
||||
)
|
||||
|
||||
|
||||
def _console_jwt() -> str:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, create_jwt
|
||||
|
||||
return create_jwt(
|
||||
user_id="test-versioning",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="test",
|
||||
secret=_TEST_JWT_SECRET,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
)
|
||||
|
||||
|
||||
_SERVER_AUTH_HEADERS = {"Authorization": f"Bearer {_server_jwt()}"}
|
||||
_CONSOLE_AUTH_HEADERS = {"Authorization": f"Bearer {_console_jwt()}"}
|
||||
|
||||
|
||||
class TestServerVersioning:
|
||||
"""Test /v1/ routes and OpenAPI endpoints on the server."""
|
||||
@@ -14,7 +45,6 @@ class TestServerVersioning:
|
||||
def client(self):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.core.auth import AuthConfig
|
||||
from turnstone.server import create_app
|
||||
|
||||
mock_mgr = MagicMock()
|
||||
@@ -26,19 +56,19 @@ class TestServerVersioning:
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
def test_v1_workstreams(self, client):
|
||||
resp = client.get("/v1/api/workstreams")
|
||||
resp = client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
assert resp.status_code == 200
|
||||
assert "workstreams" in resp.json()
|
||||
|
||||
def test_unversioned_api_404(self, client):
|
||||
resp = client.get("/api/workstreams")
|
||||
resp = client.get("/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_openapi_json(self, client):
|
||||
@@ -72,7 +102,6 @@ class TestConsoleVersioning:
|
||||
|
||||
from turnstone.console.collector import ClusterCollector
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
collector = MagicMock(spec=ClusterCollector)
|
||||
@@ -84,18 +113,18 @@ class TestConsoleVersioning:
|
||||
}
|
||||
app = create_app(
|
||||
collector=collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
def test_v1_cluster_overview(self, client):
|
||||
resp = client.get("/v1/api/cluster/overview")
|
||||
resp = client.get("/v1/api/cluster/overview", headers=_CONSOLE_AUTH_HEADERS)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_unversioned_api_404(self, client):
|
||||
resp = client.get("/api/cluster/overview")
|
||||
resp = client.get("/api/cluster/overview", headers=_CONSOLE_AUTH_HEADERS)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_openapi_json(self, client):
|
||||
|
||||
+263
-332
@@ -9,12 +9,11 @@ import pytest
|
||||
|
||||
from turnstone.core.auth import (
|
||||
WRITE_PATHS,
|
||||
AuthConfig,
|
||||
_extract_bearer,
|
||||
_extract_cookie,
|
||||
check_request,
|
||||
create_jwt,
|
||||
is_public_path,
|
||||
load_auth_config,
|
||||
make_clear_cookie,
|
||||
make_set_cookie,
|
||||
required_scope,
|
||||
@@ -199,37 +198,6 @@ class TestRequiredScope:
|
||||
assert required_scope("GET", "/api/_internal/mcp-reload") == "read"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestAuthConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAuthConfig:
|
||||
def test_check_valid_full_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full", "tok_read": "read"})
|
||||
assert cfg.check("tok_full") == "full"
|
||||
|
||||
def test_check_valid_read_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full", "tok_read": "read"})
|
||||
assert cfg.check("tok_read") == "read"
|
||||
|
||||
def test_check_invalid_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
assert cfg.check("wrong") is None
|
||||
|
||||
def test_check_none_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
assert cfg.check(None) is None
|
||||
|
||||
def test_check_empty_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
assert cfg.check("") is None
|
||||
|
||||
def test_check_no_tokens(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={})
|
||||
assert cfg.check("anything") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestExtractBearer
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -353,167 +321,157 @@ class TestMakeClearCookie:
|
||||
class TestCheckRequest:
|
||||
"""Tests for the main check_request() entry point."""
|
||||
|
||||
@pytest.fixture()
|
||||
def disabled(self):
|
||||
return AuthConfig(enabled=False)
|
||||
_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
@pytest.fixture()
|
||||
def enabled(self):
|
||||
return AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
)
|
||||
def read_jwt(self):
|
||||
return f"Bearer {create_jwt('u1', frozenset({'read'}), 'test', self._SECRET)}"
|
||||
|
||||
def test_disabled_allows_all(self, disabled):
|
||||
allowed, status, msg, _result = check_request(disabled, "POST", "/api/send", None)
|
||||
@pytest.fixture()
|
||||
def full_jwt(self):
|
||||
return f"Bearer {create_jwt('u1', frozenset({'read', 'write', 'approve'}), 'test', self._SECRET)}"
|
||||
|
||||
def test_public_path_no_token_ok(self):
|
||||
allowed, status, msg, _result = check_request("GET", "/health", None)
|
||||
assert allowed is True
|
||||
assert status == 200
|
||||
|
||||
def test_disabled_allows_no_header(self, disabled):
|
||||
allowed, status, msg, _result = check_request(disabled, "GET", "/api/workstreams", None)
|
||||
def test_public_root_no_token_ok(self):
|
||||
allowed, status, msg, _result = check_request("GET", "/", None)
|
||||
assert allowed is True
|
||||
|
||||
def test_public_path_no_token_ok(self, enabled):
|
||||
allowed, status, msg, _result = check_request(enabled, "GET", "/health", None)
|
||||
assert allowed is True
|
||||
assert status == 200
|
||||
|
||||
def test_public_root_no_token_ok(self, enabled):
|
||||
allowed, status, msg, _result = check_request(enabled, "GET", "/", None)
|
||||
def test_public_static_no_token_ok(self):
|
||||
allowed, status, msg, _result = check_request("GET", "/static/style.css", None)
|
||||
assert allowed is True
|
||||
|
||||
def test_public_static_no_token_ok(self, enabled):
|
||||
allowed, status, msg, _result = check_request(enabled, "GET", "/static/style.css", None)
|
||||
assert allowed is True
|
||||
|
||||
def test_api_no_token_401(self, enabled):
|
||||
allowed, status, msg, _result = check_request(enabled, "GET", "/api/workstreams", None)
|
||||
def test_api_no_token_401(self):
|
||||
allowed, status, msg, _result = check_request("GET", "/api/workstreams", None)
|
||||
assert allowed is False
|
||||
assert status == 401
|
||||
assert "Unauthorized" in msg
|
||||
|
||||
def test_api_invalid_token_401(self, enabled):
|
||||
def test_api_invalid_token_401(self):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "GET", "/api/workstreams", "Bearer wrong_token"
|
||||
"GET", "/api/workstreams", "Bearer wrong_token"
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 401
|
||||
|
||||
def test_api_read_token_ok(self, enabled):
|
||||
def test_api_read_token_ok(self, read_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "GET", "/api/workstreams", "Bearer tok_read"
|
||||
"GET", "/api/workstreams", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
assert status == 200
|
||||
|
||||
def test_api_full_token_ok(self, enabled):
|
||||
def test_api_full_token_ok(self, full_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "GET", "/api/workstreams", "Bearer tok_full"
|
||||
"GET", "/api/workstreams", full_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_write_read_token_403(self, enabled):
|
||||
def test_write_read_token_403(self, read_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/send", "Bearer tok_read"
|
||||
"POST", "/api/send", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
assert "Forbidden" in msg
|
||||
|
||||
def test_write_full_token_ok(self, enabled):
|
||||
def test_write_full_token_ok(self, full_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/send", "Bearer tok_full"
|
||||
"POST", "/api/send", full_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
assert status == 200
|
||||
|
||||
def test_approve_read_token_403(self, enabled):
|
||||
def test_approve_read_token_403(self, read_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/approve", "Bearer tok_read"
|
||||
"POST", "/api/approve", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_proxy_write_read_token_403(self, enabled):
|
||||
def test_proxy_write_read_token_403(self, read_jwt):
|
||||
"""Read tokens cannot escalate to write ops via proxy routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/node/node-a/api/send", "Bearer tok_read"
|
||||
"POST", "/node/node-a/api/send", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_proxy_write_trailing_slash_read_token_403(self, enabled):
|
||||
def test_proxy_write_trailing_slash_read_token_403(self, read_jwt):
|
||||
"""Trailing slash must not bypass write-role check on proxy routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/node/node-a/api/send/", "Bearer tok_read"
|
||||
"POST", "/node/node-a/api/send/", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_direct_write_trailing_slash_read_token_403(self, enabled):
|
||||
def test_direct_write_trailing_slash_read_token_403(self, read_jwt):
|
||||
"""Trailing slash must not bypass write-role check on direct routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/send/", "Bearer tok_read"
|
||||
"POST", "/api/send/", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_proxy_write_full_token_ok(self, enabled):
|
||||
def test_proxy_write_full_token_ok(self, full_jwt):
|
||||
"""Full tokens pass through proxy write routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/node/node-a/api/send", "Bearer tok_full"
|
||||
"POST", "/node/node-a/api/send", full_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_proxy_v1_write_read_token_403(self, enabled):
|
||||
def test_proxy_v1_write_read_token_403(self, read_jwt):
|
||||
"""Read tokens cannot escalate to write ops via v1 proxy routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/node/node-a/v1/api/send", "Bearer tok_read"
|
||||
"POST", "/node/node-a/v1/api/send", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_proxy_v1_write_full_token_ok(self, enabled):
|
||||
def test_proxy_v1_write_full_token_ok(self, full_jwt):
|
||||
"""Full tokens pass through v1 proxy write routes."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/node/node-a/v1/api/send", "Bearer tok_full"
|
||||
"POST", "/node/node-a/v1/api/send", full_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_proxy_v1_cluster_ws_new_read_403(self, enabled):
|
||||
def test_proxy_v1_cluster_ws_new_read_403(self, read_jwt):
|
||||
"""Read tokens cannot create workstreams via v1 proxy."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/node/node-a/v1/api/cluster/workstreams/new",
|
||||
"Bearer tok_read",
|
||||
read_jwt,
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_proxy_read_endpoint_read_token_ok(self, enabled):
|
||||
def test_proxy_read_endpoint_read_token_ok(self, read_jwt):
|
||||
"""Read tokens can access proxy read endpoints."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "GET", "/node/node-a/api/workstreams", "Bearer tok_read"
|
||||
"GET", "/node/node-a/api/workstreams", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_console_create_ws_read_token_403(self, enabled):
|
||||
def test_console_create_ws_read_token_403(self, read_jwt):
|
||||
"""Read tokens cannot create workstreams."""
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/cluster/workstreams/new", "Bearer tok_read"
|
||||
"POST", "/api/cluster/workstreams/new", read_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_approve_full_token_ok(self, enabled):
|
||||
def test_approve_full_token_ok(self, full_jwt):
|
||||
allowed, status, msg, _result = check_request(
|
||||
enabled, "POST", "/api/approve", "Bearer tok_full"
|
||||
"POST", "/api/approve", full_jwt, jwt_secret=self._SECRET
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_no_auth_header_string(self, enabled):
|
||||
allowed, status, msg, _result = check_request(enabled, "GET", "/api/dashboard", "")
|
||||
def test_no_auth_header_string(self):
|
||||
allowed, status, msg, _result = check_request("GET", "/api/dashboard", "")
|
||||
assert allowed is False
|
||||
assert status == 401
|
||||
|
||||
@@ -526,70 +484,71 @@ class TestCheckRequest:
|
||||
class TestCheckRequestWithCookie:
|
||||
"""Tests for cookie-based auth fallback in check_request."""
|
||||
|
||||
@pytest.fixture()
|
||||
def enabled(self):
|
||||
return AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
)
|
||||
_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
def test_cookie_fallback_when_no_bearer(self, enabled):
|
||||
@pytest.fixture()
|
||||
def read_jwt(self):
|
||||
return create_jwt("u1", frozenset({"read"}), "test", self._SECRET)
|
||||
|
||||
@pytest.fixture()
|
||||
def full_jwt(self):
|
||||
return create_jwt("u1", frozenset({"read", "write", "approve"}), "test", self._SECRET)
|
||||
|
||||
def test_cookie_fallback_when_no_bearer(self, read_jwt):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"GET",
|
||||
"/api/workstreams",
|
||||
None,
|
||||
cookie_header="turnstone_auth=tok_read",
|
||||
cookie_header=f"turnstone_auth={read_jwt}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is True
|
||||
assert status == 200
|
||||
|
||||
def test_bearer_takes_precedence_over_cookie(self, enabled):
|
||||
# Bearer is full, cookie is read — Bearer should win
|
||||
def test_bearer_takes_precedence_over_cookie(self, read_jwt, full_jwt):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/api/send",
|
||||
"Bearer tok_full",
|
||||
cookie_header="turnstone_auth=tok_read",
|
||||
f"Bearer {full_jwt}",
|
||||
cookie_header=f"turnstone_auth={read_jwt}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_invalid_cookie_401(self, enabled):
|
||||
def test_invalid_cookie_401(self):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"GET",
|
||||
"/api/workstreams",
|
||||
None,
|
||||
cookie_header="turnstone_auth=wrong_token",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 401
|
||||
|
||||
def test_cookie_read_on_write_403(self, enabled):
|
||||
def test_cookie_read_on_write_403(self, read_jwt):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/api/send",
|
||||
None,
|
||||
cookie_header="turnstone_auth=tok_read",
|
||||
cookie_header=f"turnstone_auth={read_jwt}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is False
|
||||
assert status == 403
|
||||
|
||||
def test_cookie_full_on_write_ok(self, enabled):
|
||||
def test_cookie_full_on_write_ok(self, full_jwt):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/api/send",
|
||||
None,
|
||||
cookie_header="turnstone_auth=tok_full",
|
||||
cookie_header=f"turnstone_auth={full_jwt}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_no_cookie_no_bearer_401(self, enabled):
|
||||
def test_no_cookie_no_bearer_401(self):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"GET",
|
||||
"/api/workstreams",
|
||||
None,
|
||||
@@ -598,18 +557,16 @@ class TestCheckRequestWithCookie:
|
||||
assert allowed is False
|
||||
assert status == 401
|
||||
|
||||
def test_login_path_public(self, enabled):
|
||||
def test_login_path_public(self):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/api/auth/login",
|
||||
None,
|
||||
)
|
||||
assert allowed is True
|
||||
|
||||
def test_logout_path_public(self, enabled):
|
||||
def test_logout_path_public(self):
|
||||
allowed, status, _, _r = check_request(
|
||||
enabled,
|
||||
"POST",
|
||||
"/api/auth/logout",
|
||||
None,
|
||||
@@ -617,139 +574,6 @@ class TestCheckRequestWithCookie:
|
||||
assert allowed is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestLoadAuthConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLoadAuthConfig:
|
||||
"""Tests for load_auth_config with mocked config + env vars."""
|
||||
|
||||
def test_default_enabled(self):
|
||||
with patch("turnstone.core.config.load_config", return_value={}):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is True
|
||||
assert cfg.tokens == {}
|
||||
|
||||
def test_explicit_disable(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={"enabled": False}),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is False
|
||||
|
||||
def test_env_disable(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_ENABLED": "0"}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is False
|
||||
|
||||
def test_config_file_tokens(self):
|
||||
mock_cfg = {
|
||||
"enabled": True,
|
||||
"tokens": [
|
||||
{"value": "tok_a", "role": "full"},
|
||||
{"value": "tok_b", "role": "read"},
|
||||
],
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value=mock_cfg),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is True
|
||||
assert cfg.tokens == {"tok_a": "full", "tok_b": "read"}
|
||||
|
||||
def test_env_var_enabled(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_ENABLED": "1"}, clear=False),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is True
|
||||
|
||||
def test_env_var_token(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_TOKEN": "tok_env"}, clear=False),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert "tok_env" in cfg.tokens
|
||||
assert cfg.tokens["tok_env"] == "full"
|
||||
|
||||
def test_config_plus_env_merge(self):
|
||||
mock_cfg = {
|
||||
"enabled": True,
|
||||
"tokens": [{"value": "tok_cfg", "role": "read"}],
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value=mock_cfg),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_TOKEN": "tok_env"}, clear=False),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.tokens["tok_cfg"] == "read"
|
||||
assert cfg.tokens["tok_env"] == "full"
|
||||
|
||||
def test_invalid_role_skipped(self):
|
||||
mock_cfg = {
|
||||
"enabled": True,
|
||||
"tokens": [
|
||||
{"value": "tok_ok", "role": "full"},
|
||||
{"value": "tok_bad", "role": "admin"},
|
||||
],
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value=mock_cfg),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert "tok_ok" in cfg.tokens
|
||||
assert "tok_bad" not in cfg.tokens
|
||||
|
||||
def test_empty_value_skipped(self):
|
||||
mock_cfg = {
|
||||
"enabled": True,
|
||||
"tokens": [{"value": "", "role": "full"}],
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value=mock_cfg),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert len(cfg.tokens) == 0
|
||||
|
||||
def test_non_dict_token_entry_skipped(self):
|
||||
mock_cfg = {
|
||||
"enabled": True,
|
||||
"tokens": ["not_a_dict", {"value": "tok_ok", "role": "full"}],
|
||||
}
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value=mock_cfg),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.tokens == {"tok_ok": "full"}
|
||||
|
||||
def test_env_enabled_true(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_ENABLED": "true"}, clear=False),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is True
|
||||
|
||||
def test_env_enabled_yes(self):
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {"TURNSTONE_AUTH_ENABLED": "yes"}, clear=False),
|
||||
):
|
||||
cfg = load_auth_config()
|
||||
assert cfg.enabled is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests — actual HTTP server with auth enabled
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -785,16 +609,22 @@ class TestServerAuth:
|
||||
mock_mgr.list_all.return_value = [mock_ws]
|
||||
mock_mgr.max_workstreams = 10
|
||||
|
||||
from turnstone.core.auth import JWT_AUD_SERVER
|
||||
|
||||
cls._jwt_secret = "test-jwt-secret-minimum-32-chars!"
|
||||
cls._read_hdr = {
|
||||
"Authorization": f"Bearer {create_jwt('u1', frozenset({'read'}), 'test', cls._jwt_secret, audience=JWT_AUD_SERVER)}"
|
||||
}
|
||||
cls._full_hdr = {
|
||||
"Authorization": f"Bearer {create_jwt('u1', frozenset({'read', 'write', 'approve'}), 'test', cls._jwt_secret, audience=JWT_AUD_SERVER)}"
|
||||
}
|
||||
app = srv_mod.create_app(
|
||||
workstreams=mock_mgr,
|
||||
global_queue=queue.Queue(),
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
),
|
||||
jwt_secret=cls._jwt_secret,
|
||||
cors_origins=["*"],
|
||||
)
|
||||
cls.client = TestClient(app, raise_server_exceptions=False)
|
||||
@@ -809,7 +639,6 @@ class TestServerAuth:
|
||||
|
||||
def test_metrics_no_token_passes_auth(self):
|
||||
resp = self.client.get("/metrics")
|
||||
# Public path — should never be 401/403
|
||||
assert resp.status_code not in (401, 403)
|
||||
|
||||
def test_root_no_token_200(self):
|
||||
@@ -826,23 +655,17 @@ class TestServerAuth:
|
||||
assert "Unauthorized" in resp.json().get("error", "")
|
||||
|
||||
def test_api_workstreams_read_token_200(self):
|
||||
resp = self.client.get(
|
||||
"/v1/api/workstreams",
|
||||
headers={"Authorization": "Bearer tok_read"},
|
||||
)
|
||||
resp = self.client.get("/v1/api/workstreams", headers=self._read_hdr)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_api_workstreams_full_token_200(self):
|
||||
resp = self.client.get(
|
||||
"/v1/api/workstreams",
|
||||
headers={"Authorization": "Bearer tok_full"},
|
||||
)
|
||||
resp = self.client.get("/v1/api/workstreams", headers=self._full_hdr)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_api_send_read_token_403(self):
|
||||
resp = self.client.post(
|
||||
"/v1/api/send",
|
||||
headers={"Authorization": "Bearer tok_read"},
|
||||
headers=self._read_hdr,
|
||||
json={"message": "hello", "ws_id": "x"},
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
@@ -851,10 +674,9 @@ class TestServerAuth:
|
||||
def test_api_send_full_token_passes_auth(self):
|
||||
resp = self.client.post(
|
||||
"/v1/api/send",
|
||||
headers={"Authorization": "Bearer tok_full"},
|
||||
headers=self._full_hdr,
|
||||
json={"message": "hello", "ws_id": "nonexistent"},
|
||||
)
|
||||
# Should get 404 (unknown workstream), not 401/403
|
||||
assert resp.status_code not in (401, 403)
|
||||
|
||||
def test_api_send_no_token_401(self):
|
||||
@@ -921,12 +743,18 @@ class TestConsoleAuth:
|
||||
"aggregate": {"total_tokens": 100},
|
||||
}
|
||||
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE
|
||||
|
||||
cls._jwt_secret = "test-jwt-secret-minimum-32-chars!"
|
||||
cls._read_hdr = {
|
||||
"Authorization": f"Bearer {create_jwt('u1', frozenset({'read'}), 'test', cls._jwt_secret, audience=JWT_AUD_CONSOLE)}"
|
||||
}
|
||||
cls._full_hdr = {
|
||||
"Authorization": f"Bearer {create_jwt('u1', frozenset({'read', 'write', 'approve'}), 'test', cls._jwt_secret, audience=JWT_AUD_CONSOLE)}"
|
||||
}
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
),
|
||||
jwt_secret=cls._jwt_secret,
|
||||
)
|
||||
cls.test_client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
@@ -947,17 +775,11 @@ class TestConsoleAuth:
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_api_overview_read_token_200(self):
|
||||
resp = self.test_client.get(
|
||||
"/v1/api/cluster/overview",
|
||||
headers={"Authorization": "Bearer tok_read"},
|
||||
)
|
||||
resp = self.test_client.get("/v1/api/cluster/overview", headers=self._read_hdr)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_api_overview_full_token_200(self):
|
||||
resp = self.test_client.get(
|
||||
"/v1/api/cluster/overview",
|
||||
headers={"Authorization": "Bearer tok_full"},
|
||||
)
|
||||
resp = self.test_client.get("/v1/api/cluster/overview", headers=self._full_hdr)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_invalid_token_401(self):
|
||||
@@ -1003,16 +825,33 @@ class TestServerLogin:
|
||||
mock_mgr.list_all.return_value = [mock_ws]
|
||||
mock_mgr.max_workstreams = 10
|
||||
|
||||
# Mock storage with a test user for password login
|
||||
from turnstone.core.auth import hash_password
|
||||
|
||||
mock_storage = MagicMock()
|
||||
mock_storage.get_user_by_username.side_effect = lambda u: (
|
||||
{
|
||||
"user_id": "uid_test",
|
||||
"username": "testuser",
|
||||
"password_hash": hash_password("testpass"),
|
||||
"display_name": "Test",
|
||||
}
|
||||
if u == "testuser"
|
||||
else None
|
||||
)
|
||||
mock_storage.list_user_roles.return_value = [
|
||||
{"role_id": "builtin-admin", "scopes": "read,write,approve"}
|
||||
]
|
||||
|
||||
cls._jwt_secret = "test-jwt-secret-minimum-32-chars!"
|
||||
app = srv_mod.create_app(
|
||||
workstreams=mock_mgr,
|
||||
global_queue=queue.Queue(),
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
),
|
||||
jwt_secret=cls._jwt_secret,
|
||||
auth_storage=mock_storage,
|
||||
)
|
||||
cls.test_client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
@@ -1020,36 +859,36 @@ class TestServerLogin:
|
||||
def teardown_class(cls):
|
||||
cls.test_client.close()
|
||||
|
||||
def test_login_valid_token_sets_cookie(self):
|
||||
def test_login_config_token_rejected(self):
|
||||
"""Config token exchange is no longer allowed."""
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"token": "tok_full"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["role"] == "full"
|
||||
cookie = resp.headers.get("set-cookie", "")
|
||||
assert "turnstone_auth=tok_full" in cookie
|
||||
assert "HttpOnly" in cookie
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_login_invalid_token_401(self):
|
||||
def test_login_invalid_credentials_401(self):
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"token": "wrong"},
|
||||
json={"username": "testuser", "password": "wrong"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_login_no_auth_required(self):
|
||||
# /v1/api/auth/login is public — shouldn't require auth itself
|
||||
def test_login_password_ok(self):
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"token": "tok_read"},
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert "jwt" in data
|
||||
|
||||
def test_cookie_auth_on_api(self):
|
||||
# Login to get cookie (TestClient tracks cookies automatically)
|
||||
login_resp = self.test_client.post("/v1/api/auth/login", json={"token": "tok_read"})
|
||||
login_resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
assert login_resp.status_code == 200
|
||||
|
||||
# Use cookie to access API — TestClient forwards cookies
|
||||
@@ -1057,7 +896,10 @@ class TestServerLogin:
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_logout_clears_cookie(self):
|
||||
self.test_client.post("/v1/api/auth/login", json={"token": "tok_read"})
|
||||
self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
|
||||
# Logout
|
||||
logout_resp = self.test_client.post("/v1/api/auth/logout")
|
||||
@@ -1081,6 +923,7 @@ class TestConsoleLogin:
|
||||
|
||||
from turnstone.console.collector import ClusterCollector
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import hash_password
|
||||
|
||||
_load_static()
|
||||
|
||||
@@ -1092,12 +935,26 @@ class TestConsoleLogin:
|
||||
"aggregate": {"total_tokens": 100},
|
||||
}
|
||||
|
||||
mock_storage = MagicMock()
|
||||
mock_storage.get_user_by_username.side_effect = lambda u: (
|
||||
{
|
||||
"user_id": "uid_test",
|
||||
"username": "testuser",
|
||||
"password_hash": hash_password("testpass"),
|
||||
"display_name": "Test",
|
||||
}
|
||||
if u == "testuser"
|
||||
else None
|
||||
)
|
||||
mock_storage.list_user_roles.return_value = [
|
||||
{"role_id": "builtin-admin", "scopes": "read,write,approve"}
|
||||
]
|
||||
|
||||
cls._jwt_secret = "test-jwt-secret-minimum-32-chars!"
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(
|
||||
enabled=True,
|
||||
tokens={"tok_full": "full", "tok_read": "read"},
|
||||
),
|
||||
jwt_secret=cls._jwt_secret,
|
||||
auth_storage=mock_storage,
|
||||
)
|
||||
cls.test_client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
@@ -1105,28 +962,34 @@ class TestConsoleLogin:
|
||||
def teardown_class(cls):
|
||||
cls.test_client.close()
|
||||
|
||||
def test_login_valid_token(self):
|
||||
def test_login_config_token_rejected(self):
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"token": "tok_read"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_login_password_ok(self):
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert "turnstone_auth" in resp.headers.get("set-cookie", "")
|
||||
|
||||
def test_login_invalid_token(self):
|
||||
resp = self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"token": "wrong"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_cookie_auth_on_api(self):
|
||||
self.test_client.post("/v1/api/auth/login", json={"token": "tok_read"})
|
||||
self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
resp = self.test_client.get("/v1/api/cluster/overview")
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_logout_then_api_fails(self):
|
||||
self.test_client.post("/v1/api/auth/login", json={"token": "tok_read"})
|
||||
self.test_client.post(
|
||||
"/v1/api/auth/login",
|
||||
json={"username": "testuser", "password": "testpass"},
|
||||
)
|
||||
self.test_client.post("/v1/api/auth/logout")
|
||||
resp = self.test_client.get("/v1/api/cluster/overview")
|
||||
assert resp.status_code == 401
|
||||
@@ -1385,25 +1248,29 @@ class TestIsSecureRequest:
|
||||
|
||||
|
||||
class TestSecretStrength:
|
||||
def test_short_secret_warns(self, caplog):
|
||||
import logging
|
||||
def test_short_secret_exits(self):
|
||||
import turnstone.core.auth as auth_mod
|
||||
|
||||
from turnstone.core.auth import _MIN_SECRET_LENGTH
|
||||
old = os.environ.get("TURNSTONE_JWT_SECRET", "")
|
||||
os.environ["TURNSTONE_JWT_SECRET"] = "short"
|
||||
try:
|
||||
with pytest.raises(SystemExit):
|
||||
auth_mod.load_jwt_secret()
|
||||
finally:
|
||||
if old:
|
||||
os.environ["TURNSTONE_JWT_SECRET"] = old
|
||||
else:
|
||||
os.environ.pop("TURNSTONE_JWT_SECRET", None)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="turnstone.core.auth"):
|
||||
import turnstone.core.auth as auth_mod
|
||||
def test_missing_secret_exits(self):
|
||||
import turnstone.core.auth as auth_mod
|
||||
|
||||
old = os.environ.get("TURNSTONE_JWT_SECRET", "")
|
||||
os.environ["TURNSTONE_JWT_SECRET"] = "short"
|
||||
try:
|
||||
secret = auth_mod.load_jwt_secret()
|
||||
assert secret == "short"
|
||||
assert any(str(_MIN_SECRET_LENGTH) in r.message for r in caplog.records)
|
||||
finally:
|
||||
if old:
|
||||
os.environ["TURNSTONE_JWT_SECRET"] = old
|
||||
else:
|
||||
os.environ.pop("TURNSTONE_JWT_SECRET", None)
|
||||
with (
|
||||
patch("turnstone.core.config.load_config", return_value={}),
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
pytest.raises(SystemExit),
|
||||
):
|
||||
auth_mod.load_jwt_secret()
|
||||
|
||||
|
||||
class TestCorsConfigurable:
|
||||
@@ -1424,7 +1291,6 @@ class TestCorsConfigurable:
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(enabled=False),
|
||||
)
|
||||
client = TestClient(app)
|
||||
resp = client.get("/health", headers={"Origin": "http://evil.com"})
|
||||
@@ -1446,7 +1312,6 @@ class TestCorsConfigurable:
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(enabled=False),
|
||||
cors_origins=["http://example.com"],
|
||||
)
|
||||
client = TestClient(app)
|
||||
@@ -1502,3 +1367,69 @@ class TestOIDCPublicPaths:
|
||||
def test_oidc_callback_is_public(self):
|
||||
assert is_public_path("/api/auth/oidc/callback") is True
|
||||
assert is_public_path("/v1/api/auth/oidc/callback") is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestRequirePermissionServiceScope — service scope bypasses permission checks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRequirePermissionServiceScope:
|
||||
"""Verify require_permission() behaviour with the service scope."""
|
||||
|
||||
def _make_request(self, auth_result):
|
||||
"""Build a mock Starlette request with the given AuthResult on state."""
|
||||
request = MagicMock()
|
||||
request.state.auth_result = auth_result
|
||||
return request
|
||||
|
||||
def test_service_scope_bypasses_permission(self):
|
||||
"""Service-scoped tokens bypass all permission checks (returns None)."""
|
||||
from turnstone.core.auth import AuthResult, require_permission
|
||||
|
||||
auth = AuthResult(
|
||||
user_id="svc-agent",
|
||||
scopes=frozenset({"service"}),
|
||||
token_source="jwt",
|
||||
)
|
||||
request = self._make_request(auth)
|
||||
result = require_permission(request, "admin.users")
|
||||
assert result is None # bypass — no 403
|
||||
|
||||
def test_without_service_scope_and_without_permission_returns_403(self):
|
||||
"""Non-service tokens without the required permission get 403."""
|
||||
from turnstone.core.auth import AuthResult, require_permission
|
||||
|
||||
auth = AuthResult(
|
||||
user_id="regular-user",
|
||||
scopes=frozenset({"read", "write"}),
|
||||
token_source="jwt",
|
||||
)
|
||||
request = self._make_request(auth)
|
||||
result = require_permission(request, "admin.users")
|
||||
assert result is not None
|
||||
assert result.status_code == 403
|
||||
|
||||
def test_without_service_scope_with_permission_returns_none(self):
|
||||
"""Non-service tokens with the required permission pass."""
|
||||
from turnstone.core.auth import AuthResult, require_permission
|
||||
|
||||
auth = AuthResult(
|
||||
user_id="admin-user",
|
||||
scopes=frozenset({"read", "write", "approve"}),
|
||||
token_source="jwt",
|
||||
permissions=frozenset({"admin.users"}),
|
||||
)
|
||||
request = self._make_request(auth)
|
||||
result = require_permission(request, "admin.users")
|
||||
assert result is None # granted — no 403
|
||||
|
||||
def test_no_auth_result_returns_401(self):
|
||||
"""Missing auth_result on request state returns 401."""
|
||||
from turnstone.core.auth import require_permission
|
||||
|
||||
request = MagicMock()
|
||||
del request.state.auth_result # ensure attribute is absent
|
||||
result = require_permission(request, "admin.users")
|
||||
assert result is not None
|
||||
assert result.status_code == 401
|
||||
|
||||
+37
-57
@@ -7,7 +7,6 @@ import time
|
||||
import pytest
|
||||
|
||||
from turnstone.core.auth import (
|
||||
AuthConfig,
|
||||
AuthResult,
|
||||
_authenticate_token,
|
||||
check_request,
|
||||
@@ -203,24 +202,10 @@ class TestRequiredScope:
|
||||
|
||||
|
||||
class TestAuthenticateToken:
|
||||
def test_config_token_read(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_read": "read"})
|
||||
result = _authenticate_token("tok_read", cfg)
|
||||
assert result is not None
|
||||
assert result.scopes == frozenset({"read"})
|
||||
assert result.token_source == "config"
|
||||
|
||||
def test_config_token_full(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
result = _authenticate_token("tok_full", cfg)
|
||||
assert result is not None
|
||||
assert result.scopes == frozenset({"read", "write", "approve"})
|
||||
|
||||
def test_jwt_token(self):
|
||||
secret = "test-secret-key-for-jwt-min-32b!"
|
||||
jwt_tok = create_jwt("user1", frozenset({"read", "write"}), "db", secret)
|
||||
cfg = AuthConfig(enabled=True)
|
||||
result = _authenticate_token(jwt_tok, cfg, jwt_secret=secret)
|
||||
result = _authenticate_token(jwt_tok, jwt_secret=secret)
|
||||
assert result is not None
|
||||
assert result.user_id == "user1"
|
||||
assert result.token_source == "db"
|
||||
@@ -243,8 +228,7 @@ class TestAuthenticateToken:
|
||||
}
|
||||
return None
|
||||
|
||||
cfg = AuthConfig(enabled=True)
|
||||
result = _authenticate_token(raw, cfg, storage=MockStorage())
|
||||
result = _authenticate_token(raw, storage=MockStorage())
|
||||
assert result is not None
|
||||
assert result.user_id == "user1"
|
||||
assert result.has_scope("write")
|
||||
@@ -266,13 +250,11 @@ class TestAuthenticateToken:
|
||||
"expires": "2020-01-02T00:00:00",
|
||||
}
|
||||
|
||||
cfg = AuthConfig(enabled=True)
|
||||
result = _authenticate_token(raw, cfg, storage=MockStorage())
|
||||
result = _authenticate_token(raw, storage=MockStorage())
|
||||
assert result is None
|
||||
|
||||
def test_unknown_token(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok": "full"})
|
||||
result = _authenticate_token("unknown", cfg)
|
||||
result = _authenticate_token("unknown")
|
||||
assert result is None
|
||||
|
||||
|
||||
@@ -282,76 +264,74 @@ class TestAuthenticateToken:
|
||||
|
||||
|
||||
class TestCheckRequestScopes:
|
||||
def test_config_read_on_write_403(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_read": "read"})
|
||||
allowed, status, msg, _ = check_request(cfg, "POST", "/api/send", "Bearer tok_read")
|
||||
_SECRET = "test-secret-key-for-jwt-min-32b!"
|
||||
|
||||
def test_jwt_read_on_write_403(self):
|
||||
jwt_tok = create_jwt("u1", frozenset({"read"}), "test", self._SECRET)
|
||||
allowed, status, msg, _ = check_request(
|
||||
"POST",
|
||||
"/api/send",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert not allowed
|
||||
assert status == 403
|
||||
assert "write" in msg
|
||||
|
||||
def test_config_read_on_approve_403(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_read": "read"})
|
||||
allowed, status, msg, _ = check_request(cfg, "POST", "/api/approve", "Bearer tok_read")
|
||||
def test_jwt_read_on_approve_403(self):
|
||||
jwt_tok = create_jwt("u1", frozenset({"read"}), "test", self._SECRET)
|
||||
allowed, status, msg, _ = check_request(
|
||||
"POST",
|
||||
"/api/approve",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert not allowed
|
||||
assert status == 403
|
||||
assert "approve" in msg
|
||||
|
||||
def test_config_full_on_approve_ok(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
allowed, status, msg, result = check_request(cfg, "POST", "/api/approve", "Bearer tok_full")
|
||||
def test_jwt_full_on_approve_ok(self):
|
||||
jwt_tok = create_jwt("u1", frozenset({"read", "write", "approve"}), "test", self._SECRET)
|
||||
allowed, status, msg, result = check_request(
|
||||
"POST",
|
||||
"/api/approve",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed
|
||||
assert result is not None
|
||||
assert result.has_scope("approve")
|
||||
|
||||
def test_jwt_with_scopes(self):
|
||||
secret = "test-secret-key-for-jwt-min-32b!"
|
||||
jwt_tok = create_jwt("u1", frozenset({"read", "write"}), "db", secret)
|
||||
cfg = AuthConfig(enabled=True)
|
||||
jwt_tok = create_jwt("u1", frozenset({"read", "write"}), "db", self._SECRET)
|
||||
allowed, status, msg, result = check_request(
|
||||
cfg,
|
||||
"POST",
|
||||
"/api/send",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=secret,
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert allowed
|
||||
assert result is not None
|
||||
assert result.user_id == "u1"
|
||||
|
||||
def test_jwt_insufficient_scope(self):
|
||||
secret = "test-secret-key-for-jwt-min-32b!"
|
||||
jwt_tok = create_jwt("u1", frozenset({"read"}), "db", secret)
|
||||
cfg = AuthConfig(enabled=True)
|
||||
jwt_tok = create_jwt("u1", frozenset({"read"}), "db", self._SECRET)
|
||||
allowed, status, msg, _ = check_request(
|
||||
cfg,
|
||||
"POST",
|
||||
"/api/send",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=secret,
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert not allowed
|
||||
assert status == 403
|
||||
|
||||
def test_admin_path_requires_approve(self):
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_read": "read"})
|
||||
jwt_tok = create_jwt("u1", frozenset({"read"}), "test", self._SECRET)
|
||||
allowed, status, msg, _ = check_request(
|
||||
cfg,
|
||||
"GET",
|
||||
"/v1/api/admin/users",
|
||||
"Bearer tok_read",
|
||||
f"Bearer {jwt_tok}",
|
||||
jwt_secret=self._SECRET,
|
||||
)
|
||||
assert not allowed
|
||||
assert status == 403
|
||||
|
||||
def test_backward_compat_role_full(self):
|
||||
"""Config tokens with role='full' get all scopes."""
|
||||
cfg = AuthConfig(enabled=True, tokens={"tok_full": "full"})
|
||||
allowed, _, _, result = check_request(
|
||||
cfg,
|
||||
"GET",
|
||||
"/v1/api/admin/users",
|
||||
"Bearer tok_full",
|
||||
)
|
||||
assert allowed
|
||||
assert result is not None
|
||||
assert result.has_scope("approve")
|
||||
|
||||
+48
-1
@@ -19,6 +19,7 @@ from turnstone.bootstrap import (
|
||||
_tool_generate_secret,
|
||||
_tool_read_file,
|
||||
_tool_validate_api_key,
|
||||
_tool_write_compose,
|
||||
_tool_write_file,
|
||||
execute_tool,
|
||||
)
|
||||
@@ -103,6 +104,52 @@ class TestWriteFile:
|
||||
assert (tmp_path / "changed.txt").read_text() == "new\n"
|
||||
|
||||
|
||||
class TestWriteCompose:
|
||||
def test_writes_compose_file(self, tmp_path: Path) -> None:
|
||||
with patch("builtins.input", return_value="y"):
|
||||
result = _tool_write_compose(tmp_path, {})
|
||||
assert "written successfully" in result
|
||||
assert "ghcr.io" in result
|
||||
content = (tmp_path / "compose.yaml").read_text()
|
||||
assert "ghcr.io/turnstonelabs/turnstone" in content
|
||||
assert "TURNSTONE_IMAGE_TAG" in content
|
||||
|
||||
def test_user_declines(self, tmp_path: Path) -> None:
|
||||
with patch("builtins.input", return_value="n"):
|
||||
result = _tool_write_compose(tmp_path, {})
|
||||
assert "declined" in result
|
||||
assert not (tmp_path / "compose.yaml").exists()
|
||||
|
||||
def test_identical_content_skipped(self, tmp_path: Path) -> None:
|
||||
# Write it once
|
||||
with patch("builtins.input", return_value="y"):
|
||||
_tool_write_compose(tmp_path, {})
|
||||
# Second call should skip
|
||||
result = _tool_write_compose(tmp_path, {})
|
||||
assert "already exists" in result
|
||||
|
||||
def test_no_build_blocks(self, tmp_path: Path) -> None:
|
||||
with patch("builtins.input", return_value="y"):
|
||||
_tool_write_compose(tmp_path, {})
|
||||
content = (tmp_path / "compose.yaml").read_text()
|
||||
assert "build:" not in content
|
||||
assert "dockerfile:" not in content.lower()
|
||||
|
||||
def test_overwrites_different_content(self, tmp_path: Path) -> None:
|
||||
(tmp_path / "compose.yaml").write_text("old content\n")
|
||||
with patch("builtins.input", return_value="y"):
|
||||
result = _tool_write_compose(tmp_path, {})
|
||||
assert "written successfully" in result
|
||||
content = (tmp_path / "compose.yaml").read_text()
|
||||
assert "ghcr.io" in content
|
||||
|
||||
def test_no_local_image_references(self, tmp_path: Path) -> None:
|
||||
with patch("builtins.input", return_value="y"):
|
||||
_tool_write_compose(tmp_path, {})
|
||||
content = (tmp_path / "compose.yaml").read_text()
|
||||
assert "turnstone:local" not in content
|
||||
|
||||
|
||||
class TestGenerateSecret:
|
||||
def test_default_length(self) -> None:
|
||||
secret = _tool_generate_secret({})
|
||||
@@ -620,7 +667,7 @@ class TestConstants:
|
||||
assert func["parameters"]["type"] == "object"
|
||||
|
||||
def test_tool_count(self) -> None:
|
||||
assert len(TOOLS) == 7
|
||||
assert len(TOOLS) == 8
|
||||
|
||||
def test_all_tools_have_implementations(self) -> None:
|
||||
from turnstone.bootstrap import TOOL_FUNCTIONS
|
||||
|
||||
@@ -8,6 +8,13 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# discord.utils.escape_markdown passes 'count' as positional to re.sub,
|
||||
# which is deprecated in Python 3.13+. This is a discord.py bug (fixed
|
||||
# in newer releases); suppress here to keep the test output clean.
|
||||
pytestmark = pytest.mark.filterwarnings(
|
||||
"ignore:.*'count' is passed as positional argument:DeprecationWarning"
|
||||
)
|
||||
|
||||
discord = pytest.importorskip("discord")
|
||||
|
||||
|
||||
@@ -886,6 +893,192 @@ class TestFormatToolResult:
|
||||
assert result.count("```") == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Media embed detection and rendering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTryParseMedia:
|
||||
"""Tests for try_parse_media in _formatter.py."""
|
||||
|
||||
def test_stream_url_detected(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = json.dumps({"stream_url": "http://jf:8096/Videos/abc/stream", "container": "mp4"})
|
||||
result = try_parse_media(data)
|
||||
assert result is not None
|
||||
assert result["stream_url"] == "http://jf:8096/Videos/abc/stream"
|
||||
|
||||
def test_media_details_detected(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = json.dumps({"id": "abc", "name": "Test Movie", "type": "Movie", "year": 2024})
|
||||
result = try_parse_media(data)
|
||||
assert result is not None
|
||||
assert result["name"] == "Test Movie"
|
||||
|
||||
def test_search_results_detected(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = json.dumps({"results": [{"id": "1", "name": "Hit"}], "total_count": 1})
|
||||
result = try_parse_media(data)
|
||||
assert result is not None
|
||||
assert len(result["results"]) == 1
|
||||
|
||||
def test_sessions_detected(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = json.dumps({"sessions": [{"id": "s1", "user_name": "ptrck"}]})
|
||||
result = try_parse_media(data)
|
||||
assert result is not None
|
||||
|
||||
def test_empty_results_returns_none(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
assert try_parse_media(json.dumps({"results": []})) is None
|
||||
|
||||
def test_plain_text_returns_none(self):
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
assert try_parse_media("just a string") is None
|
||||
|
||||
def test_non_dict_json_returns_none(self):
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
assert try_parse_media("[1, 2, 3]") is None
|
||||
|
||||
def test_unrelated_dict_returns_none(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
assert try_parse_media(json.dumps({"foo": "bar"})) is None
|
||||
|
||||
|
||||
class TestIsSafeImageUrl:
|
||||
"""Tests for _is_safe_image_url in _formatter.py."""
|
||||
|
||||
def test_http_url(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("http://jellyfin:8096/Items/abc/Images/Primary") is True
|
||||
|
||||
def test_https_url(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("https://jellyfin.example.com/Items/abc/Images/Primary") is True
|
||||
|
||||
def test_ftp_rejected(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("ftp://evil.com/image.jpg") is False
|
||||
|
||||
def test_file_rejected(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("file:///etc/passwd") is False
|
||||
|
||||
def test_userinfo_rejected(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("http://user:pass@jellyfin:8096/image") is False
|
||||
|
||||
def test_empty_rejected(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("") is False
|
||||
|
||||
def test_private_ip_allowed(self):
|
||||
from turnstone.channels._formatter import _is_safe_image_url
|
||||
|
||||
assert _is_safe_image_url("http://192.168.0.6:8096/Items/abc/Images/Primary") is True
|
||||
|
||||
|
||||
class TestBuildMediaEmbed:
|
||||
"""Tests for try_build_media_embed and embed builders."""
|
||||
|
||||
def test_single_item_embed_uses_web_url_not_stream_url(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = {
|
||||
"name": "Test Movie",
|
||||
"type": "Movie",
|
||||
"year": 2024,
|
||||
"stream_url": "http://jf:8096/Videos/abc/stream?api_key=SECRET",
|
||||
"web_url": "http://jf:8096/web/#/details?id=abc",
|
||||
"overview": "A test movie.",
|
||||
}
|
||||
parsed = try_parse_media(json.dumps(data))
|
||||
assert parsed is not None
|
||||
|
||||
from turnstone.channels._formatter import _build_single_media_embed
|
||||
|
||||
embed = _build_single_media_embed(parsed, "mcp__mediamcp__get_stream_url")
|
||||
# web_url should be the embed URL, never stream_url
|
||||
assert embed.url == "http://jf:8096/web/#/details?id=abc"
|
||||
assert "SECRET" not in str(embed.to_dict())
|
||||
|
||||
def test_search_results_embed_format(self):
|
||||
import json
|
||||
|
||||
from turnstone.channels._formatter import try_parse_media
|
||||
|
||||
data = {
|
||||
"results": [
|
||||
{"name": "Movie A", "year": 2020, "type": "Movie", "runtime_minutes": 120},
|
||||
{"name": "Movie B", "year": 2021, "type": "Movie"},
|
||||
],
|
||||
"total_count": 2,
|
||||
}
|
||||
parsed = try_parse_media(json.dumps(data))
|
||||
|
||||
from turnstone.channels._formatter import _build_search_results_embed
|
||||
|
||||
embed = _build_search_results_embed(parsed)
|
||||
assert "Movie A" in embed.description
|
||||
assert "Movie B" in embed.description
|
||||
assert "2 of 2" in embed.footer.text
|
||||
|
||||
def test_build_media_embed_returns_none_for_plain_text(self):
|
||||
from turnstone.channels._formatter import try_build_media_embed
|
||||
|
||||
http = MagicMock()
|
||||
result = _run(try_build_media_embed("tool", "plain text", http=http))
|
||||
assert result is None
|
||||
|
||||
def test_season_episode_string_values(self):
|
||||
"""Season/episode numbers as strings should not raise."""
|
||||
|
||||
from turnstone.channels._formatter import _build_search_results_embed
|
||||
|
||||
data = {
|
||||
"results": [
|
||||
{
|
||||
"name": "Pilot",
|
||||
"type": "Episode",
|
||||
"series_name": "Show",
|
||||
"season_number": "1",
|
||||
"episode_number": "1",
|
||||
},
|
||||
],
|
||||
"total_count": 1,
|
||||
}
|
||||
embed = _build_search_results_embed(data)
|
||||
assert "S01E01" in embed.description
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Thinking indicator lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1103,6 +1296,7 @@ class TestToolResultEvent:
|
||||
bot._tool_info_msgs = {}
|
||||
bot._pending_approval_msgs = {}
|
||||
bot._notify_reply_channels = {}
|
||||
bot._http_client = MagicMock()
|
||||
bot._should_auto_approve = MagicMock(return_value=False)
|
||||
bot._on_ws_event = TurnstoneBot._on_ws_event.__get__(bot, TurnstoneBot)
|
||||
return bot
|
||||
|
||||
@@ -105,7 +105,8 @@ class TestDelete:
|
||||
assert store.get("tools.timeout") == defn.default
|
||||
|
||||
def test_returns_false_for_non_existent(self, store):
|
||||
assert store.delete("tools.timeout") is False
|
||||
result = store.delete("tools.timeout")
|
||||
assert result is False
|
||||
|
||||
def test_rejects_unknown_key(self, store):
|
||||
with pytest.raises(ValueError, match="Unknown setting"):
|
||||
|
||||
+35
-24
@@ -9,6 +9,24 @@ import pytest
|
||||
|
||||
from turnstone.console.collector import ClusterCollector, NodeSnapshot
|
||||
|
||||
# Shared test auth — JWT-based
|
||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
|
||||
def _test_jwt() -> str:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, create_jwt
|
||||
|
||||
return create_jwt(
|
||||
user_id="test-console",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="test",
|
||||
secret=_TEST_JWT_SECRET,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
)
|
||||
|
||||
|
||||
_TEST_AUTH_HEADERS = {"Authorization": f"Bearer {_test_jwt()}"}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mock storage for collector tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -21,7 +39,7 @@ class MockStorage:
|
||||
self.services: list[dict[str, str]] = []
|
||||
|
||||
def list_services(self, service_type: str, max_age_seconds: int = 120) -> list[dict[str, str]]:
|
||||
return [s for s in self.services if True] # all services match
|
||||
return list(self.services)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -711,13 +729,11 @@ class TestConsoleHTTPEndpoints:
|
||||
|
||||
_load_static()
|
||||
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
@@ -953,12 +969,11 @@ class TestConsoleWorkstreamCreation:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
|
||||
# Set up a mock proxy_client (lifespan doesn't run in TestClient)
|
||||
@@ -974,7 +989,7 @@ class TestConsoleWorkstreamCreation:
|
||||
mock_proxy.post = mock_post
|
||||
app.state.proxy_client = mock_proxy
|
||||
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
yield client, mock_post
|
||||
client.close()
|
||||
|
||||
@@ -1151,14 +1166,13 @@ class TestConsoleProxy:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
@@ -1322,14 +1336,13 @@ class TestConsoleVersionEndpoints:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
app = create_app(
|
||||
collector=mock_collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
@@ -1364,7 +1377,6 @@ class TestSharedStatic:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
collector = MagicMock(spec=ClusterCollector)
|
||||
@@ -1376,9 +1388,9 @@ class TestSharedStatic:
|
||||
}
|
||||
app = create_app(
|
||||
collector=collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
@@ -1481,7 +1493,6 @@ class TestProxySharedStatic:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
collector = MagicMock(spec=ClusterCollector)
|
||||
@@ -1494,9 +1505,9 @@ class TestProxySharedStatic:
|
||||
collector.get_node_detail.return_value = None
|
||||
app = create_app(
|
||||
collector=collector,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
client = TestClient(app, raise_server_exceptions=False, headers=_TEST_AUTH_HEADERS)
|
||||
resp = client.get("/node/unknown/shared/base.css")
|
||||
assert resp.status_code == 404
|
||||
client.close()
|
||||
@@ -1815,14 +1826,14 @@ class TestProxyAuthHeaders:
|
||||
# Should use ServiceTokenManager, not mint a user JWT
|
||||
assert headers["Authorization"] == f"Bearer {mgr.token}"
|
||||
|
||||
def test_fallback_static_token(self):
|
||||
"""No auth_result, no ServiceTokenManager → uses static proxy_auth_token."""
|
||||
def test_no_mgr_no_user_returns_empty(self):
|
||||
"""No auth_result, no ServiceTokenManager → empty headers."""
|
||||
from turnstone.console.server import _proxy_auth_headers
|
||||
|
||||
req = self._make_request(proxy_auth_token="static-tok-123")
|
||||
req = self._make_request()
|
||||
headers = _proxy_auth_headers(req)
|
||||
|
||||
assert headers == {"Authorization": "Bearer static-tok-123"}
|
||||
assert headers == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -13,6 +13,24 @@ from turnstone.console.collector import ClusterCollector
|
||||
from turnstone.console.router import ConsoleRouter, NodeRef
|
||||
from turnstone.core.hash_ring import NoAvailableNodeError
|
||||
|
||||
# Shared test auth — JWT-based
|
||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
|
||||
def _test_jwt() -> str:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, create_jwt
|
||||
|
||||
return create_jwt(
|
||||
user_id="test-routing",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="test",
|
||||
secret=_TEST_JWT_SECRET,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
)
|
||||
|
||||
|
||||
_TEST_AUTH_HEADERS: dict[str, str] = {"Authorization": f"Bearer {_test_jwt()}"}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -42,12 +60,11 @@ def _make_app(
|
||||
router: Any = None,
|
||||
) -> Any:
|
||||
from turnstone.console.server import _load_static, create_app
|
||||
from turnstone.core.auth import AuthConfig
|
||||
|
||||
_load_static()
|
||||
return create_app(
|
||||
collector=collector or _make_mock_collector(),
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
router=router,
|
||||
)
|
||||
|
||||
@@ -100,6 +117,7 @@ class TestRouteCreate:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test-ws"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
@@ -109,6 +127,7 @@ class TestRouteCreate:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test-ws"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
@@ -126,6 +145,7 @@ class TestRouteCreate:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"resume_ws": "old_ws_id"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
@@ -150,6 +170,7 @@ class TestRouteCreate:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"target_node": "node-c"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
@@ -203,6 +224,7 @@ class TestRouteCreate503Retry:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test-ws"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
@@ -233,6 +255,7 @@ class TestRouteProxy:
|
||||
resp = client.post(
|
||||
"/v1/api/route/send",
|
||||
json={"ws_id": "abc123", "message": "hello"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
# Verify upstream URL was /v1/api/send (not /v1/api/route/send)
|
||||
@@ -245,6 +268,7 @@ class TestRouteProxy:
|
||||
resp = client.post(
|
||||
"/v1/api/route/approve",
|
||||
json={"ws_id": "abc123", "approved": True},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@@ -252,6 +276,7 @@ class TestRouteProxy:
|
||||
resp = client.post(
|
||||
"/v1/api/route/cancel",
|
||||
json={"ws_id": "abc123"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@@ -259,6 +284,7 @@ class TestRouteProxy:
|
||||
resp = client.post(
|
||||
"/v1/api/route/command",
|
||||
json={"ws_id": "abc123", "command": "status"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@@ -266,6 +292,7 @@ class TestRouteProxy:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/close",
|
||||
json={"ws_id": "abc123"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@@ -288,14 +315,14 @@ class TestRouteLookup:
|
||||
client.close()
|
||||
|
||||
def test_route_lookup(self, client):
|
||||
resp = client.get("/v1/api/route?ws_id=abc123")
|
||||
resp = client.get("/v1/api/route?ws_id=abc123", headers=_TEST_AUTH_HEADERS)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["node_url"] == "http://a:8080"
|
||||
assert data["node_id"] == "node-a"
|
||||
|
||||
def test_route_lookup_missing_ws_id(self, client):
|
||||
resp = client.get("/v1/api/route")
|
||||
resp = client.get("/v1/api/route", headers=_TEST_AUTH_HEADERS)
|
||||
assert resp.status_code == 400
|
||||
assert "ws_id" in resp.json()["error"]
|
||||
|
||||
@@ -329,6 +356,7 @@ class TestRouteNotReady:
|
||||
resp = client_no_router.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
@@ -336,6 +364,7 @@ class TestRouteNotReady:
|
||||
resp = client_empty_cache.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
@@ -343,22 +372,24 @@ class TestRouteNotReady:
|
||||
resp = client_no_router.post(
|
||||
"/v1/api/route/send",
|
||||
json={"ws_id": "abc", "message": "hello"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_route_lookup_no_router_503(self, client_no_router):
|
||||
resp = client_no_router.get("/v1/api/route?ws_id=abc")
|
||||
resp = client_no_router.get("/v1/api/route?ws_id=abc", headers=_TEST_AUTH_HEADERS)
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_route_proxy_empty_cache_503(self, client_empty_cache):
|
||||
resp = client_empty_cache.post(
|
||||
"/v1/api/route/send",
|
||||
json={"ws_id": "abc", "message": "hello"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_route_lookup_empty_cache_503(self, client_empty_cache):
|
||||
resp = client_empty_cache.get("/v1/api/route?ws_id=abc")
|
||||
resp = client_empty_cache.get("/v1/api/route?ws_id=abc", headers=_TEST_AUTH_HEADERS)
|
||||
assert resp.status_code == 503
|
||||
|
||||
|
||||
@@ -384,6 +415,7 @@ class TestRouteNoNode:
|
||||
resp = client.post(
|
||||
"/v1/api/route/workstreams/new",
|
||||
json={"name": "test"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
assert "No available node" in resp.json()["error"]
|
||||
@@ -392,9 +424,10 @@ class TestRouteNoNode:
|
||||
resp = client.post(
|
||||
"/v1/api/route/send",
|
||||
json={"ws_id": "abc", "message": "hello"},
|
||||
headers=_TEST_AUTH_HEADERS,
|
||||
)
|
||||
assert resp.status_code == 503
|
||||
|
||||
def test_route_lookup_no_node_503(self, client):
|
||||
resp = client.get("/v1/api/route?ws_id=abc")
|
||||
resp = client.get("/v1/api/route?ws_id=abc", headers=_TEST_AUTH_HEADERS)
|
||||
assert resp.status_code == 503
|
||||
|
||||
+11
-11
@@ -154,7 +154,7 @@ class TestSingleEdit:
|
||||
)
|
||||
assert result["needs_approval"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "applied 1 edit" in msg
|
||||
with open(path) as f:
|
||||
assert f.read() == "foo\nbar\nbaz\n"
|
||||
@@ -172,7 +172,7 @@ class TestSingleEdit:
|
||||
assert result["needs_approval"]
|
||||
assert "deletion" in result["preview"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
with open(sample_file) as f:
|
||||
assert f.read() == "line1\nline2\nline4\nline5\n"
|
||||
|
||||
@@ -196,7 +196,7 @@ class TestBatchEdit:
|
||||
assert result["needs_approval"]
|
||||
assert "2 edits" in result["header"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "applied 2 edits" in msg
|
||||
with open(sample_file) as f:
|
||||
assert f.read() == "first\nline2\nline3\nline4\nlast\n"
|
||||
@@ -216,7 +216,7 @@ class TestBatchEdit:
|
||||
)
|
||||
assert result["needs_approval"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "applied 3 edits" in msg
|
||||
with open(sample_file) as f:
|
||||
assert f.read() == "line1\nsecond\nthird\nfourth\nline5\n"
|
||||
@@ -238,7 +238,7 @@ class TestBatchEdit:
|
||||
)
|
||||
assert result["needs_approval"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "overlap" in msg.lower()
|
||||
# File should be untouched
|
||||
with open(path) as f:
|
||||
@@ -305,7 +305,7 @@ class TestBatchEdit:
|
||||
)
|
||||
assert result["needs_approval"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "applied 2 edits" in msg
|
||||
with open(path) as f:
|
||||
assert f.read() == "first_foo\nbar\nsecond_foo\nbaz\n"
|
||||
@@ -324,7 +324,7 @@ class TestBatchEdit:
|
||||
)
|
||||
assert result["needs_approval"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
with open(sample_file) as f:
|
||||
assert f.read() == "line1\nline3\nline5\n"
|
||||
|
||||
@@ -344,7 +344,7 @@ class TestBatchEdit:
|
||||
# Single edit — no "(N edits)" count in header
|
||||
assert "edits)" not in result["header"]
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "applied 1 edit" in msg
|
||||
with open(sample_file) as f:
|
||||
assert f.read() == "line1\nline2\nmiddle\nline4\nline5\n"
|
||||
@@ -414,7 +414,7 @@ class TestExecEdgeCases:
|
||||
with open(sample_file, "w") as f:
|
||||
f.write("completely different content\n")
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "no longer found" in msg
|
||||
|
||||
def test_file_deleted_between_prepare_and_exec(self, session, sample_file):
|
||||
@@ -431,7 +431,7 @@ class TestExecEdgeCases:
|
||||
|
||||
os.unlink(sample_file)
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "Error" in msg
|
||||
|
||||
def test_batch_file_changed_partial_match(self, session, sample_file):
|
||||
@@ -453,7 +453,7 @@ class TestExecEdgeCases:
|
||||
with open(sample_file, "w") as f:
|
||||
f.write("line1\nline2\nline3\nline4\n")
|
||||
|
||||
call_id, msg = session._exec_edit_file(result)
|
||||
_, msg = session._exec_edit_file(result)
|
||||
assert "no longer found" in msg
|
||||
# line1 should NOT have been edited (atomic failure)
|
||||
with open(sample_file) as f:
|
||||
|
||||
@@ -57,6 +57,7 @@ class _InjectAuthMiddleware(BaseHTTPMiddleware):
|
||||
"admin.users",
|
||||
"admin.orgs",
|
||||
"admin.policies",
|
||||
"admin.prompt_policies",
|
||||
"admin.skills",
|
||||
"admin.usage",
|
||||
"admin.audit",
|
||||
|
||||
@@ -41,6 +41,21 @@ class TestHashRingBuckets:
|
||||
# Empty list returns 0
|
||||
assert storage.assign_buckets([], "node-x") == 0
|
||||
|
||||
def test_assign_large_list_exceeds_chunk_size(self, storage):
|
||||
"""Regression: lists larger than chunk_size must not hit param limits."""
|
||||
n = 1200 # exceeds SQLite chunk_size (500) and exercises multi-chunk path
|
||||
storage.seed_ring_buckets([(i, "node-a") for i in range(n)])
|
||||
count = storage.assign_buckets(list(range(n)), "node-b")
|
||||
assert count == n
|
||||
rows = storage.list_ring_buckets()
|
||||
assert all(r["node_id"] == "node-b" for r in rows)
|
||||
|
||||
def test_assign_deduplicates_input(self, storage):
|
||||
"""Duplicates in the input list should not inflate rowcount."""
|
||||
storage.seed_ring_buckets([(0, "node-a"), (1, "node-a")])
|
||||
count = storage.assign_buckets([0, 1, 0, 1, 0], "node-b")
|
||||
assert count == 2
|
||||
|
||||
|
||||
class TestBucketStats:
|
||||
def test_increment_creates_row(self, storage):
|
||||
|
||||
@@ -37,3 +37,55 @@ class TestStripHtml:
|
||||
def test_self_closing_tags(self):
|
||||
result = strip_html("hello<br/>world")
|
||||
assert result == "helloworld"
|
||||
|
||||
# -- invisible element stripping -----------------------------------------
|
||||
|
||||
def test_strips_script_content(self):
|
||||
html = "<p>before</p><script>var x = 1;</script><p>after</p>"
|
||||
result = strip_html(html)
|
||||
assert "var x" not in result
|
||||
assert "before" in result
|
||||
assert "after" in result
|
||||
|
||||
def test_strips_style_content(self):
|
||||
html = "<style>.foo { color: red; }</style><p>visible</p>"
|
||||
result = strip_html(html)
|
||||
assert "color" not in result
|
||||
assert "visible" in result
|
||||
|
||||
def test_strips_template_content(self):
|
||||
html = "<template><div>hidden</div></template><p>shown</p>"
|
||||
result = strip_html(html)
|
||||
assert "hidden" not in result
|
||||
assert "shown" in result
|
||||
|
||||
def test_strips_noscript_content(self):
|
||||
html = "<noscript>Enable JS</noscript><p>content</p>"
|
||||
result = strip_html(html)
|
||||
assert "Enable JS" not in result
|
||||
assert "content" in result
|
||||
|
||||
def test_strips_multiple_script_blocks(self):
|
||||
html = "<script>a()</script><p>middle</p><script>b()</script>"
|
||||
result = strip_html(html)
|
||||
assert "a()" not in result
|
||||
assert "b()" not in result
|
||||
assert "middle" in result
|
||||
|
||||
def test_strips_multiline_script(self):
|
||||
html = "<script>\nfunction foo() {\n return 1;\n}\n</script><p>ok</p>"
|
||||
result = strip_html(html)
|
||||
assert "function" not in result
|
||||
assert "ok" in result
|
||||
|
||||
def test_strips_script_case_insensitive(self):
|
||||
html = "<SCRIPT>code()</SCRIPT><p>text</p>"
|
||||
result = strip_html(html)
|
||||
assert "code()" not in result
|
||||
assert "text" in result
|
||||
|
||||
def test_strips_script_with_attributes(self):
|
||||
html = '<script type="text/javascript" src="app.js">init();</script><p>done</p>'
|
||||
result = strip_html(html)
|
||||
assert "init()" not in result
|
||||
assert "done" in result
|
||||
|
||||
+409
-3
@@ -3,7 +3,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import time
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
@@ -141,7 +143,7 @@ class TestMcpToOpenai:
|
||||
assert result["type"] == "function"
|
||||
func = result["function"]
|
||||
assert func["name"] == "mcp__github__search_repos"
|
||||
assert "[MCP: github]" in func["description"]
|
||||
assert func["description"] == "Search GitHub repos"
|
||||
assert func["parameters"]["type"] == "object"
|
||||
assert "query" in func["parameters"]["properties"]
|
||||
|
||||
@@ -164,7 +166,7 @@ class TestMcpToOpenai:
|
||||
tool.description = ""
|
||||
tool.inputSchema = {"type": "object", "properties": {}}
|
||||
result = _mcp_to_openai("test", tool)
|
||||
assert result["function"]["description"] == "[MCP: test] "
|
||||
assert result["function"]["description"] == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -304,7 +306,7 @@ class TestMCPClientManager:
|
||||
def test_call_tool_sync_disconnected_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._tool_map["mcp__dead__ping"] = ("dead", "ping")
|
||||
# No session registered for "dead"
|
||||
# No session registered for "dead", no config/loop → reconnect fails
|
||||
with pytest.raises(RuntimeError, match="not connected"):
|
||||
mgr.call_tool_sync("mcp__dead__ping", {})
|
||||
|
||||
@@ -1553,3 +1555,407 @@ class TestSafeCloseStack:
|
||||
await MCPClientManager._safe_close_stack(stack)
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 1: Cancel orphaned futures on timeout
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFutureCancellation:
|
||||
"""Verify future.cancel() is called when sync bridge methods time out."""
|
||||
|
||||
def _make_manager_with_session(self) -> MCPClientManager:
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
# Prevent auto-spec from creating async coroutines that trigger warnings
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mock_session.read_resource = MagicMock(return_value="sentinel")
|
||||
mock_session.get_prompt = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__search"] = ("test", "search")
|
||||
mgr._resource_map["file:///a.txt"] = ("test", "file:///a.txt")
|
||||
mgr._prompt_map["mcp__test__review"] = ("test", "review")
|
||||
return mgr
|
||||
|
||||
def test_call_tool_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__search", {"query": "x"}, timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_read_resource_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.read_resource_sync("file:///a.txt", timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_get_prompt_sync_cancels_future_on_timeout(self):
|
||||
mgr = self._make_manager_with_session()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.get_prompt_sync("mcp__test__review", timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
def test_refresh_sync_cancels_future_on_timeout(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._loop = MagicMock()
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch.object(mgr, "_refresh_all", return_value=MagicMock()),
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError, match="timed out"),
|
||||
):
|
||||
mgr.refresh_sync(timeout=1)
|
||||
mock_future.cancel.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 2: Per-server circuit breaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCircuitBreaker:
|
||||
"""Verify per-server circuit breaker behavior."""
|
||||
|
||||
def test_circuit_stays_closed_below_threshold(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._cb_record_failure("srv")
|
||||
is_open, _ = mgr._cb_check("srv")
|
||||
assert not is_open
|
||||
|
||||
def test_circuit_opens_at_threshold(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
is_open, cooldown_expired = mgr._cb_check("srv")
|
||||
assert is_open
|
||||
assert not cooldown_expired # just opened, cooldown not expired
|
||||
|
||||
def test_circuit_half_open_after_cooldown(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
# Simulate cooldown expiry
|
||||
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
||||
is_open, cooldown_expired = mgr._cb_check("srv")
|
||||
assert is_open
|
||||
assert cooldown_expired
|
||||
|
||||
def test_circuit_resets_on_success(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
assert "srv" in mgr._circuit_open_until
|
||||
mgr._cb_record_success("srv")
|
||||
is_open, _ = mgr._cb_check("srv")
|
||||
assert not is_open
|
||||
assert mgr._consecutive_failures.get("srv") is None
|
||||
|
||||
def test_success_decays_trip_count(self):
|
||||
"""Success decays trip_count by 1 so flapping servers escalate backoff."""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._circuit_trip_count["srv"] = 3
|
||||
mgr._cb_record_success("srv")
|
||||
assert mgr._circuit_trip_count["srv"] == 2
|
||||
mgr._cb_record_success("srv")
|
||||
assert mgr._circuit_trip_count["srv"] == 1
|
||||
mgr._cb_record_success("srv")
|
||||
assert "srv" not in mgr._circuit_trip_count
|
||||
|
||||
def test_cooldown_is_exponential(self):
|
||||
mgr = MCPClientManager({})
|
||||
# First trip (trip_count starts at 0)
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline1 = mgr._circuit_open_until["srv"]
|
||||
base1 = deadline1 - time.monotonic()
|
||||
# Reset circuit but keep trip_count at 1 (set by first trip)
|
||||
mgr._cb_record_success("srv")
|
||||
# trip_count decayed from 1 to 0 — manually set to 1 for test
|
||||
mgr._circuit_trip_count["srv"] = 1
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline2 = mgr._circuit_open_until["srv"]
|
||||
base2 = deadline2 - time.monotonic()
|
||||
# Second trip should have longer cooldown (roughly 2x, within jitter)
|
||||
assert base2 > base1 * 1.5
|
||||
|
||||
def test_cooldown_capped_at_max(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._circuit_trip_count["srv"] = 100 # very high trip count
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
deadline = mgr._circuit_open_until["srv"]
|
||||
cooldown = deadline - time.monotonic()
|
||||
# Should not exceed max (300s) + 10% jitter = 330s
|
||||
assert cooldown <= mgr._CB_MAX_COOLDOWN * 1.11
|
||||
|
||||
def test_cb_gate_rejects_when_open(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
with pytest.raises(RuntimeError, match="circuit open"):
|
||||
mgr._cb_gate("srv")
|
||||
|
||||
def test_cb_gate_allows_after_cooldown(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
||||
# Should not raise
|
||||
mgr._cb_gate("srv")
|
||||
# Deadline should be removed (half-open probe allowed)
|
||||
assert "srv" not in mgr._circuit_open_until
|
||||
|
||||
def test_cb_clear_removes_all_state(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
mgr._cb_clear("srv")
|
||||
assert "srv" not in mgr._consecutive_failures
|
||||
assert "srv" not in mgr._circuit_open_until
|
||||
assert "srv" not in mgr._circuit_trip_count
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
|
||||
@pytest.mark.filterwarnings("ignore:coroutine.*was never awaited:RuntimeWarning")
|
||||
def test_call_tool_sync_records_failure_on_timeout(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(TimeoutError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=1)
|
||||
assert mgr._consecutive_failures.get("test", 0) == 1
|
||||
|
||||
def test_call_tool_sync_records_success(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
# Pre-set a failure
|
||||
mgr._consecutive_failures["test"] = 2
|
||||
mock_result = MagicMock()
|
||||
mock_result.content = []
|
||||
mock_result.isError = False
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.return_value = mock_result
|
||||
with patch("asyncio.run_coroutine_threadsafe", return_value=mock_future):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
assert mgr._consecutive_failures.get("test") is None
|
||||
|
||||
def test_connection_error_evicts_session(self):
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = BrokenPipeError("dead")
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(BrokenPipeError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
assert "test" not in mgr._sessions
|
||||
|
||||
def test_independent_circuits_per_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("a")
|
||||
is_open_a, _ = mgr._cb_check("a")
|
||||
is_open_b, _ = mgr._cb_check("b")
|
||||
assert is_open_a
|
||||
assert not is_open_b
|
||||
|
||||
def test_mcp_error_does_not_trip_circuit(self):
|
||||
"""Protocol errors (McpError) should not count as transport failures."""
|
||||
from mcp import McpError
|
||||
from mcp.types import ErrorData
|
||||
|
||||
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = MagicMock(return_value="sentinel")
|
||||
mgr._sessions["test"] = mock_session
|
||||
mgr._loop = MagicMock()
|
||||
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
||||
mock_future = MagicMock()
|
||||
mock_future.result.side_effect = McpError(ErrorData(code=-32601, message="tool not found"))
|
||||
with (
|
||||
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
||||
pytest.raises(McpError),
|
||||
):
|
||||
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
||||
# Circuit should NOT have recorded a failure
|
||||
assert mgr._consecutive_failures.get("test", 0) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 3: Safe transport stream pre-close
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSafeTransportStreams:
|
||||
"""Verify stream references are stored and pre-closed."""
|
||||
|
||||
def test_pre_close_streams_closes_both(self):
|
||||
mgr = MCPClientManager({})
|
||||
stream_a = MagicMock()
|
||||
stream_b = MagicMock()
|
||||
mgr._server_streams["srv"] = (stream_a, stream_b)
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("srv")
|
||||
|
||||
asyncio.run(_run())
|
||||
stream_a.aclose.assert_called_once()
|
||||
stream_b.aclose.assert_called_once()
|
||||
assert "srv" not in mgr._server_streams
|
||||
|
||||
def test_pre_close_streams_ignores_missing(self):
|
||||
mgr = MCPClientManager({})
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("nonexistent")
|
||||
|
||||
asyncio.run(_run()) # should not raise
|
||||
|
||||
def test_pre_close_streams_suppresses_errors(self):
|
||||
mgr = MCPClientManager({})
|
||||
stream_a = MagicMock()
|
||||
stream_a.aclose.side_effect = RuntimeError("boom")
|
||||
stream_b = MagicMock()
|
||||
mgr._server_streams["srv"] = (stream_a, stream_b)
|
||||
|
||||
async def _run():
|
||||
await mgr._pre_close_streams("srv")
|
||||
|
||||
asyncio.run(_run()) # should not raise despite stream_a error
|
||||
stream_b.aclose.assert_called_once()
|
||||
|
||||
def test_shutdown_clears_stream_refs(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._server_streams["srv"] = (MagicMock(), MagicMock())
|
||||
mgr.shutdown()
|
||||
assert len(mgr._server_streams) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 4: Notification debounce
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestNotificationDebounce:
|
||||
"""Verify notification-triggered refreshes are debounced."""
|
||||
|
||||
def test_debounce_within_window(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._last_notification_refresh["srv"] = time.monotonic()
|
||||
# We can't easily call _on_notification (it's a closure), so test
|
||||
# the debounce logic directly via the timestamp check
|
||||
now = time.monotonic()
|
||||
last = mgr._last_notification_refresh.get("srv", 0.0)
|
||||
assert now - last < mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
def test_debounce_passes_after_window(self):
|
||||
mgr = MCPClientManager({})
|
||||
# Set timestamp well in the past
|
||||
mgr._last_notification_refresh["srv"] = time.monotonic() - 10
|
||||
now = time.monotonic()
|
||||
last = mgr._last_notification_refresh.get("srv", 0.0)
|
||||
assert now - last >= mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
def test_debounce_is_per_server(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._last_notification_refresh["srv_a"] = time.monotonic()
|
||||
# srv_b has no timestamp — should pass debounce
|
||||
now = time.monotonic()
|
||||
last_b = mgr._last_notification_refresh.get("srv_b", 0.0)
|
||||
assert now - last_b >= mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix 5: Periodic refresh backoff
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPeriodicRefreshBackoff:
|
||||
"""Verify periodic refresh backoff and auto-reconnect."""
|
||||
|
||||
def test_backoff_set_on_failure(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._refresh_failures["srv"] = 1
|
||||
# Simulate what _periodic_refresh does on failure
|
||||
failures = mgr._refresh_failures.get("srv", 0) + 1
|
||||
mgr._refresh_failures["srv"] = failures
|
||||
backoff = min(mgr._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)), mgr._REFRESH_BACKOFF_MAX)
|
||||
mgr._refresh_backoff_until["srv"] = time.monotonic() + backoff
|
||||
assert mgr._refresh_backoff_until["srv"] > time.monotonic()
|
||||
assert failures == 2
|
||||
|
||||
def test_backoff_doubles(self):
|
||||
mgr = MCPClientManager({})
|
||||
b1 = min(mgr._REFRESH_BACKOFF_BASE * (2**0), mgr._REFRESH_BACKOFF_MAX)
|
||||
b2 = min(mgr._REFRESH_BACKOFF_BASE * (2**1), mgr._REFRESH_BACKOFF_MAX)
|
||||
b3 = min(mgr._REFRESH_BACKOFF_BASE * (2**2), mgr._REFRESH_BACKOFF_MAX)
|
||||
assert b1 == 60
|
||||
assert b2 == 120
|
||||
assert b3 == 240
|
||||
|
||||
def test_backoff_capped(self):
|
||||
mgr = MCPClientManager({})
|
||||
b = min(mgr._REFRESH_BACKOFF_BASE * (2**20), mgr._REFRESH_BACKOFF_MAX)
|
||||
assert b == mgr._REFRESH_BACKOFF_MAX
|
||||
|
||||
def test_backoff_clears_on_success(self):
|
||||
mgr = MCPClientManager({})
|
||||
mgr._refresh_failures["srv"] = 3
|
||||
mgr._refresh_backoff_until["srv"] = time.monotonic() + 1000
|
||||
# Simulate success
|
||||
mgr._refresh_failures.pop("srv", None)
|
||||
mgr._refresh_backoff_until.pop("srv", None)
|
||||
assert "srv" not in mgr._refresh_failures
|
||||
assert "srv" not in mgr._refresh_backoff_until
|
||||
|
||||
def test_server_status_includes_circuit_info(self):
|
||||
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
||||
status = mgr.get_server_status("srv")
|
||||
assert "circuit_open" in status
|
||||
assert "consecutive_failures" in status
|
||||
assert status["circuit_open"] is False
|
||||
assert status["consecutive_failures"] == 0
|
||||
|
||||
def test_server_status_shows_open_circuit(self):
|
||||
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
||||
for _ in range(3):
|
||||
mgr._cb_record_failure("srv")
|
||||
status = mgr.get_server_status("srv")
|
||||
assert status["circuit_open"] is True
|
||||
assert status["consecutive_failures"] == 3
|
||||
|
||||
@@ -762,8 +762,12 @@ class TestSessionAgentModel:
|
||||
def test_agent_model_resolved(self) -> None:
|
||||
reg = ModelRegistry(
|
||||
models={
|
||||
"main": ModelConfig("main", "http://m/v1", "k", "main-model"),
|
||||
"agent": ModelConfig("agent", "http://a/v1", "k", "agent-model"),
|
||||
"main": ModelConfig(
|
||||
"main", "http://m/v1", "k", "main-model", provider="openai-compatible"
|
||||
),
|
||||
"agent": ModelConfig(
|
||||
"agent", "http://a/v1", "k", "agent-model", provider="openai-compatible"
|
||||
),
|
||||
},
|
||||
default="main",
|
||||
agent_model="agent",
|
||||
|
||||
+38
-48
@@ -8,8 +8,26 @@ import pytest
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from turnstone.channels._http import create_channel_app
|
||||
from turnstone.core.auth import JWT_AUD_CHANNEL, create_jwt
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
_JWT_SECRET = "a" * 32
|
||||
|
||||
|
||||
def _make_jwt() -> str:
|
||||
"""Create a valid JWT for channel auth."""
|
||||
return create_jwt(
|
||||
user_id="system",
|
||||
scopes=frozenset({"write"}),
|
||||
source="service",
|
||||
secret=_JWT_SECRET,
|
||||
audience=JWT_AUD_CHANNEL,
|
||||
)
|
||||
|
||||
|
||||
def _auth_headers() -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {_make_jwt()}"}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage(tmp_path):
|
||||
@@ -33,22 +51,22 @@ def no_auth_client(storage, mock_adapter):
|
||||
|
||||
@pytest.fixture
|
||||
def client(storage, mock_adapter):
|
||||
"""Default client with static auth token configured."""
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, auth_token="test-secret-token")
|
||||
"""Default client with JWT auth configured."""
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, jwt_secret=_JWT_SECRET)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def authed_client(storage, mock_adapter):
|
||||
"""Alias — same as client, for auth-specific test clarity."""
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, auth_token="test-secret-token")
|
||||
"""Alias -- same as client, for auth-specific test clarity."""
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, jwt_secret=_JWT_SECRET)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def jwt_client(storage, mock_adapter):
|
||||
"""Client with JWT auth configured."""
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, jwt_secret="a" * 32)
|
||||
app = create_channel_app({"discord": mock_adapter}, storage, jwt_secret=_JWT_SECRET)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@@ -58,9 +76,6 @@ class TestNotifyEndpoint:
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["status"] == "ok"
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {"Authorization": "Bearer test-secret-token"}
|
||||
|
||||
def test_direct_discord_target(self, client, mock_adapter):
|
||||
resp = client.post(
|
||||
"/v1/api/notify",
|
||||
@@ -68,7 +83,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"channel_type": "discord", "channel_id": "123456"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
results = resp.json()["results"]
|
||||
@@ -85,7 +100,7 @@ class TestNotifyEndpoint:
|
||||
"message": "Hello!",
|
||||
"title": "Alert",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
mock_adapter.send.assert_called_once_with("123456", "**Alert**\nHello!")
|
||||
@@ -101,7 +116,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"username": "testuser"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
results = resp.json()["results"]
|
||||
@@ -116,7 +131,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"username": "nobody"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
error = resp.json()["error"]
|
||||
@@ -132,10 +147,10 @@ class TestNotifyEndpoint:
|
||||
"target": {"username": "testuser"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers={"Authorization": "Bearer test-secret-token"},
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
# Generic message — must not differentiate "not found" vs "no channels"
|
||||
# Generic message -- must not differentiate "not found" vs "no channels"
|
||||
error = resp.json()["error"]
|
||||
assert "testuser" not in error
|
||||
assert "not found or has no linked channels" in error
|
||||
@@ -144,7 +159,7 @@ class TestNotifyEndpoint:
|
||||
resp = client.post(
|
||||
"/v1/api/notify",
|
||||
json={"target": {"username": "x"}},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
@@ -152,7 +167,7 @@ class TestNotifyEndpoint:
|
||||
resp = client.post(
|
||||
"/v1/api/notify",
|
||||
json={"message": "Hello!"},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
@@ -163,7 +178,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"invalid": "field"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
@@ -175,7 +190,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"channel_type": "email", "channel_id": "test@example.com"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
results = resp.json()["results"]
|
||||
@@ -189,7 +204,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"channel_type": "discord", "channel_id": "123456"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
results = resp.json()["results"]
|
||||
@@ -201,7 +216,7 @@ class TestNotifyEndpoint:
|
||||
content=b"not json",
|
||||
headers={
|
||||
"content-type": "application/json",
|
||||
"Authorization": "Bearer test-secret-token",
|
||||
"Authorization": f"Bearer {_make_jwt()}",
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
@@ -214,7 +229,7 @@ class TestNotifyEndpoint:
|
||||
"target": {"channel_type": "discord", "channel_id": "123"},
|
||||
"message": " ",
|
||||
},
|
||||
headers=self._headers(),
|
||||
headers=_auth_headers(),
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
@@ -256,30 +271,9 @@ class TestNotifyAuth:
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
def test_accept_valid_static_token(self, authed_client, mock_adapter):
|
||||
"""Requests with correct static token are accepted."""
|
||||
resp = authed_client.post(
|
||||
"/v1/api/notify",
|
||||
json={
|
||||
"target": {"channel_type": "discord", "channel_id": "123"},
|
||||
"message": "Hello!",
|
||||
},
|
||||
headers={"Authorization": "Bearer test-secret-token"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["results"][0]["status"] == "sent"
|
||||
|
||||
def test_accept_valid_jwt(self, jwt_client, mock_adapter):
|
||||
"""Requests with a valid JWT for the channel audience are accepted."""
|
||||
from turnstone.core.auth import JWT_AUD_CHANNEL, create_jwt
|
||||
|
||||
token = create_jwt(
|
||||
user_id="system",
|
||||
scopes=frozenset({"write"}),
|
||||
source="service",
|
||||
secret="a" * 32,
|
||||
audience=JWT_AUD_CHANNEL,
|
||||
)
|
||||
token = _make_jwt()
|
||||
resp = jwt_client.post(
|
||||
"/v1/api/notify",
|
||||
json={
|
||||
@@ -292,13 +286,11 @@ class TestNotifyAuth:
|
||||
|
||||
def test_reject_jwt_wrong_audience(self, jwt_client):
|
||||
"""JWTs with wrong audience are rejected."""
|
||||
from turnstone.core.auth import create_jwt
|
||||
|
||||
token = create_jwt(
|
||||
user_id="system",
|
||||
scopes=frozenset({"write"}),
|
||||
source="service",
|
||||
secret="a" * 32,
|
||||
secret=_JWT_SECRET,
|
||||
audience="turnstone-server", # wrong audience
|
||||
)
|
||||
resp = jwt_client.post(
|
||||
@@ -313,8 +305,6 @@ class TestNotifyAuth:
|
||||
|
||||
def test_reject_jwt_wrong_secret(self, jwt_client):
|
||||
"""JWTs signed with wrong secret are rejected."""
|
||||
from turnstone.core.auth import JWT_AUD_CHANNEL, create_jwt
|
||||
|
||||
token = create_jwt(
|
||||
user_id="system",
|
||||
scopes=frozenset({"write"}),
|
||||
|
||||
@@ -403,7 +403,7 @@ class TestMCPTemplates:
|
||||
|
||||
|
||||
class TestResumeDeletedTemplate:
|
||||
def test_resume_with_deleted_template_degrades_gracefully(self, tmp_db, capsys):
|
||||
def test_resume_with_deleted_template_degrades_gracefully(self, tmp_db, caplog):
|
||||
from turnstone.core.memory import save_message
|
||||
from turnstone.core.storage import get_storage
|
||||
|
||||
@@ -430,8 +430,7 @@ class TestResumeDeletedTemplate:
|
||||
content = _sys_content(session2)
|
||||
assert "EPHEMERAL_CONTENT" not in content
|
||||
# Warning should be logged via structlog
|
||||
captured = capsys.readouterr()
|
||||
assert "not_found" in captured.out or "not_found" in captured.err
|
||||
assert "not_found" in caplog.text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+584
-31
@@ -9,9 +9,18 @@ from unittest.mock import MagicMock, PropertyMock, patch
|
||||
import pytest
|
||||
|
||||
from turnstone.core.providers._openai import OpenAIProvider
|
||||
from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider
|
||||
from turnstone.core.providers._openai_common import (
|
||||
apply_cache_retention,
|
||||
apply_temperature_and_effort,
|
||||
apply_tool_search,
|
||||
format_citations,
|
||||
sanitize_messages,
|
||||
)
|
||||
from turnstone.core.providers._protocol import (
|
||||
CompletionResult,
|
||||
LLMProvider,
|
||||
ModelCapabilities,
|
||||
StreamChunk,
|
||||
ToolCallDelta,
|
||||
UsageInfo,
|
||||
@@ -133,38 +142,38 @@ def _anthropic_event(
|
||||
|
||||
|
||||
class TestOpenAIProvider:
|
||||
"""Tests for the OpenAI-compatible provider adapter."""
|
||||
"""Tests for the OpenAI Chat Completions provider adapter."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
self.provider = OpenAIProvider()
|
||||
|
||||
def test_provider_name(self) -> None:
|
||||
assert self.provider.provider_name == "openai"
|
||||
assert self.provider.provider_name == "openai-compatible"
|
||||
|
||||
# -- _sanitize_messages ---------------------------------------------------
|
||||
|
||||
def test_sanitize_messages_none_content_no_tool_calls(self) -> None:
|
||||
msgs = [{"role": "assistant", "content": None}]
|
||||
assert self.provider._sanitize_messages(msgs) == [{"role": "assistant", "content": ""}]
|
||||
assert sanitize_messages(msgs) == [{"role": "assistant", "content": ""}]
|
||||
|
||||
def test_sanitize_messages_none_content_with_tool_calls(self) -> None:
|
||||
msgs = [{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]}]
|
||||
result = self.provider._sanitize_messages(msgs)
|
||||
result = sanitize_messages(msgs)
|
||||
assert result[0]["content"] is None
|
||||
assert result[0]["tool_calls"] == [{"id": "1"}]
|
||||
|
||||
def test_sanitize_messages_empty_string_passthrough(self) -> None:
|
||||
msgs = [{"role": "assistant", "content": ""}]
|
||||
assert self.provider._sanitize_messages(msgs) == msgs
|
||||
assert sanitize_messages(msgs) == msgs
|
||||
|
||||
def test_sanitize_messages_non_assistant_unchanged(self) -> None:
|
||||
msgs = [{"role": "user", "content": None}]
|
||||
result = self.provider._sanitize_messages(msgs)
|
||||
result = sanitize_messages(msgs)
|
||||
assert result[0]["content"] is None
|
||||
|
||||
def test_sanitize_messages_does_not_mutate_original(self) -> None:
|
||||
original = {"role": "assistant", "content": None}
|
||||
self.provider._sanitize_messages([original])
|
||||
sanitize_messages([original])
|
||||
assert original["content"] is None
|
||||
|
||||
# -- convert_tools --------------------------------------------------------
|
||||
@@ -992,10 +1001,10 @@ class TestProviderFactory:
|
||||
"""Tests for create_provider and create_client factory functions."""
|
||||
|
||||
def test_create_provider_openai(self) -> None:
|
||||
from turnstone.core.providers import create_provider
|
||||
from turnstone.core.providers import OpenAIResponsesProvider, create_provider
|
||||
|
||||
provider = create_provider("openai")
|
||||
assert isinstance(provider, OpenAIProvider)
|
||||
assert isinstance(provider, OpenAIResponsesProvider)
|
||||
assert provider.provider_name == "openai"
|
||||
|
||||
def test_create_provider_anthropic(self) -> None:
|
||||
@@ -1040,6 +1049,24 @@ class TestProviderFactory:
|
||||
|
||||
assert not isinstance(NotAProvider(), LLMProvider)
|
||||
|
||||
def test_create_provider_openai_compatible(self) -> None:
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
provider = create_provider("openai-compatible")
|
||||
assert isinstance(provider, OpenAIChatCompletionsProvider)
|
||||
assert provider.provider_name == "openai-compatible"
|
||||
|
||||
def test_create_provider_openai_vs_compatible_distinct(self) -> None:
|
||||
from turnstone.core.providers import OpenAIResponsesProvider, create_provider
|
||||
|
||||
openai_prov = create_provider("openai")
|
||||
compat = create_provider("openai-compatible")
|
||||
assert openai_prov is not compat
|
||||
assert isinstance(openai_prov, OpenAIResponsesProvider)
|
||||
assert isinstance(compat, OpenAIChatCompletionsProvider)
|
||||
assert openai_prov.provider_name == "openai"
|
||||
assert compat.provider_name == "openai-compatible"
|
||||
|
||||
def test_create_provider_returns_singleton(self) -> None:
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
@@ -1111,7 +1138,7 @@ class TestOpenAIParameterGating:
|
||||
"""Unknown/local models should NOT receive top-level reasoning_effort."""
|
||||
caps = self.provider.get_capabilities("my-local-model")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
assert "reasoning_effort" not in kwargs
|
||||
assert kwargs["temperature"] == 0.7
|
||||
|
||||
@@ -1119,7 +1146,7 @@ class TestOpenAIParameterGating:
|
||||
"""GPT-5 base: no temperature, reasoning_effort sent."""
|
||||
caps = self.provider.get_capabilities("gpt-5")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
assert "temperature" not in kwargs
|
||||
assert kwargs["reasoning_effort"] == "high"
|
||||
|
||||
@@ -1127,7 +1154,7 @@ class TestOpenAIParameterGating:
|
||||
"""GPT-5.1: temperature only when reasoning_effort='none'."""
|
||||
caps = self.provider.get_capabilities("gpt-5.1")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="none")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="none")
|
||||
assert kwargs["temperature"] == 0.7
|
||||
assert "reasoning_effort" not in kwargs # "none" is skipped
|
||||
|
||||
@@ -1135,7 +1162,7 @@ class TestOpenAIParameterGating:
|
||||
"""GPT-5.1: no temperature when reasoning is active."""
|
||||
caps = self.provider.get_capabilities("gpt-5.1")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
assert "temperature" not in kwargs
|
||||
assert kwargs["reasoning_effort"] == "high"
|
||||
|
||||
@@ -1143,7 +1170,7 @@ class TestOpenAIParameterGating:
|
||||
"""O-series: no temperature, no reasoning_effort."""
|
||||
caps = self.provider.get_capabilities("o3")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
assert "temperature" not in kwargs
|
||||
assert "reasoning_effort" not in kwargs
|
||||
|
||||
@@ -1151,7 +1178,7 @@ class TestOpenAIParameterGating:
|
||||
"""GPT-5 pro only supports 'high'; unsupported values fall back to default."""
|
||||
caps = self.provider.get_capabilities("gpt-5-pro")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="medium")
|
||||
assert "temperature" not in kwargs
|
||||
assert kwargs["reasoning_effort"] == "high" # fell back to default
|
||||
|
||||
@@ -1159,7 +1186,7 @@ class TestOpenAIParameterGating:
|
||||
"""GPT-5 pro accepts 'high' directly."""
|
||||
caps = self.provider.get_capabilities("gpt-5-pro")
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="high")
|
||||
assert kwargs["reasoning_effort"] == "high"
|
||||
|
||||
def test_gpt54_1m_context_and_effort(self) -> None:
|
||||
@@ -1167,11 +1194,11 @@ class TestOpenAIParameterGating:
|
||||
caps = self.provider.get_capabilities("gpt-5.4")
|
||||
assert caps.context_window == 1050000
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="none")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="none")
|
||||
assert kwargs["temperature"] == 0.7
|
||||
assert "reasoning_effort" not in kwargs
|
||||
kwargs2: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs2, caps, temperature=0.7, reasoning_effort="xhigh")
|
||||
apply_temperature_and_effort(kwargs2, caps, temperature=0.7, reasoning_effort="xhigh")
|
||||
assert "temperature" not in kwargs2
|
||||
assert kwargs2["reasoning_effort"] == "xhigh"
|
||||
|
||||
@@ -1180,7 +1207,7 @@ class TestOpenAIParameterGating:
|
||||
caps = self.provider.get_capabilities("gpt-5.4-pro")
|
||||
assert caps.context_window == 1050000
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_model_params(kwargs, caps, temperature=0.7, reasoning_effort="low")
|
||||
apply_temperature_and_effort(kwargs, caps, temperature=0.7, reasoning_effort="low")
|
||||
assert "temperature" not in kwargs
|
||||
assert kwargs["reasoning_effort"] == "medium" # fell back from unsupported "low"
|
||||
|
||||
@@ -1827,7 +1854,7 @@ class TestOpenAIWebSearch:
|
||||
ann.url_citation = citation
|
||||
|
||||
content = "Some search result text."
|
||||
result = OpenAIProvider._format_citations(content, [ann])
|
||||
result = format_citations(content, [ann])
|
||||
assert "Sources:" in result
|
||||
assert "[Example Page](https://example.com)" in result
|
||||
|
||||
@@ -1842,7 +1869,7 @@ class TestOpenAIWebSearch:
|
||||
ann2.url_citation = MagicMock(title="Page Again", url="https://example.com")
|
||||
|
||||
content = "Text."
|
||||
result = OpenAIProvider._format_citations(content, [ann1, ann2])
|
||||
result = format_citations(content, [ann1, ann2])
|
||||
assert result.count("example.com") == 1
|
||||
|
||||
def test_format_citations_skips_non_url_citation(self) -> None:
|
||||
@@ -1851,7 +1878,7 @@ class TestOpenAIWebSearch:
|
||||
ann.type = "something_else"
|
||||
|
||||
content = "Text."
|
||||
result = OpenAIProvider._format_citations(content, [ann])
|
||||
result = format_citations(content, [ann])
|
||||
assert "Sources:" not in result
|
||||
|
||||
def test_format_citations_empty_title(self) -> None:
|
||||
@@ -1860,7 +1887,7 @@ class TestOpenAIWebSearch:
|
||||
ann.type = "url_citation"
|
||||
ann.url_citation = MagicMock(title="", url="https://example.com")
|
||||
|
||||
result = OpenAIProvider._format_citations("Text.", [ann])
|
||||
result = format_citations("Text.", [ann])
|
||||
assert "https://example.com" in result
|
||||
# Should not have markdown link format when title is empty
|
||||
assert "[](https://example.com)" not in result
|
||||
@@ -1871,7 +1898,7 @@ class TestOpenAIWebSearch:
|
||||
ann.type = "url_citation"
|
||||
ann.url_citation = None
|
||||
|
||||
result = OpenAIProvider._format_citations("Text.", [ann])
|
||||
result = format_citations("Text.", [ann])
|
||||
assert "Sources:" not in result
|
||||
|
||||
def test_apply_web_search_with_no_tools(self) -> None:
|
||||
@@ -2303,8 +2330,8 @@ class TestAnthropicToolSearch:
|
||||
# MCP tool should be deferred
|
||||
assert result[1]["defer_loading"] is True
|
||||
# Search tool should be appended
|
||||
assert result[-1]["type"] == "tool_search_tool_bm25_20251119"
|
||||
assert result[-1]["name"] == "tool_search"
|
||||
assert result[-1]["type"] == "tool_search_tool_bm25"
|
||||
assert result[-1]["name"] == "tool_search_tool_bm25"
|
||||
|
||||
def test_inject_tool_search_no_op_without_deferred(self, provider):
|
||||
caps = provider.get_capabilities("claude-opus-4-6-20260101")
|
||||
@@ -2345,7 +2372,7 @@ class TestOpenAIToolSearch:
|
||||
},
|
||||
]
|
||||
deferred = frozenset(["mcp__slack__send"])
|
||||
result = provider._apply_tool_search(caps, tools, deferred)
|
||||
result = apply_tool_search(caps, tools, deferred)
|
||||
assert result is not None
|
||||
# bash not deferred
|
||||
assert result[0].get("defer_loading") is None or result[0].get("defer_loading") is False
|
||||
@@ -2357,7 +2384,7 @@ class TestOpenAIToolSearch:
|
||||
tools = [
|
||||
{"type": "function", "function": {"name": "bash", "description": "Run commands"}},
|
||||
]
|
||||
result = provider._apply_tool_search(caps, tools, None)
|
||||
result = apply_tool_search(caps, tools, None)
|
||||
assert result == tools
|
||||
|
||||
def test_apply_tool_search_no_op_on_unsupported_model(self, provider):
|
||||
@@ -2366,7 +2393,7 @@ class TestOpenAIToolSearch:
|
||||
{"type": "function", "function": {"name": "bash", "description": "Run commands"}},
|
||||
]
|
||||
deferred = frozenset(["some_tool"])
|
||||
result = provider._apply_tool_search(caps, tools, deferred)
|
||||
result = apply_tool_search(caps, tools, deferred)
|
||||
assert result == tools
|
||||
|
||||
|
||||
@@ -2699,14 +2726,14 @@ class TestOpenAIPromptCaching:
|
||||
"""GPT-5.x models get prompt_cache_retention=24h."""
|
||||
for model in ("gpt-5", "gpt-5.1", "gpt-5.2", "gpt-5.4", "gpt-5-mini", "gpt-5-pro"):
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_cache_retention(kwargs, model)
|
||||
apply_cache_retention(kwargs, model)
|
||||
assert kwargs.get("prompt_cache_retention") == "24h", f"Failed for {model}"
|
||||
|
||||
def test_cache_retention_not_set_for_non_gpt5(self) -> None:
|
||||
"""Non-GPT-5 models do not get cache retention."""
|
||||
for model in ("o3", "o4-mini", "local-model", "gpt-4o"):
|
||||
kwargs: dict[str, Any] = {}
|
||||
self.provider._apply_cache_retention(kwargs, model)
|
||||
apply_cache_retention(kwargs, model)
|
||||
assert "prompt_cache_retention" not in kwargs, f"Unexpected retention for {model}"
|
||||
|
||||
def test_streaming_cached_tokens_from_usage(self) -> None:
|
||||
@@ -2845,3 +2872,529 @@ class TestMetricsCacheTokens:
|
||||
assert 'turnstone_tokens_total{type="cache_creation"} 800' in text
|
||||
assert 'turnstone_tokens_total{type="cache_read"} 200' in text
|
||||
assert 'turnstone_tokens_total{type="prompt"} 1000' in text
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# TestOpenAIResponsesProvider — Responses API provider
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestOpenAIResponsesProvider:
|
||||
"""Tests for the OpenAI Responses API provider."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def test_provider_name(self) -> None:
|
||||
assert self.provider.provider_name == "openai"
|
||||
|
||||
def test_get_capabilities(self) -> None:
|
||||
caps = self.provider.get_capabilities("gpt-5.4")
|
||||
assert caps.context_window == 1050000
|
||||
assert caps.supports_tool_search is True
|
||||
|
||||
|
||||
class TestResponsesMessageConversion:
|
||||
"""Tests for _convert_messages — Chat Completions format to Responses API."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def test_system_message_to_instructions(self) -> None:
|
||||
messages = [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
instructions, items = self.provider._convert_messages(messages)
|
||||
assert instructions == "You are helpful."
|
||||
assert len(items) == 1
|
||||
assert items[0]["role"] == "user"
|
||||
assert items[0]["content"] == "Hello"
|
||||
|
||||
def test_multiple_system_messages_concatenated(self) -> None:
|
||||
messages = [
|
||||
{"role": "system", "content": "Rule 1"},
|
||||
{"role": "developer", "content": "Rule 2"},
|
||||
{"role": "user", "content": "Hi"},
|
||||
]
|
||||
instructions, items = self.provider._convert_messages(messages)
|
||||
assert instructions == "Rule 1\n\nRule 2"
|
||||
assert len(items) == 1
|
||||
|
||||
def test_assistant_text_message(self) -> None:
|
||||
messages = [
|
||||
{"role": "assistant", "content": "Hello back"},
|
||||
]
|
||||
_, items = self.provider._convert_messages(messages)
|
||||
assert len(items) == 1
|
||||
assert items[0]["type"] == "message"
|
||||
assert items[0]["role"] == "assistant"
|
||||
assert items[0]["content"] == "Hello back"
|
||||
|
||||
def test_assistant_tool_calls(self) -> None:
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {"name": "read_file", "arguments": '{"path": "/tmp"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
_, items = self.provider._convert_messages(messages)
|
||||
assert len(items) == 1
|
||||
assert items[0]["type"] == "function_call"
|
||||
assert items[0]["call_id"] == "call_1"
|
||||
assert items[0]["name"] == "read_file"
|
||||
assert items[0]["arguments"] == '{"path": "/tmp"}'
|
||||
|
||||
def test_tool_result(self) -> None:
|
||||
messages = [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "file contents"},
|
||||
]
|
||||
_, items = self.provider._convert_messages(messages)
|
||||
assert len(items) == 1
|
||||
assert items[0]["type"] == "function_call_output"
|
||||
assert items[0]["call_id"] == "call_1"
|
||||
assert items[0]["output"] == "file contents"
|
||||
|
||||
def test_provider_content_ignored_with_store_false(self) -> None:
|
||||
"""With store=False, provider_content is ignored — rebuild from content."""
|
||||
provider_items = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Hi"}],
|
||||
},
|
||||
{"type": "function_call", "call_id": "c1", "name": "f", "arguments": "{}"},
|
||||
]
|
||||
messages = [
|
||||
{"role": "assistant", "content": "Hi", "_provider_content": provider_items},
|
||||
]
|
||||
_, items = self.provider._convert_messages(messages)
|
||||
# Should rebuild from content, not passthrough provider_content
|
||||
assert len(items) == 1
|
||||
assert items[0]["type"] == "message"
|
||||
assert items[0]["content"] == "Hi"
|
||||
|
||||
def test_no_system_returns_none_instructions(self) -> None:
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
instructions, _ = self.provider._convert_messages(messages)
|
||||
assert instructions is None
|
||||
|
||||
def test_assistant_with_content_and_tool_calls(self) -> None:
|
||||
"""Assistant message with both text and tool calls emits separate items."""
|
||||
messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I'll read that file",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {"name": "read_file", "arguments": '{"path": "/tmp"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
_, items = self.provider._convert_messages(messages)
|
||||
assert len(items) == 2
|
||||
assert items[0]["type"] == "message"
|
||||
assert items[0]["content"] == "I'll read that file"
|
||||
assert items[1]["type"] == "function_call"
|
||||
assert items[1]["name"] == "read_file"
|
||||
|
||||
|
||||
class TestResponsesToolConversion:
|
||||
"""Tests for _convert_tools — Chat Completions tool format to Responses API."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def test_function_tool_conversion(self) -> None:
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a file",
|
||||
"parameters": {"type": "object", "properties": {"path": {"type": "string"}}},
|
||||
},
|
||||
}
|
||||
]
|
||||
caps = ModelCapabilities()
|
||||
result = self.provider._convert_tools(tools, caps)
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0]["type"] == "function"
|
||||
assert result[0]["name"] == "read_file"
|
||||
assert result[0]["description"] == "Read a file"
|
||||
assert result[0]["strict"] is False
|
||||
|
||||
def test_web_search_replaced_with_native(self) -> None:
|
||||
tools = [
|
||||
{"type": "function", "function": {"name": "web_search", "description": "Search"}},
|
||||
{"type": "function", "function": {"name": "read_file", "description": "Read"}},
|
||||
]
|
||||
caps = ModelCapabilities(supports_web_search=True)
|
||||
result = self.provider._convert_tools(tools, caps)
|
||||
assert result is not None
|
||||
names = [t.get("name", t.get("type")) for t in result]
|
||||
assert "web_search" in names # native web_search tool
|
||||
assert "read_file" in names
|
||||
|
||||
def test_none_tools_returns_none(self) -> None:
|
||||
caps = ModelCapabilities()
|
||||
assert self.provider._convert_tools(None, caps) is None
|
||||
|
||||
def test_defer_loading_preserved(self) -> None:
|
||||
tools = [
|
||||
{"type": "function", "function": {"name": "f"}, "defer_loading": True},
|
||||
]
|
||||
caps = ModelCapabilities()
|
||||
result = self.provider._convert_tools(tools, caps)
|
||||
assert result is not None
|
||||
assert result[0].get("defer_loading") is True
|
||||
|
||||
|
||||
class TestResponsesParamBuilding:
|
||||
"""Tests for _build_kwargs — parameter construction for Responses API."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def test_reasoning_effort_as_dict(self) -> None:
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5.4",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="high",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert kwargs["reasoning"] == {"effort": "high"}
|
||||
assert "reasoning_effort" not in kwargs
|
||||
|
||||
def test_no_reasoning_when_none_effort(self) -> None:
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5.4",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="none",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert "reasoning" not in kwargs
|
||||
|
||||
def test_store_is_false(self) -> None:
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5.4",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="medium",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert kwargs["store"] is False
|
||||
|
||||
def test_cache_retention_for_gpt5(self) -> None:
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5.4",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="medium",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert kwargs["prompt_cache_retention"] == "24h"
|
||||
|
||||
def test_instructions_from_system_messages(self) -> None:
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5.4",
|
||||
messages=[
|
||||
{"role": "system", "content": "Be helpful"},
|
||||
{"role": "user", "content": "Hi"},
|
||||
],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="none",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert kwargs["instructions"] == "Be helpful"
|
||||
|
||||
def test_web_search_injected_with_no_tools(self) -> None:
|
||||
"""Search-capable models get web_search tool even when tools=None."""
|
||||
kwargs = self.provider._build_kwargs(
|
||||
model="gpt-5-search-api",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=None,
|
||||
max_tokens=4096,
|
||||
temperature=0.5,
|
||||
reasoning_effort="none",
|
||||
deferred_names=None,
|
||||
)
|
||||
assert "tools" in kwargs
|
||||
tool_types = [t.get("type") for t in kwargs["tools"]]
|
||||
assert "web_search" in tool_types
|
||||
|
||||
|
||||
class TestResponsesCitationFormat:
|
||||
"""Test format_citations handles Responses API flat annotation format."""
|
||||
|
||||
def test_responses_api_flat_annotation(self) -> None:
|
||||
"""Responses API annotations have title/url directly on the object."""
|
||||
|
||||
class FlatAnnotation:
|
||||
type = "url_citation"
|
||||
url_citation = None # Not present in Responses API
|
||||
title = "Example"
|
||||
url = "https://example.com"
|
||||
|
||||
result = format_citations("Text.", [FlatAnnotation()])
|
||||
assert "Sources:" in result
|
||||
assert "[Example](https://example.com)" in result
|
||||
|
||||
|
||||
class TestResponsesStreaming:
|
||||
"""Tests for Responses API streaming event handling."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def _make_event(self, event_type: str, **attrs: Any) -> MagicMock:
|
||||
event = MagicMock()
|
||||
event.type = event_type
|
||||
for k, v in attrs.items():
|
||||
setattr(event, k, v)
|
||||
return event
|
||||
|
||||
def test_text_delta(self) -> None:
|
||||
events = [
|
||||
self._make_event("response.output_text.delta", delta="Hello"),
|
||||
self._make_event("response.output_text.delta", delta=" world"),
|
||||
self._make_event(
|
||||
"response.completed",
|
||||
response=MagicMock(
|
||||
status="completed",
|
||||
usage=None,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = list(self.provider._iter_stream(iter(events)))
|
||||
text_chunks = [c for c in chunks if c.content_delta]
|
||||
assert len(text_chunks) == 2
|
||||
assert text_chunks[0].content_delta == "Hello"
|
||||
assert text_chunks[0].is_first is True
|
||||
assert text_chunks[1].content_delta == " world"
|
||||
|
||||
def test_reasoning_delta(self) -> None:
|
||||
events = [
|
||||
self._make_event("response.reasoning_text.delta", delta="thinking..."),
|
||||
self._make_event(
|
||||
"response.completed",
|
||||
response=MagicMock(
|
||||
status="completed",
|
||||
usage=None,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = list(self.provider._iter_stream(iter(events)))
|
||||
reasoning = [c for c in chunks if c.reasoning_delta]
|
||||
assert len(reasoning) == 1
|
||||
assert reasoning[0].reasoning_delta == "thinking..."
|
||||
assert reasoning[0].is_first is True
|
||||
|
||||
def test_tool_call_streaming(self) -> None:
|
||||
item = MagicMock()
|
||||
item.type = "function_call"
|
||||
item.id = "fc_abc123"
|
||||
item.call_id = "call_1"
|
||||
item.name = "read_file"
|
||||
|
||||
events = [
|
||||
self._make_event("response.output_item.added", item=item),
|
||||
self._make_event(
|
||||
"response.function_call_arguments.delta",
|
||||
item_id="fc_abc123",
|
||||
delta='{"path":',
|
||||
),
|
||||
self._make_event(
|
||||
"response.function_call_arguments.delta",
|
||||
item_id="fc_abc123",
|
||||
delta='"/tmp"}',
|
||||
),
|
||||
self._make_event(
|
||||
"response.completed",
|
||||
response=MagicMock(
|
||||
status="completed",
|
||||
usage=None,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = list(self.provider._iter_stream(iter(events)))
|
||||
tc_chunks = [c for c in chunks if c.tool_call_deltas]
|
||||
assert len(tc_chunks) == 3
|
||||
# First chunk: tool call added with name
|
||||
assert tc_chunks[0].tool_call_deltas[0].name == "read_file"
|
||||
assert tc_chunks[0].tool_call_deltas[0].id == "call_1"
|
||||
# Argument deltas
|
||||
assert tc_chunks[1].tool_call_deltas[0].arguments_delta == '{"path":'
|
||||
assert tc_chunks[2].tool_call_deltas[0].arguments_delta == '"/tmp"}'
|
||||
|
||||
def test_completed_event_with_usage(self) -> None:
|
||||
usage = MagicMock()
|
||||
usage.input_tokens = 100
|
||||
usage.output_tokens = 50
|
||||
usage.total_tokens = 150
|
||||
usage.input_tokens_details = MagicMock(cached_tokens=80)
|
||||
# Ensure Chat Completions attributes are not present
|
||||
del usage.prompt_tokens
|
||||
del usage.completion_tokens
|
||||
del usage.prompt_tokens_details
|
||||
|
||||
events = [
|
||||
self._make_event(
|
||||
"response.completed",
|
||||
response=MagicMock(
|
||||
status="completed",
|
||||
usage=usage,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = list(self.provider._iter_stream(iter(events)))
|
||||
final = [c for c in chunks if c.finish_reason]
|
||||
assert len(final) == 1
|
||||
assert final[0].finish_reason == "stop"
|
||||
assert final[0].usage is not None
|
||||
assert final[0].usage.prompt_tokens == 100
|
||||
assert final[0].usage.completion_tokens == 50
|
||||
assert final[0].usage.cache_read_tokens == 80
|
||||
|
||||
def test_web_search_events(self) -> None:
|
||||
events = [
|
||||
self._make_event("response.web_search_call.searching"),
|
||||
self._make_event("response.web_search_call.completed"),
|
||||
self._make_event(
|
||||
"response.completed",
|
||||
response=MagicMock(
|
||||
status="completed",
|
||||
usage=None,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = list(self.provider._iter_stream(iter(events)))
|
||||
info = [c for c in chunks if c.info_delta]
|
||||
assert len(info) == 2
|
||||
assert "Searching" in info[0].info_delta
|
||||
assert "complete" in info[1].info_delta
|
||||
|
||||
|
||||
class TestResponsesCompletion:
|
||||
"""Tests for non-streaming Responses API completion."""
|
||||
|
||||
def setup_method(self) -> None:
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
self.provider = OpenAIResponsesProvider()
|
||||
|
||||
def _make_response(
|
||||
self,
|
||||
text: str = "Hello",
|
||||
tool_calls: list[dict[str, Any]] | None = None,
|
||||
status: str = "completed",
|
||||
) -> MagicMock:
|
||||
resp = MagicMock()
|
||||
resp.status = status
|
||||
resp.usage = MagicMock()
|
||||
resp.usage.input_tokens = 10
|
||||
resp.usage.output_tokens = 5
|
||||
resp.usage.total_tokens = 15
|
||||
resp.usage.input_tokens_details = MagicMock(cached_tokens=0)
|
||||
# Remove Chat Completions attributes
|
||||
del resp.usage.prompt_tokens
|
||||
del resp.usage.completion_tokens
|
||||
del resp.usage.prompt_tokens_details
|
||||
|
||||
output: list[Any] = []
|
||||
if text:
|
||||
msg = MagicMock()
|
||||
msg.type = "message"
|
||||
text_part = MagicMock()
|
||||
text_part.type = "output_text"
|
||||
text_part.text = text
|
||||
text_part.annotations = []
|
||||
msg.content = [text_part]
|
||||
msg.model_dump.return_value = {
|
||||
"type": "message",
|
||||
"content": [{"type": "output_text", "text": text}],
|
||||
}
|
||||
output.append(msg)
|
||||
if tool_calls:
|
||||
for tc in tool_calls:
|
||||
item = MagicMock()
|
||||
item.type = "function_call"
|
||||
item.call_id = tc["id"]
|
||||
item.name = tc["name"]
|
||||
item.arguments = tc["arguments"]
|
||||
item.model_dump.return_value = {
|
||||
"type": "function_call",
|
||||
"call_id": tc["id"],
|
||||
"name": tc["name"],
|
||||
"arguments": tc["arguments"],
|
||||
}
|
||||
output.append(item)
|
||||
resp.output = output
|
||||
return resp
|
||||
|
||||
def test_basic_text_completion(self) -> None:
|
||||
resp = self._make_response(text="Hello world")
|
||||
result = self.provider._parse_response(resp)
|
||||
assert result.content == "Hello world"
|
||||
assert result.tool_calls is None
|
||||
assert result.finish_reason == "stop"
|
||||
|
||||
def test_completion_with_tool_calls(self) -> None:
|
||||
resp = self._make_response(
|
||||
text="",
|
||||
tool_calls=[{"id": "call_1", "name": "read_file", "arguments": '{"path": "/tmp"}'}],
|
||||
)
|
||||
result = self.provider._parse_response(resp)
|
||||
assert result.tool_calls is not None
|
||||
assert len(result.tool_calls) == 1
|
||||
assert result.tool_calls[0]["id"] == "call_1"
|
||||
assert result.tool_calls[0]["function"]["name"] == "read_file"
|
||||
|
||||
def test_provider_blocks_captured(self) -> None:
|
||||
resp = self._make_response(text="Hello")
|
||||
result = self.provider._parse_response(resp)
|
||||
assert len(result.provider_blocks) > 0
|
||||
assert result.provider_blocks[0]["type"] == "message"
|
||||
|
||||
def test_incomplete_status_maps_to_length(self) -> None:
|
||||
resp = self._make_response(text="Partial", status="incomplete")
|
||||
result = self.provider._parse_response(resp)
|
||||
assert result.finish_reason == "length"
|
||||
|
||||
def test_usage_extraction(self) -> None:
|
||||
resp = self._make_response(text="Hi")
|
||||
result = self.provider._parse_response(resp)
|
||||
assert result.usage is not None
|
||||
assert result.usage.prompt_tokens == 10
|
||||
assert result.usage.completion_tokens == 5
|
||||
|
||||
@@ -64,6 +64,7 @@ class _InjectAuthMiddleware(BaseHTTPMiddleware):
|
||||
"admin.roles",
|
||||
"admin.orgs",
|
||||
"admin.policies",
|
||||
"admin.prompt_policies",
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
+34
-11
@@ -138,6 +138,8 @@ def tmp_db():
|
||||
|
||||
def _make_session(client, model_id, tmp_db, **kwargs) -> tuple[ChatSession, RecordingUI]:
|
||||
"""Create a ChatSession with RecordingUI and sensible test defaults."""
|
||||
from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider
|
||||
|
||||
ui = RecordingUI()
|
||||
defaults = dict(
|
||||
client=client,
|
||||
@@ -151,6 +153,8 @@ def _make_session(client, model_id, tmp_db, **kwargs) -> tuple[ChatSession, Reco
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
session = ChatSession(**defaults)
|
||||
# Mock-based tests use Chat Completions format (client.chat.completions)
|
||||
session._provider = OpenAIChatCompletionsProvider()
|
||||
session.auto_approve = True
|
||||
return session, ui
|
||||
|
||||
@@ -580,6 +584,24 @@ class TestSessionConfig:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_TEST_JWT_SECRET = "test-jwt-secret-minimum-32-chars!"
|
||||
|
||||
|
||||
def _server_jwt() -> str:
|
||||
from turnstone.core.auth import JWT_AUD_SERVER, create_jwt
|
||||
|
||||
return create_jwt(
|
||||
user_id="test-server-live",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="test",
|
||||
secret=_TEST_JWT_SECRET,
|
||||
audience=JWT_AUD_SERVER,
|
||||
)
|
||||
|
||||
|
||||
_SERVER_AUTH_HEADERS = {"Authorization": f"Bearer {_server_jwt()}"}
|
||||
|
||||
|
||||
class TestServerHealthMetrics:
|
||||
"""Verify /health and /metrics endpoints using a Starlette TestClient.
|
||||
|
||||
@@ -596,7 +618,6 @@ class TestServerHealthMetrics:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
import turnstone.server as srv_mod
|
||||
from turnstone.core.auth import AuthConfig
|
||||
from turnstone.core.metrics import MetricsCollector
|
||||
from turnstone.core.workstream import WorkstreamState
|
||||
|
||||
@@ -631,7 +652,7 @@ class TestServerHealthMetrics:
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
)
|
||||
cls.client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
@@ -727,8 +748,8 @@ class TestServerHealthMetrics:
|
||||
assert 'le="+Inf"' in body
|
||||
|
||||
def test_unknown_endpoint_returns_404(self):
|
||||
status, _, _ = self._get("/does-not-exist")
|
||||
assert status == 404
|
||||
resp = self.client.get("/does-not-exist", headers=_SERVER_AUTH_HEADERS)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_health_contains_backend_field(self):
|
||||
_, _, body = self._get("/health")
|
||||
@@ -772,7 +793,6 @@ class TestServerRateLimiting:
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
import turnstone.server as srv_mod
|
||||
from turnstone.core.auth import AuthConfig
|
||||
from turnstone.core.metrics import MetricsCollector
|
||||
from turnstone.core.ratelimit import RateLimiter
|
||||
from turnstone.core.workstream import WorkstreamState
|
||||
@@ -808,7 +828,7 @@ class TestServerRateLimiting:
|
||||
global_listeners=[],
|
||||
global_listeners_lock=threading.Lock(),
|
||||
skip_permissions=False,
|
||||
auth_config=AuthConfig(),
|
||||
jwt_secret=_TEST_JWT_SECRET,
|
||||
rate_limiter=RateLimiter(enabled=True, rate=2.0, burst=3),
|
||||
)
|
||||
cls.client = TestClient(app, raise_server_exceptions=False)
|
||||
@@ -830,16 +850,19 @@ class TestServerRateLimiting:
|
||||
"""After exhausting burst on a non-exempt endpoint, get 429."""
|
||||
# Exhaust burst on a non-exempt endpoint
|
||||
for _ in range(5):
|
||||
self._get("/v1/api/workstreams")
|
||||
self.client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
# At least one should be 429
|
||||
statuses = [self._get("/v1/api/workstreams").status_code for _ in range(3)]
|
||||
statuses = [
|
||||
self.client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS).status_code
|
||||
for _ in range(3)
|
||||
]
|
||||
assert 429 in statuses
|
||||
|
||||
def test_429_includes_retry_after(self):
|
||||
"""429 response includes Retry-After header."""
|
||||
# Burn through burst
|
||||
for _ in range(10):
|
||||
resp = self._get("/v1/api/workstreams")
|
||||
resp = self.client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
if resp.status_code == 429:
|
||||
assert "retry-after" in resp.headers
|
||||
data = resp.json()
|
||||
@@ -851,7 +874,7 @@ class TestServerRateLimiting:
|
||||
"""Health endpoint is always accessible regardless of rate limit."""
|
||||
# Burn through bucket on non-exempt path
|
||||
for _ in range(10):
|
||||
self._get("/v1/api/workstreams")
|
||||
self.client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
# Health should still work
|
||||
resp = self._get("/health")
|
||||
assert resp.status_code == 200
|
||||
@@ -859,6 +882,6 @@ class TestServerRateLimiting:
|
||||
def test_metrics_exempt_from_ratelimit(self):
|
||||
"""Metrics endpoint is always accessible regardless of rate limit."""
|
||||
for _ in range(10):
|
||||
self._get("/v1/api/workstreams")
|
||||
self.client.get("/v1/api/workstreams", headers=_SERVER_AUTH_HEADERS)
|
||||
resp = self._get("/metrics")
|
||||
assert resp.status_code == 200
|
||||
|
||||
+78
-20
@@ -640,13 +640,11 @@ class TestExecReadImage:
|
||||
self._make_png(str(img))
|
||||
|
||||
session = _make_session()
|
||||
# Mock provider to report vision support
|
||||
mock_caps = MagicMock()
|
||||
mock_caps.supports_vision = True
|
||||
session._provider.get_capabilities = MagicMock(return_value=mock_caps)
|
||||
|
||||
item = {"call_id": "c1", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
with patch.object(session._provider, "get_capabilities", return_value=mock_caps):
|
||||
item = {"call_id": "c1", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
|
||||
assert call_id == "c1"
|
||||
assert isinstance(output, list)
|
||||
@@ -669,10 +667,9 @@ class TestExecReadImage:
|
||||
session = _make_session()
|
||||
mock_caps = MagicMock()
|
||||
mock_caps.supports_vision = False
|
||||
session._provider.get_capabilities = MagicMock(return_value=mock_caps)
|
||||
|
||||
item = {"call_id": "c2", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
with patch.object(session._provider, "get_capabilities", return_value=mock_caps):
|
||||
item = {"call_id": "c2", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
|
||||
assert call_id == "c2"
|
||||
assert isinstance(output, str)
|
||||
@@ -689,10 +686,9 @@ class TestExecReadImage:
|
||||
session = _make_session()
|
||||
mock_caps = MagicMock()
|
||||
mock_caps.supports_vision = True
|
||||
session._provider.get_capabilities = MagicMock(return_value=mock_caps)
|
||||
|
||||
item = {"call_id": "c3", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
with patch.object(session._provider, "get_capabilities", return_value=mock_caps):
|
||||
item = {"call_id": "c3", "path": str(img), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
|
||||
assert call_id == "c3"
|
||||
assert isinstance(output, str)
|
||||
@@ -703,10 +699,14 @@ class TestExecReadImage:
|
||||
session = _make_session()
|
||||
mock_caps = MagicMock()
|
||||
mock_caps.supports_vision = True
|
||||
session._provider.get_capabilities = MagicMock(return_value=mock_caps)
|
||||
|
||||
item = {"call_id": "c4", "path": str(tmp_path / "nope.png"), "offset": None, "limit": None}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
with patch.object(session._provider, "get_capabilities", return_value=mock_caps):
|
||||
item = {
|
||||
"call_id": "c4",
|
||||
"path": str(tmp_path / "nope.png"),
|
||||
"offset": None,
|
||||
"limit": None,
|
||||
}
|
||||
call_id, output = session._exec_read_file(item)
|
||||
assert isinstance(output, str)
|
||||
assert "not found" in output
|
||||
|
||||
@@ -742,9 +742,10 @@ class TestGetCapabilitiesOverride:
|
||||
default="qwen-vl",
|
||||
)
|
||||
session = _make_session(registry=registry, model_alias="qwen-vl")
|
||||
# Ensure provider returns a real ModelCapabilities (not MagicMock)
|
||||
session._provider.get_capabilities = MagicMock(return_value=ModelCapabilities())
|
||||
caps = session._get_capabilities()
|
||||
# Ensure provider returns a real ModelCapabilities (not MagicMock).
|
||||
# Use patch.object so the singleton provider is restored after the test.
|
||||
with patch.object(session._provider, "get_capabilities", return_value=ModelCapabilities()):
|
||||
caps = session._get_capabilities()
|
||||
assert caps.supports_vision is True
|
||||
|
||||
def test_no_override_uses_provider_default(self, tmp_db):
|
||||
@@ -759,6 +760,8 @@ class TestTitleRetry:
|
||||
"""_generate_title resets _title_generated on failure."""
|
||||
|
||||
def test_title_generated_reset_on_failure(self, tmp_db):
|
||||
from turnstone.core.providers._protocol import ModelCapabilities
|
||||
|
||||
session = _make_session()
|
||||
session._title_generated = True
|
||||
session.messages = [
|
||||
@@ -767,6 +770,7 @@ class TestTitleRetry:
|
||||
]
|
||||
# Mock provider to raise
|
||||
session._provider = MagicMock()
|
||||
session._provider.get_capabilities.return_value = ModelCapabilities()
|
||||
session._provider.create_completion.side_effect = RuntimeError("API error")
|
||||
|
||||
session._generate_title()
|
||||
@@ -774,6 +778,8 @@ class TestTitleRetry:
|
||||
assert session._title_generated is False
|
||||
|
||||
def test_title_generated_stays_true_on_success(self, tmp_db):
|
||||
from turnstone.core.providers._protocol import ModelCapabilities
|
||||
|
||||
session = _make_session()
|
||||
session._title_generated = True
|
||||
session.messages = [
|
||||
@@ -783,6 +789,7 @@ class TestTitleRetry:
|
||||
result = MagicMock()
|
||||
result.content = "Test Title"
|
||||
session._provider = MagicMock()
|
||||
session._provider.get_capabilities.return_value = ModelCapabilities()
|
||||
session._provider.create_completion.return_value = result
|
||||
|
||||
with patch("turnstone.core.session.update_workstream_title"):
|
||||
@@ -793,6 +800,8 @@ class TestTitleRetry:
|
||||
|
||||
def test_title_skipped_after_resume_changes_ws_id(self, tmp_db):
|
||||
"""If ws_id changes (via resume) during title generation, discard the result."""
|
||||
from turnstone.core.providers._protocol import ModelCapabilities
|
||||
|
||||
session = _make_session()
|
||||
session._title_generated = True
|
||||
session.messages = [
|
||||
@@ -803,6 +812,7 @@ class TestTitleRetry:
|
||||
result = MagicMock()
|
||||
result.content = "Test Title"
|
||||
session._provider = MagicMock()
|
||||
session._provider.get_capabilities.return_value = ModelCapabilities()
|
||||
session._provider.create_completion.return_value = result
|
||||
|
||||
# Simulate resume() changing ws_id while title generation is in flight
|
||||
@@ -912,8 +922,10 @@ class TestAgentOutputGuard:
|
||||
def test_agent_loop_calls_evaluate_output(self):
|
||||
"""_run_agent passes tool output through _evaluate_output when output_guard is enabled."""
|
||||
from turnstone.core.judge import JudgeConfig
|
||||
from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider
|
||||
|
||||
session = _make_session(judge_config=JudgeConfig(output_guard=True))
|
||||
session._provider = OpenAIChatCompletionsProvider()
|
||||
|
||||
with patch.object(session, "_evaluate_output", wraps=lambda cid, o, fn: o) as mock_eval:
|
||||
# Simulate _run_agent getting a tool call response then a text response
|
||||
@@ -971,8 +983,10 @@ class TestAgentOutputGuard:
|
||||
def test_agent_loop_skips_guard_when_disabled(self):
|
||||
"""_run_agent does not call _evaluate_output when output_guard is disabled."""
|
||||
from turnstone.core.judge import JudgeConfig
|
||||
from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider
|
||||
|
||||
session = _make_session(judge_config=JudgeConfig(output_guard=False))
|
||||
session._provider = OpenAIChatCompletionsProvider()
|
||||
|
||||
with patch.object(session, "_evaluate_output") as mock_eval:
|
||||
call_count = [0]
|
||||
@@ -1018,3 +1032,47 @@ class TestAgentOutputGuard:
|
||||
)
|
||||
|
||||
mock_eval.assert_not_called()
|
||||
|
||||
|
||||
class TestProviderExtraParams:
|
||||
"""Tests for _provider_extra_params — local-only chat_template_kwargs."""
|
||||
|
||||
def _session_with_provider(self, provider_name: str, tmp_db) -> ChatSession:
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
session = _make_session(reasoning_effort="medium")
|
||||
session._provider = create_provider(provider_name)
|
||||
return session
|
||||
|
||||
def test_openai_compatible_returns_chat_template_kwargs(self, tmp_db):
|
||||
session = self._session_with_provider("openai-compatible", tmp_db)
|
||||
result = session._provider_extra_params()
|
||||
assert result is not None
|
||||
assert "chat_template_kwargs" in result
|
||||
assert result["chat_template_kwargs"]["reasoning_effort"] == "medium"
|
||||
|
||||
def test_openai_commercial_returns_none(self, tmp_db):
|
||||
session = self._session_with_provider("openai", tmp_db)
|
||||
result = session._provider_extra_params()
|
||||
assert result is None
|
||||
|
||||
def test_anthropic_returns_none(self, tmp_db):
|
||||
session = self._session_with_provider("anthropic", tmp_db)
|
||||
result = session._provider_extra_params()
|
||||
assert result is None
|
||||
|
||||
def test_reasoning_effort_override(self, tmp_db):
|
||||
session = self._session_with_provider("openai-compatible", tmp_db)
|
||||
result = session._provider_extra_params(reasoning_effort="high")
|
||||
assert result is not None
|
||||
assert result["chat_template_kwargs"]["reasoning_effort"] == "high"
|
||||
|
||||
def test_explicit_openai_provider_overrides_session(self, tmp_db):
|
||||
"""Passing an explicit commercial OpenAI provider returns None even
|
||||
when the session's own provider is openai-compatible."""
|
||||
from turnstone.core.providers import create_provider
|
||||
|
||||
session = self._session_with_provider("openai-compatible", tmp_db)
|
||||
openai_prov = create_provider("openai")
|
||||
result = session._provider_extra_params(provider=openai_prov)
|
||||
assert result is None
|
||||
|
||||
@@ -335,8 +335,6 @@ class TestSaveMessageUpdatesWorkstream:
|
||||
def test_updated_timestamp_bumped(self, tmp_db):
|
||||
register_workstream("s1")
|
||||
save_message("s1", "user", "first")
|
||||
rows = list_workstreams_with_history()
|
||||
_original_updated = rows[0][4]
|
||||
|
||||
import time
|
||||
|
||||
|
||||
@@ -0,0 +1,479 @@
|
||||
"""Tests for skill resource materialization to disk.
|
||||
|
||||
Verifies that skill-bundled resources (scripts, references, assets) stored
|
||||
in the ``skill_resources`` table are written to a temp directory when a
|
||||
skill is loaded, exposed via ``SKILL_RESOURCES_DIR`` env var and ``PATH``,
|
||||
and cleaned up on skill change or session close.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import stat
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from turnstone.core.session import ChatSession
|
||||
from turnstone.core.storage._registry import get_storage
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers (mirrors test_skills.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class NullUI:
|
||||
"""UI adapter that discards all output."""
|
||||
|
||||
def on_thinking_start(self):
|
||||
pass
|
||||
|
||||
def on_thinking_stop(self):
|
||||
pass
|
||||
|
||||
def on_reasoning_token(self, text):
|
||||
pass
|
||||
|
||||
def on_content_token(self, text):
|
||||
pass
|
||||
|
||||
def on_stream_end(self):
|
||||
pass
|
||||
|
||||
def approve_tools(self, items):
|
||||
return True, None
|
||||
|
||||
def on_tool_result(self, call_id, name, output, **kwargs):
|
||||
pass
|
||||
|
||||
def on_tool_output_chunk(self, call_id, chunk):
|
||||
pass
|
||||
|
||||
def on_status(self, usage, context_window, effort):
|
||||
pass
|
||||
|
||||
def on_plan_review(self, content):
|
||||
return ""
|
||||
|
||||
def on_info(self, message):
|
||||
pass
|
||||
|
||||
def on_error(self, message):
|
||||
pass
|
||||
|
||||
def on_state_change(self, state):
|
||||
pass
|
||||
|
||||
def on_rename(self, name):
|
||||
pass
|
||||
|
||||
def on_output_warning(self, call_id, assessment):
|
||||
pass
|
||||
|
||||
|
||||
def _make_session(**kwargs: Any) -> ChatSession:
|
||||
defaults: dict[str, Any] = dict(
|
||||
client=MagicMock(),
|
||||
model="test-model",
|
||||
ui=NullUI(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
max_tokens=4096,
|
||||
tool_timeout=30,
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return ChatSession(**defaults)
|
||||
|
||||
|
||||
def _create_skill(db: Any, skill_id: str, name: str, content: str, **kw: Any) -> None:
|
||||
db.create_prompt_template(
|
||||
template_id=skill_id,
|
||||
name=name,
|
||||
category=kw.get("category", "general"),
|
||||
content=content,
|
||||
variables=kw.get("variables", "[]"),
|
||||
is_default=kw.get("is_default", False),
|
||||
org_id="",
|
||||
created_by="test",
|
||||
origin="manual",
|
||||
mcp_server="",
|
||||
readonly=False,
|
||||
description="",
|
||||
tags="[]",
|
||||
source_url="",
|
||||
version="1.0.0",
|
||||
author="",
|
||||
activation=kw.get("activation", "named"),
|
||||
token_estimate=0,
|
||||
model="",
|
||||
auto_approve=False,
|
||||
temperature=None,
|
||||
reasoning_effort="",
|
||||
max_tokens=None,
|
||||
token_budget=0,
|
||||
agent_max_turns=None,
|
||||
notify_on_complete="{}",
|
||||
enabled=True,
|
||||
allowed_tools="[]",
|
||||
priority=0,
|
||||
)
|
||||
|
||||
|
||||
def _sys_content(session: ChatSession) -> str:
|
||||
msgs = [m for m in session.system_messages if m["role"] == "system"]
|
||||
assert msgs
|
||||
return msgs[0]["content"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMaterializeResources:
|
||||
def test_materialize_creates_files(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "test-skill", "Use the scripts.")
|
||||
db.create_skill_resource("r1", "s1", "scripts/helper.py", "print('hello')")
|
||||
db.create_skill_resource("r2", "s1", "references/api.md", "# API")
|
||||
|
||||
session = _make_session(skill="test-skill")
|
||||
assert session._skill_resources_dir is not None
|
||||
base = session._skill_resources_dir
|
||||
assert os.path.isdir(base)
|
||||
|
||||
helper = os.path.join(base, "scripts", "helper.py")
|
||||
assert os.path.isfile(helper)
|
||||
with open(helper) as f:
|
||||
assert f.read() == "print('hello')"
|
||||
|
||||
api_md = os.path.join(base, "references", "api.md")
|
||||
assert os.path.isfile(api_md)
|
||||
with open(api_md) as f:
|
||||
assert f.read() == "# API"
|
||||
|
||||
session.close()
|
||||
|
||||
def test_scripts_executable(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "exec-skill", "Run scripts/run.sh")
|
||||
db.create_skill_resource("r1", "s1", "scripts/run.sh", "#!/bin/bash\necho hi")
|
||||
|
||||
session = _make_session(skill="exec-skill")
|
||||
base = session._skill_resources_dir
|
||||
run_sh = os.path.join(base, "scripts", "run.sh")
|
||||
mode = os.stat(run_sh).st_mode
|
||||
assert mode & stat.S_IXUSR # owner execute
|
||||
session.close()
|
||||
|
||||
def test_non_scripts_not_executable(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "ref-skill", "Read references/guide.md")
|
||||
db.create_skill_resource("r1", "s1", "references/guide.md", "# Guide")
|
||||
|
||||
session = _make_session(skill="ref-skill")
|
||||
base = session._skill_resources_dir
|
||||
guide = os.path.join(base, "references", "guide.md")
|
||||
mode = os.stat(guide).st_mode
|
||||
assert not (mode & stat.S_IXUSR) # not executable
|
||||
session.close()
|
||||
|
||||
def test_cleanup_on_close(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "cleanup-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "scripts/a.py", "code")
|
||||
|
||||
session = _make_session(skill="cleanup-skill")
|
||||
base = session._skill_resources_dir
|
||||
assert os.path.isdir(base)
|
||||
|
||||
session.close()
|
||||
assert not os.path.exists(base)
|
||||
assert session._skill_resources_dir is None
|
||||
|
||||
def test_cleanup_on_skill_switch(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "skill-a", "Skill A")
|
||||
db.create_skill_resource("r1", "s1", "scripts/a.py", "code_a")
|
||||
_create_skill(db, "s2", "skill-b", "Skill B")
|
||||
db.create_skill_resource("r2", "s2", "scripts/b.py", "code_b")
|
||||
|
||||
session = _make_session(skill="skill-a")
|
||||
dir_a = session._skill_resources_dir
|
||||
assert os.path.isfile(os.path.join(dir_a, "scripts", "a.py"))
|
||||
|
||||
session.set_skill("skill-b")
|
||||
dir_b = session._skill_resources_dir
|
||||
assert dir_b != dir_a
|
||||
assert not os.path.exists(dir_a)
|
||||
assert os.path.isfile(os.path.join(dir_b, "scripts", "b.py"))
|
||||
|
||||
session.close()
|
||||
|
||||
def test_cleanup_on_skill_clear(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "clear-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "scripts/x.py", "code")
|
||||
|
||||
session = _make_session(skill="clear-skill")
|
||||
base = session._skill_resources_dir
|
||||
assert os.path.isdir(base)
|
||||
|
||||
session.set_skill(None)
|
||||
assert not os.path.exists(base)
|
||||
assert session._skill_resources_dir is None
|
||||
|
||||
session.close()
|
||||
|
||||
def test_empty_resources_no_dir(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "no-res-skill", "content")
|
||||
# No resources added
|
||||
|
||||
session = _make_session(skill="no-res-skill")
|
||||
assert session._skill_resources_dir is None
|
||||
session.close()
|
||||
|
||||
def test_no_skill_no_dir(self, tmp_db):
|
||||
session = _make_session()
|
||||
assert session._skill_resources_dir is None
|
||||
session.close()
|
||||
|
||||
def test_path_traversal_rejected(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "traversal-skill", "content")
|
||||
# Inject a malicious path directly into storage
|
||||
db.create_skill_resource("r1", "s1", "../etc/passwd", "bad content")
|
||||
db.create_skill_resource("r2", "s1", "scripts/good.py", "good content")
|
||||
|
||||
session = _make_session(skill="traversal-skill")
|
||||
base = session._skill_resources_dir
|
||||
# The traversal path must not be written inside the resources dir
|
||||
assert not os.path.exists(os.path.join(base, "etc"))
|
||||
# The good resource should still be materialized
|
||||
assert os.path.isfile(os.path.join(base, "scripts", "good.py"))
|
||||
session.close()
|
||||
|
||||
|
||||
class TestSkillResourceEnv:
|
||||
def test_env_with_resources(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "env-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "scripts/tool.py", "code")
|
||||
|
||||
session = _make_session(skill="env-skill")
|
||||
env = session._skill_resource_env()
|
||||
assert env["SKILL_RESOURCES_DIR"] == session._skill_resources_dir
|
||||
assert "PATH" in env
|
||||
scripts_dir = os.path.join(session._skill_resources_dir, "scripts")
|
||||
assert env["PATH"].startswith(scripts_dir + ":")
|
||||
session.close()
|
||||
|
||||
def test_env_without_scripts_dir(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "no-scripts-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "references/doc.md", "# Doc")
|
||||
|
||||
session = _make_session(skill="no-scripts-skill")
|
||||
env = session._skill_resource_env()
|
||||
assert "SKILL_RESOURCES_DIR" in env
|
||||
# No scripts/ subdir so PATH should not be overridden
|
||||
assert "PATH" not in env
|
||||
session.close()
|
||||
|
||||
def test_env_empty_when_no_resources(self, tmp_db):
|
||||
session = _make_session()
|
||||
assert session._skill_resource_env() == {}
|
||||
session.close()
|
||||
|
||||
|
||||
class TestSystemMessageHint:
|
||||
def test_hint_present_when_resources_exist(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "hint-skill", "Use the bundled scripts.")
|
||||
db.create_skill_resource("r1", "s1", "scripts/run.py", "code")
|
||||
|
||||
session = _make_session(skill="hint-skill")
|
||||
content = _sys_content(session)
|
||||
assert "$SKILL_RESOURCES_DIR" in content
|
||||
assert "scripts/ are on PATH" in content
|
||||
session.close()
|
||||
|
||||
def test_no_hint_when_no_resources(self, tmp_db):
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "plain-skill", "No resources here.")
|
||||
|
||||
session = _make_session(skill="plain-skill")
|
||||
content = _sys_content(session)
|
||||
assert "SKILL_RESOURCES_DIR" not in content
|
||||
session.close()
|
||||
|
||||
|
||||
class TestMaterializeEdgeCases:
|
||||
def test_all_resources_rejected_no_dir(self, tmp_db):
|
||||
"""When every resource fails path validation, no temp dir is left."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "all-bad", "content")
|
||||
db.create_skill_resource("r1", "s1", "../escape", "bad")
|
||||
db.create_skill_resource("r2", "s1", "/absolute", "bad")
|
||||
|
||||
session = _make_session(skill="all-bad")
|
||||
assert session._skill_resources_dir is None
|
||||
session.close()
|
||||
|
||||
def test_dot_path_rejected(self, tmp_db):
|
||||
"""A bare '.' path is rejected rather than crashing."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "dot-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", ".", "bad")
|
||||
db.create_skill_resource("r2", "s1", "scripts/ok.py", "good")
|
||||
|
||||
session = _make_session(skill="dot-skill")
|
||||
base = session._skill_resources_dir
|
||||
assert os.path.isfile(os.path.join(base, "scripts", "ok.py"))
|
||||
session.close()
|
||||
|
||||
def test_empty_path_rejected(self, tmp_db):
|
||||
"""An empty string path is rejected."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "empty-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "", "bad")
|
||||
db.create_skill_resource("r2", "s1", "scripts/ok.py", "good")
|
||||
|
||||
session = _make_session(skill="empty-skill")
|
||||
assert session._skill_resources_dir is not None
|
||||
session.close()
|
||||
|
||||
def test_nested_traversal_rejected(self, tmp_db):
|
||||
"""Traversal hidden inside a valid prefix is still caught."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "nested-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "scripts/../../../etc/passwd", "bad")
|
||||
db.create_skill_resource("r2", "s1", "scripts/ok.py", "good")
|
||||
|
||||
session = _make_session(skill="nested-skill")
|
||||
base = session._skill_resources_dir
|
||||
assert not os.path.exists(os.path.join(base, "etc"))
|
||||
assert os.path.isfile(os.path.join(base, "scripts", "ok.py"))
|
||||
session.close()
|
||||
|
||||
def test_double_close_idempotent(self, tmp_db):
|
||||
"""Calling close() twice does not raise."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "double-skill", "content")
|
||||
db.create_skill_resource("r1", "s1", "scripts/x.py", "code")
|
||||
|
||||
session = _make_session(skill="double-skill")
|
||||
session.close()
|
||||
session.close() # must not raise
|
||||
|
||||
|
||||
class TestPreflightValidation:
|
||||
def test_missing_resource_warns(self, tmp_db):
|
||||
"""Skill content references a script not in resources."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "warn-skill", "Run scripts/missing.py to start.")
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="warn-skill")
|
||||
ui.on_info.assert_called_once()
|
||||
msg = ui.on_info.call_args[0][0]
|
||||
assert "scripts/missing.py" in msg
|
||||
assert "warn-skill" in msg
|
||||
session.close()
|
||||
|
||||
def test_all_resources_present_no_warn(self, tmp_db):
|
||||
"""No warning when all referenced paths are bundled."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "ok-skill", "Run scripts/helper.py for help.")
|
||||
db.create_skill_resource("r1", "s1", "scripts/helper.py", "print('hi')")
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="ok-skill")
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
|
||||
def test_no_references_no_warn(self, tmp_db):
|
||||
"""Skill content with no resource paths triggers no validation warning."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "plain-skill", "Just a plain skill with no paths.")
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="plain-skill")
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
|
||||
def test_multiple_missing_warns_once(self, tmp_db):
|
||||
"""Multiple missing resources produce a single warning listing all."""
|
||||
db = get_storage()
|
||||
_create_skill(
|
||||
db,
|
||||
"s1",
|
||||
"multi-skill",
|
||||
"Use scripts/a.py and scripts/b.sh to process references/guide.md",
|
||||
)
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="multi-skill")
|
||||
ui.on_info.assert_called_once()
|
||||
msg = ui.on_info.call_args[0][0]
|
||||
assert "3 resource(s)" in msg
|
||||
assert "scripts/a.py" in msg
|
||||
assert "scripts/b.sh" in msg
|
||||
assert "references/guide.md" in msg
|
||||
session.close()
|
||||
|
||||
def test_validation_skipped_no_skill(self, tmp_db):
|
||||
"""No crash or warning when no skill is active."""
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui)
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
|
||||
def test_json_extension_not_truncated(self, tmp_db):
|
||||
"""assets/config.json should match as .json, not .js."""
|
||||
db = get_storage()
|
||||
_create_skill(db, "s1", "json-skill", "Load assets/config.json for settings.")
|
||||
db.create_skill_resource("r1", "s1", "assets/config.json", "{}")
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="json-skill")
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
|
||||
def test_compound_prefix_not_matched(self, tmp_db):
|
||||
"""'myscripts/tool.py' should not match as 'scripts/tool.py'."""
|
||||
db = get_storage()
|
||||
_create_skill(
|
||||
db,
|
||||
"s1",
|
||||
"compound-skill",
|
||||
"The myscripts/tool.py file is unrelated.",
|
||||
)
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="compound-skill")
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
|
||||
def test_extension_suffix_not_matched(self, tmp_db):
|
||||
"""'scripts/tool.python' should not match as 'scripts/tool.py'."""
|
||||
db = get_storage()
|
||||
_create_skill(
|
||||
db,
|
||||
"s1",
|
||||
"suffix-skill",
|
||||
"Run scripts/tool.python to start.",
|
||||
)
|
||||
|
||||
ui = NullUI()
|
||||
ui.on_info = MagicMock()
|
||||
session = _make_session(ui=ui, skill="suffix-skill")
|
||||
ui.on_info.assert_not_called()
|
||||
session.close()
|
||||
@@ -1,8 +1,17 @@
|
||||
"""Tests for the storage backend registry."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
|
||||
from turnstone.core.storage import get_storage, init_storage, reset_storage
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
|
||||
from turnstone.core.storage import (
|
||||
StorageUnavailableError,
|
||||
get_storage,
|
||||
init_storage,
|
||||
reset_storage,
|
||||
)
|
||||
from turnstone.core.storage._postgresql import PostgreSQLBackend
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
|
||||
@@ -48,7 +57,72 @@ class TestResetStorage:
|
||||
s1 = get_storage()
|
||||
reset_storage()
|
||||
# After reset, get_storage() auto-inits a new instance
|
||||
monkeypatch_not_needed = True # noqa: F841
|
||||
init_storage("sqlite", path=str(tmp_path / "test2.db"), run_migrations=False)
|
||||
s2 = get_storage()
|
||||
assert s1 is not s2
|
||||
|
||||
|
||||
class TestConnUnavailableLogging:
|
||||
"""Test that _conn() deduplicates DB unavailable/restored logging."""
|
||||
|
||||
def _make_backend(self, tmp_path):
|
||||
"""Create a minimal SQLite backend for testing _conn()."""
|
||||
from turnstone.core.storage._sqlite import SQLiteBackend
|
||||
|
||||
return SQLiteBackend(str(tmp_path / "test.db"), create_tables=True)
|
||||
|
||||
def test_logs_unavailable_once(self, tmp_path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
backend = self._make_backend(tmp_path)
|
||||
with patch.object(backend, "_engine") as mock_engine:
|
||||
mock_engine.connect.side_effect = sa.exc.OperationalError(
|
||||
"conn", {}, Exception("refused")
|
||||
)
|
||||
for _ in range(3):
|
||||
with pytest.raises(StorageUnavailableError), backend._conn():
|
||||
pass # pragma: no cover
|
||||
unavailable_msgs = [r for r in caplog.records if "database.unavailable" in r.message]
|
||||
assert len(unavailable_msgs) == 1
|
||||
|
||||
def test_logs_restored_on_recovery(self, tmp_path, caplog: pytest.LogCaptureFixture) -> None:
|
||||
import logging
|
||||
|
||||
caplog.set_level(logging.INFO)
|
||||
backend = self._make_backend(tmp_path)
|
||||
# Simulate outage
|
||||
with patch.object(backend, "_engine") as mock_engine:
|
||||
mock_engine.connect.side_effect = sa.exc.OperationalError(
|
||||
"conn", {}, Exception("refused")
|
||||
)
|
||||
with pytest.raises(StorageUnavailableError), backend._conn():
|
||||
pass # pragma: no cover
|
||||
assert backend._db_unavailable is True
|
||||
# Real connection — should log restored
|
||||
caplog.clear()
|
||||
with backend._conn():
|
||||
pass
|
||||
restored_msgs = [r for r in caplog.records if "database.connection_restored" in r.message]
|
||||
assert len(restored_msgs) == 1
|
||||
assert backend._db_unavailable is False
|
||||
|
||||
def test_postgresql_conn_raises_storage_unavailable(self) -> None:
|
||||
import threading
|
||||
|
||||
backend = PostgreSQLBackend.__new__(PostgreSQLBackend)
|
||||
backend._db_unavailable = False
|
||||
backend._db_unavailable_lock = threading.Lock()
|
||||
|
||||
def _raise_op_error():
|
||||
raise sa.exc.OperationalError("conn", {}, Exception("refused"))
|
||||
|
||||
mock_engine = type(
|
||||
"E",
|
||||
(),
|
||||
{
|
||||
"connect": staticmethod(_raise_op_error),
|
||||
"url": sa.engine.make_url("postgresql://user:pass@localhost/db"),
|
||||
},
|
||||
)()
|
||||
backend._engine = mock_engine
|
||||
with pytest.raises(StorageUnavailableError), backend._conn():
|
||||
pass # pragma: no cover
|
||||
assert backend._db_unavailable is True
|
||||
|
||||
+109
-2
@@ -55,8 +55,8 @@ def _make_app(tls_manager):
|
||||
async def _grant_access(request, call_next): # type: ignore[no-untyped-def]
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="",
|
||||
scopes=frozenset({"approve"}),
|
||||
token_source="config",
|
||||
scopes=frozenset({"approve", "service"}),
|
||||
token_source="test",
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
@@ -124,6 +124,113 @@ def test_delete_cert_not_found(tls_manager):
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
# ── Auth enforcement ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_app_no_auth(tls_manager):
|
||||
"""Create app without auth middleware — simulates unauthenticated requests."""
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import Route
|
||||
|
||||
from turnstone.console.server import (
|
||||
tls_ca_cert,
|
||||
tls_ca_status,
|
||||
tls_delete_cert,
|
||||
tls_list_certs,
|
||||
tls_renew_cert,
|
||||
)
|
||||
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/ca", tls_ca_status),
|
||||
Route("/ca.pem", tls_ca_cert),
|
||||
Route("/certs", tls_list_certs),
|
||||
Route("/certs/{domain}/renew", tls_renew_cert, methods=["POST"]),
|
||||
Route("/certs/{domain}", tls_delete_cert, methods=["DELETE"]),
|
||||
],
|
||||
)
|
||||
app.state.tls_manager = tls_manager
|
||||
return app
|
||||
|
||||
|
||||
def _make_app_read_only(tls_manager):
|
||||
"""Create app with read-only auth — should be rejected by admin endpoints."""
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.routing import Route
|
||||
|
||||
from turnstone.console.server import (
|
||||
tls_ca_cert,
|
||||
tls_ca_status,
|
||||
tls_delete_cert,
|
||||
tls_list_certs,
|
||||
tls_renew_cert,
|
||||
)
|
||||
from turnstone.core.auth import AuthResult
|
||||
|
||||
async def _grant_read(request, call_next): # type: ignore[no-untyped-def]
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="viewer",
|
||||
scopes=frozenset({"read"}),
|
||||
token_source="jwt",
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
app = Starlette(
|
||||
routes=[
|
||||
Route("/ca", tls_ca_status),
|
||||
Route("/ca.pem", tls_ca_cert),
|
||||
Route("/certs", tls_list_certs),
|
||||
Route("/certs/{domain}/renew", tls_renew_cert, methods=["POST"]),
|
||||
Route("/certs/{domain}", tls_delete_cert, methods=["DELETE"]),
|
||||
],
|
||||
middleware=[Middleware(BaseHTTPMiddleware, dispatch=_grant_read)],
|
||||
)
|
||||
app.state.tls_manager = tls_manager
|
||||
return app
|
||||
|
||||
|
||||
def test_unauthenticated_list_certs_401(tls_manager):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
client = TestClient(_make_app_no_auth(tls_manager))
|
||||
resp = client.get("/certs")
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_unauthenticated_renew_401(tls_manager):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
client = TestClient(_make_app_no_auth(tls_manager))
|
||||
resp = client.post("/certs/test.internal/renew")
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_unauthenticated_delete_401(tls_manager):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
client = TestClient(_make_app_no_auth(tls_manager))
|
||||
resp = client.delete("/certs/test.internal")
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
def test_read_only_renew_403(tls_manager):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
client = TestClient(_make_app_read_only(tls_manager))
|
||||
resp = client.post("/certs/test.internal/renew")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
def test_read_only_delete_403(tls_manager):
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
client = TestClient(_make_app_read_only(tls_manager))
|
||||
resp = client.delete("/certs/test.internal")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ── CLI bootstrap ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -145,7 +145,7 @@ async def test_tls_ca_cert_endpoint(tls_manager):
|
||||
|
||||
async def _grant_access(request, call_next): # type: ignore[no-untyped-def]
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="", scopes=frozenset({"approve"}), token_source="config"
|
||||
user_id="", scopes=frozenset({"approve", "service"}), token_source="test"
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
@@ -190,7 +190,7 @@ async def test_tls_endpoints_disabled():
|
||||
|
||||
async def _grant_access(request, call_next): # type: ignore[no-untyped-def]
|
||||
request.state.auth_result = AuthResult(
|
||||
user_id="", scopes=frozenset({"approve"}), token_source="config"
|
||||
user_id="", scopes=frozenset({"approve", "service"}), token_source="test"
|
||||
)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
"""Tests for capacity-aware tool output truncation and context overflow recovery."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.session import ChatSession
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session(tmp_db, mock_openai_client):
|
||||
"""Create a ChatSession with defaults for truncation testing."""
|
||||
return ChatSession(
|
||||
client=mock_openai_client,
|
||||
model="test-model",
|
||||
ui=MagicMock(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
tool_timeout=10,
|
||||
context_window=10_000,
|
||||
max_tokens=1_000,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _truncate_output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTruncateOutput:
|
||||
def test_no_truncation_when_under_limit(self, session):
|
||||
result = session._truncate_output("short text")
|
||||
assert result == "short text"
|
||||
|
||||
def test_truncates_to_tool_truncation_limit(self, session):
|
||||
session.tool_truncation = 100
|
||||
big = "x" * 500
|
||||
result = session._truncate_output(big)
|
||||
assert len(result) <= 200 # head + tail + marker
|
||||
assert "chars truncated" in result
|
||||
|
||||
def test_budget_aware_truncation(self, session):
|
||||
session.tool_truncation = 100_000
|
||||
session._chars_per_token = 4.0
|
||||
# Budget of 50 tokens = 200 chars
|
||||
big = "x" * 1000
|
||||
result = session._truncate_output(big, remaining_budget_tokens=50)
|
||||
assert len(result) <= 400 # head + tail + marker
|
||||
assert "chars truncated" in result
|
||||
|
||||
def test_budget_takes_precedence_when_smaller(self, session):
|
||||
session.tool_truncation = 10_000
|
||||
session._chars_per_token = 4.0
|
||||
# Budget of 25 tokens = 100 chars, smaller than tool_truncation
|
||||
big = "x" * 500
|
||||
result = session._truncate_output(big, remaining_budget_tokens=25)
|
||||
assert "chars truncated" in result
|
||||
|
||||
def test_zero_budget_returns_placeholder(self, session):
|
||||
big = "x" * 1000
|
||||
result = session._truncate_output(big, remaining_budget_tokens=0)
|
||||
assert "exceeded context budget" in result
|
||||
assert len(result) < 100
|
||||
|
||||
def test_negative_budget_returns_placeholder(self, session):
|
||||
big = "x" * 1000
|
||||
result = session._truncate_output(big, remaining_budget_tokens=-10)
|
||||
assert "exceeded context budget" in result
|
||||
|
||||
def test_none_budget_uses_fixed_limit(self, session):
|
||||
session.tool_truncation = 100
|
||||
big = "x" * 500
|
||||
result = session._truncate_output(big, remaining_budget_tokens=None)
|
||||
assert "100 char limit" in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _remaining_token_budget
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRemainingTokenBudget:
|
||||
def test_empty_session(self, session):
|
||||
session._system_tokens = 500
|
||||
session._msg_tokens = []
|
||||
budget = session._remaining_token_budget()
|
||||
# 10000 - 500 - 0 - 1000 - 500 (5%) = 8000
|
||||
assert budget == 8000
|
||||
|
||||
def test_partially_full(self, session):
|
||||
session._system_tokens = 500
|
||||
session._msg_tokens = [2000, 3000]
|
||||
budget = session._remaining_token_budget()
|
||||
# 10000 - 500 - 5000 - 1000 - 500 = 3000
|
||||
assert budget == 3000
|
||||
|
||||
def test_overfull_returns_zero(self, session):
|
||||
session._system_tokens = 500
|
||||
session._msg_tokens = [9000]
|
||||
assert session._remaining_token_budget() == 0
|
||||
|
||||
def test_exactly_full_returns_zero(self, session):
|
||||
session._system_tokens = 500
|
||||
session._msg_tokens = [8000]
|
||||
assert session._remaining_token_budget() == 0
|
||||
|
||||
def test_max_tokens_equals_context_window(self, tmp_db, mock_openai_client):
|
||||
"""Regression: max_tokens >= context_window must not zero the budget."""
|
||||
s = ChatSession(
|
||||
client=mock_openai_client,
|
||||
model="test-model",
|
||||
ui=MagicMock(),
|
||||
instructions=None,
|
||||
temperature=0.5,
|
||||
tool_timeout=10,
|
||||
context_window=32_768,
|
||||
max_tokens=32_768,
|
||||
)
|
||||
s._system_tokens = 500
|
||||
s._msg_tokens = [1000]
|
||||
budget = s._remaining_token_budget()
|
||||
# response_reserve = min(32768, 32768//4) = 8192
|
||||
# safety = 32768 * 0.05 = 1638
|
||||
# budget = 32768 - 500 - 1000 - 8192 - 1638 = 21438
|
||||
assert budget > 20_000
|
||||
# Tool output should NOT be collapsed to a placeholder
|
||||
big = "x" * 5000
|
||||
result = s._truncate_output(big, remaining_budget_tokens=budget)
|
||||
assert result == big # 5000 chars fits easily in 21K+ token budget
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Context overflow recovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestContextOverflowRecovery:
|
||||
"""Test that context-length errors trigger compact-and-retry."""
|
||||
|
||||
def test_openai_context_length_error_triggers_compact(self, session):
|
||||
session.messages = [{"role": "user", "content": "hi"}]
|
||||
session._msg_tokens = [1]
|
||||
|
||||
call_count = 0
|
||||
|
||||
def mock_create_stream(msgs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise Exception("maximum context length exceeded")
|
||||
return iter([])
|
||||
|
||||
compact_mock = MagicMock()
|
||||
with (
|
||||
patch.object(session, "_create_stream_with_retry", side_effect=mock_create_stream),
|
||||
patch.object(session, "_compact_messages", compact_mock),
|
||||
patch.object(
|
||||
session, "_stream_response", return_value={"role": "assistant", "content": "ok"}
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
patch.object(session, "_update_token_table"),
|
||||
patch.object(session, "_print_status_line"),
|
||||
patch.object(session, "_emit_state"),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
session.send("hello")
|
||||
|
||||
compact_mock.assert_called_once_with(auto=True)
|
||||
assert call_count == 2
|
||||
|
||||
def test_anthropic_prompt_too_long_triggers_compact(self, session):
|
||||
session.messages = [{"role": "user", "content": "hi"}]
|
||||
session._msg_tokens = [1]
|
||||
|
||||
call_count = 0
|
||||
|
||||
def mock_create_stream(msgs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise Exception("prompt is too long: 250000 tokens > 200000 maximum")
|
||||
return iter([])
|
||||
|
||||
compact_mock = MagicMock()
|
||||
with (
|
||||
patch.object(session, "_create_stream_with_retry", side_effect=mock_create_stream),
|
||||
patch.object(session, "_compact_messages", compact_mock),
|
||||
patch.object(
|
||||
session, "_stream_response", return_value={"role": "assistant", "content": "ok"}
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
patch.object(session, "_update_token_table"),
|
||||
patch.object(session, "_print_status_line"),
|
||||
patch.object(session, "_emit_state"),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
):
|
||||
session.send("hello")
|
||||
|
||||
compact_mock.assert_called_once_with(auto=True)
|
||||
|
||||
def test_non_context_error_propagates(self, session):
|
||||
session.messages = [{"role": "user", "content": "hi"}]
|
||||
session._msg_tokens = [1]
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
side_effect=Exception("authentication failed"),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
patch.object(session, "_emit_state"),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
pytest.raises(Exception, match="authentication failed"),
|
||||
):
|
||||
session.send("hello")
|
||||
|
||||
def test_compact_failure_raises_original_error(self, session):
|
||||
session.messages = [{"role": "user", "content": "hi"}]
|
||||
session._msg_tokens = [1]
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
side_effect=Exception("maximum context length exceeded"),
|
||||
),
|
||||
patch.object(session, "_compact_messages", side_effect=RuntimeError("compact failed")),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
patch.object(session, "_emit_state"),
|
||||
patch("turnstone.core.session.save_message"),
|
||||
pytest.raises(Exception, match="maximum context length exceeded"),
|
||||
):
|
||||
session.send("hello")
|
||||
+120
-73
@@ -130,54 +130,54 @@ class TestWorkstream:
|
||||
class TestManagerCreation:
|
||||
def test_create_first_sets_active(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.active_id == ws.id
|
||||
assert mgr.get_active() is ws
|
||||
|
||||
def test_create_second_does_not_change_active(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
_ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.active_id == ws1.id
|
||||
|
||||
def test_create_assigns_session(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert isinstance(ws.session, FakeSession)
|
||||
|
||||
def test_create_assigns_ui(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert isinstance(ws.ui, FakeUI)
|
||||
assert ws.ui.ws_id == ws.id
|
||||
|
||||
def test_create_custom_name(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(name="research", ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(name="research", ui_factory=FakeUI)
|
||||
assert ws.name == "research"
|
||||
|
||||
def test_create_default_name(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert ws.name.startswith("ws-")
|
||||
|
||||
def test_create_max_workstreams_all_active(self):
|
||||
mgr = WorkstreamManager(_fake_factory, max_workstreams=3)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws3 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
ws3 = mgr.create(ui_factory=FakeUI)
|
||||
# Mark all as non-idle so eviction cannot help
|
||||
mgr.set_state(ws1.id, WorkstreamState.THINKING)
|
||||
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
||||
mgr.set_state(ws3.id, WorkstreamState.ATTENTION)
|
||||
with pytest.raises(RuntimeError, match="All 3 workstreams are active"):
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
|
||||
|
||||
class TestManagerLookup:
|
||||
def test_get_existing(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.get(ws.id) is ws
|
||||
|
||||
def test_get_nonexistent(self):
|
||||
@@ -186,16 +186,16 @@ class TestManagerLookup:
|
||||
|
||||
def test_list_all_creation_order(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(name="a", ui_factory=lambda wid: FakeUI(wid))
|
||||
_ws2 = mgr.create(name="b", ui_factory=lambda wid: FakeUI(wid))
|
||||
_ws3 = mgr.create(name="c", ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(name="a", ui_factory=FakeUI)
|
||||
mgr.create(name="b", ui_factory=FakeUI)
|
||||
mgr.create(name="c", ui_factory=FakeUI)
|
||||
result = mgr.list_all()
|
||||
assert [w.name for w in result] == ["a", "b", "c"]
|
||||
|
||||
def test_index_of(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.index_of(ws1.id) == 1
|
||||
assert mgr.index_of(ws2.id) == 2
|
||||
assert mgr.index_of("nonexistent") == 0
|
||||
@@ -203,9 +203,9 @@ class TestManagerLookup:
|
||||
def test_count(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
assert mgr.count == 0
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.count == 1
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.count == 2
|
||||
|
||||
|
||||
@@ -217,8 +217,8 @@ class TestManagerLookup:
|
||||
class TestManagerSwitching:
|
||||
def test_switch_by_id(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.active_id == ws1.id
|
||||
|
||||
result = mgr.switch(ws2.id)
|
||||
@@ -227,13 +227,13 @@ class TestManagerSwitching:
|
||||
|
||||
def test_switch_nonexistent_returns_none(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.switch("bad-id") is None
|
||||
|
||||
def test_switch_by_index(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
result = mgr.switch_by_index(2)
|
||||
assert result is ws2
|
||||
@@ -241,7 +241,7 @@ class TestManagerSwitching:
|
||||
|
||||
def test_switch_by_index_out_of_range(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.switch_by_index(0) is None
|
||||
assert mgr.switch_by_index(5) is None
|
||||
|
||||
@@ -254,29 +254,32 @@ class TestManagerSwitching:
|
||||
class TestManagerClose:
|
||||
def test_close_removes_workstream(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
assert mgr.close(ws2.id) is True
|
||||
closed = mgr.close(ws2.id)
|
||||
assert closed is True
|
||||
assert mgr.count == 1
|
||||
assert mgr.get(ws2.id) is None
|
||||
|
||||
def test_close_last_returns_false(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
assert mgr.close(ws.id) is False
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
closed = mgr.close(ws.id)
|
||||
assert closed is False
|
||||
assert mgr.count == 1
|
||||
|
||||
def test_close_nonexistent_returns_false(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
assert mgr.close("nonexistent") is False
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
closed = mgr.close("nonexistent")
|
||||
assert closed is False
|
||||
|
||||
def test_close_active_switches_to_first(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
mgr.switch(ws2.id)
|
||||
|
||||
mgr.close(ws2.id)
|
||||
@@ -284,9 +287,9 @@ class TestManagerClose:
|
||||
|
||||
def test_close_updates_order(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(name="a", ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(name="b", ui_factory=lambda wid: FakeUI(wid))
|
||||
_ws3 = mgr.create(name="c", ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(name="a", ui_factory=FakeUI)
|
||||
ws2 = mgr.create(name="b", ui_factory=FakeUI)
|
||||
mgr.create(name="c", ui_factory=FakeUI)
|
||||
|
||||
mgr.close(ws2.id)
|
||||
names = [w.name for w in mgr.list_all()]
|
||||
@@ -295,7 +298,7 @@ class TestManagerClose:
|
||||
def test_close_unblocks_approval_event(self):
|
||||
"""Closing a workstream whose UI has a pending approval should unblock it."""
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
|
||||
# Create a workstream with a WebUI-like approval mechanism
|
||||
from turnstone.server import WebUI
|
||||
@@ -310,7 +313,7 @@ class TestManagerClose:
|
||||
def test_close_unblocks_plan_event(self):
|
||||
"""Closing a workstream with pending plan review should unblock it."""
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
_ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
|
||||
from turnstone.server import WebUI
|
||||
|
||||
@@ -331,13 +334,13 @@ class TestManagerEviction:
|
||||
def test_evict_oldest_idle_on_create(self):
|
||||
"""At capacity with idle workstreams, create() succeeds by evicting the oldest idle."""
|
||||
mgr = WorkstreamManager(_fake_factory, max_workstreams=3)
|
||||
ws1 = mgr.create(name="oldest", ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(name="middle", ui_factory=lambda wid: FakeUI(wid))
|
||||
_ws3 = mgr.create(name="newest", ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(name="oldest", ui_factory=FakeUI)
|
||||
ws2 = mgr.create(name="middle", ui_factory=FakeUI)
|
||||
mgr.create(name="newest", ui_factory=FakeUI)
|
||||
# All three are IDLE. Mark ws2 as RUNNING so it won't be evicted.
|
||||
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
||||
# ws1 is oldest idle, ws3 is newer idle. Creating should evict ws1.
|
||||
ws4 = mgr.create(name="four", ui_factory=lambda wid: FakeUI(wid))
|
||||
ws4 = mgr.create(name="four", ui_factory=FakeUI)
|
||||
assert mgr.count == 3
|
||||
assert mgr.get(ws1.id) is None, "oldest idle should have been evicted"
|
||||
assert mgr.get(ws4.id) is ws4
|
||||
@@ -349,35 +352,35 @@ class TestManagerEviction:
|
||||
def test_create_fails_when_all_active(self):
|
||||
"""At capacity with ALL non-idle workstreams, create() raises RuntimeError."""
|
||||
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
mgr.set_state(ws1.id, WorkstreamState.THINKING)
|
||||
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
||||
with pytest.raises(RuntimeError, match="All 2 workstreams are active"):
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
|
||||
def test_configurable_max(self):
|
||||
"""Constructor accepts max_workstreams param and respects it."""
|
||||
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
||||
ws1 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws2 = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws1 = mgr.create(ui_factory=FakeUI)
|
||||
ws2 = mgr.create(ui_factory=FakeUI)
|
||||
mgr.set_state(ws1.id, WorkstreamState.RUNNING)
|
||||
mgr.set_state(ws2.id, WorkstreamState.RUNNING)
|
||||
with pytest.raises(RuntimeError):
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.count == 2
|
||||
|
||||
def test_eviction_counter(self):
|
||||
"""eviction_count increments on each auto-eviction."""
|
||||
mgr = WorkstreamManager(_fake_factory, max_workstreams=2)
|
||||
assert mgr.eviction_count == 0
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
# Both IDLE — create should evict the oldest
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.eviction_count == 1
|
||||
# Again — evict another idle one
|
||||
mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
mgr.create(ui_factory=FakeUI)
|
||||
assert mgr.eviction_count == 2
|
||||
assert mgr.count == 2
|
||||
|
||||
@@ -390,7 +393,7 @@ class TestManagerEviction:
|
||||
class TestManagerState:
|
||||
def test_set_state(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
assert ws.state == WorkstreamState.IDLE
|
||||
|
||||
mgr.set_state(ws.id, WorkstreamState.THINKING)
|
||||
@@ -398,7 +401,7 @@ class TestManagerState:
|
||||
|
||||
def test_set_state_with_error(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
mgr.set_state(ws.id, WorkstreamState.ERROR, error_msg="API timeout")
|
||||
assert ws.state == WorkstreamState.ERROR
|
||||
@@ -410,7 +413,7 @@ class TestManagerState:
|
||||
|
||||
def test_on_state_change_callback(self):
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
changes = []
|
||||
mgr._on_state_change = lambda wid, state: changes.append((wid, state))
|
||||
@@ -433,7 +436,7 @@ class TestManagerThreadSafety:
|
||||
|
||||
def do_create():
|
||||
try:
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
# Mark as non-idle immediately so auto-eviction cannot reclaim it
|
||||
mgr.set_state(ws.id, WorkstreamState.RUNNING)
|
||||
created.append(ws.id)
|
||||
@@ -456,7 +459,7 @@ class TestManagerThreadSafety:
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ids = []
|
||||
for _ in range(5):
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
ids.append(ws.id)
|
||||
|
||||
def do_switch(wid):
|
||||
@@ -476,10 +479,10 @@ class TestManagerThreadSafety:
|
||||
"""close() and list_all() running concurrently should not crash."""
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
# Keep one alive to prevent closing the last
|
||||
anchor = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
anchor = mgr.create(ui_factory=FakeUI)
|
||||
targets = []
|
||||
for _ in range(5):
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
targets.append(ws.id)
|
||||
|
||||
def do_close():
|
||||
@@ -713,6 +716,41 @@ class TestWebUI:
|
||||
assert ui._plan_result == "approved"
|
||||
t.join()
|
||||
|
||||
def test_pending_plan_review_stored_and_replayed(self):
|
||||
"""Plan review state is stored for SSE reconnection replay."""
|
||||
from turnstone.server import WebUI
|
||||
|
||||
ui = WebUI(ws_id="test")
|
||||
assert ui._pending_plan_review is None
|
||||
|
||||
# Simulate on_plan_review in a background thread (it blocks)
|
||||
def review():
|
||||
ui.on_plan_review("Here is the plan")
|
||||
|
||||
t = threading.Thread(target=review)
|
||||
t.start()
|
||||
time.sleep(0.1)
|
||||
|
||||
# While blocking, pending state should be set
|
||||
assert ui._pending_plan_review is not None
|
||||
assert ui._pending_plan_review["type"] == "plan_review"
|
||||
assert ui._pending_plan_review["content"] == "Here is the plan"
|
||||
|
||||
# Resolve — pending state should be cleared
|
||||
ui.resolve_plan("looks good")
|
||||
t.join(timeout=2)
|
||||
assert ui._pending_plan_review is None
|
||||
assert ui._plan_result == "looks good"
|
||||
|
||||
def test_pending_plan_review_cleared_on_resolve_before_wait_returns(self):
|
||||
"""resolve_plan clears pending state immediately, not just after wait."""
|
||||
from turnstone.server import WebUI
|
||||
|
||||
ui = WebUI(ws_id="test")
|
||||
ui._pending_plan_review = {"type": "plan_review", "content": "test"}
|
||||
ui.resolve_plan("ok")
|
||||
assert ui._pending_plan_review is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# WebUI SSE fan-out
|
||||
@@ -730,13 +768,23 @@ class TestWebUIFanOut:
|
||||
ui._enqueue({"type": "content", "text": "hello"}) # should not raise
|
||||
|
||||
def test_enqueue_single_listener(self):
|
||||
"""Single listener receives the event."""
|
||||
"""Single listener receives the event with ws_id stamped."""
|
||||
from turnstone.server import WebUI
|
||||
|
||||
ui = WebUI(ws_id="test")
|
||||
q = ui._register_listener()
|
||||
ui._enqueue({"type": "content", "text": "hello"})
|
||||
assert q.get_nowait() == {"type": "content", "text": "hello"}
|
||||
assert q.get_nowait() == {"type": "content", "text": "hello", "ws_id": "test"}
|
||||
|
||||
def test_enqueue_does_not_mutate_input(self):
|
||||
"""_enqueue must not mutate the caller's dict."""
|
||||
from turnstone.server import WebUI
|
||||
|
||||
ui = WebUI(ws_id="test")
|
||||
ui._register_listener()
|
||||
original = {"type": "content", "text": "hello"}
|
||||
ui._enqueue(original)
|
||||
assert "ws_id" not in original
|
||||
|
||||
def test_enqueue_multiple_listeners(self):
|
||||
"""All registered listeners receive an identical copy."""
|
||||
@@ -747,12 +795,12 @@ class TestWebUIFanOut:
|
||||
q2 = ui._register_listener()
|
||||
q3 = ui._register_listener()
|
||||
|
||||
event = {"type": "content", "text": "world"}
|
||||
ui._enqueue(event)
|
||||
ui._enqueue({"type": "content", "text": "world"})
|
||||
|
||||
assert q1.get_nowait() == event
|
||||
assert q2.get_nowait() == event
|
||||
assert q3.get_nowait() == event
|
||||
expected = {"type": "content", "text": "world", "ws_id": "test"}
|
||||
assert q1.get_nowait() == expected
|
||||
assert q2.get_nowait() == expected
|
||||
assert q3.get_nowait() == expected
|
||||
|
||||
def test_unregister_stops_delivery(self):
|
||||
"""After unregister, the queue receives no further events."""
|
||||
@@ -784,11 +832,10 @@ class TestWebUIFanOut:
|
||||
assert fast.qsize() == 0
|
||||
|
||||
# Enqueue via fan-out — slow drops (full), fast receives
|
||||
event = {"type": "content", "text": "overflow"}
|
||||
ui._enqueue(event)
|
||||
ui._enqueue({"type": "content", "text": "overflow"})
|
||||
assert slow.qsize() == 500 # still full, overflow dropped
|
||||
assert fast.qsize() == 1
|
||||
assert fast.get_nowait() == event
|
||||
assert fast.get_nowait() == {"type": "content", "text": "overflow", "ws_id": "test"}
|
||||
|
||||
def test_unregister_idempotent(self):
|
||||
"""Double unregister does not raise."""
|
||||
@@ -834,7 +881,7 @@ class TestStateTransitions:
|
||||
def test_full_lifecycle(self):
|
||||
"""Verify the expected state transition sequence."""
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
# Simulate the state transitions that ChatSession.send() would emit
|
||||
mgr.set_state(ws.id, WorkstreamState.THINKING)
|
||||
@@ -855,7 +902,7 @@ class TestStateTransitions:
|
||||
def test_error_recovery(self):
|
||||
"""After an error, sending again should transition back to thinking."""
|
||||
mgr = WorkstreamManager(_fake_factory)
|
||||
ws = mgr.create(ui_factory=lambda wid: FakeUI(wid))
|
||||
ws = mgr.create(ui_factory=FakeUI)
|
||||
|
||||
mgr.set_state(ws.id, WorkstreamState.ERROR, "API failed")
|
||||
assert ws.state == WorkstreamState.ERROR
|
||||
|
||||
@@ -56,11 +56,9 @@
|
||||
# --- Auth (node, console) ---
|
||||
|
||||
[auth]
|
||||
# enabled = true # env: TURNSTONE_AUTH_ENABLED
|
||||
# Auth is always enabled. JWT secret is required.
|
||||
# jwt_secret = "" # HS256 signing secret (min 32 bytes recommended)
|
||||
# env: TURNSTONE_JWT_SECRET
|
||||
# token = "" # Static config token for full access
|
||||
# env: TURNSTONE_AUTH_TOKEN
|
||||
|
||||
# --- Logging (turnstone, node, console) ---
|
||||
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
"""turnstone - Multi-node AI orchestration platform with tool use, agent routing, and cluster simulation."""
|
||||
|
||||
__version__ = "0.9.8"
|
||||
__version__ = "1.1.1"
|
||||
|
||||
+13
-18
@@ -265,9 +265,19 @@ def _cmd_tls_list(args: argparse.Namespace) -> None:
|
||||
|
||||
url = f"{console_url}/v1/api/admin/tls/certs"
|
||||
headers = {}
|
||||
token = getattr(args, "auth_token", "") or _get_config_token()
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
# Prefer JWT via ServiceTokenManager when JWT secret is available
|
||||
jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip()
|
||||
if jwt_secret:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, ServiceTokenManager
|
||||
|
||||
mgr = ServiceTokenManager(
|
||||
user_id="admin-cli",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="cli",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {mgr.token}"
|
||||
resp = httpx.get(url, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
@@ -283,20 +293,6 @@ def _cmd_tls_list(args: argparse.Namespace) -> None:
|
||||
print(f"{c['domain']:<30s} {c['issued_at']:<22s} {c['expires_at']:<22s}")
|
||||
|
||||
|
||||
def _get_config_token() -> str:
|
||||
"""Try to load auth token from config.toml or environment."""
|
||||
token = os.environ.get("TURNSTONE_AUTH_TOKEN", "")
|
||||
if token:
|
||||
return token
|
||||
try:
|
||||
from turnstone.core.config import load_config
|
||||
|
||||
cfg = load_config("auth")
|
||||
return str(cfg.get("token", ""))
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
|
||||
def _discover_console_url() -> str:
|
||||
"""Discover console URL from the services table."""
|
||||
from turnstone.core.storage import get_storage
|
||||
@@ -381,7 +377,6 @@ def main() -> None:
|
||||
|
||||
p_tlslist = sub.add_parser("tls-list", help="List issued certificates")
|
||||
p_tlslist.add_argument("--console-url", default="", help="Console URL")
|
||||
p_tlslist.add_argument("--auth-token", default="", help="Auth token for admin API")
|
||||
|
||||
args = parser.parse_args()
|
||||
if not args.command:
|
||||
|
||||
+103
-44
@@ -2,14 +2,15 @@
|
||||
|
||||
Entry point: turnstone-bootstrap
|
||||
|
||||
Walks users through configuring a single-node or multi-node Turnstone
|
||||
deployment via a conversational AI assistant. Generates .env files,
|
||||
docker-compose overrides, and post-start setup scripts.
|
||||
Walks users through configuring a Turnstone deployment via a conversational
|
||||
AI assistant. Generates compose.yaml, .env files, and post-start setup
|
||||
scripts.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import getpass
|
||||
import importlib.resources
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
@@ -60,8 +61,6 @@ Turnstone is a multi-node AI orchestration platform. A deployment consists of:
|
||||
## Deployment Profiles (compose.yaml)
|
||||
- **Default** (no flag): console only (infrastructure, good for running external servers)
|
||||
- **Production** (`--profile production`): 1 server + console + PostgreSQL + channel (single node)
|
||||
- **Cluster** (`--profile cluster`): 10-node server fleet + PostgreSQL + channel + console (multi-node)
|
||||
- **ddgCluster** (`--profile ddgCluster`): Cluster + DuckDuckGo Search MCP sidecar (web search via MCP, no API key needed)
|
||||
|
||||
## Environment Variables (.env)
|
||||
The compose.yaml reads these from a `.env` file:
|
||||
@@ -78,14 +77,13 @@ For commercial providers (OpenAI, Anthropic-via-proxy), use the real key.
|
||||
|
||||
### Database
|
||||
- `DB_BACKEND` — `sqlite` (default) or `postgresql`
|
||||
- `DATABASE_URL` — PostgreSQL connection string (production/cluster only)
|
||||
- `DATABASE_URL` — PostgreSQL connection string (production only)
|
||||
- `POSTGRES_USER` — PostgreSQL username (default: turnstone)
|
||||
- `POSTGRES_PASSWORD` — PostgreSQL password (required for production/cluster)
|
||||
- `POSTGRES_PASSWORD` — PostgreSQL password (required for production)
|
||||
|
||||
### Authentication
|
||||
- `TURNSTONE_AUTH_ENABLED` — Enable auth (`true`/empty)
|
||||
- `TURNSTONE_JWT_SECRET` — JWT signing secret (required if auth enabled)
|
||||
- `TURNSTONE_AUTH_TOKEN` — Static bearer token for inter-service auth
|
||||
### Authentication (always enabled)
|
||||
- `TURNSTONE_JWT_SECRET` — JWT signing secret (required). All services must share the same secret. \
|
||||
Generate with: `python -c "import secrets; print(secrets.token_hex(32))"`
|
||||
|
||||
### OIDC SSO (optional)
|
||||
- `TURNSTONE_OIDC_ISSUER` — OIDC issuer URL (e.g., https://accounts.google.com). Setting this + CLIENT_ID + CLIENT_SECRET enables SSO.
|
||||
@@ -105,16 +103,15 @@ For commercial providers (OpenAI, Anthropic-via-proxy), use the real key.
|
||||
- `TURNSTONE_DISCORD_TOKEN` — Discord bot token
|
||||
- `TURNSTONE_DISCORD_GUILD` — Restrict to single guild ID
|
||||
|
||||
### MCP Integration (optional)
|
||||
- `MCP_CONFIG` — Path to MCP server config inside the container \
|
||||
(e.g., `/etc/turnstone/mcp-ddg.json`). When set, servers connect to configured MCP servers on startup.
|
||||
- The `ddgCluster` profile runs a DuckDuckGo Search MCP sidecar (Python) that provides \
|
||||
`duckduckgo_web_search` and `duckduckgo_fetch_content` tools to every node. No API key required. \
|
||||
The sidecar uses MCP streamable-http transport with DNS rebinding protection disabled \
|
||||
(required for Docker internal networking) and binds to 0.0.0.0:3000 via FastMCP settings. \
|
||||
Safe search is disabled by default.
|
||||
### Docker Image
|
||||
- `TURNSTONE_IMAGE_TAG` — Docker image tag (default: `latest`). \
|
||||
Set this to pin the image version (e.g., `1.1.0`, `stable`, `experimental`).
|
||||
|
||||
### Cluster
|
||||
### MCP Integration (optional)
|
||||
- `MCP_CONFIG` — Path to MCP server config inside the container. \
|
||||
When set, servers connect to configured MCP servers on startup.
|
||||
|
||||
### Other
|
||||
- `APPROVAL_TIMEOUT` — Tool approval timeout in seconds (default: 3600)
|
||||
|
||||
## Auth Setup Flow
|
||||
@@ -155,27 +152,28 @@ Categories like "engineering", "analysis", etc.
|
||||
## Your Task
|
||||
Walk the user through setting up their deployment step by step:
|
||||
|
||||
1. **First**: Call `check_docker` and `read_file` on `.env` to detect existing state.
|
||||
2. **Deployment mode**: Ask if they want single-node (`--profile production`) or multi-node \
|
||||
(`--profile cluster`). Explain trade-offs.
|
||||
3. **LLM provider for the deployment**: Which LLM backend their Turnstone will use \
|
||||
1. **First**: Call `check_docker`, `read_file` on `.env`, and `read_file` on `compose.yaml` \
|
||||
to detect existing state. If `compose.yaml` does not exist, call `write_compose` to \
|
||||
extract the bundled production compose file. This is essential — without it, \
|
||||
`docker compose` will fail.
|
||||
2. **LLM provider for the deployment**: Which LLM backend their Turnstone will use \
|
||||
(may differ from this wizard's model). Ask for base URL, API key, model name.
|
||||
4. **Database**: SQLite (dev/simple) vs PostgreSQL (production/cluster). \
|
||||
PostgreSQL is required for cluster mode.
|
||||
5. **Security**: Recommend enabling auth for any non-local deployment. \
|
||||
Use `generate_secret` for JWT secret, auth token, and Postgres password. \
|
||||
3. **Database**: SQLite (dev/simple) vs PostgreSQL (production). \
|
||||
PostgreSQL is recommended for production use.
|
||||
4. **Security**: Auth is always enabled and requires `TURNSTONE_JWT_SECRET`. \
|
||||
Use `generate_secret` for JWT secret and Postgres password. \
|
||||
Always set `TURNSTONE_JWT_SECRET` in the .env. \
|
||||
Ask for initial admin username and password. \
|
||||
If the user's deployment will use an external identity provider (Okta, Azure AD, Google, etc.), \
|
||||
offer to configure OIDC SSO. Ask for the issuer URL, client ID, and client secret. \
|
||||
Optionally configure role mapping and OIDC-only mode.
|
||||
6. **Ports**: Check defaults with `check_port`, suggest alternatives if conflicts.
|
||||
7. **Optional features**: Discord integration, web search (Tavily key), \
|
||||
DuckDuckGo Search MCP (for cluster — uses `ddgCluster` profile with \
|
||||
`MCP_CONFIG=/etc/turnstone/mcp-ddg.json`, no API key needed).
|
||||
8. **Generate .env**: Call `write_file` with the complete `.env` content.
|
||||
9. **Generate setup.sh**: Call `write_file` with a post-start script that creates the admin \
|
||||
5. **Ports**: Check defaults with `check_port`, suggest alternatives if conflicts.
|
||||
6. **Optional features**: Discord integration, web search (Tavily key).
|
||||
7. **Generate .env**: Call `write_file` with the complete `.env` content. \
|
||||
Include `TURNSTONE_IMAGE_TAG` set to the version matching the installed package.
|
||||
8. **Generate setup.sh**: Call `write_file` with a post-start script that creates the admin \
|
||||
user and any roles/policies/skills the user wants.
|
||||
10. **Finish**: Call the `finish` tool with a summary of what was configured and the \
|
||||
9. **Finish**: Call the `finish` tool with a summary of what was configured and the \
|
||||
exact commands to run next (e.g., `docker compose --profile production up -d` then `./setup.sh`).
|
||||
|
||||
## Rules
|
||||
@@ -183,14 +181,9 @@ exact commands to run next (e.g., `docker compose --profile production up -d` th
|
||||
- NEVER echo API keys or passwords back to the user in your text responses.
|
||||
- ALWAYS use `generate_secret` for passwords and secrets — never invent them.
|
||||
- When writing files, use `write_file` — the user will see a preview and confirm.
|
||||
- If `compose.yaml` is missing, call `write_compose` before anything else. \
|
||||
The compose file uses pre-built images from ghcr.io — no local Docker build is needed.
|
||||
- If an existing .env is detected, summarize what's configured and ask what to change.
|
||||
- For cluster mode, the compose.yaml has a fixed 10-node fleet — no override needed.
|
||||
- For cluster + DuckDuckGo Search, use `--profile ddgCluster` instead of `--profile cluster`. \
|
||||
Set `MCP_CONFIG=/etc/turnstone/mcp-ddg.json` in `.env`. No API key needed. \
|
||||
The DuckDuckGo MCP sidecar starts automatically and all cluster nodes connect to it. \
|
||||
Note: the MCP SDK's DNS rebinding protection must be disabled for Docker-internal networking \
|
||||
(the compose.yaml handles this), and the server must bind to 0.0.0.0 (not 127.0.0.1) to be \
|
||||
reachable from other containers.
|
||||
- The `DATABASE_URL` for docker compose internal networking uses the hostname `postgres` \
|
||||
(e.g., `postgresql+psycopg://turnstone:<password>@postgres:5432/turnstone`).
|
||||
- For local LLM backends (vLLM, llama.cpp, Ollama, etc.), set `OPENAI_API_KEY=dummy` in the \
|
||||
@@ -342,6 +335,23 @@ TOOLS: list[dict[str, Any]] = [
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "write_compose",
|
||||
"description": (
|
||||
"Write the production Docker Compose file to the project directory. "
|
||||
"This extracts the compose.yaml bundled with Turnstone, which uses "
|
||||
"pre-built images from ghcr.io (no local Docker build required). "
|
||||
"The user will be shown a preview and asked to confirm."
|
||||
),
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
@@ -427,7 +437,7 @@ def _tool_write_file(project_dir: Path, args: dict[str, Any]) -> str:
|
||||
if existing == content:
|
||||
return f"File already exists with identical content: {args['path']}"
|
||||
except (OSError, UnicodeDecodeError):
|
||||
pass
|
||||
pass # best-effort duplicate check
|
||||
|
||||
line_count = content.count("\n") + (1 if content and not content.endswith("\n") else 0)
|
||||
|
||||
@@ -561,6 +571,54 @@ def _tool_check_docker(args: dict[str, Any]) -> str:
|
||||
return "\n".join(results)
|
||||
|
||||
|
||||
def _tool_write_compose(project_dir: Path, args: dict[str, Any]) -> str:
|
||||
"""Extract the bundled production compose.yaml to the project directory."""
|
||||
dest = project_dir / "compose.yaml"
|
||||
|
||||
# Read the bundled template
|
||||
try:
|
||||
ref = importlib.resources.files("turnstone.deploy").joinpath("compose.yaml")
|
||||
content = ref.read_text(encoding="utf-8")
|
||||
except Exception as exc:
|
||||
return f"Error: could not read bundled compose template: {exc}"
|
||||
|
||||
# Skip if identical
|
||||
if dest.exists():
|
||||
try:
|
||||
existing = dest.read_text(encoding="utf-8")
|
||||
if existing == content:
|
||||
return "compose.yaml already exists with identical content."
|
||||
except (OSError, UnicodeDecodeError):
|
||||
pass # best-effort duplicate check
|
||||
|
||||
line_count = content.count("\n") + (1 if content and not content.endswith("\n") else 0)
|
||||
|
||||
# Show preview
|
||||
print(f"\n{YELLOW} Writing compose.yaml ({line_count} lines){RESET}")
|
||||
print(f"{DIM}{'─' * 50}{RESET}")
|
||||
for line in content.split("\n")[:30]:
|
||||
print(f" {DIM}{line}{RESET}")
|
||||
if line_count > 30:
|
||||
print(f" {DIM}... ({line_count - 30} more lines){RESET}")
|
||||
print(f"{DIM}{'─' * 50}{RESET}")
|
||||
|
||||
try:
|
||||
choice = input(f"{BOLD}Write this file? [Y/n]{RESET} ").strip().lower()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
return "User cancelled the write."
|
||||
if choice in ("n", "no"):
|
||||
return "User declined to write compose.yaml."
|
||||
|
||||
dest.write_text(content, encoding="utf-8")
|
||||
|
||||
return (
|
||||
f"compose.yaml written successfully. "
|
||||
f"It uses ghcr.io/turnstonelabs/turnstone images. "
|
||||
f"Add TURNSTONE_IMAGE_TAG={__version__} to .env to pin the image "
|
||||
f"to the currently installed version, or omit it to use 'latest'."
|
||||
)
|
||||
|
||||
|
||||
class _FinishError(Exception):
|
||||
"""Raised by the finish tool to signal the wizard is done."""
|
||||
|
||||
@@ -581,11 +639,12 @@ TOOL_FUNCTIONS: dict[str, Any] = {
|
||||
"check_port": _tool_check_port,
|
||||
"validate_api_key": _tool_validate_api_key,
|
||||
"check_docker": _tool_check_docker,
|
||||
"write_compose": _tool_write_compose,
|
||||
"finish": _tool_finish,
|
||||
}
|
||||
|
||||
# Tools that need the project_dir argument
|
||||
_PROJECT_DIR_TOOLS = frozenset({"read_file", "write_file"})
|
||||
_PROJECT_DIR_TOOLS = frozenset({"read_file", "write_file", "write_compose"})
|
||||
|
||||
|
||||
def execute_tool(name: str, args: dict[str, Any], project_dir: Path) -> str:
|
||||
|
||||
@@ -1,12 +1,17 @@
|
||||
"""Message formatting utilities for channel adapters.
|
||||
|
||||
Handles chunking long messages for platforms with character limits, formatting
|
||||
tool-approval requests, and plan-review prompts.
|
||||
tool-approval requests, plan-review prompts, and rich media embeds for
|
||||
platforms that support them (e.g. Discord).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
||||
|
||||
def chunk_message(text: str, max_length: int = 2000) -> list[str]:
|
||||
@@ -164,3 +169,298 @@ def truncate(text: str, max_length: int = 200) -> str:
|
||||
if len(text) <= max_length:
|
||||
return text
|
||||
return text[: max_length - 1] + "\u2026"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Rich media embed helpers (Discord)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def try_parse_media(output: str) -> dict[str, Any] | None:
|
||||
"""Attempt to parse tool output as a media result.
|
||||
|
||||
Returns the parsed dict when the output looks like structured media
|
||||
(single item, search results, or session list), otherwise ``None``.
|
||||
"""
|
||||
try:
|
||||
data = json.loads(output)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
# Single item with stream URL or detailed metadata.
|
||||
if "stream_url" in data or ("name" in data and "type" in data and "id" in data):
|
||||
return data
|
||||
# Search results.
|
||||
if "results" in data and isinstance(data["results"], list) and data["results"]:
|
||||
return data
|
||||
# Active sessions.
|
||||
if "sessions" in data and isinstance(data["sessions"], list):
|
||||
return data
|
||||
return None
|
||||
|
||||
|
||||
_BLOCKED_HOSTNAMES = frozenset({"localhost", "metadata.google.internal"})
|
||||
|
||||
|
||||
def _is_safe_image_url(url: str) -> bool:
|
||||
"""Validate that *url* uses http(s), has no embedded credentials, and does
|
||||
not target loopback or cloud metadata endpoints.
|
||||
|
||||
Private/LAN IPs are intentionally allowed (media servers are typically
|
||||
on the local network).
|
||||
"""
|
||||
import ipaddress
|
||||
from urllib.parse import urlparse
|
||||
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
return False
|
||||
if parsed.username or parsed.password:
|
||||
return False
|
||||
hostname = parsed.hostname
|
||||
if not hostname:
|
||||
return False
|
||||
if hostname in _BLOCKED_HOSTNAMES:
|
||||
return False
|
||||
try:
|
||||
ip = ipaddress.ip_address(hostname)
|
||||
if ip.is_loopback or ip.is_link_local:
|
||||
return False
|
||||
except ValueError:
|
||||
pass # Not an IP literal — hostname is fine
|
||||
return True
|
||||
|
||||
|
||||
async def _fetch_thumbnail(
|
||||
http: httpx.AsyncClient,
|
||||
url: str,
|
||||
*,
|
||||
timeout: float = 5.0,
|
||||
max_bytes: int = 2 * 1024 * 1024,
|
||||
) -> tuple[bytes, str] | None:
|
||||
"""Fetch a thumbnail image, returning ``(bytes, filename)`` or ``None``.
|
||||
|
||||
Never raises — a failed image fetch must not break tool result
|
||||
rendering. Private/LAN URLs are intentionally allowed (media servers
|
||||
are typically on the local network), but scheme is restricted to
|
||||
http(s) and userinfo is rejected.
|
||||
"""
|
||||
if not _is_safe_image_url(url):
|
||||
return None
|
||||
try:
|
||||
async with http.stream("GET", url, timeout=timeout) as resp:
|
||||
if resp.status_code != 200:
|
||||
return None
|
||||
cl = resp.headers.get("content-length")
|
||||
if cl and cl.isdigit() and int(cl) > max_bytes:
|
||||
return None
|
||||
content_type = resp.headers.get("content-type", "image/jpeg").lower()
|
||||
if not content_type.startswith("image/"):
|
||||
return None
|
||||
ext = "jpg"
|
||||
if "png" in content_type:
|
||||
ext = "png"
|
||||
elif "webp" in content_type:
|
||||
ext = "webp"
|
||||
data = bytearray()
|
||||
async for chunk in resp.aiter_bytes():
|
||||
data.extend(chunk)
|
||||
if len(data) > max_bytes:
|
||||
return None
|
||||
return bytes(data), f"poster.{ext}"
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
|
||||
|
||||
async def try_build_media_embed(
|
||||
tool_name: str,
|
||||
output: str,
|
||||
*,
|
||||
http: httpx.AsyncClient,
|
||||
) -> tuple[Any, Any | None] | None:
|
||||
"""Attempt to build a rich Discord embed from media tool output.
|
||||
|
||||
Returns ``(embed, optional_file)`` if the output is parseable as media,
|
||||
or ``None`` to fall through to the default code-block formatter.
|
||||
|
||||
The ``discord`` library is imported lazily since this module is shared
|
||||
across adapters and ``discord.py`` is an optional dependency.
|
||||
"""
|
||||
data = try_parse_media(output)
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
import io
|
||||
|
||||
import discord
|
||||
|
||||
# Dispatch on result shape.
|
||||
if "results" in data and isinstance(data["results"], list):
|
||||
embed = _build_search_results_embed(data)
|
||||
elif "sessions" in data and isinstance(data["sessions"], list):
|
||||
embed = _build_sessions_embed(data)
|
||||
else:
|
||||
embed = _build_single_media_embed(data, tool_name)
|
||||
|
||||
# Proxy thumbnail image.
|
||||
thumbnail_url = data.get("thumbnail_url") or data.get("image_url")
|
||||
if not thumbnail_url and data.get("results"):
|
||||
first = data["results"][0]
|
||||
thumbnail_url = first.get("thumbnail_url") or first.get("image_url")
|
||||
|
||||
file: discord.File | None = None
|
||||
if thumbnail_url:
|
||||
fetched = await _fetch_thumbnail(http, thumbnail_url)
|
||||
if fetched:
|
||||
image_bytes, filename = fetched
|
||||
file = discord.File(io.BytesIO(image_bytes), filename=filename)
|
||||
embed.set_thumbnail(url=f"attachment://{filename}")
|
||||
|
||||
return embed, file
|
||||
|
||||
|
||||
# -- Private embed builders ------------------------------------------------
|
||||
|
||||
|
||||
def _build_single_media_embed(data: dict[str, Any], tool_name: str) -> Any:
|
||||
"""Build a Discord embed for a single media item."""
|
||||
import discord
|
||||
|
||||
title = data.get("name", "Unknown")
|
||||
if data.get("year"):
|
||||
title += f" ({data['year']})"
|
||||
|
||||
embed = discord.Embed(
|
||||
title=title,
|
||||
url=data.get("web_url"), # safe link — NOT stream_url
|
||||
description=truncate(data.get("overview", ""), 200),
|
||||
color=discord.Color.teal(),
|
||||
)
|
||||
|
||||
# Metadata fields (inline).
|
||||
meta_parts: list[str] = []
|
||||
if data.get("type"):
|
||||
meta_parts.append(data["type"])
|
||||
if data.get("official_rating"):
|
||||
meta_parts.append(data["official_rating"])
|
||||
if data.get("runtime_minutes"):
|
||||
hours = int(data["runtime_minutes"] // 60)
|
||||
mins = int(data["runtime_minutes"] % 60)
|
||||
meta_parts.append(f"{hours}h {mins}m" if hours else f"{mins}m")
|
||||
if meta_parts:
|
||||
embed.add_field(name="Info", value=" \u00b7 ".join(meta_parts), inline=True)
|
||||
|
||||
if data.get("genres"):
|
||||
embed.add_field(name="Genres", value=", ".join(data["genres"][:5]), inline=True)
|
||||
|
||||
if data.get("community_rating"):
|
||||
embed.add_field(
|
||||
name="Rating",
|
||||
value=f"{data['community_rating']:.1f}/10",
|
||||
inline=True,
|
||||
)
|
||||
|
||||
# Extract server name from tool_name (mcp__servername__toolname).
|
||||
parts = tool_name.split("__")
|
||||
if len(parts) >= 3:
|
||||
embed.set_footer(text=parts[1])
|
||||
|
||||
return embed
|
||||
|
||||
|
||||
def _build_search_results_embed(data: dict[str, Any]) -> Any:
|
||||
"""Build a Discord embed for a list of search results."""
|
||||
import discord
|
||||
|
||||
results = data.get("results", [])
|
||||
total = data.get("total_count", len(results))
|
||||
|
||||
lines: list[str] = []
|
||||
char_count = 0
|
||||
for i, r in enumerate(results[:10], 1):
|
||||
line = f"**{i}.** {r.get('name', '?')}"
|
||||
if r.get("year"):
|
||||
line += f" ({r['year']})"
|
||||
meta: list[str] = []
|
||||
if r.get("type"):
|
||||
meta.append(r["type"])
|
||||
if r.get("series_name"):
|
||||
meta.append(r["series_name"])
|
||||
if r.get("season_number") is not None and r.get("episode_number") is not None:
|
||||
meta.append(f"S{int(r['season_number']):02d}E{int(r['episode_number']):02d}")
|
||||
if r.get("runtime_minutes"):
|
||||
mins = r["runtime_minutes"]
|
||||
meta.append(f"{int(mins // 60)}h {int(mins % 60)}m" if mins >= 60 else f"{int(mins)}m")
|
||||
if meta:
|
||||
line += " \u00b7 " + " \u00b7 ".join(meta)
|
||||
if char_count + len(line) + 1 > 4000:
|
||||
break
|
||||
lines.append(line)
|
||||
char_count += len(line) + 1
|
||||
|
||||
embed = discord.Embed(
|
||||
title="Search results",
|
||||
description="\n".join(lines),
|
||||
color=discord.Color.teal(),
|
||||
)
|
||||
embed.set_footer(text=f"showing {len(lines)} of {total}")
|
||||
return embed
|
||||
|
||||
|
||||
def _build_sessions_embed(data: dict[str, Any]) -> Any:
|
||||
"""Build a Discord embed for active playback sessions."""
|
||||
import discord
|
||||
|
||||
sessions = data.get("sessions", [])
|
||||
if not sessions:
|
||||
embed = discord.Embed(
|
||||
title="Now Playing",
|
||||
description="No active sessions.",
|
||||
color=discord.Color.light_grey(),
|
||||
)
|
||||
return embed
|
||||
|
||||
lines: list[str] = []
|
||||
has_active = False
|
||||
for s in sessions:
|
||||
np = s.get("now_playing")
|
||||
device = s.get("device_name", "Unknown device")
|
||||
user = s.get("user_name", "")
|
||||
if np:
|
||||
has_active = True
|
||||
title = np.get("name", "Unknown")
|
||||
if np.get("year"):
|
||||
title += f" ({np['year']})"
|
||||
ps = s.get("play_state", {}) or {}
|
||||
pos = ps.get("position_seconds")
|
||||
runtime_min = np.get("runtime_minutes")
|
||||
time_str = ""
|
||||
if pos is not None and runtime_min:
|
||||
total_sec = int(runtime_min * 60)
|
||||
pos_i = int(pos)
|
||||
time_str = (
|
||||
f" {pos_i // 3600}:{pos_i % 3600 // 60:02d}:{pos_i % 60:02d}"
|
||||
f" / {total_sec // 3600}:{total_sec % 3600 // 60:02d}:{total_sec % 60:02d}"
|
||||
)
|
||||
paused = ps.get("is_paused", False)
|
||||
icon = "\u23f8" if paused else "\u25b6"
|
||||
line = f"**{title}** on {device}\n{icon}{time_str}"
|
||||
if user:
|
||||
line += f" \u00b7 {user}"
|
||||
lines.append(line)
|
||||
else:
|
||||
line = f"*{device}* \u2014 idle"
|
||||
if user:
|
||||
line += f" ({user})"
|
||||
lines.append(line)
|
||||
|
||||
embed = discord.Embed(
|
||||
title="Now Playing",
|
||||
description="\n\n".join(lines),
|
||||
color=discord.Color.green() if has_active else discord.Color.light_grey(),
|
||||
)
|
||||
return embed
|
||||
|
||||
@@ -37,10 +37,9 @@ async def _handle_health(request: Request) -> JSONResponse:
|
||||
|
||||
def _check_auth(request: Request) -> JSONResponse | None:
|
||||
"""Validate the request's Authorization header. Returns an error response or None."""
|
||||
auth_token: str = getattr(request.app.state, "auth_token", "")
|
||||
jwt_secret: str = getattr(request.app.state, "jwt_secret", "")
|
||||
|
||||
if not auth_token and not jwt_secret:
|
||||
if not jwt_secret:
|
||||
log.warning("notify.auth_not_configured")
|
||||
return JSONResponse({"error": "authentication not configured"}, status_code=401)
|
||||
|
||||
@@ -50,15 +49,8 @@ def _check_auth(request: Request) -> JSONResponse | None:
|
||||
|
||||
token = header[7:]
|
||||
|
||||
# Static token check
|
||||
if auth_token:
|
||||
import hmac
|
||||
|
||||
if hmac.compare_digest(token, auth_token):
|
||||
return None
|
||||
|
||||
# JWT check
|
||||
if jwt_secret and "." in token:
|
||||
# JWT validation
|
||||
if "." in token:
|
||||
from turnstone.core.auth import JWT_AUD_CHANNEL, validate_jwt
|
||||
|
||||
result = validate_jwt(token, jwt_secret, audience=JWT_AUD_CHANNEL)
|
||||
@@ -178,7 +170,6 @@ def create_channel_app(
|
||||
adapters: dict[str, ChannelAdapter],
|
||||
storage: StorageBackend,
|
||||
*,
|
||||
auth_token: str = "",
|
||||
jwt_secret: str = "",
|
||||
) -> Starlette:
|
||||
"""Create the channel gateway HTTP application."""
|
||||
@@ -195,7 +186,6 @@ def create_channel_app(
|
||||
)
|
||||
app.state.adapters = adapters
|
||||
app.state.storage = storage
|
||||
app.state.auth_token = auth_token
|
||||
app.state.jwt_secret = jwt_secret
|
||||
return app
|
||||
|
||||
|
||||
@@ -244,7 +244,7 @@ class ChannelRouter:
|
||||
self._node_urls[ws_id] = node_url.rstrip("/")
|
||||
return self._node_urls[ws_id]
|
||||
except Exception:
|
||||
pass
|
||||
log.debug("Console route lookup failed for ws %s", ws_id, exc_info=True)
|
||||
return self._server_url
|
||||
|
||||
# -- user resolution -----------------------------------------------------
|
||||
|
||||
@@ -72,13 +72,6 @@ def main() -> None:
|
||||
parser.add_argument("--ssl-keyfile", default=None, help="SSL private key file")
|
||||
parser.add_argument("--ssl-ca-certs", default=None, help="SSL CA certs for client verification")
|
||||
|
||||
# -- Auth ----------------------------------------------------------------
|
||||
parser.add_argument(
|
||||
"--auth-token",
|
||||
default=os.environ.get("TURNSTONE_CHANNEL_AUTH_TOKEN", ""),
|
||||
help="Static auth token for /v1/api/notify (default: $TURNSTONE_CHANNEL_AUTH_TOKEN)",
|
||||
)
|
||||
|
||||
# -- Workstream defaults -------------------------------------------------
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
@@ -121,7 +114,6 @@ def main() -> None:
|
||||
)
|
||||
|
||||
# -- Auth config ---------------------------------------------------------
|
||||
auth_token = os.environ.get("TURNSTONE_AUTH_TOKEN", "") or args.auth_token
|
||||
jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip()
|
||||
|
||||
# Prefer auto-rotating service JWTs when jwt_secret is available.
|
||||
@@ -132,7 +124,7 @@ def main() -> None:
|
||||
if jwt_secret:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, JWT_AUD_SERVER, ServiceTokenManager
|
||||
|
||||
_scopes = frozenset({"read", "write", "approve"})
|
||||
_scopes = frozenset({"read", "write", "approve", "service"})
|
||||
_console_mgr = ServiceTokenManager(
|
||||
user_id="channel-gateway",
|
||||
scopes=_scopes,
|
||||
@@ -151,7 +143,6 @@ def main() -> None:
|
||||
)
|
||||
_console_token_factory = lambda: _console_mgr.token # noqa: E731
|
||||
_server_token_factory = lambda: _server_mgr.token # noqa: E731
|
||||
auth_token = "" # don't also pass static token
|
||||
|
||||
server_url: str = args.server_url
|
||||
console_url: str = args.console_url
|
||||
@@ -242,7 +233,6 @@ def main() -> None:
|
||||
config,
|
||||
server_url,
|
||||
storage,
|
||||
api_token=auth_token,
|
||||
console_url=console_url,
|
||||
console_token_factory=_console_token_factory,
|
||||
server_token_factory=_server_token_factory,
|
||||
@@ -253,7 +243,6 @@ def main() -> None:
|
||||
channel_app = create_channel_app(
|
||||
adapters, # type: ignore[arg-type]
|
||||
storage,
|
||||
auth_token=auth_token,
|
||||
jwt_secret=jwt_secret,
|
||||
)
|
||||
|
||||
@@ -293,10 +282,14 @@ def main() -> None:
|
||||
|
||||
async def _heartbeat_loop() -> None:
|
||||
"""Periodically update service heartbeat."""
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30)
|
||||
try:
|
||||
await asyncio.to_thread(storage.heartbeat_service, "channel", service_id)
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.exception("channel.heartbeat_failed")
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ import contextlib
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -180,6 +180,7 @@ class TurnstoneBot:
|
||||
server_token_factory=server_token_factory,
|
||||
)
|
||||
|
||||
self._commands_synced: bool = False
|
||||
self._subscribed_ws: set[str] = set()
|
||||
self._sse_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
self._streaming: dict[str, StreamingMessage] = {}
|
||||
@@ -204,12 +205,17 @@ class TurnstoneBot:
|
||||
# response message can be re-tracked for multi-turn DM conversations.
|
||||
self._notify_reply_channels: dict[str, tuple[discord.abc.Messageable, str]] = {}
|
||||
|
||||
# Shared HTTP client for SSE connections (long-lived, no timeout).
|
||||
# Shared HTTP client for SSE connections.
|
||||
# Read timeout detects half-open connections (server sends ping=5s
|
||||
# keepalives, so 90s is very conservative).
|
||||
# Token factory provides auto-rotating JWTs; static token is fallback.
|
||||
headers: dict[str, str] = {}
|
||||
if api_token and not server_token_factory:
|
||||
headers["Authorization"] = f"Bearer {api_token}"
|
||||
self._http_client = httpx.AsyncClient(headers=headers, timeout=None)
|
||||
self._http_client = httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(connect=10.0, read=90.0, write=10.0, pool=10.0),
|
||||
)
|
||||
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
@@ -230,6 +236,10 @@ class TurnstoneBot:
|
||||
async def on_ready() -> None:
|
||||
await self._on_ready()
|
||||
|
||||
@self._bot.event
|
||||
async def on_resumed() -> None:
|
||||
await self._on_resumed()
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
async def _setup_hook(self) -> None:
|
||||
@@ -247,23 +257,57 @@ class TurnstoneBot:
|
||||
log.info("discord.setup_hook_complete")
|
||||
|
||||
async def _on_ready(self) -> None:
|
||||
"""Sync slash commands and recover existing routes."""
|
||||
"""Sync slash commands (once) and recover existing routes."""
|
||||
import discord
|
||||
|
||||
bot = self._bot
|
||||
log.info("discord.ready", user=str(bot.user), guild_count=len(bot.guilds))
|
||||
|
||||
if self.config.guild_id:
|
||||
guild = discord.Object(id=self.config.guild_id)
|
||||
bot.tree.copy_global_to(guild=guild)
|
||||
await bot.tree.sync(guild=guild)
|
||||
log.info("discord.commands_synced", guild_id=self.config.guild_id)
|
||||
else:
|
||||
await bot.tree.sync()
|
||||
log.info("discord.commands_synced_global")
|
||||
if not self._commands_synced:
|
||||
if self.config.guild_id:
|
||||
guild = discord.Object(id=self.config.guild_id)
|
||||
bot.tree.copy_global_to(guild=guild)
|
||||
await bot.tree.sync(guild=guild)
|
||||
log.info("discord.commands_synced", guild_id=self.config.guild_id)
|
||||
else:
|
||||
await bot.tree.sync()
|
||||
log.info("discord.commands_synced_global")
|
||||
self._commands_synced = True
|
||||
|
||||
self._purge_dead_sse_tasks("ready")
|
||||
await self._recover_routes()
|
||||
|
||||
async def _on_resumed(self) -> None:
|
||||
"""Recover dead SSE tasks after a gateway session resume.
|
||||
|
||||
Unlike ``on_ready``, ``on_resumed`` fires when discord.py resumes
|
||||
an existing session after a brief disconnect — ``on_ready`` is NOT
|
||||
called in that case. Any SSE listener tasks that died during the
|
||||
blip need to be cleaned up and re-subscribed.
|
||||
"""
|
||||
self._purge_dead_sse_tasks("resumed")
|
||||
await self._recover_routes()
|
||||
|
||||
def _purge_dead_sse_tasks(self, trigger: str) -> None:
|
||||
"""Remove completed/failed SSE tasks so they can be re-subscribed."""
|
||||
dead = [ws_id for ws_id, task in self._sse_tasks.items() if task.done()]
|
||||
for ws_id in dead:
|
||||
task = self._sse_tasks.pop(ws_id)
|
||||
self._subscribed_ws.discard(ws_id)
|
||||
# Retrieve exception to suppress "Task exception was never
|
||||
# retrieved" warnings and log the underlying failure.
|
||||
if not task.cancelled():
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
log.warning(
|
||||
"discord.sse_task_failed",
|
||||
trigger=trigger,
|
||||
ws_id=ws_id,
|
||||
error=str(exc),
|
||||
)
|
||||
if dead:
|
||||
log.info("discord.purged_dead_tasks", trigger=trigger, count=len(dead), ws_ids=dead)
|
||||
|
||||
async def _recover_routes(self) -> None:
|
||||
"""Re-subscribe to event channels for existing discord routes.
|
||||
|
||||
@@ -361,14 +405,15 @@ class TurnstoneBot:
|
||||
"""
|
||||
import httpx_sse
|
||||
|
||||
# When routing through the console, connect SSE directly to the
|
||||
# assigned server node (node_url from the create response).
|
||||
node_base = await self.router.get_node_url(ws_id)
|
||||
url = f"{node_base}/v1/api/events"
|
||||
delay = _SSE_RECONNECT_DELAY
|
||||
url = "" # set before loop so exception handlers can reference it
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Re-resolve node URL on each attempt so reconnects pick up
|
||||
# changes after bot restarts or router cache expiry.
|
||||
node_base = await self.router.get_node_url(ws_id)
|
||||
url = f"{node_base}/v1/api/events"
|
||||
# Refresh auth header per-connection (token may have rotated)
|
||||
sse_headers: dict[str, str] | None = None
|
||||
if self._token_factory is not None:
|
||||
@@ -392,7 +437,13 @@ class TurnstoneBot:
|
||||
ws_id=ws_id,
|
||||
status=status,
|
||||
)
|
||||
# Fall through to backoff/retry for transient errors.
|
||||
# Don't try to parse a non-SSE error body —
|
||||
# fall through to backoff/retry below.
|
||||
raise httpx.HTTPStatusError(
|
||||
f"SSE upstream {status}",
|
||||
request=event_source.response.request,
|
||||
response=event_source.response,
|
||||
)
|
||||
delay = _SSE_RECONNECT_DELAY # reset on successful connect
|
||||
async for sse in event_source.aiter_sse():
|
||||
if sse.event == "message" or not sse.event:
|
||||
@@ -406,12 +457,27 @@ class TurnstoneBot:
|
||||
)
|
||||
continue
|
||||
event = ServerEvent.from_dict(data)
|
||||
await self._on_ws_event(ws_id, thread, event)
|
||||
try:
|
||||
await self._on_ws_event(ws_id, thread, event)
|
||||
except Exception:
|
||||
# Discord API failures (rate limits, outages)
|
||||
# must not kill the SSE connection.
|
||||
log.warning(
|
||||
"discord.event_dispatch_failed",
|
||||
ws_id=ws_id,
|
||||
exc_info=True,
|
||||
)
|
||||
except httpx.HTTPStatusError:
|
||||
pass # already logged above; fall through to backoff
|
||||
except httpx.RemoteProtocolError:
|
||||
# Server closed connection (normal on stream_end or shutdown).
|
||||
log.debug("discord.sse_remote_closed", ws_id=ws_id)
|
||||
except asyncio.CancelledError:
|
||||
return # unsubscribe or shutdown
|
||||
except httpx.ReadTimeout:
|
||||
# No data received within read timeout — likely a half-open
|
||||
# connection. Reconnect to recover.
|
||||
log.info("discord.sse_read_timeout", ws_id=ws_id)
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout) as exc:
|
||||
log.warning(
|
||||
"discord.sse_connect_failed",
|
||||
@@ -492,15 +558,16 @@ class TurnstoneBot:
|
||||
# authorize this?" while the running embed says "this tool is
|
||||
# executing." Both can coexist in the thread.
|
||||
for it in event.items:
|
||||
name = it.get("func_name") or it.get("approval_label") or "tool"
|
||||
raw_name = it.get("func_name") or it.get("approval_label") or "tool"
|
||||
display_name = discord.utils.escape_markdown(raw_name)
|
||||
raw_preview = it.get("preview", "")
|
||||
# Sanitize preview: escape backticks to prevent markdown
|
||||
# breakout and strip @-mentions.
|
||||
# Escape backticks to prevent markdown breakout and
|
||||
# strip @-mentions.
|
||||
raw_preview = raw_preview.replace("`", "\\`")
|
||||
raw_preview = discord.utils.escape_mentions(raw_preview)
|
||||
preview = truncate(raw_preview, max_length=120) or None
|
||||
embed = discord.Embed(
|
||||
title=name,
|
||||
title=display_name,
|
||||
description=preview,
|
||||
color=discord.Color.light_grey(),
|
||||
)
|
||||
@@ -515,8 +582,9 @@ class TurnstoneBot:
|
||||
else:
|
||||
msg = await thread.send(embed=embed)
|
||||
call_id = it.get("call_id", "")
|
||||
# Store raw (unescaped) name for matching against ToolResultEvent.name
|
||||
self._tool_info_msgs.setdefault(ws_id, []).append(
|
||||
(call_id, name, preview or "", msg)
|
||||
(call_id, raw_name, preview or "", msg)
|
||||
)
|
||||
|
||||
# If no items consumed the thinking message (empty event), clean up.
|
||||
@@ -548,7 +616,7 @@ class TurnstoneBot:
|
||||
status = "Error" if event.is_error else "Done"
|
||||
status_color = discord.Color.red() if event.is_error else discord.Color.dark_grey()
|
||||
status_embed = discord.Embed(
|
||||
title=f"{event.name} \u2014 {status}",
|
||||
title=f"{discord.utils.escape_markdown(event.name)} \u2014 {status}",
|
||||
description=matched_preview or None,
|
||||
color=status_color,
|
||||
)
|
||||
@@ -558,14 +626,40 @@ class TurnstoneBot:
|
||||
log.debug("discord.tool_info_status_edit_failed", ws_id=ws_id)
|
||||
|
||||
# Send the result as a separate message.
|
||||
desc = format_tool_result(event.output)
|
||||
color = discord.Color.red() if event.is_error else discord.Color.dark_grey()
|
||||
result_embed = discord.Embed(
|
||||
title=event.name,
|
||||
description=desc,
|
||||
color=color,
|
||||
)
|
||||
await thread.send(embed=result_embed)
|
||||
if not event.is_error:
|
||||
from turnstone.channels._formatter import try_build_media_embed
|
||||
|
||||
media_result = None
|
||||
try:
|
||||
media_result = await try_build_media_embed(
|
||||
event.name,
|
||||
event.output,
|
||||
http=self._http_client,
|
||||
)
|
||||
except Exception:
|
||||
log.debug("discord.media_embed_failed", ws_id=ws_id, tool=event.name)
|
||||
if media_result is not None:
|
||||
embed, file = media_result
|
||||
kwargs: dict[str, Any] = {"embed": embed}
|
||||
if file is not None:
|
||||
kwargs["file"] = file
|
||||
await thread.send(**kwargs)
|
||||
else:
|
||||
desc = format_tool_result(event.output)
|
||||
result_embed = discord.Embed(
|
||||
title=event.name,
|
||||
description=desc,
|
||||
color=discord.Color.dark_grey(),
|
||||
)
|
||||
await thread.send(embed=result_embed)
|
||||
else:
|
||||
desc = format_tool_result(event.output)
|
||||
result_embed = discord.Embed(
|
||||
title=event.name,
|
||||
description=desc,
|
||||
color=discord.Color.red(),
|
||||
)
|
||||
await thread.send(embed=result_embed)
|
||||
|
||||
elif isinstance(event, ApproveRequestEvent):
|
||||
# Evaluate admin tool policies before auto-approve.
|
||||
|
||||
+18
-11
@@ -7,6 +7,7 @@ model auto-detection, workstream management, and the main() REPL entry point.
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import readline
|
||||
import sys
|
||||
@@ -165,7 +166,7 @@ class TerminalUI(SessionUI):
|
||||
it for it in items if it.get("needs_approval") and not it.get("error")
|
||||
]
|
||||
except Exception:
|
||||
pass # Best-effort — no policy enforcement on error
|
||||
logging.getLogger(__name__).debug("Policy evaluation unavailable", exc_info=True)
|
||||
|
||||
with self._print_lock:
|
||||
# Print all headers, previews, and heuristic verdicts
|
||||
@@ -176,7 +177,8 @@ class TerminalUI(SessionUI):
|
||||
else:
|
||||
sys.stdout.write(f" {yellow(item['header'])}\n")
|
||||
if item.get("preview"):
|
||||
sys.stdout.write(item["preview"] + "\n")
|
||||
styled = dim(item["preview"]) if not item.get("error") else red(item["preview"])
|
||||
sys.stdout.write(styled + "\n")
|
||||
verdict = item.get("_heuristic_verdict")
|
||||
if verdict:
|
||||
risk = verdict.get("risk_level", "medium")
|
||||
@@ -633,7 +635,7 @@ def _handle_ws_command(
|
||||
# ─── Cluster commands ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _handle_cluster_command(cmd_line: str, console_url: str | None, auth_token: str = "") -> None:
|
||||
def _handle_cluster_command(cmd_line: str, console_url: str | None) -> None:
|
||||
"""Handle /cluster subcommands querying the turnstone-console API."""
|
||||
import httpx
|
||||
|
||||
@@ -642,8 +644,18 @@ def _handle_cluster_command(cmd_line: str, console_url: str | None, auth_token:
|
||||
return
|
||||
|
||||
headers: dict[str, str] = {}
|
||||
if auth_token:
|
||||
headers["Authorization"] = f"Bearer {auth_token}"
|
||||
jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip()
|
||||
if jwt_secret:
|
||||
from turnstone.core.auth import JWT_AUD_CONSOLE, ServiceTokenManager
|
||||
|
||||
_cluster_token_mgr = ServiceTokenManager(
|
||||
user_id="cli",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="cli",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_CONSOLE,
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {_cluster_token_mgr.token}"
|
||||
|
||||
parts = cmd_line.strip().split()
|
||||
sub = parts[1] if len(parts) > 1 else "status"
|
||||
@@ -968,11 +980,6 @@ def main() -> None:
|
||||
default=None,
|
||||
help="Turnstone console URL for /cluster commands (e.g., http://localhost:8090)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auth-token",
|
||||
default=os.environ.get("TURNSTONE_AUTH_TOKEN", ""),
|
||||
help="Bearer token for authenticating to turnstone services (default: $TURNSTONE_AUTH_TOKEN)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mcp-config",
|
||||
default=None,
|
||||
@@ -1247,7 +1254,7 @@ def main() -> None:
|
||||
continue
|
||||
|
||||
if user_input.startswith("/cluster"):
|
||||
_handle_cluster_command(user_input, args.console_url, args.auth_token)
|
||||
_handle_cluster_command(user_input, args.console_url)
|
||||
continue
|
||||
|
||||
active = manager.get_active()
|
||||
|
||||
@@ -59,7 +59,6 @@ class ClusterCollector:
|
||||
storage: StorageBackend,
|
||||
discovery_interval: float = 60.0,
|
||||
http_timeout: float = 30.0,
|
||||
auth_token: str = "",
|
||||
token_manager: ServiceTokenManager | None = None,
|
||||
tls_verify: Any = True,
|
||||
tls_cert: tuple[str, str] | None = None,
|
||||
@@ -74,10 +73,6 @@ class ClusterCollector:
|
||||
self._console_metrics = console_metrics
|
||||
self._tls_verify = tls_verify
|
||||
self._tls_cert = tls_cert
|
||||
# Static auth header — only used when no token_manager is present.
|
||||
self._static_auth: dict[str, str] | None = None
|
||||
if auth_token and token_manager is None:
|
||||
self._static_auth = {"Authorization": f"Bearer {auth_token}"}
|
||||
|
||||
self._lock = threading.Lock()
|
||||
self._nodes: dict[str, NodeSnapshot] = {}
|
||||
@@ -165,7 +160,7 @@ class ClusterCollector:
|
||||
"""Build auth headers for the current SSE connection."""
|
||||
if self._token_manager is not None:
|
||||
return {"Authorization": f"Bearer {self._token_manager.token}"}
|
||||
return dict(self._static_auth) if self._static_auth else {}
|
||||
return {}
|
||||
|
||||
# -- SSE manager ---------------------------------------------------------
|
||||
|
||||
@@ -294,9 +289,13 @@ class ClusterCollector:
|
||||
|
||||
def _discovery_loop(self) -> None:
|
||||
"""Periodically scan the service registry for active nodes."""
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
self._discover_nodes()
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.exception("Node discovery error")
|
||||
time.sleep(self._discovery_interval)
|
||||
|
||||
@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any
|
||||
import structlog
|
||||
|
||||
from turnstone.core.hash_ring import RING_SIZE, RingNode, bucket_of
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.console.collector import ClusterCollector
|
||||
@@ -64,6 +65,7 @@ class Rebalancer:
|
||||
lock_ttl: int = 120,
|
||||
eager_migrate: bool = False,
|
||||
api_token: str = "",
|
||||
token_manager: Any = None,
|
||||
) -> None:
|
||||
self._storage = storage
|
||||
self._router = router
|
||||
@@ -75,6 +77,7 @@ class Rebalancer:
|
||||
self._lock_ttl = lock_ttl
|
||||
self._eager_migrate = eager_migrate
|
||||
self._api_token = api_token
|
||||
self._token_manager = token_manager
|
||||
self._stop_event = threading.Event()
|
||||
self._trigger_event = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
@@ -147,6 +150,8 @@ class Rebalancer:
|
||||
result = self.rebalance_once(trigger=trigger)
|
||||
self._last_result = result
|
||||
self._record_result_metrics(result)
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.exception("rebalancer.error")
|
||||
finally:
|
||||
@@ -530,7 +535,9 @@ class Rebalancer:
|
||||
return 0
|
||||
|
||||
headers: dict[str, str] = {}
|
||||
if self._api_token:
|
||||
if self._token_manager is not None:
|
||||
headers["Authorization"] = f"Bearer {self._token_manager.token}"
|
||||
elif self._api_token:
|
||||
headers["Authorization"] = f"Bearer {self._api_token}"
|
||||
|
||||
migrated = 0
|
||||
|
||||
@@ -90,9 +90,13 @@ class TaskScheduler:
|
||||
|
||||
def _loop(self) -> None:
|
||||
"""Main scheduler loop — tick then sleep."""
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
self._tick()
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.exception("scheduler.tick_error")
|
||||
self._stop_event.wait(self._check_interval)
|
||||
|
||||
+39
-71
@@ -162,19 +162,11 @@ def _proxy_auth_headers(request: Request) -> dict[str, str]:
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
# Fallback: service identity (no user context).
|
||||
# When auth is disabled on the console, auth_result is None, so all proxied
|
||||
# requests use the full-privilege service identity. This is safe only when
|
||||
# the upstream server also has auth disabled.
|
||||
# Fallback: service identity via ServiceTokenManager.
|
||||
mgr = getattr(request.app.state, "proxy_token_mgr", None)
|
||||
if mgr is not None:
|
||||
return dict(mgr.bearer_header)
|
||||
|
||||
# Fall back to static proxy_auth_token (e.g. from --auth-token)
|
||||
static_token = getattr(request.app.state, "proxy_auth_token", "")
|
||||
if static_token:
|
||||
return {"Authorization": f"Bearer {static_token}"}
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
@@ -295,7 +287,7 @@ async def cluster_events_sse(request: Request) -> Response:
|
||||
)
|
||||
yield {"data": json.dumps(event)}
|
||||
except queue.Empty:
|
||||
pass
|
||||
pass # poll timeout, retry
|
||||
if await request.is_disconnected():
|
||||
break
|
||||
finally:
|
||||
@@ -1175,10 +1167,14 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
import asyncio
|
||||
|
||||
async def _console_heartbeat() -> None:
|
||||
from turnstone.core.storage._registry import StorageUnavailableError
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30)
|
||||
try:
|
||||
storage.heartbeat_service("console", "console")
|
||||
except StorageUnavailableError:
|
||||
pass # already logged by storage layer
|
||||
except Exception:
|
||||
log.warning("console.heartbeat_failed", exc_info=True)
|
||||
|
||||
@@ -2129,6 +2125,7 @@ _VALID_PERMISSIONS = frozenset(
|
||||
"admin.roles",
|
||||
"admin.orgs",
|
||||
"admin.policies",
|
||||
"admin.prompt_policies",
|
||||
"admin.skills",
|
||||
"admin.audit",
|
||||
"admin.usage",
|
||||
@@ -5489,8 +5486,14 @@ async def admin_detect_model(request: Request) -> JSONResponse:
|
||||
base_url = row.get("base_url", "")
|
||||
|
||||
# For commercial endpoints an api_key is required
|
||||
_normalized = (base_url if "://" in base_url else f"https://{base_url}") if base_url else ""
|
||||
_hostname = (urllib.parse.urlparse(_normalized).hostname or "") if _normalized else ""
|
||||
if not api_key and (
|
||||
not base_url or "api.openai.com" in base_url or "api.anthropic.com" in base_url
|
||||
not base_url
|
||||
or _hostname == "api.openai.com"
|
||||
or _hostname.endswith(".openai.com")
|
||||
or _hostname == "api.anthropic.com"
|
||||
or _hostname.endswith(".anthropic.com")
|
||||
):
|
||||
return JSONResponse({"error": "api_key is required"}, status_code=400)
|
||||
|
||||
@@ -5924,10 +5927,8 @@ def _seed_config_from_env(config_store: Any, storage: Any) -> None:
|
||||
def create_app(
|
||||
*,
|
||||
collector: ClusterCollector,
|
||||
auth_config: Any,
|
||||
jwt_secret: str = "",
|
||||
auth_storage: Any = None,
|
||||
proxy_auth_token: str = "",
|
||||
proxy_token_mgr: Any = None,
|
||||
cors_origins: list[str] | None = None,
|
||||
tls_manager: Any = None,
|
||||
@@ -6265,10 +6266,8 @@ def create_app(
|
||||
lifespan=_lifespan,
|
||||
)
|
||||
app.state.collector = collector
|
||||
app.state.auth_config = auth_config
|
||||
app.state.jwt_secret = jwt_secret
|
||||
app.state.auth_storage = auth_storage
|
||||
app.state.proxy_auth_token = proxy_auth_token
|
||||
app.state.proxy_token_mgr = proxy_token_mgr
|
||||
app.state.console_url = console_url
|
||||
app.state.tls_manager = tls_manager
|
||||
@@ -6301,7 +6300,7 @@ def create_app(
|
||||
scheduler = TaskScheduler(
|
||||
collector=collector,
|
||||
storage=auth_storage,
|
||||
api_token=proxy_auth_token,
|
||||
api_token="",
|
||||
token_manager=proxy_token_mgr,
|
||||
)
|
||||
app.state.scheduler = scheduler
|
||||
@@ -6351,12 +6350,6 @@ def main() -> None:
|
||||
from turnstone.core.log import add_log_args
|
||||
|
||||
add_log_args(parser)
|
||||
parser.add_argument(
|
||||
"--auth-token",
|
||||
default=os.environ.get("TURNSTONE_AUTH_TOKEN", ""),
|
||||
help="Bearer token for polling turnstone-server nodes (default: $TURNSTONE_AUTH_TOKEN)",
|
||||
)
|
||||
|
||||
from turnstone.core.config import add_config_arg, apply_config
|
||||
|
||||
add_config_arg(parser)
|
||||
@@ -6367,10 +6360,9 @@ def main() -> None:
|
||||
|
||||
configure_logging_from_args(args, "console")
|
||||
|
||||
from turnstone.core.auth import load_auth_config, load_jwt_secret
|
||||
from turnstone.core.auth import load_jwt_secret
|
||||
|
||||
auth_config = load_auth_config()
|
||||
jwt_secret = load_jwt_secret() if auth_config.enabled else ""
|
||||
jwt_secret = load_jwt_secret()
|
||||
|
||||
# Initialize storage early — the collector needs it for service discovery.
|
||||
auth_storage = None
|
||||
@@ -6399,38 +6391,23 @@ def main() -> None:
|
||||
)
|
||||
raise SystemExit(1)
|
||||
|
||||
# If no explicit auth token is provided, use a ServiceTokenManager
|
||||
# so collector JWTs auto-rotate. A shared JWT secret is required for
|
||||
# multi-service deployments — ephemeral secrets differ per process.
|
||||
collector_token = args.auth_token
|
||||
collector_token_mgr = None
|
||||
if not collector_token:
|
||||
_jwt_secret = os.environ.get("TURNSTONE_JWT_SECRET", "")
|
||||
if not _jwt_secret:
|
||||
log.error(
|
||||
"TURNSTONE_JWT_SECRET is not set and no --auth-token provided. "
|
||||
"The console cannot authenticate to server nodes. Set TURNSTONE_JWT_SECRET "
|
||||
"to a shared secret (at least 32 characters) or pass --auth-token."
|
||||
)
|
||||
raise SystemExit(1)
|
||||
from turnstone.core.auth import JWT_AUD_SERVER, ServiceTokenManager
|
||||
from turnstone.core.auth import JWT_AUD_SERVER, ServiceTokenManager
|
||||
|
||||
collector_token_mgr = ServiceTokenManager(
|
||||
user_id="console-collector",
|
||||
scopes=frozenset({"read"}),
|
||||
source="console",
|
||||
secret=_jwt_secret,
|
||||
audience=JWT_AUD_SERVER,
|
||||
expiry_hours=1,
|
||||
)
|
||||
log.info("console.collector_token_manager_created")
|
||||
collector_token_mgr = ServiceTokenManager(
|
||||
user_id="console-collector",
|
||||
scopes=frozenset({"read"}),
|
||||
source="console",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_SERVER,
|
||||
expiry_hours=1,
|
||||
)
|
||||
log.info("console.collector_token_manager_created")
|
||||
|
||||
router = ConsoleRouter(storage=auth_storage)
|
||||
console_metrics = ConsoleMetrics()
|
||||
|
||||
collector = ClusterCollector(
|
||||
storage=auth_storage,
|
||||
auth_token=collector_token if collector_token_mgr is None else "",
|
||||
token_manager=collector_token_mgr,
|
||||
router=router,
|
||||
console_metrics=console_metrics,
|
||||
@@ -6439,22 +6416,15 @@ def main() -> None:
|
||||
|
||||
_load_static()
|
||||
|
||||
# If no explicit auth token is provided, use a ServiceTokenManager
|
||||
# so proxy JWTs auto-rotate.
|
||||
proxy_token = args.auth_token
|
||||
proxy_token_mgr = None
|
||||
if not proxy_token and jwt_secret:
|
||||
from turnstone.core.auth import JWT_AUD_SERVER, ServiceTokenManager
|
||||
|
||||
proxy_token_mgr = ServiceTokenManager(
|
||||
user_id="console-proxy",
|
||||
scopes=frozenset({"read", "write", "approve"}),
|
||||
source="console",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_SERVER,
|
||||
expiry_hours=1,
|
||||
)
|
||||
log.info("console.proxy_token_manager_created")
|
||||
proxy_token_mgr = ServiceTokenManager(
|
||||
user_id="console-proxy",
|
||||
scopes=frozenset({"read", "write", "approve", "service"}),
|
||||
source="console",
|
||||
secret=jwt_secret,
|
||||
audience=JWT_AUD_SERVER,
|
||||
expiry_hours=1,
|
||||
)
|
||||
log.info("console.proxy_token_manager_created")
|
||||
|
||||
from turnstone.core.web_helpers import parse_cors_origins
|
||||
|
||||
@@ -6541,7 +6511,8 @@ def main() -> None:
|
||||
threshold=_rcs.get("rebalancer.threshold", 0.10),
|
||||
vnodes_per_unit=_rcs.get("ring.vnodes_per_unit", 150),
|
||||
eager_migrate=_rcs.get("rebalancer.eager_migrate", False),
|
||||
api_token=proxy_token if proxy_token_mgr is None else "",
|
||||
api_token="",
|
||||
token_manager=proxy_token_mgr,
|
||||
)
|
||||
log.info("rebalancer.configured")
|
||||
except Exception:
|
||||
@@ -6549,10 +6520,8 @@ def main() -> None:
|
||||
|
||||
app = create_app(
|
||||
collector=collector,
|
||||
auth_config=auth_config,
|
||||
jwt_secret=jwt_secret,
|
||||
auth_storage=auth_storage,
|
||||
proxy_auth_token=proxy_token if proxy_token_mgr is None else "",
|
||||
proxy_token_mgr=proxy_token_mgr,
|
||||
cors_origins=cors_origins,
|
||||
tls_manager=tls_mgr,
|
||||
@@ -6563,8 +6532,7 @@ def main() -> None:
|
||||
)
|
||||
|
||||
log.info("Console starting on %s", console_url)
|
||||
if auth_config.enabled:
|
||||
log.info("Auth: enabled (%d config token(s))", len(auth_config.tokens))
|
||||
log.info("Auth: enabled (JWT)")
|
||||
print("Press Ctrl+C to stop.")
|
||||
|
||||
import uvicorn
|
||||
|
||||
@@ -94,6 +94,29 @@ function showAdmin() {
|
||||
|
||||
// Mobile: ensure sidebar starts hidden + inert; desktop: ensure it's accessible
|
||||
var sidebar = document.getElementById("admin-sidebar");
|
||||
|
||||
// Inject close header for mobile drawer (once)
|
||||
if (!document.getElementById("admin-sidebar-close")) {
|
||||
var closeHeader = document.createElement("div");
|
||||
closeHeader.id = "admin-sidebar-close";
|
||||
closeHeader.className = "admin-sidebar-close";
|
||||
var label = document.createElement("span");
|
||||
label.textContent = "Navigation";
|
||||
var closeBtn = document.createElement("button");
|
||||
closeBtn.setAttribute("aria-label", "Close navigation");
|
||||
closeBtn.textContent = "\u00d7";
|
||||
closeBtn.addEventListener("click", function () {
|
||||
if (_mobileSidebarOpen) {
|
||||
_toggleMobileSidebar();
|
||||
var mt = document.getElementById("admin-mobile-toggle");
|
||||
if (mt) mt.focus();
|
||||
}
|
||||
});
|
||||
closeHeader.appendChild(label);
|
||||
closeHeader.appendChild(closeBtn);
|
||||
sidebar.insertBefore(closeHeader, sidebar.firstChild);
|
||||
}
|
||||
|
||||
if (window.innerWidth <= 700) {
|
||||
_mobileSidebarOpen = false;
|
||||
sidebar.classList.add("collapsed");
|
||||
@@ -147,15 +170,20 @@ function _injectMobileToggle(tab) {
|
||||
toggle.id = "admin-mobile-toggle";
|
||||
toggle.className = "admin-mobile-toggle";
|
||||
toggle.setAttribute("aria-label", "Open navigation");
|
||||
toggle.setAttribute("aria-expanded", "false");
|
||||
toggle.onclick = function () {
|
||||
_mobileSidebarOpen = false;
|
||||
_toggleMobileSidebar();
|
||||
};
|
||||
}
|
||||
var panel = document.getElementById("admin-" + tab);
|
||||
if (panel) {
|
||||
var toolbar = panel.querySelector(".admin-toolbar");
|
||||
if (toolbar) toolbar.insertBefore(toggle, toolbar.firstChild);
|
||||
if (!panel) return;
|
||||
var toolbar = panel.querySelector(".admin-toolbar");
|
||||
if (toolbar) {
|
||||
if (!toolbar.contains(toggle))
|
||||
toolbar.insertBefore(toggle, toolbar.firstChild);
|
||||
} else {
|
||||
// Panel has no toolbar — prepend toggle directly so it remains accessible
|
||||
if (!panel.contains(toggle)) panel.insertBefore(toggle, panel.firstChild);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -169,6 +197,20 @@ function _toggleMobileSidebar() {
|
||||
else sidebar.setAttribute("inert", "");
|
||||
var backdrop = document.getElementById("admin-sidebar-backdrop");
|
||||
if (backdrop) backdrop.classList.toggle("visible", _mobileSidebarOpen);
|
||||
// Update hamburger aria-label to reflect current state
|
||||
var mt = document.getElementById("admin-mobile-toggle");
|
||||
if (mt) {
|
||||
mt.setAttribute(
|
||||
"aria-label",
|
||||
_mobileSidebarOpen ? "Close navigation" : "Open navigation",
|
||||
);
|
||||
mt.setAttribute("aria-expanded", _mobileSidebarOpen ? "true" : "false");
|
||||
}
|
||||
// Move focus into drawer on open; callers handle focus-return on close
|
||||
if (_mobileSidebarOpen) {
|
||||
var closeBtn = sidebar.querySelector(".admin-sidebar-close button");
|
||||
if (closeBtn) closeBtn.focus();
|
||||
}
|
||||
}
|
||||
|
||||
function switchAdminTab(tab) {
|
||||
@@ -238,6 +280,15 @@ function switchAdminTab(tab) {
|
||||
// On mobile, auto-close sidebar after tab selection
|
||||
if (window.innerWidth <= 700 && _mobileSidebarOpen) {
|
||||
_toggleMobileSidebar();
|
||||
// Move focus to the newly active panel instead of leaving it in the inert sidebar
|
||||
var panel = document.getElementById("admin-" + tab);
|
||||
var focusTarget =
|
||||
panel &&
|
||||
panel.querySelector("h2, .section-header, button:not([disabled])");
|
||||
if (focusTarget) {
|
||||
focusTarget.setAttribute("tabindex", "-1");
|
||||
focusTarget.focus();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2002,18 +2053,20 @@ document.addEventListener("keydown", function (e) {
|
||||
if (!sidebar) return;
|
||||
var isMobile = window.innerWidth <= 700;
|
||||
var backdrop = document.getElementById("admin-sidebar-backdrop");
|
||||
if (isMobile && !_mobileSidebarOpen) {
|
||||
if (!isMobile) {
|
||||
// Crossed into desktop: close drawer cleanly if it was open
|
||||
if (_mobileSidebarOpen) _toggleMobileSidebar();
|
||||
sidebar.removeAttribute("aria-hidden");
|
||||
sidebar.removeAttribute("inert");
|
||||
sidebar.classList.remove("collapsed", "open");
|
||||
if (backdrop) backdrop.classList.remove("visible");
|
||||
} else if (!_mobileSidebarOpen) {
|
||||
// Mobile with drawer closed: ensure collapsed state
|
||||
sidebar.setAttribute("aria-hidden", "true");
|
||||
sidebar.setAttribute("inert", "");
|
||||
sidebar.classList.add("collapsed");
|
||||
sidebar.classList.remove("open");
|
||||
if (backdrop) backdrop.classList.remove("visible");
|
||||
} else if (!isMobile) {
|
||||
sidebar.removeAttribute("aria-hidden");
|
||||
sidebar.removeAttribute("inert");
|
||||
sidebar.classList.remove("collapsed", "open");
|
||||
if (backdrop) backdrop.classList.remove("visible");
|
||||
_mobileSidebarOpen = false;
|
||||
}
|
||||
}, 150);
|
||||
});
|
||||
|
||||
@@ -720,7 +720,9 @@ function buildNodeRow(node) {
|
||||
function toggleGroup(prefix) {
|
||||
expandedGroups[prefix] = !expandedGroups[prefix];
|
||||
var body = document.querySelector(
|
||||
'.node-group-body[data-prefix="' + prefix.replace(/"/g, '\\"') + '"]',
|
||||
'.node-group-body[data-prefix="' +
|
||||
prefix.replace(/\\/g, "\\\\").replace(/"/g, '\\"') +
|
||||
'"]',
|
||||
);
|
||||
if (!body) return;
|
||||
var isExpanded = expandedGroups[prefix];
|
||||
|
||||
@@ -2185,8 +2185,7 @@ function searchSkillDiscover() {
|
||||
var searchBtn = document.getElementById("skill-discover-search-btn");
|
||||
if (searchBtn) searchBtn.disabled = true;
|
||||
|
||||
var url = "/v1/api/admin/skills/discover?limit=20";
|
||||
if (q) url += "&q=" + encodeURIComponent(q);
|
||||
var url = "/v1/api/admin/skills/discover?limit=20&q=" + encodeURIComponent(q);
|
||||
|
||||
authFetch(url)
|
||||
.then(function (r) {
|
||||
|
||||
@@ -768,7 +768,8 @@
|
||||
color: var(--fg-dim);
|
||||
padding: 12px 16px 4px;
|
||||
}
|
||||
.admin-sidebar-group:first-child .admin-sidebar-group-label {
|
||||
.admin-sidebar-group:first-child .admin-sidebar-group-label,
|
||||
.admin-sidebar-close + .admin-sidebar-group .admin-sidebar-group-label {
|
||||
padding-top: 4px;
|
||||
}
|
||||
|
||||
@@ -818,13 +819,16 @@
|
||||
z-index: 499;
|
||||
opacity: 0;
|
||||
pointer-events: none;
|
||||
transition: opacity 0.25s ease;
|
||||
transition: opacity 0.25s cubic-bezier(0.4, 0, 0.2, 1);
|
||||
}
|
||||
.admin-sidebar-backdrop.visible {
|
||||
opacity: 1;
|
||||
pointer-events: auto;
|
||||
}
|
||||
|
||||
/* Close header — hidden on desktop, shown via mobile media query */
|
||||
.admin-sidebar-close { display: none; }
|
||||
|
||||
/* Mobile menu toggle — visible only on mobile, lives in toolbars */
|
||||
.admin-mobile-toggle {
|
||||
display: none;
|
||||
@@ -832,8 +836,8 @@
|
||||
border: 1px solid var(--border);
|
||||
border-radius: var(--radius-sm);
|
||||
color: var(--fg-dim);
|
||||
width: 28px;
|
||||
height: 28px;
|
||||
min-width: 44px;
|
||||
min-height: 44px;
|
||||
cursor: pointer;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
@@ -850,6 +854,10 @@
|
||||
box-shadow: 0 4px 0 currentColor, 0 8px 0 currentColor;
|
||||
}
|
||||
.admin-mobile-toggle:hover { color: var(--fg); }
|
||||
.admin-mobile-toggle:focus-visible {
|
||||
outline: 2px solid var(--accent);
|
||||
outline-offset: 2px;
|
||||
}
|
||||
@media (max-width: 700px) {
|
||||
.admin-mobile-toggle { display: flex; }
|
||||
}
|
||||
@@ -912,7 +920,7 @@
|
||||
.admin-row {
|
||||
display: grid;
|
||||
padding: 8px 12px;
|
||||
align-items: center;
|
||||
align-items: start;
|
||||
border-radius: var(--radius-sm);
|
||||
transition: background 0.1s;
|
||||
}
|
||||
@@ -1464,18 +1472,62 @@ h3.skill-spec-heading { font-size: inherit; margin-block: 0; }
|
||||
right: 0;
|
||||
bottom: 0;
|
||||
left: auto;
|
||||
width: 220px;
|
||||
width: 260px;
|
||||
max-width: 80vw;
|
||||
z-index: 500;
|
||||
background: var(--bg-surface);
|
||||
border-left: 1px solid var(--border-strong);
|
||||
border-right: none;
|
||||
box-shadow: -4px 0 24px rgba(0, 0, 0, 0.35);
|
||||
transform: translateX(100%);
|
||||
transition: transform 0.25s ease;
|
||||
padding-top: 48px;
|
||||
transition: transform 0.25s cubic-bezier(0.4, 0, 0.2, 1);
|
||||
padding-top: 0;
|
||||
overflow-y: auto;
|
||||
-webkit-overflow-scrolling: touch;
|
||||
}
|
||||
.admin-sidebar.open { transform: translateX(0); }
|
||||
.admin-sidebar.collapsed { transform: translateX(100%); width: 220px; }
|
||||
.admin-sidebar.collapsed { transform: translateX(100%); }
|
||||
.admin-content { padding-right: 0; }
|
||||
|
||||
/* Close button at top of mobile drawer */
|
||||
.admin-sidebar-close {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
padding: 12px 16px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
font-family: var(--font-display);
|
||||
font-size: 11px;
|
||||
font-weight: 600;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.08em;
|
||||
color: var(--fg-dim);
|
||||
}
|
||||
.admin-sidebar-close button {
|
||||
background: none;
|
||||
border: none;
|
||||
color: var(--fg-dim);
|
||||
font-size: 20px;
|
||||
line-height: 1;
|
||||
cursor: pointer;
|
||||
padding: 10px;
|
||||
min-width: 44px;
|
||||
min-height: 44px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
border-radius: var(--radius-sm);
|
||||
}
|
||||
.admin-sidebar-close button:hover { color: var(--fg); }
|
||||
.admin-sidebar-close button:focus-visible {
|
||||
outline: 2px solid var(--accent);
|
||||
outline-offset: 2px;
|
||||
}
|
||||
|
||||
/* Flip active indicator to left border on mobile (drawer is on right edge) */
|
||||
.admin-nav { border-right: none; border-left: 2px solid transparent; }
|
||||
.admin-nav:hover { border-right-color: transparent; border-left-color: var(--border-strong); }
|
||||
.admin-nav.active { border-right-color: transparent; border-left-color: var(--accent); }
|
||||
}
|
||||
|
||||
/* ==========================================================================
|
||||
@@ -1868,7 +1920,7 @@ h3.skill-spec-heading { font-size: inherit; margin-block: 0; }
|
||||
grid-template-columns: 200px 1fr auto;
|
||||
gap: 8px 16px;
|
||||
padding: 8px 12px;
|
||||
align-items: center;
|
||||
align-items: start;
|
||||
border-bottom: 1px solid var(--border);
|
||||
}
|
||||
.settings-row:hover { background: var(--row-alt, rgba(255,255,255,0.015)); }
|
||||
@@ -1913,6 +1965,7 @@ h3.skill-spec-heading { font-size: inherit; margin-block: 0; }
|
||||
}
|
||||
|
||||
/* Bool toggle */
|
||||
.settings-input .settings-toggle { margin-top: 2px; }
|
||||
.settings-toggle {
|
||||
position: relative;
|
||||
display: inline-block;
|
||||
|
||||
+32
-131
@@ -1,15 +1,12 @@
|
||||
"""Bearer token authentication and authorization for turnstone HTTP servers.
|
||||
|
||||
Supports three token types:
|
||||
Supports two token types:
|
||||
|
||||
1. **Config-file tokens** — static tokens in ``config.toml`` or the
|
||||
``TURNSTONE_AUTH_TOKEN`` env var. Validated in-memory via
|
||||
``hmac.compare_digest``. Map to scopes via their role.
|
||||
2. **API tokens** — database-backed, prefixed ``ts_``, stored as SHA-256
|
||||
1. **API tokens** — database-backed, prefixed ``ts_``, stored as SHA-256
|
||||
hashes. Exchanged for JWTs via ``/api/auth/login``.
|
||||
3. **JWTs** — short-lived session tokens issued after API token validation.
|
||||
Validated locally via shared HMAC-SHA256 secret. Contain user_id and
|
||||
scopes in claims.
|
||||
2. **JWTs** — short-lived session tokens issued after login or by
|
||||
:class:`ServiceTokenManager`. Validated locally via shared HMAC-SHA256
|
||||
secret. Contain user_id and scopes in claims.
|
||||
|
||||
Public paths (``/``, ``/static/*``, ``/shared/*``, ``/health``, ``/metrics``,
|
||||
``/openapi.json``, ``/docs``, ``/api/auth/login``, ``/api/auth/logout``) are
|
||||
@@ -19,7 +16,6 @@ always accessible without authentication.
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
@@ -28,7 +24,7 @@ import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -56,7 +52,7 @@ JWT_AUD_CONSOLE = "turnstone-console"
|
||||
JWT_AUD_CHANNEL = "turnstone-channel"
|
||||
_MIN_SECRET_LENGTH = 32 # 256 bits minimum for HMAC-SHA256
|
||||
|
||||
VALID_SCOPES: frozenset[str] = frozenset({"read", "write", "approve"})
|
||||
VALID_SCOPES: frozenset[str] = frozenset({"read", "write", "approve", "service"})
|
||||
|
||||
_USERNAME_RE = re.compile(r"^[a-zA-Z0-9._-]+$")
|
||||
USERNAME_MAX_LEN = 64
|
||||
@@ -72,16 +68,12 @@ def is_valid_username(username: str) -> bool:
|
||||
|
||||
|
||||
# Hierarchical: each scope implies all lower scopes.
|
||||
# "service" is a superset that grants full access + bypasses RBAC permission checks.
|
||||
SCOPE_HIERARCHY: dict[str, frozenset[str]] = {
|
||||
"read": frozenset({"read"}),
|
||||
"write": frozenset({"read", "write"}),
|
||||
"approve": frozenset({"read", "write", "approve"}),
|
||||
}
|
||||
|
||||
# Map old role names to scope sets.
|
||||
_ROLE_TO_SCOPES: dict[str, frozenset[str]] = {
|
||||
"read": frozenset({"read"}),
|
||||
"full": frozenset({"read", "write", "approve"}),
|
||||
"service": frozenset({"read", "write", "approve", "service"}),
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -106,7 +98,7 @@ def _permissions_to_scopes(permissions: set[str]) -> frozenset[str]:
|
||||
scopes.add("read")
|
||||
return frozenset(scopes)
|
||||
for perm in permissions:
|
||||
if perm in VALID_SCOPES:
|
||||
if perm in VALID_SCOPES and perm != "service":
|
||||
scopes.update(SCOPE_HIERARCHY.get(perm, {perm}))
|
||||
# Any admin.* permission requires access to admin endpoints → approve scope
|
||||
if any(p.startswith("admin.") for p in permissions):
|
||||
@@ -120,15 +112,14 @@ def require_permission(request: Request, permission: str) -> JSONResponse | None
|
||||
"""Return a 403 JSONResponse if the user lacks *permission*, else None.
|
||||
|
||||
Call from admin handlers after the middleware scope check passes.
|
||||
Config-file tokens (no user_id) are treated as full-access.
|
||||
Service tokens (scope ``service``) bypass permission checks.
|
||||
"""
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
auth_result: AuthResult | None = getattr(getattr(request, "state", None), "auth_result", None)
|
||||
if auth_result is None:
|
||||
return JSONResponse({"error": "Unauthorized"}, status_code=401)
|
||||
# Config-file tokens (no user_id) are treated as full-access
|
||||
if not auth_result.user_id:
|
||||
if auth_result.has_scope("service"):
|
||||
return None
|
||||
if auth_result.has_permission(permission):
|
||||
return None
|
||||
@@ -200,9 +191,9 @@ def _strip_version_prefix(path: str) -> str:
|
||||
class AuthResult:
|
||||
"""Result of successful authentication."""
|
||||
|
||||
user_id: str # empty string for config-file tokens
|
||||
user_id: str
|
||||
scopes: frozenset[str]
|
||||
token_source: str # "config", "jwt", "database"
|
||||
token_source: str # "jwt", "database", "password", or service origin (e.g. "console", "cli")
|
||||
permissions: frozenset[str] = frozenset()
|
||||
|
||||
def has_scope(self, scope: str) -> bool:
|
||||
@@ -214,28 +205,6 @@ class AuthResult:
|
||||
return permission in self.permissions
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AuthConfig (unchanged from before — static config-file tokens)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuthConfig:
|
||||
"""Auth configuration loaded once at startup (not modified after creation)."""
|
||||
|
||||
enabled: bool = False
|
||||
tokens: dict[str, str] = field(default_factory=dict) # token_value → role
|
||||
|
||||
def check(self, token: str | None) -> str | None:
|
||||
"""Return the role for a valid config token, or *None*."""
|
||||
if not token:
|
||||
return None
|
||||
for known_token, role in self.tokens.items():
|
||||
if hmac.compare_digest(token, known_token):
|
||||
return role
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token generation and hashing
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -303,7 +272,11 @@ def parse_scopes(scopes_str: str) -> frozenset[str]:
|
||||
|
||||
|
||||
def load_jwt_secret() -> str:
|
||||
"""Load JWT signing secret from env or config, or auto-generate."""
|
||||
"""Load JWT signing secret from env or config.
|
||||
|
||||
Raises :class:`SystemExit` if no secret is configured. A JWT secret
|
||||
is required for auth, inter-service communication, and session tokens.
|
||||
"""
|
||||
secret = os.environ.get("TURNSTONE_JWT_SECRET", "").strip()
|
||||
if not secret:
|
||||
from turnstone.core.config import load_config
|
||||
@@ -312,18 +285,19 @@ def load_jwt_secret() -> str:
|
||||
secret = str(auth_cfg.get("jwt_secret", "")).strip()
|
||||
|
||||
if not secret:
|
||||
# Auto-generate an ephemeral secret
|
||||
secret = secrets.token_hex(32)
|
||||
log.warning(
|
||||
"No JWT secret configured — using ephemeral secret (tokens will not survive restart)"
|
||||
log.error(
|
||||
"TURNSTONE_JWT_SECRET is required but not set. "
|
||||
'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"'
|
||||
)
|
||||
return secret
|
||||
raise SystemExit(1)
|
||||
|
||||
if len(secret) < _MIN_SECRET_LENGTH:
|
||||
log.warning(
|
||||
"JWT secret is shorter than %d characters — consider using a stronger secret",
|
||||
log.error(
|
||||
"JWT secret must be at least %d characters. "
|
||||
'Generate one with: python -c "import secrets; print(secrets.token_hex(32))"',
|
||||
_MIN_SECRET_LENGTH,
|
||||
)
|
||||
raise SystemExit(1)
|
||||
return secret
|
||||
|
||||
|
||||
@@ -397,62 +371,6 @@ def validate_jwt(token: str, secret: str, audience: str = "") -> AuthResult | No
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_auth_config() -> AuthConfig:
|
||||
"""Build :class:`AuthConfig` from ``config.toml`` ``[auth]`` + env vars.
|
||||
|
||||
Auth is **enabled by default**. Set ``[auth] enabled = false`` or
|
||||
``TURNSTONE_AUTH_ENABLED=0`` to disable.
|
||||
|
||||
Config format::
|
||||
|
||||
[auth]
|
||||
enabled = false # opt out
|
||||
|
||||
[[auth.tokens]]
|
||||
value = "tok_abc123"
|
||||
role = "full"
|
||||
|
||||
Environment variables:
|
||||
|
||||
- ``TURNSTONE_AUTH_ENABLED=0`` — disables auth
|
||||
- ``TURNSTONE_AUTH_ENABLED=1`` — enables auth (default)
|
||||
- ``TURNSTONE_AUTH_TOKEN=<token>`` — registers a single full-access token
|
||||
"""
|
||||
from turnstone.core.config import load_config
|
||||
|
||||
auth_cfg = load_config("auth")
|
||||
enabled = bool(auth_cfg.get("enabled", True))
|
||||
tokens: dict[str, str] = {}
|
||||
|
||||
# Tokens from config file (TOML array-of-tables)
|
||||
for entry in auth_cfg.get("tokens", []):
|
||||
value = entry.get("value", "") if isinstance(entry, dict) else ""
|
||||
role = entry.get("role", "read") if isinstance(entry, dict) else ""
|
||||
if value and role in ("read", "full"):
|
||||
tokens[value] = role
|
||||
|
||||
# Environment variable overrides
|
||||
env_enabled = os.environ.get("TURNSTONE_AUTH_ENABLED", "").strip().lower()
|
||||
if env_enabled in ("1", "true", "yes"):
|
||||
enabled = True
|
||||
elif env_enabled in ("0", "false", "no"):
|
||||
enabled = False
|
||||
|
||||
env_token = os.environ.get("TURNSTONE_AUTH_TOKEN", "").strip()
|
||||
if env_token:
|
||||
tokens[env_token] = "full"
|
||||
|
||||
if enabled and not tokens:
|
||||
log.info("Auth enabled (no config tokens — use /api/auth/setup or turnstone-admin)")
|
||||
|
||||
return AuthConfig(enabled=enabled, tokens=tokens)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -529,7 +447,6 @@ def _extract_proxied_path(normalized: str) -> str | None:
|
||||
|
||||
|
||||
def check_request(
|
||||
auth_config: AuthConfig,
|
||||
method: str,
|
||||
path: str,
|
||||
auth_header: str | None,
|
||||
@@ -539,20 +456,16 @@ def check_request(
|
||||
jwt_audience: str = "",
|
||||
storage: Any = None,
|
||||
) -> tuple[bool, int, str, AuthResult | None]:
|
||||
"""Validate a request against the auth config.
|
||||
"""Validate a request.
|
||||
|
||||
Checks ``Authorization: Bearer <token>`` first, then falls back to the
|
||||
``turnstone_auth`` cookie. Token types are auto-detected:
|
||||
|
||||
- Contains ``.`` → JWT (validated with *jwt_secret*)
|
||||
- Starts with ``ts_`` → API token (looked up in *storage* by hash)
|
||||
- Otherwise → config-file token (hmac check)
|
||||
|
||||
Returns ``(allowed, status_code, message, auth_result)``.
|
||||
"""
|
||||
if not auth_config.enabled:
|
||||
return True, 200, "", None
|
||||
|
||||
if is_public_path(path):
|
||||
return True, 200, "", None
|
||||
|
||||
@@ -566,7 +479,7 @@ def check_request(
|
||||
|
||||
# Authenticate
|
||||
result = _authenticate_token(
|
||||
raw_token, auth_config, jwt_secret=jwt_secret, jwt_audience=jwt_audience, storage=storage
|
||||
raw_token, jwt_secret=jwt_secret, jwt_audience=jwt_audience, storage=storage
|
||||
)
|
||||
if result is None:
|
||||
return False, 401, "Unauthorized: missing or invalid token", None
|
||||
@@ -581,7 +494,6 @@ def check_request(
|
||||
|
||||
def _authenticate_token(
|
||||
token: str,
|
||||
auth_config: AuthConfig,
|
||||
*,
|
||||
jwt_secret: str = "",
|
||||
jwt_audience: str = "",
|
||||
@@ -601,12 +513,6 @@ def _authenticate_token(
|
||||
if token.startswith(TOKEN_PREFIX) and storage is not None:
|
||||
return _authenticate_api_token(token, storage)
|
||||
|
||||
# 3. Config-file token (hmac comparison)
|
||||
role = auth_config.check(token)
|
||||
if role is not None:
|
||||
scopes = _ROLE_TO_SCOPES.get(role, frozenset({"read"}))
|
||||
return AuthResult(user_id="", scopes=scopes, token_source="config")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@@ -852,7 +758,6 @@ class AuthMiddleware:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
auth_config = request.app.state.auth_config
|
||||
jwt_secret = getattr(request.app.state, "jwt_secret", "")
|
||||
storage = getattr(request.app.state, "auth_storage", None)
|
||||
method = request.method
|
||||
@@ -860,7 +765,6 @@ class AuthMiddleware:
|
||||
auth_header = request.headers.get("Authorization")
|
||||
cookie_header = request.headers.get("Cookie")
|
||||
allowed, status, msg, auth_result = check_request(
|
||||
auth_config,
|
||||
method,
|
||||
path,
|
||||
auth_header,
|
||||
@@ -904,7 +808,6 @@ async def handle_auth_login(request: Request, audience: str) -> Response:
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
return JSONResponse({"error": "Invalid JSON body"}, status_code=400)
|
||||
|
||||
auth_config = request.app.state.auth_config
|
||||
jwt_secret = getattr(request.app.state, "jwt_secret", "")
|
||||
storage = getattr(request.app.state, "auth_storage", None)
|
||||
login_limiter: LoginRateLimiter | None = getattr(request.app.state, "login_limiter", None)
|
||||
@@ -955,7 +858,6 @@ async def handle_auth_login(request: Request, audience: str) -> Response:
|
||||
elif body.get("token"):
|
||||
result = _authenticate_token(
|
||||
body["token"],
|
||||
auth_config,
|
||||
jwt_secret=jwt_secret,
|
||||
jwt_audience=audience,
|
||||
storage=storage,
|
||||
@@ -1011,7 +913,6 @@ async def handle_auth_status(request: Request) -> Response:
|
||||
"""Shared ``GET /api/auth/status`` handler — login UI state detection."""
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
auth_config = request.app.state.auth_config
|
||||
storage = getattr(request.app.state, "auth_storage", None)
|
||||
|
||||
has_users = False
|
||||
@@ -1027,9 +928,9 @@ async def handle_auth_status(request: Request) -> Response:
|
||||
oidc_enabled = bool(oidc_config and oidc_config.enabled)
|
||||
|
||||
resp: dict[str, Any] = {
|
||||
"auth_enabled": auth_config.enabled,
|
||||
"auth_enabled": True,
|
||||
"has_users": has_users,
|
||||
"setup_required": auth_config.enabled and not has_users,
|
||||
"setup_required": not has_users,
|
||||
}
|
||||
if oidc_enabled and oidc_config is not None:
|
||||
resp["oidc_enabled"] = True
|
||||
@@ -1291,7 +1192,7 @@ async def handle_oidc_callback(request: Request, audience: str) -> Response:
|
||||
jwks_data = await fetch_jwks(oidc_config.jwks_uri)
|
||||
request.app.state.jwks_data = jwks_data
|
||||
except OIDCError:
|
||||
pass
|
||||
log.warning("JWKS fetch failed from %s", oidc_config.jwks_uri, exc_info=True)
|
||||
if jwks_data is None:
|
||||
return RedirectResponse("/?oidc_error=OIDC+temporarily+unavailable", status_code=302)
|
||||
|
||||
|
||||
@@ -66,7 +66,6 @@ _EXPLICIT_SCRUB: frozenset[str] = frozenset(
|
||||
"ANTHROPIC_API_KEY",
|
||||
"TAVILY_API_KEY",
|
||||
"TURNSTONE_JWT_SECRET",
|
||||
"TURNSTONE_AUTH_TOKEN",
|
||||
"TURNSTONE_DISCORD_TOKEN",
|
||||
"TURNSTONE_GITHUB_TOKEN",
|
||||
"TURNSTONE_OIDC_CLIENT_SECRET",
|
||||
|
||||
@@ -1382,7 +1382,7 @@ class IntentJudge:
|
||||
confidence = float(data.get("confidence", 0.5))
|
||||
confidence = max(0.0, min(1.0, confidence))
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
pass # keeps default 0.5
|
||||
|
||||
evidence = data.get("evidence", [])
|
||||
if isinstance(evidence, str):
|
||||
@@ -1415,7 +1415,7 @@ class IntentJudge:
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
pass # falls through to strategy 2
|
||||
|
||||
# Strategy 2: Markdown code block
|
||||
md_match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
|
||||
@@ -1425,7 +1425,7 @@ class IntentJudge:
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
pass # falls through to strategy 3
|
||||
|
||||
# Strategy 3: Find first { and matching }
|
||||
start = text.find("{")
|
||||
@@ -1442,7 +1442,7 @@ class IntentJudge:
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
pass # falls through to regex extraction
|
||||
break
|
||||
|
||||
# Strategy 4: Regex field extraction (last resort)
|
||||
|
||||
+331
-11
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
import mcp.types as mcp_types
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp import ClientSession, McpError, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
|
||||
@@ -67,7 +67,7 @@ def _mcp_to_openai(server_name: str, tool: Any) -> dict[str, Any]:
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": f"mcp__{server_name}__{tool.name}",
|
||||
"description": f"[MCP: {server_name}] {description}",
|
||||
"description": description,
|
||||
"parameters": input_schema,
|
||||
},
|
||||
}
|
||||
@@ -151,6 +151,22 @@ class MCPClientManager:
|
||||
self._refresh_interval = refresh_interval
|
||||
self._refresh_task: asyncio.Task[None] | None = None
|
||||
|
||||
# Circuit breaker (per-server) — prevents repeated calls to broken servers
|
||||
self._consecutive_failures: dict[str, int] = {}
|
||||
self._circuit_open_until: dict[str, float] = {} # monotonic timestamp
|
||||
self._circuit_trip_count: dict[str, int] = {} # backoff exponent
|
||||
|
||||
# Safe transport stream refs (pre-close before stack teardown to avoid
|
||||
# the anyio cancel-scope CPU busy-loop — MCP SDK #2147)
|
||||
self._server_streams: dict[str, tuple[Any, Any]] = {}
|
||||
|
||||
# Notification debounce (per-server)
|
||||
self._last_notification_refresh: dict[str, float] = {}
|
||||
|
||||
# Periodic refresh backoff (per-server)
|
||||
self._refresh_failures: dict[str, int] = {}
|
||||
self._refresh_backoff_until: dict[str, float] = {} # monotonic timestamp
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
@@ -181,6 +197,7 @@ class MCPClientManager:
|
||||
except Exception as exc:
|
||||
log.warning("Failed to connect MCP server '%s'", name, exc_info=True)
|
||||
self._set_error(name, f"{type(exc).__name__}: {exc}")
|
||||
self._cb_record_failure(name)
|
||||
|
||||
self._connected.set()
|
||||
|
||||
@@ -203,6 +220,94 @@ class MCPClientManager:
|
||||
_CONNECT_TIMEOUT = 30 # seconds — prevents hung connections on broken remotes
|
||||
_TCP_PROBE_TIMEOUT = 5 # seconds — fast TCP pre-flight for HTTP transports
|
||||
|
||||
# Circuit breaker constants
|
||||
_CB_FAILURE_THRESHOLD = 3
|
||||
_CB_BASE_COOLDOWN = 30.0 # seconds
|
||||
_CB_MAX_COOLDOWN = 300.0 # 5 minutes
|
||||
|
||||
# Notification debounce
|
||||
_NOTIFICATION_DEBOUNCE = 5.0 # seconds between refreshes per server
|
||||
|
||||
# Periodic refresh backoff
|
||||
_REFRESH_BACKOFF_BASE = 60.0 # seconds
|
||||
_REFRESH_BACKOFF_MAX = 3600.0 # 1 hour
|
||||
|
||||
# -- circuit breaker (per-server) -----------------------------------------
|
||||
|
||||
def _cb_check(self, name: str) -> tuple[bool, bool]:
|
||||
"""Check circuit breaker state for *name*.
|
||||
|
||||
Returns ``(is_open, cooldown_expired)``. When the circuit is closed
|
||||
both values are False. When open, *cooldown_expired* indicates
|
||||
whether a probe attempt is allowed.
|
||||
"""
|
||||
deadline = self._circuit_open_until.get(name)
|
||||
if deadline is None:
|
||||
return False, False
|
||||
now = time.monotonic()
|
||||
if now >= deadline:
|
||||
return True, True # half-open: allow one probe
|
||||
return True, False # still in cooldown
|
||||
|
||||
def _cb_record_failure(self, name: str) -> None:
|
||||
"""Record a failure against *name*, potentially opening the circuit."""
|
||||
count = self._consecutive_failures.get(name, 0) + 1
|
||||
self._consecutive_failures[name] = count
|
||||
# Guard: don't extend an already-open deadline. Additional failures
|
||||
# while open still accumulate in _consecutive_failures, so the circuit
|
||||
# re-opens immediately after the next half-open probe fails (count is
|
||||
# already >= threshold).
|
||||
if count >= self._CB_FAILURE_THRESHOLD and name not in self._circuit_open_until:
|
||||
trips = self._circuit_trip_count.get(name, 0)
|
||||
cooldown = min(self._CB_BASE_COOLDOWN * (2**trips), self._CB_MAX_COOLDOWN)
|
||||
# Per-server jitter seeded from server name (varies across process
|
||||
# restarts via PYTHONHASHSEED, which is desirable — each cluster
|
||||
# node gets different jitter to avoid thundering herd).
|
||||
jitter = random.Random(hash(name)).random() * cooldown * 0.1
|
||||
self._circuit_open_until[name] = time.monotonic() + cooldown + jitter
|
||||
self._circuit_trip_count[name] = trips + 1
|
||||
log.warning(
|
||||
"MCP circuit open for '%s': %d consecutive failures, cooldown %.0fs",
|
||||
name,
|
||||
count,
|
||||
cooldown + jitter,
|
||||
)
|
||||
|
||||
def _cb_record_success(self, name: str) -> None:
|
||||
"""Record a successful operation for *name*, decaying circuit state.
|
||||
|
||||
Decays trip count by 1 rather than resetting to 0, so a chronically
|
||||
flapping server escalates its backoff over time instead of always
|
||||
restarting at the minimum cooldown.
|
||||
"""
|
||||
self._consecutive_failures.pop(name, None)
|
||||
self._circuit_open_until.pop(name, None)
|
||||
trips = self._circuit_trip_count.get(name, 0)
|
||||
if trips > 1:
|
||||
self._circuit_trip_count[name] = trips - 1
|
||||
else:
|
||||
self._circuit_trip_count.pop(name, None)
|
||||
|
||||
def _cb_clear(self, name: str) -> None:
|
||||
"""Remove all circuit breaker state for *name*."""
|
||||
self._consecutive_failures.pop(name, None)
|
||||
self._circuit_open_until.pop(name, None)
|
||||
self._circuit_trip_count.pop(name, None)
|
||||
|
||||
# -- safe transport helpers ------------------------------------------------
|
||||
|
||||
async def _pre_close_streams(self, name: str) -> None:
|
||||
"""Close MCP transport streams before stack teardown.
|
||||
|
||||
Pre-closing unblocks anyio transport tasks stuck on zero-buffer
|
||||
``send()`` calls, preventing the CPU busy-loop from SDK #2147.
|
||||
"""
|
||||
streams = self._server_streams.pop(name, None)
|
||||
if streams:
|
||||
for s in streams:
|
||||
with contextlib.suppress(Exception):
|
||||
await s.aclose()
|
||||
|
||||
async def _tcp_probe(self, name: str, url: str) -> None:
|
||||
"""Fast TCP connect check before entering the MCP transport context.
|
||||
|
||||
@@ -251,6 +356,16 @@ class MCPClientManager:
|
||||
log.error("MCP server name '%s' contains '__' (reserved delimiter), skipping", name)
|
||||
return
|
||||
|
||||
# Guard: tear down stale session/stack so we don't leak. Checks both
|
||||
# _sessions and _per_server_stacks because transport errors in the sync
|
||||
# dispatch methods evict the session but leave the stack behind.
|
||||
if name in self._sessions or name in self._per_server_stacks:
|
||||
self._sessions.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
old_stack = self._per_server_stacks.pop(name, None)
|
||||
if old_stack:
|
||||
await self._safe_close_stack(old_stack)
|
||||
|
||||
# Per-server exit stack for clean per-server lifecycle management
|
||||
stack = AsyncExitStack()
|
||||
await stack.__aenter__()
|
||||
@@ -271,6 +386,9 @@ class MCPClientManager:
|
||||
),
|
||||
timeout=self._CONNECT_TIMEOUT,
|
||||
)
|
||||
# Stash stream refs so _pre_close_streams can unblock anyio
|
||||
# transport tasks before the cancel scope fires (SDK #2147).
|
||||
self._server_streams[name] = (read, write)
|
||||
else:
|
||||
# Default: stdio transport
|
||||
command = cfg.get("command", "")
|
||||
@@ -287,24 +405,29 @@ class MCPClientManager:
|
||||
env=env,
|
||||
)
|
||||
read, write = await stack.enter_async_context(stdio_client(params))
|
||||
self._server_streams[name] = (read, write)
|
||||
except asyncio.CancelledError:
|
||||
# Stray CancelledError from broken anyio cancel scope — treat as
|
||||
# Stray CancelledError from broken anyio cancel scope -- treat as
|
||||
# connection failure. But if the task is genuinely being cancelled
|
||||
# (shutdown), re-raise so we don't block teardown.
|
||||
task = asyncio.current_task()
|
||||
if task is not None and task.cancelling():
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
log.warning("MCP server '%s' connection failed (anyio cancel)", name)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"Connection failed for '{name}'") from None
|
||||
except TimeoutError:
|
||||
log.warning(
|
||||
"MCP server '%s' connection timed out after %ds", name, self._CONNECT_TIMEOUT
|
||||
)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"Connection timed out after {self._CONNECT_TIMEOUT}s") from None
|
||||
except Exception:
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
|
||||
@@ -316,15 +439,30 @@ class MCPClientManager:
|
||||
if not isinstance(msg, mcp_types.ServerNotification):
|
||||
return
|
||||
root = msg.root
|
||||
|
||||
# Debounce: skip if we refreshed this server very recently
|
||||
now = time.monotonic()
|
||||
last = self._last_notification_refresh.get(name, 0.0)
|
||||
if now - last < self._NOTIFICATION_DEBOUNCE:
|
||||
log.debug(
|
||||
"Debouncing notification from '%s' (%.1fs since last refresh)",
|
||||
name,
|
||||
now - last,
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
if isinstance(root, mcp_types.ToolListChangedNotification):
|
||||
log.info("Received tools/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_tools(name)
|
||||
elif isinstance(root, mcp_types.ResourceListChangedNotification):
|
||||
log.info("Received resources/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_resources(name)
|
||||
elif isinstance(root, mcp_types.PromptListChangedNotification):
|
||||
log.info("Received prompts/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_prompts(name)
|
||||
self._last_error.pop(name, None)
|
||||
except Exception as exc:
|
||||
@@ -336,6 +474,7 @@ class MCPClientManager:
|
||||
ClientSession(read, write, message_handler=_on_notification) # type: ignore[arg-type]
|
||||
)
|
||||
except Exception:
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
|
||||
@@ -346,16 +485,20 @@ class MCPClientManager:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
task = asyncio.current_task()
|
||||
if task is not None and task.cancelling():
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"MCP handshake failed for '{name}'") from None
|
||||
except TimeoutError:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise TimeoutError(f"MCP handshake timed out after {self._CONNECT_TIMEOUT}s") from None
|
||||
except Exception:
|
||||
self._per_server_stacks.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
await self._safe_close_stack(stack)
|
||||
raise
|
||||
self._sessions[name] = session
|
||||
@@ -548,12 +691,14 @@ class MCPClientManager:
|
||||
if cfg:
|
||||
log.info("Reconnecting MCP server '%s'", name)
|
||||
await self._connect_one(name, cfg)
|
||||
self._cb_record_success(name)
|
||||
new_names = [
|
||||
t["function"]["name"] for t in self._per_server_tools.get(name, [])
|
||||
]
|
||||
results[name] = (new_names, [])
|
||||
continue
|
||||
added, removed = await self._refresh_server(name)
|
||||
self._cb_record_success(name)
|
||||
results[name] = (added, removed)
|
||||
except Exception as exc:
|
||||
log.warning("Refresh failed for MCP server '%s'", name, exc_info=True)
|
||||
@@ -577,10 +722,18 @@ class MCPClientManager:
|
||||
"""
|
||||
assert self._loop is not None
|
||||
future = asyncio.run_coroutine_threadsafe(self._refresh_all(server_name), self._loop)
|
||||
return future.result(timeout=timeout)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
raise TimeoutError(f"MCP refresh timed out after {timeout}s") from None
|
||||
|
||||
async def _periodic_refresh(self) -> None:
|
||||
"""Periodically refresh servers that lack push notifications."""
|
||||
"""Periodically refresh servers that lack push notifications.
|
||||
|
||||
Applies per-server exponential backoff on failure and attempts
|
||||
reconnection for disconnected servers.
|
||||
"""
|
||||
# Stagger start using a launch-time seed so cluster nodes don't
|
||||
# all hit MCP servers simultaneously.
|
||||
seed = random.Random(time.monotonic_ns() ^ os.getpid()).random()
|
||||
@@ -588,8 +741,42 @@ class MCPClientManager:
|
||||
await asyncio.sleep(initial_delay)
|
||||
while True:
|
||||
for name in list(self._server_configs):
|
||||
now = time.monotonic()
|
||||
|
||||
# Check per-server backoff
|
||||
backoff_until = self._refresh_backoff_until.get(name, 0.0)
|
||||
if now < backoff_until:
|
||||
continue # still in backoff
|
||||
|
||||
if name not in self._sessions:
|
||||
continue # not connected — skip (reconnect on manual refresh)
|
||||
# Attempt reconnection for disconnected servers
|
||||
cfg = self._server_configs.get(name)
|
||||
if cfg:
|
||||
try:
|
||||
log.info("Periodic reconnect attempt for '%s'", name)
|
||||
await self._connect_one(name, cfg)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_record_success(name)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
failures = self._refresh_failures.get(name, 0) + 1
|
||||
self._refresh_failures[name] = failures
|
||||
backoff = min(
|
||||
self._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)),
|
||||
self._REFRESH_BACKOFF_MAX,
|
||||
)
|
||||
self._refresh_backoff_until[name] = time.monotonic() + backoff
|
||||
log.warning(
|
||||
"Periodic reconnect failed for '%s' (attempt %d, backoff %.0fs)",
|
||||
name,
|
||||
failures,
|
||||
backoff,
|
||||
)
|
||||
self._set_error(name, f"Reconnect failed: {exc}")
|
||||
continue
|
||||
|
||||
try:
|
||||
if not self._supports_list_changed.get(name, False):
|
||||
await self._refresh_server_tools(name)
|
||||
@@ -598,9 +785,26 @@ class MCPClientManager:
|
||||
if not self._supports_prompt_list_changed.get(name, False):
|
||||
await self._refresh_server_prompts(name)
|
||||
self._last_error.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
except Exception as exc:
|
||||
log.warning("Periodic refresh failed for '%s'", name, exc_info=True)
|
||||
failures = self._refresh_failures.get(name, 0) + 1
|
||||
self._refresh_failures[name] = failures
|
||||
backoff = min(
|
||||
self._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)),
|
||||
self._REFRESH_BACKOFF_MAX,
|
||||
)
|
||||
self._refresh_backoff_until[name] = time.monotonic() + backoff
|
||||
log.warning(
|
||||
"Periodic refresh failed for '%s' (attempt %d, backoff %.0fs)",
|
||||
name,
|
||||
failures,
|
||||
backoff,
|
||||
)
|
||||
self._set_error(name, f"Periodic refresh failed: {exc}")
|
||||
# Note: per-server backoff (max 1h) is only meaningful when
|
||||
# refresh_interval is shorter than _REFRESH_BACKOFF_MAX. With
|
||||
# the default 4h interval this sleep already bounds retry frequency.
|
||||
await asyncio.sleep(self._refresh_interval)
|
||||
|
||||
# -- resource refresh ----------------------------------------------------
|
||||
@@ -940,6 +1144,9 @@ class MCPClientManager:
|
||||
if self._loop and self._per_server_stacks:
|
||||
|
||||
async def _close_all_stacks() -> None:
|
||||
# Pre-close streams to prevent anyio CPU busy-loop during teardown
|
||||
for srv_name in list(self._server_streams):
|
||||
await self._pre_close_streams(srv_name)
|
||||
for stack in self._per_server_stacks.values():
|
||||
await self._safe_close_stack(stack)
|
||||
|
||||
@@ -985,6 +1192,14 @@ class MCPClientManager:
|
||||
self._listeners.clear()
|
||||
self._resource_listeners.clear()
|
||||
self._prompt_listeners.clear()
|
||||
# Clear resilience state
|
||||
self._consecutive_failures.clear()
|
||||
self._circuit_open_until.clear()
|
||||
self._circuit_trip_count.clear()
|
||||
self._server_streams.clear()
|
||||
self._last_notification_refresh.clear()
|
||||
self._refresh_failures.clear()
|
||||
self._refresh_backoff_until.clear()
|
||||
|
||||
log.info("MCP client shut down")
|
||||
|
||||
@@ -1049,6 +1264,7 @@ class MCPClientManager:
|
||||
async def _remove() -> None:
|
||||
# Close session + transport via per-server stack
|
||||
self._sessions.pop(name, None)
|
||||
await self._pre_close_streams(name)
|
||||
stack = self._per_server_stacks.pop(name, None)
|
||||
if stack is not None:
|
||||
await self._safe_close_stack(stack)
|
||||
@@ -1062,6 +1278,10 @@ class MCPClientManager:
|
||||
self._supports_prompts.pop(name, None)
|
||||
self._supports_prompt_list_changed.pop(name, None)
|
||||
self._last_error.pop(name, None)
|
||||
self._last_notification_refresh.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_clear(name)
|
||||
# Rebuild merged state (serialized with notification handlers)
|
||||
self._rebuild_tools()
|
||||
self._rebuild_resources()
|
||||
@@ -1075,6 +1295,7 @@ class MCPClientManager:
|
||||
else:
|
||||
# No event loop (tests / pre-start) — mutate directly
|
||||
self._sessions.pop(name, None)
|
||||
self._server_streams.pop(name, None)
|
||||
self._per_server_tools.pop(name, None)
|
||||
self._per_server_resources.pop(name, None)
|
||||
self._per_server_prompts.pop(name, None)
|
||||
@@ -1084,6 +1305,10 @@ class MCPClientManager:
|
||||
self._supports_prompts.pop(name, None)
|
||||
self._supports_prompt_list_changed.pop(name, None)
|
||||
self._last_error.pop(name, None)
|
||||
self._last_notification_refresh.pop(name, None)
|
||||
self._refresh_failures.pop(name, None)
|
||||
self._refresh_backoff_until.pop(name, None)
|
||||
self._cb_clear(name)
|
||||
self._rebuild_tools()
|
||||
self._rebuild_resources()
|
||||
self._rebuild_prompts()
|
||||
@@ -1107,6 +1332,8 @@ class MCPClientManager:
|
||||
connected = name in self._sessions
|
||||
cfg = self._server_configs.get(name, {})
|
||||
transport = cfg.get("type", "stdio")
|
||||
cb_deadline = self._circuit_open_until.get(name)
|
||||
cb_open = cb_deadline is not None and time.monotonic() < cb_deadline
|
||||
return {
|
||||
"connected": connected,
|
||||
"tools": len(self._per_server_tools.get(name, [])) if connected else 0,
|
||||
@@ -1116,6 +1343,8 @@ class MCPClientManager:
|
||||
"transport": transport,
|
||||
"command": cfg.get("command", "") if transport == "stdio" else "",
|
||||
"url": cfg.get("url", "") if transport != "stdio" else "",
|
||||
"circuit_open": cb_open,
|
||||
"consecutive_failures": self._consecutive_failures.get(name, 0),
|
||||
}
|
||||
|
||||
def get_all_server_status(self) -> dict[str, dict[str, Any]]:
|
||||
@@ -1245,6 +1474,55 @@ class MCPClientManager:
|
||||
|
||||
# -- tool invocation -----------------------------------------------------
|
||||
|
||||
def _cb_gate(self, server_name: str) -> None:
|
||||
"""Check circuit breaker before dispatching to *server_name*.
|
||||
|
||||
Raises ``RuntimeError`` if the circuit is open and cooldown has not
|
||||
expired. When the cooldown has expired (half-open), clears the
|
||||
deadline so the probe attempt is allowed through.
|
||||
"""
|
||||
is_open, cooldown_expired = self._cb_check(server_name)
|
||||
if is_open and not cooldown_expired:
|
||||
remaining = self._circuit_open_until.get(server_name, 0) - time.monotonic()
|
||||
raise RuntimeError(
|
||||
f"MCP server '{server_name}' circuit open "
|
||||
f"(cooldown {remaining:.0f}s remaining). "
|
||||
f"Use '/mcp refresh {server_name}' to retry manually."
|
||||
)
|
||||
if cooldown_expired:
|
||||
# Remove deadline so concurrent callers aren't rejected while the
|
||||
# probe is in-flight. This intentionally allows multiple callers
|
||||
# through rather than a single probe: reconnects serialize on the
|
||||
# event loop via _connect_one's guard, and if the server is truly
|
||||
# broken the first failure re-trips the circuit immediately.
|
||||
self._circuit_open_until.pop(server_name, None)
|
||||
|
||||
def _cb_auto_reconnect(self, server_name: str) -> Any:
|
||||
"""Attempt reconnection for a disconnected server during half-open probe.
|
||||
|
||||
Returns the new session on success, or raises on failure.
|
||||
"""
|
||||
cfg = self._server_configs.get(server_name)
|
||||
if not cfg or self._loop is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
reconnect_future = asyncio.run_coroutine_threadsafe(
|
||||
self._connect_one(server_name, cfg), self._loop
|
||||
)
|
||||
try:
|
||||
reconnect_future.result(timeout=self._CONNECT_TIMEOUT)
|
||||
except concurrent.futures.TimeoutError:
|
||||
reconnect_future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect timed out") from None
|
||||
except Exception as exc:
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect failed: {exc}") from None
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
self._cb_record_failure(server_name)
|
||||
raise RuntimeError(f"MCP server '{server_name}' reconnect produced no session")
|
||||
return session
|
||||
|
||||
def call_tool_sync(
|
||||
self,
|
||||
func_name: str,
|
||||
@@ -1254,15 +1532,19 @@ class MCPClientManager:
|
||||
"""Execute an MCP tool call synchronously (blocks the calling thread).
|
||||
|
||||
Dispatches an async ``tools/call`` to the background event loop and
|
||||
waits for the result.
|
||||
waits for the result. Includes circuit-breaker gating and automatic
|
||||
reconnection for servers recovering from failure.
|
||||
"""
|
||||
mapping = self._tool_map.get(func_name)
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP tool: {func_name}")
|
||||
server_name, original_name = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
@@ -1271,7 +1553,19 @@ class MCPClientManager:
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP tool call timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
# Protocol errors (McpError) come from a healthy connection that
|
||||
# rejected the request — only transport errors trip the breaker.
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
# Extract text from the content array
|
||||
texts: list[str] = []
|
||||
@@ -1320,16 +1614,29 @@ class MCPClientManager:
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP resource: {uri}")
|
||||
server_name, _ = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(session.read_resource(uri), self._loop)
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP resource read timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
parts: list[str] = []
|
||||
for item in result.contents:
|
||||
@@ -1357,9 +1664,12 @@ class MCPClientManager:
|
||||
if mapping is None:
|
||||
raise ValueError(f"Unknown MCP prompt: {prefixed_name}")
|
||||
server_name, original_name = mapping
|
||||
|
||||
self._cb_gate(server_name)
|
||||
|
||||
session = self._sessions.get(server_name)
|
||||
if session is None:
|
||||
raise RuntimeError(f"MCP server '{server_name}' is not connected")
|
||||
session = self._cb_auto_reconnect(server_name)
|
||||
assert self._loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
@@ -1368,7 +1678,17 @@ class MCPClientManager:
|
||||
try:
|
||||
result = future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError:
|
||||
future.cancel()
|
||||
self._cb_record_failure(server_name)
|
||||
raise TimeoutError(f"MCP prompt retrieval timed out after {timeout}s") from None
|
||||
except Exception as exc:
|
||||
if not isinstance(exc, McpError):
|
||||
self._cb_record_failure(server_name)
|
||||
if isinstance(exc, (BrokenPipeError, ConnectionResetError, EOFError)):
|
||||
self._sessions.pop(server_name, None)
|
||||
raise
|
||||
|
||||
self._cb_record_success(server_name)
|
||||
|
||||
messages: list[dict[str, Any]] = []
|
||||
for msg in result.messages:
|
||||
|
||||
@@ -199,6 +199,18 @@ def _resolve_env_vars(value: str) -> str:
|
||||
return re.sub(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}", _replace, value)
|
||||
|
||||
|
||||
def _resolve_openai_provider(provider: str, base_url: str) -> str:
|
||||
"""Distinguish commercial OpenAI from local OpenAI-compatible servers.
|
||||
|
||||
When ``provider`` is ``"openai"`` but the ``base_url`` does not point to
|
||||
``api.openai.com``, the model is on a local server (vLLM, llama.cpp, etc.)
|
||||
and should use the Chat Completions provider (``"openai-compatible"``).
|
||||
"""
|
||||
if provider == "openai" and base_url and "api.openai.com" not in base_url:
|
||||
return "openai-compatible"
|
||||
return provider
|
||||
|
||||
|
||||
def load_model_registry(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
@@ -241,15 +253,16 @@ def load_model_registry(
|
||||
if isinstance(parsed, dict):
|
||||
caps = parsed
|
||||
except (_json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
row_provider = row.get("provider", "openai")
|
||||
pass # falls back to empty capabilities
|
||||
row_base_url = _resolve_env_vars(row.get("base_url", ""))
|
||||
row_provider = _resolve_openai_provider(row.get("provider", "openai"), row_base_url)
|
||||
row_model = row["model"]
|
||||
# 0 = auto-detect: inherit CLI-detected context_window,
|
||||
# same fallback chain as config.toml models
|
||||
row_ctx = row.get("context_window", 0) or context_window
|
||||
configs[alias] = ModelConfig(
|
||||
alias=alias,
|
||||
base_url=_resolve_env_vars(row.get("base_url", "")),
|
||||
base_url=row_base_url,
|
||||
api_key=_resolve_env_vars(row.get("api_key", "")),
|
||||
model=row_model,
|
||||
context_window=row_ctx,
|
||||
@@ -268,13 +281,14 @@ def load_model_registry(
|
||||
if not model_name:
|
||||
log.warning("Model entry '%s' has no model name, skipping", alias)
|
||||
continue
|
||||
entry_base_url = _resolve_env_vars(entry.get("base_url", base_url))
|
||||
configs[alias] = ModelConfig(
|
||||
alias=alias,
|
||||
base_url=entry.get("base_url", base_url),
|
||||
api_key=entry.get("api_key", api_key),
|
||||
base_url=entry_base_url,
|
||||
api_key=_resolve_env_vars(entry.get("api_key", api_key)),
|
||||
model=model_name,
|
||||
context_window=entry.get("context_window", context_window),
|
||||
provider=entry.get("provider", "openai"),
|
||||
provider=_resolve_openai_provider(entry.get("provider", "openai"), entry_base_url),
|
||||
capabilities=entry.get("capabilities", {})
|
||||
if isinstance(entry.get("capabilities"), dict)
|
||||
else {},
|
||||
@@ -290,7 +304,7 @@ def load_model_registry(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
context_window=context_window,
|
||||
provider=provider,
|
||||
provider=_resolve_openai_provider(provider, base_url),
|
||||
)
|
||||
|
||||
# Determine default alias
|
||||
@@ -546,7 +560,11 @@ def _detect_openai_compat(
|
||||
result["context_window"] = known["context_window"]
|
||||
|
||||
# Server type heuristics
|
||||
if base_url and "api.openai.com" in base_url:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
_normalized = (base_url if "://" in base_url else f"https://{base_url}") if base_url else ""
|
||||
_hostname = urlparse(_normalized).hostname or "" if _normalized else ""
|
||||
if base_url and (_hostname == "api.openai.com" or _hostname.endswith(".openai.com")):
|
||||
result["server_type"] = "openai"
|
||||
elif meta is not None and "n_ctx_train" in meta:
|
||||
result["server_type"] = "llama.cpp"
|
||||
|
||||
@@ -6,6 +6,8 @@ import threading
|
||||
from typing import Any
|
||||
|
||||
from turnstone.core.providers._openai import OpenAIProvider
|
||||
from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
from turnstone.core.providers._protocol import (
|
||||
CompletionResult,
|
||||
LLMProvider,
|
||||
@@ -19,7 +21,9 @@ __all__ = [
|
||||
"CompletionResult",
|
||||
"LLMProvider",
|
||||
"ModelCapabilities",
|
||||
"OpenAIChatCompletionsProvider",
|
||||
"OpenAIProvider",
|
||||
"OpenAIResponsesProvider",
|
||||
"StreamChunk",
|
||||
"ToolCallDelta",
|
||||
"UsageInfo",
|
||||
@@ -31,15 +35,18 @@ __all__ = [
|
||||
|
||||
# Singleton instances (stateless, safe to share)
|
||||
_provider_lock = threading.Lock()
|
||||
_openai_provider = OpenAIProvider()
|
||||
_openai_provider = OpenAIResponsesProvider()
|
||||
_openai_compat_provider = OpenAIChatCompletionsProvider()
|
||||
_anthropic_provider: LLMProvider | None = None
|
||||
|
||||
|
||||
def create_provider(provider_name: str) -> LLMProvider:
|
||||
"""Return a provider adapter for the given provider name. Thread-safe."""
|
||||
global _anthropic_provider # noqa: PLW0603
|
||||
if provider_name in ("openai", "openai-compatible"):
|
||||
if provider_name == "openai":
|
||||
return _openai_provider
|
||||
if provider_name == "openai-compatible":
|
||||
return _openai_compat_provider
|
||||
if provider_name == "anthropic":
|
||||
with _provider_lock:
|
||||
if _anthropic_provider is None:
|
||||
@@ -99,9 +106,9 @@ def lookup_model_capabilities(provider: str, model: str) -> dict[str, Any] | Non
|
||||
def list_known_models(provider: str) -> list[str]:
|
||||
"""Return the model name prefixes in the static capability table."""
|
||||
if provider == "openai":
|
||||
from turnstone.core.providers._openai import _OPENAI_CAPABILITIES
|
||||
from turnstone.core.providers._openai_common import OPENAI_CAPABILITIES
|
||||
|
||||
return sorted(_OPENAI_CAPABILITIES.keys())
|
||||
return sorted(OPENAI_CAPABILITIES.keys())
|
||||
if provider == "anthropic":
|
||||
from turnstone.core.providers._anthropic import _ANTHROPIC_CAPABILITIES
|
||||
|
||||
|
||||
@@ -69,7 +69,7 @@ def _merge_consecutive(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
_WEB_SEARCH_TOOL_TYPE = "web_search_20250305"
|
||||
|
||||
# Tool search: server-side BM25 tool discovery for deferred tools
|
||||
_TOOL_SEARCH_TOOL_TYPE = "tool_search_tool_bm25_20251119"
|
||||
_TOOL_SEARCH_TOOL_TYPE = "tool_search_tool_bm25"
|
||||
|
||||
# -- model capabilities -------------------------------------------------------
|
||||
|
||||
@@ -205,7 +205,7 @@ class AnthropicProvider:
|
||||
result.append({**tool, "defer_loading": True})
|
||||
else:
|
||||
result.append(tool)
|
||||
result.append({"type": _TOOL_SEARCH_TOOL_TYPE, "name": "tool_search"})
|
||||
result.append({"type": _TOOL_SEARCH_TOOL_TYPE, "name": _TOOL_SEARCH_TOOL_TYPE})
|
||||
return result
|
||||
|
||||
# -- shared param logic --------------------------------------------------
|
||||
@@ -543,9 +543,13 @@ class AnthropicProvider:
|
||||
if extra_params and "thinking_budget_tokens" in extra_params:
|
||||
budget = extra_params["thinking_budget_tokens"]
|
||||
if budget > 0:
|
||||
# Budget must leave room for the response
|
||||
# Budget must be strictly less than max_tokens (API requirement).
|
||||
# If max_tokens is too small to fit even a minimal thinking
|
||||
# budget alongside the response, disable thinking entirely.
|
||||
if budget >= max_tokens:
|
||||
budget = max(1024, max_tokens - 1024)
|
||||
budget = max_tokens - 1024
|
||||
if budget < 1:
|
||||
return {}
|
||||
return {"thinking": {"type": "enabled", "budget_tokens": budget}}
|
||||
return {}
|
||||
|
||||
@@ -704,7 +708,7 @@ class AnthropicProvider:
|
||||
parsed = json.loads(info["input_json"])
|
||||
query = parsed.get("query", "")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
pass # best-effort query extraction for status
|
||||
sc.info_delta = f"[Searching: {query}]" if query else "[Searching...]"
|
||||
|
||||
elif event_type == "message_delta":
|
||||
|
||||
@@ -1,577 +1,30 @@
|
||||
"""OpenAI-compatible provider — wraps current behavior with zero semantic change.
|
||||
"""Re-export shim for backwards compatibility.
|
||||
|
||||
Handles OpenAI, vLLM, llama.cpp, and any server that speaks the
|
||||
OpenAI Chat Completions API.
|
||||
The OpenAI provider family is split into:
|
||||
- ``_openai_chat.py`` — Chat Completions API (local model servers)
|
||||
- ``_openai_responses.py`` — Responses API (commercial OpenAI)
|
||||
- ``_openai_common.py`` — shared capability table, helpers
|
||||
|
||||
``OpenAIProvider`` is preserved as an alias for ``OpenAIChatCompletionsProvider``
|
||||
so existing code that imports it directly continues to work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
import structlog
|
||||
|
||||
from turnstone.core.providers._protocol import (
|
||||
CompletionResult,
|
||||
ModelCapabilities,
|
||||
StreamChunk,
|
||||
ToolCallDelta,
|
||||
UsageInfo,
|
||||
_lookup_capabilities,
|
||||
from turnstone.core.providers._openai_chat import (
|
||||
OpenAIChatCompletionsProvider,
|
||||
)
|
||||
from turnstone.core.providers._openai_chat import (
|
||||
OpenAIChatCompletionsProvider as OpenAIProvider,
|
||||
)
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
# Backwards-compatible aliases for the capability tables
|
||||
from turnstone.core.providers._openai_common import (
|
||||
OPENAI_CAPABILITIES as _OPENAI_CAPABILITIES, # noqa: F401
|
||||
)
|
||||
from turnstone.core.providers._openai_common import OPENAI_DEFAULT as _OPENAI_DEFAULT # noqa: F401
|
||||
from turnstone.core.providers._openai_responses import OpenAIResponsesProvider
|
||||
|
||||
# -- model capabilities -------------------------------------------------------
|
||||
|
||||
_OPENAI_CAPABILITIES: dict[str, ModelCapabilities] = {
|
||||
# GPT-5 base — NO temperature support
|
||||
"gpt-5": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
"gpt-5-mini": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
"gpt-5-nano": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5 pro — high reasoning only, extended output
|
||||
"gpt-5-pro": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=272000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("high",),
|
||||
default_reasoning_effort="high",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.1 — temperature OK when reasoning_effort=none (default)
|
||||
"gpt-5.1": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.2 — adds xhigh
|
||||
"gpt-5.2": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.2 pro — always-reasoning variant
|
||||
"gpt-5.2-pro": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("medium", "high", "xhigh"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.3 — same capabilities as 5.2 (matches gpt-5.3-chat-latest, codex)
|
||||
"gpt-5.3": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.4 — 1M context window, native tool search
|
||||
"gpt-5.4": ModelCapabilities(
|
||||
context_window=1050000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_tool_search=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.4 pro — always-reasoning, 1M context, native tool search
|
||||
"gpt-5.4-pro": ModelCapabilities(
|
||||
context_window=1050000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("medium", "high", "xhigh"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_tool_search=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
# O-series reasoning models
|
||||
"o1": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o1-mini": ModelCapabilities(
|
||||
context_window=128000,
|
||||
max_output_tokens=65536,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3-mini": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3-pro": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o4-mini": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
# Search models — always search on every request, no reasoning_effort
|
||||
"gpt-5-search-api": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
supports_web_search=True,
|
||||
reasoning_effort_values=(),
|
||||
supports_vision=True,
|
||||
),
|
||||
}
|
||||
|
||||
# Default for unknown models (local servers: vLLM, llama.cpp, etc.)
|
||||
_OPENAI_DEFAULT = ModelCapabilities()
|
||||
|
||||
|
||||
class OpenAIProvider:
|
||||
"""Provider for OpenAI-compatible APIs (OpenAI, vLLM, llama.cpp, etc.)."""
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai"
|
||||
|
||||
def get_capabilities(self, model: str) -> ModelCapabilities:
|
||||
return _lookup_capabilities(model, _OPENAI_CAPABILITIES, _OPENAI_DEFAULT)
|
||||
|
||||
# -- shared param logic --------------------------------------------------
|
||||
|
||||
def _apply_model_params(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
caps: ModelCapabilities,
|
||||
temperature: float,
|
||||
reasoning_effort: str,
|
||||
) -> None:
|
||||
"""Conditionally add temperature and reasoning_effort to *kwargs*.
|
||||
|
||||
- Models with ``supports_temperature=False`` (GPT-5 base, O-series)
|
||||
never receive temperature.
|
||||
- Models that list ``"none"`` in their effort values (GPT-5.1/5.2)
|
||||
only receive temperature when reasoning is inactive.
|
||||
- ``reasoning_effort`` is forwarded as a first-class API parameter
|
||||
only for models that declare supported effort values.
|
||||
"""
|
||||
if caps.supports_temperature:
|
||||
# GPT-5.1/5.2: temperature only valid when reasoning_effort is "none"
|
||||
if "none" in caps.reasoning_effort_values and reasoning_effort not in (
|
||||
"none",
|
||||
"",
|
||||
):
|
||||
pass # Skip temperature when reasoning is active
|
||||
else:
|
||||
kwargs["temperature"] = temperature
|
||||
if caps.reasoning_effort_values and reasoning_effort and reasoning_effort != "none":
|
||||
# Validate against supported values; fall back to model default
|
||||
if reasoning_effort in caps.reasoning_effort_values:
|
||||
kwargs["reasoning_effort"] = reasoning_effort
|
||||
elif caps.default_reasoning_effort and caps.default_reasoning_effort != "none":
|
||||
kwargs["reasoning_effort"] = caps.default_reasoning_effort
|
||||
|
||||
# -- web search ----------------------------------------------------------
|
||||
|
||||
def _apply_web_search(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
caps: ModelCapabilities,
|
||||
tools: list[dict[str, Any]] | None,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Inject ``web_search_options`` for search models.
|
||||
|
||||
For models with ``supports_web_search``, the web search function tool
|
||||
is removed (the model searches automatically) and ``web_search_options``
|
||||
is added to the request kwargs.
|
||||
|
||||
Returns the (possibly filtered) tools list.
|
||||
"""
|
||||
if not caps.supports_web_search:
|
||||
return tools
|
||||
# Remove web_search function tool — model has built-in search
|
||||
if tools:
|
||||
tools = [t for t in tools if t.get("function", {}).get("name") != "web_search"]
|
||||
if not tools:
|
||||
tools = None
|
||||
kwargs["web_search_options"] = {}
|
||||
return tools
|
||||
|
||||
# -- prompt cache retention -----------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _apply_cache_retention(kwargs: dict[str, Any], model: str) -> None:
|
||||
"""Enable 24-hour extended prompt cache retention for GPT-5.x models.
|
||||
|
||||
OpenAI caching is automatic (no code changes for basic caching), but
|
||||
the default TTL is only 5-10 minutes. Extended retention keeps cached
|
||||
KV tensors for up to 24 hours at no additional cost, which is valuable
|
||||
for workstreams with bursty activity patterns.
|
||||
"""
|
||||
# GPT-5, GPT-5.1, GPT-5.2, GPT-5.3, GPT-5.4 and variants
|
||||
if model.startswith("gpt-5"):
|
||||
kwargs["prompt_cache_retention"] = "24h"
|
||||
|
||||
# -- tool search ---------------------------------------------------------
|
||||
|
||||
def _apply_tool_search(
|
||||
self,
|
||||
caps: ModelCapabilities,
|
||||
tools: list[dict[str, Any]] | None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Mark deferred tools with ``defer_loading: true`` for native search.
|
||||
|
||||
For GPT-5.4+ models that support tool search, OpenAI's API handles
|
||||
discovery automatically — no explicit search tool is needed.
|
||||
"""
|
||||
if not caps.supports_tool_search or not deferred_names or not tools:
|
||||
return tools
|
||||
result = []
|
||||
for tool in tools:
|
||||
name = tool.get("function", {}).get("name", "")
|
||||
if name in deferred_names:
|
||||
result.append({**tool, "defer_loading": True})
|
||||
else:
|
||||
result.append(tool)
|
||||
return result
|
||||
|
||||
# -- message sanitisation ------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_messages(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Ensure assistant messages always have ``content`` or ``tool_calls``.
|
||||
|
||||
OpenAI-compatible APIs reject assistant messages that have neither.
|
||||
This is a defensive catch-all; the upstream layers should already
|
||||
guarantee well-formed messages.
|
||||
"""
|
||||
out: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if (
|
||||
msg.get("role") == "assistant"
|
||||
and msg.get("content") is None
|
||||
and not msg.get("tool_calls")
|
||||
):
|
||||
msg = {**msg, "content": ""}
|
||||
out.append(msg)
|
||||
return out
|
||||
|
||||
# -- streaming -----------------------------------------------------------
|
||||
|
||||
def create_streaming(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
caps = self.get_capabilities(model)
|
||||
messages = self._sanitize_messages(messages)
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
caps.token_param: max_tokens,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
self._apply_model_params(kwargs, caps, temperature, reasoning_effort)
|
||||
self._apply_cache_retention(kwargs, model)
|
||||
tools = self._apply_web_search(kwargs, caps, tools)
|
||||
tools = self._apply_tool_search(caps, tools, deferred_names)
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if extra_params:
|
||||
kwargs["extra_body"] = extra_params
|
||||
|
||||
log.debug(
|
||||
"openai.request",
|
||||
model=model,
|
||||
stream=True,
|
||||
max_tokens=max_tokens,
|
||||
message_count=len(messages),
|
||||
tool_count=len(tools) if tools else 0,
|
||||
)
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
if cancel_ref is not None:
|
||||
cancel_ref.append(stream)
|
||||
return self._iter_stream(stream)
|
||||
|
||||
def _iter_stream(self, stream: Any) -> Iterator[StreamChunk]:
|
||||
"""Convert OpenAI stream chunks to normalized StreamChunks."""
|
||||
first = True
|
||||
annotations: list[Any] = []
|
||||
content_len = 0
|
||||
tool_call_count = 0
|
||||
last_finish_reason: str | None = None
|
||||
completion_tokens: int | None = None
|
||||
for chunk in stream:
|
||||
sc = StreamChunk()
|
||||
|
||||
# Finish reason
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
sc.finish_reason = chunk.choices[0].finish_reason
|
||||
last_finish_reason = sc.finish_reason
|
||||
|
||||
# Usage from final chunk
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
u = chunk.usage
|
||||
pt = getattr(u, "prompt_tokens", None)
|
||||
ct = getattr(u, "completion_tokens", None)
|
||||
tt = getattr(u, "total_tokens", None)
|
||||
completion_tokens = ct
|
||||
if pt is not None and ct is not None:
|
||||
# Extract cached_tokens from prompt_tokens_details.
|
||||
# OpenAI caching is automatic with no write premium, so
|
||||
# cache_creation_tokens is always 0 (only Anthropic reports it).
|
||||
ptd = getattr(u, "prompt_tokens_details", None)
|
||||
cached = getattr(ptd, "cached_tokens", 0) if ptd else 0
|
||||
sc.usage = UsageInfo(
|
||||
prompt_tokens=pt,
|
||||
completion_tokens=ct,
|
||||
total_tokens=tt or (pt + ct),
|
||||
cache_read_tokens=cached or 0,
|
||||
)
|
||||
|
||||
if not chunk.choices:
|
||||
if sc.usage:
|
||||
yield sc
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
|
||||
# Reasoning field (vLLM --reasoning-parser, llama.cpp)
|
||||
rc = getattr(delta, "reasoning", None) or getattr(delta, "reasoning_content", None)
|
||||
if rc:
|
||||
sc.reasoning_delta = rc
|
||||
|
||||
# Content
|
||||
if delta.content:
|
||||
sc.content_delta = delta.content
|
||||
content_len += len(delta.content)
|
||||
|
||||
# Tool calls
|
||||
if delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tcd = ToolCallDelta(index=tc_delta.index)
|
||||
if tc_delta.id:
|
||||
tcd.id = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tcd.name = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tcd.arguments_delta = tc_delta.function.arguments
|
||||
sc.tool_call_deltas.append(tcd)
|
||||
tool_call_count += 1
|
||||
|
||||
# Accumulate url_citation annotations from search models
|
||||
delta_anns = getattr(delta, "annotations", None)
|
||||
if delta_anns:
|
||||
annotations.extend(delta_anns)
|
||||
|
||||
has_content = sc.content_delta or sc.reasoning_delta or sc.tool_call_deltas
|
||||
if has_content and first:
|
||||
sc.is_first = True
|
||||
first = False
|
||||
|
||||
if has_content or sc.finish_reason or sc.usage:
|
||||
yield sc
|
||||
|
||||
log.debug(
|
||||
"openai.response",
|
||||
stream=True,
|
||||
finish_reason=last_finish_reason,
|
||||
content_length=content_len,
|
||||
tool_call_deltas=tool_call_count,
|
||||
completion_tokens=completion_tokens,
|
||||
)
|
||||
|
||||
# Emit accumulated citations as a final info chunk
|
||||
if annotations:
|
||||
citation_text = self._format_citations("", annotations).strip()
|
||||
if citation_text:
|
||||
yield StreamChunk(info_delta=citation_text)
|
||||
|
||||
# -- non-streaming -------------------------------------------------------
|
||||
|
||||
def create_completion(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
) -> CompletionResult:
|
||||
caps = self.get_capabilities(model)
|
||||
messages = self._sanitize_messages(messages)
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
caps.token_param: max_tokens,
|
||||
"stream": False,
|
||||
}
|
||||
self._apply_model_params(kwargs, caps, temperature, reasoning_effort)
|
||||
self._apply_cache_retention(kwargs, model)
|
||||
tools = self._apply_web_search(kwargs, caps, tools)
|
||||
tools = self._apply_tool_search(caps, tools, deferred_names)
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if extra_params:
|
||||
kwargs["extra_body"] = extra_params
|
||||
|
||||
log.debug(
|
||||
"openai.request",
|
||||
model=model,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
message_count=len(messages),
|
||||
tool_count=len(tools) if tools else 0,
|
||||
)
|
||||
response = client.chat.completions.create(**kwargs)
|
||||
choice = response.choices[0]
|
||||
msg = choice.message
|
||||
|
||||
tool_calls = None
|
||||
if msg.tool_calls:
|
||||
tool_calls = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in msg.tool_calls
|
||||
]
|
||||
|
||||
# Extract url_citation annotations from web search models
|
||||
content = msg.content or ""
|
||||
annotations = getattr(msg, "annotations", None)
|
||||
if annotations:
|
||||
content = self._format_citations(content, annotations)
|
||||
|
||||
usage = None
|
||||
if hasattr(response, "usage") and response.usage:
|
||||
u = response.usage
|
||||
ptd = getattr(u, "prompt_tokens_details", None)
|
||||
cached = getattr(ptd, "cached_tokens", 0) if ptd else 0
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=u.prompt_tokens,
|
||||
completion_tokens=u.completion_tokens,
|
||||
total_tokens=getattr(u, "total_tokens", None)
|
||||
or (u.prompt_tokens + u.completion_tokens),
|
||||
cache_read_tokens=cached or 0,
|
||||
)
|
||||
|
||||
result = CompletionResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=choice.finish_reason or "stop",
|
||||
usage=usage,
|
||||
)
|
||||
log.debug(
|
||||
"openai.response",
|
||||
stream=False,
|
||||
finish_reason=result.finish_reason,
|
||||
content_length=len(content),
|
||||
tool_call_count=len(tool_calls) if tool_calls else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else None,
|
||||
)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _format_citations(content: str, annotations: list[Any]) -> str:
|
||||
"""Append url_citation sources as footnotes at the end of the content."""
|
||||
seen_urls: set[str] = set()
|
||||
sources: list[str] = []
|
||||
for ann in annotations:
|
||||
ann_type = getattr(ann, "type", None)
|
||||
if ann_type == "url_citation":
|
||||
citation = getattr(ann, "url_citation", None)
|
||||
if citation:
|
||||
title = getattr(citation, "title", "")
|
||||
url = getattr(citation, "url", "")
|
||||
if url and url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
sources.append(f"[{title}]({url})" if title else url)
|
||||
if sources:
|
||||
content += "\n\nSources:\n" + "\n".join(f"- {s}" for s in sources)
|
||||
return content
|
||||
|
||||
# -- tool conversion -----------------------------------------------------
|
||||
|
||||
def convert_tools(
|
||||
self,
|
||||
tools: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
return tools # Already in OpenAI format
|
||||
|
||||
# -- retryable errors ----------------------------------------------------
|
||||
|
||||
@property
|
||||
def retryable_error_names(self) -> frozenset[str]:
|
||||
return frozenset(
|
||||
{
|
||||
"APIError",
|
||||
"APIConnectionError",
|
||||
"RateLimitError",
|
||||
"Timeout",
|
||||
"APITimeoutError",
|
||||
}
|
||||
)
|
||||
__all__ = [
|
||||
"OpenAIChatCompletionsProvider",
|
||||
"OpenAIProvider",
|
||||
"OpenAIResponsesProvider",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
"""Chat Completions provider — for local model servers (vLLM, llama.cpp, SGLang).
|
||||
|
||||
Wraps the OpenAI Chat Completions API (``/v1/chat/completions``).
|
||||
Commercial OpenAI models should use ``OpenAIResponsesProvider`` instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
import structlog
|
||||
|
||||
from turnstone.core.providers._openai_common import (
|
||||
RETRYABLE_ERROR_NAMES,
|
||||
apply_cache_retention,
|
||||
apply_temperature_and_effort,
|
||||
apply_tool_search,
|
||||
extract_usage,
|
||||
format_citations,
|
||||
lookup_openai_capabilities,
|
||||
sanitize_messages,
|
||||
)
|
||||
from turnstone.core.providers._protocol import (
|
||||
CompletionResult,
|
||||
ModelCapabilities,
|
||||
StreamChunk,
|
||||
ToolCallDelta,
|
||||
)
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
class OpenAIChatCompletionsProvider:
|
||||
"""Provider for local OpenAI-compatible servers (vLLM, llama.cpp, SGLang).
|
||||
|
||||
Uses the Chat Completions API (``/v1/chat/completions``).
|
||||
"""
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-compatible"
|
||||
|
||||
def get_capabilities(self, model: str) -> ModelCapabilities:
|
||||
return lookup_openai_capabilities(model)
|
||||
|
||||
# -- web search ----------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _apply_web_search(
|
||||
kwargs: dict[str, Any],
|
||||
caps: ModelCapabilities,
|
||||
tools: list[dict[str, Any]] | None,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Inject ``web_search_options`` for search models.
|
||||
|
||||
For models with ``supports_web_search``, the web search function tool
|
||||
is removed (the model searches automatically) and ``web_search_options``
|
||||
is added to the request kwargs.
|
||||
|
||||
Returns the (possibly filtered) tools list.
|
||||
"""
|
||||
if not caps.supports_web_search:
|
||||
return tools
|
||||
if tools:
|
||||
tools = [t for t in tools if t.get("function", {}).get("name") != "web_search"]
|
||||
if not tools:
|
||||
tools = None
|
||||
kwargs["web_search_options"] = {}
|
||||
return tools
|
||||
|
||||
# -- streaming -----------------------------------------------------------
|
||||
|
||||
def create_streaming(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
caps = self.get_capabilities(model)
|
||||
messages = sanitize_messages(messages)
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
caps.token_param: max_tokens,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
apply_temperature_and_effort(kwargs, caps, temperature, reasoning_effort)
|
||||
apply_cache_retention(kwargs, model)
|
||||
tools = self._apply_web_search(kwargs, caps, tools)
|
||||
tools = apply_tool_search(caps, tools, deferred_names)
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if extra_params:
|
||||
kwargs["extra_body"] = extra_params
|
||||
|
||||
log.debug(
|
||||
"openai.chat.request",
|
||||
model=model,
|
||||
stream=True,
|
||||
max_tokens=max_tokens,
|
||||
message_count=len(messages),
|
||||
tool_count=len(tools) if tools else 0,
|
||||
)
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
if cancel_ref is not None:
|
||||
cancel_ref.append(stream)
|
||||
return self._iter_stream(stream)
|
||||
|
||||
def _iter_stream(self, stream: Any) -> Iterator[StreamChunk]:
|
||||
"""Convert OpenAI Chat Completions stream chunks to StreamChunks."""
|
||||
first = True
|
||||
annotations: list[Any] = []
|
||||
content_len = 0
|
||||
tool_call_count = 0
|
||||
last_finish_reason: str | None = None
|
||||
completion_tokens: int | None = None
|
||||
for chunk in stream:
|
||||
sc = StreamChunk()
|
||||
|
||||
# Finish reason
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
sc.finish_reason = chunk.choices[0].finish_reason
|
||||
last_finish_reason = sc.finish_reason
|
||||
|
||||
# Usage from final chunk
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
sc.usage = extract_usage(chunk.usage)
|
||||
if sc.usage:
|
||||
completion_tokens = sc.usage.completion_tokens
|
||||
|
||||
if not chunk.choices:
|
||||
if sc.usage:
|
||||
yield sc
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
|
||||
# Reasoning field (vLLM --reasoning-parser, llama.cpp)
|
||||
rc = getattr(delta, "reasoning", None) or getattr(delta, "reasoning_content", None)
|
||||
if rc:
|
||||
sc.reasoning_delta = rc
|
||||
|
||||
# Content
|
||||
if delta.content:
|
||||
sc.content_delta = delta.content
|
||||
content_len += len(delta.content)
|
||||
|
||||
# Tool calls
|
||||
if delta.tool_calls:
|
||||
for tc_delta in delta.tool_calls:
|
||||
tcd = ToolCallDelta(index=tc_delta.index)
|
||||
if tc_delta.id:
|
||||
tcd.id = tc_delta.id
|
||||
if tc_delta.function:
|
||||
if tc_delta.function.name:
|
||||
tcd.name = tc_delta.function.name
|
||||
if tc_delta.function.arguments:
|
||||
tcd.arguments_delta = tc_delta.function.arguments
|
||||
sc.tool_call_deltas.append(tcd)
|
||||
tool_call_count += 1
|
||||
|
||||
# Accumulate url_citation annotations from search models
|
||||
delta_anns = getattr(delta, "annotations", None)
|
||||
if delta_anns:
|
||||
annotations.extend(delta_anns)
|
||||
|
||||
has_content = sc.content_delta or sc.reasoning_delta or sc.tool_call_deltas
|
||||
if has_content and first:
|
||||
sc.is_first = True
|
||||
first = False
|
||||
|
||||
if has_content or sc.finish_reason or sc.usage:
|
||||
yield sc
|
||||
|
||||
log.debug(
|
||||
"openai.chat.response",
|
||||
stream=True,
|
||||
finish_reason=last_finish_reason,
|
||||
content_length=content_len,
|
||||
tool_call_deltas=tool_call_count,
|
||||
completion_tokens=completion_tokens,
|
||||
)
|
||||
|
||||
# Emit accumulated citations as a final info chunk
|
||||
if annotations:
|
||||
citation_text = format_citations("", annotations).strip()
|
||||
if citation_text:
|
||||
yield StreamChunk(info_delta=citation_text)
|
||||
|
||||
# -- non-streaming -------------------------------------------------------
|
||||
|
||||
def create_completion(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
) -> CompletionResult:
|
||||
caps = self.get_capabilities(model)
|
||||
messages = sanitize_messages(messages)
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
caps.token_param: max_tokens,
|
||||
"stream": False,
|
||||
}
|
||||
apply_temperature_and_effort(kwargs, caps, temperature, reasoning_effort)
|
||||
apply_cache_retention(kwargs, model)
|
||||
tools = self._apply_web_search(kwargs, caps, tools)
|
||||
tools = apply_tool_search(caps, tools, deferred_names)
|
||||
if tools:
|
||||
kwargs["tools"] = tools
|
||||
if extra_params:
|
||||
kwargs["extra_body"] = extra_params
|
||||
|
||||
log.debug(
|
||||
"openai.chat.request",
|
||||
model=model,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
message_count=len(messages),
|
||||
tool_count=len(tools) if tools else 0,
|
||||
)
|
||||
response = client.chat.completions.create(**kwargs)
|
||||
choice = response.choices[0]
|
||||
msg = choice.message
|
||||
|
||||
tool_calls = None
|
||||
if msg.tool_calls:
|
||||
tool_calls = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in msg.tool_calls
|
||||
]
|
||||
|
||||
# Extract url_citation annotations from web search models
|
||||
content = msg.content or ""
|
||||
annotations = getattr(msg, "annotations", None)
|
||||
if annotations:
|
||||
content = format_citations(content, annotations)
|
||||
|
||||
usage = extract_usage(getattr(response, "usage", None))
|
||||
|
||||
result = CompletionResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=choice.finish_reason or "stop",
|
||||
usage=usage,
|
||||
)
|
||||
log.debug(
|
||||
"openai.chat.response",
|
||||
stream=False,
|
||||
finish_reason=result.finish_reason,
|
||||
content_length=len(content),
|
||||
tool_call_count=len(tool_calls) if tool_calls else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else None,
|
||||
)
|
||||
return result
|
||||
|
||||
# -- tool conversion -----------------------------------------------------
|
||||
|
||||
def convert_tools(
|
||||
self,
|
||||
tools: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
return tools # Already in OpenAI Chat Completions format
|
||||
|
||||
# -- retryable errors ----------------------------------------------------
|
||||
|
||||
@property
|
||||
def retryable_error_names(self) -> frozenset[str]:
|
||||
return RETRYABLE_ERROR_NAMES
|
||||
@@ -0,0 +1,378 @@
|
||||
"""Shared helpers for OpenAI-family providers (Chat Completions & Responses).
|
||||
|
||||
Capability table, temperature/reasoning gating, cache retention, citation
|
||||
formatting, and message sanitisation live here so both
|
||||
``OpenAIChatCompletionsProvider`` and ``OpenAIResponsesProvider`` stay DRY.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from turnstone.core.providers._protocol import (
|
||||
ModelCapabilities,
|
||||
UsageInfo,
|
||||
_lookup_capabilities,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model capability table
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
OPENAI_CAPABILITIES: dict[str, ModelCapabilities] = {
|
||||
# GPT-5 base — NO temperature support
|
||||
"gpt-5": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
"gpt-5-mini": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
"gpt-5-nano": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("minimal", "low", "medium", "high"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5 pro — high reasoning only, extended output
|
||||
"gpt-5-pro": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=272000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("high",),
|
||||
default_reasoning_effort="high",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.1 — temperature OK when reasoning_effort=none (default)
|
||||
"gpt-5.1": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.2 — adds xhigh
|
||||
"gpt-5.2": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.2 pro — always-reasoning variant
|
||||
"gpt-5.2-pro": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("medium", "high", "xhigh"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.3 — same capabilities as 5.2 (matches gpt-5.3-chat-latest, codex)
|
||||
"gpt-5.3": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.4 — 1M context window, native tool search
|
||||
"gpt-5.4": ModelCapabilities(
|
||||
context_window=1050000,
|
||||
max_output_tokens=128000,
|
||||
reasoning_effort_values=("none", "low", "medium", "high", "xhigh"),
|
||||
default_reasoning_effort="none",
|
||||
supports_tool_search=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
# GPT-5.4 pro — always-reasoning, 1M context, native tool search
|
||||
"gpt-5.4-pro": ModelCapabilities(
|
||||
context_window=1050000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
reasoning_effort_values=("medium", "high", "xhigh"),
|
||||
default_reasoning_effort="medium",
|
||||
supports_tool_search=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
# O-series reasoning models
|
||||
"o1": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o1-mini": ModelCapabilities(
|
||||
context_window=128000,
|
||||
max_output_tokens=65536,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3-mini": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o3-pro": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_streaming=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
"o4-mini": ModelCapabilities(
|
||||
context_window=200000,
|
||||
max_output_tokens=100000,
|
||||
supports_temperature=False,
|
||||
supports_vision=True,
|
||||
),
|
||||
# Search models — always search on every request, no reasoning_effort
|
||||
"gpt-5-search-api": ModelCapabilities(
|
||||
context_window=400000,
|
||||
max_output_tokens=128000,
|
||||
supports_temperature=False,
|
||||
supports_web_search=True,
|
||||
reasoning_effort_values=(),
|
||||
supports_vision=True,
|
||||
),
|
||||
}
|
||||
|
||||
# Default for unknown models (local servers: vLLM, llama.cpp, etc.)
|
||||
OPENAI_DEFAULT = ModelCapabilities()
|
||||
|
||||
|
||||
def lookup_openai_capabilities(model: str) -> ModelCapabilities:
|
||||
"""Find capabilities for *model* by longest prefix match."""
|
||||
return _lookup_capabilities(model, OPENAI_CAPABILITIES, OPENAI_DEFAULT)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Temperature and reasoning effort gating
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def apply_temperature(
|
||||
kwargs: dict[str, Any],
|
||||
caps: ModelCapabilities,
|
||||
temperature: float,
|
||||
reasoning_effort: str,
|
||||
) -> None:
|
||||
"""Conditionally add temperature to *kwargs*.
|
||||
|
||||
- Models with ``supports_temperature=False`` (GPT-5 base, O-series)
|
||||
never receive temperature.
|
||||
- Models that list ``"none"`` in their effort values (GPT-5.1/5.2)
|
||||
only receive temperature when reasoning is inactive.
|
||||
"""
|
||||
if not caps.supports_temperature:
|
||||
return
|
||||
if "none" in caps.reasoning_effort_values and reasoning_effort not in ("none", ""):
|
||||
return # Skip temperature when reasoning is active
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
|
||||
def resolve_reasoning_effort(caps: ModelCapabilities, reasoning_effort: str) -> str | None:
|
||||
"""Return the validated reasoning effort value, or ``None`` to omit.
|
||||
|
||||
Validates against supported values and falls back to model default.
|
||||
"""
|
||||
if not caps.reasoning_effort_values or not reasoning_effort or reasoning_effort == "none":
|
||||
return None
|
||||
if reasoning_effort in caps.reasoning_effort_values:
|
||||
return reasoning_effort
|
||||
if caps.default_reasoning_effort and caps.default_reasoning_effort != "none":
|
||||
return caps.default_reasoning_effort
|
||||
return None
|
||||
|
||||
|
||||
def apply_temperature_and_effort(
|
||||
kwargs: dict[str, Any],
|
||||
caps: ModelCapabilities,
|
||||
temperature: float,
|
||||
reasoning_effort: str,
|
||||
) -> None:
|
||||
"""Conditionally add temperature and reasoning_effort to *kwargs*.
|
||||
|
||||
Chat Completions API version — reasoning effort is a flat parameter.
|
||||
"""
|
||||
apply_temperature(kwargs, caps, temperature, reasoning_effort)
|
||||
effort = resolve_reasoning_effort(caps, reasoning_effort)
|
||||
if effort:
|
||||
kwargs["reasoning_effort"] = effort
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache retention
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def apply_cache_retention(kwargs: dict[str, Any], model: str) -> None:
|
||||
"""Enable 24-hour extended prompt cache retention for GPT-5.x models.
|
||||
|
||||
OpenAI caching is automatic (no code changes for basic caching), but
|
||||
the default TTL is only 5-10 minutes. Extended retention keeps cached
|
||||
KV tensors for up to 24 hours at no additional cost, which is valuable
|
||||
for workstreams with bursty activity patterns.
|
||||
"""
|
||||
if model.startswith("gpt-5"):
|
||||
kwargs["prompt_cache_retention"] = "24h"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool search (native deferred loading)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def apply_tool_search(
|
||||
caps: ModelCapabilities,
|
||||
tools: list[dict[str, Any]] | None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Mark deferred tools with ``defer_loading: true`` for native search.
|
||||
|
||||
For GPT-5.4+ models that support tool search, OpenAI's API handles
|
||||
discovery automatically — no explicit search tool is needed.
|
||||
"""
|
||||
if not caps.supports_tool_search or not deferred_names or not tools:
|
||||
return tools
|
||||
result = []
|
||||
for tool in tools:
|
||||
name = tool.get("function", {}).get("name", "")
|
||||
if name in deferred_names:
|
||||
result.append({**tool, "defer_loading": True})
|
||||
else:
|
||||
result.append(tool)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Citation formatting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def format_citations(content: str, annotations: list[Any]) -> str:
|
||||
"""Append url_citation sources as footnotes at the end of the content."""
|
||||
seen_urls: set[str] = set()
|
||||
sources: list[str] = []
|
||||
for ann in annotations:
|
||||
ann_type = getattr(ann, "type", None)
|
||||
if ann_type == "url_citation":
|
||||
title: str = ""
|
||||
url: str = ""
|
||||
citation = getattr(ann, "url_citation", None)
|
||||
if citation is not None:
|
||||
# Chat Completions API: nested url_citation object
|
||||
title = getattr(citation, "title", "") or ""
|
||||
url = getattr(citation, "url", "") or ""
|
||||
elif hasattr(ann, "url") and isinstance(getattr(ann, "url", None), str):
|
||||
# Responses API: attributes directly on the annotation
|
||||
title = getattr(ann, "title", "") or ""
|
||||
url = getattr(ann, "url", "") or ""
|
||||
if url and url not in seen_urls:
|
||||
seen_urls.add(url)
|
||||
sources.append(f"[{title}]({url})" if title else url)
|
||||
if sources:
|
||||
content += "\n\nSources:\n" + "\n".join(f"- {s}" for s in sources)
|
||||
return content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Message sanitisation (Chat Completions specific but shared for compat)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def sanitize_messages(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Ensure assistant messages always have ``content`` or ``tool_calls``.
|
||||
|
||||
OpenAI-compatible APIs reject assistant messages that have neither.
|
||||
This is a defensive catch-all; the upstream layers should already
|
||||
guarantee well-formed messages.
|
||||
"""
|
||||
out: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if (
|
||||
msg.get("role") == "assistant"
|
||||
and msg.get("content") is None
|
||||
and not msg.get("tool_calls")
|
||||
):
|
||||
msg = {**msg, "content": ""}
|
||||
out.append(msg)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Usage extraction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def extract_usage(usage_obj: Any) -> UsageInfo | None:
|
||||
"""Normalize usage from either Chat Completions or Responses API.
|
||||
|
||||
Chat Completions uses ``prompt_tokens`` / ``completion_tokens``.
|
||||
Responses API uses ``input_tokens`` / ``output_tokens``.
|
||||
We check for each in order, preferring the real SDK attribute names.
|
||||
"""
|
||||
if usage_obj is None:
|
||||
return None
|
||||
|
||||
# Token counts — prefer Chat Completions names, fall back to Responses API
|
||||
pt = getattr(usage_obj, "prompt_tokens", None)
|
||||
if not isinstance(pt, int):
|
||||
pt = getattr(usage_obj, "input_tokens", None)
|
||||
ct = getattr(usage_obj, "completion_tokens", None)
|
||||
if not isinstance(ct, int):
|
||||
ct = getattr(usage_obj, "output_tokens", None)
|
||||
tt = getattr(usage_obj, "total_tokens", None)
|
||||
if not isinstance(pt, int) or not isinstance(ct, int):
|
||||
return None
|
||||
|
||||
# Cache tokens — Chat Completions: prompt_tokens_details.cached_tokens,
|
||||
# Responses API: input_tokens_details.cached_tokens
|
||||
ptd = getattr(usage_obj, "prompt_tokens_details", None)
|
||||
if ptd is None:
|
||||
ptd = getattr(usage_obj, "input_tokens_details", None)
|
||||
cached = getattr(ptd, "cached_tokens", 0) if ptd is not None else 0
|
||||
|
||||
return UsageInfo(
|
||||
prompt_tokens=pt,
|
||||
completion_tokens=ct,
|
||||
total_tokens=tt if isinstance(tt, int) else (pt + ct),
|
||||
cache_read_tokens=cached if isinstance(cached, int) else 0,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Retryable error names (shared across both OpenAI providers)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
RETRYABLE_ERROR_NAMES: frozenset[str] = frozenset(
|
||||
{
|
||||
"APIError",
|
||||
"APIConnectionError",
|
||||
"RateLimitError",
|
||||
"Timeout",
|
||||
"APITimeoutError",
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,556 @@
|
||||
"""Responses API provider — for commercial OpenAI models (GPT-5.x, O-series).
|
||||
|
||||
Uses the OpenAI Responses API (``/v1/responses``) which natively supports
|
||||
reasoning, tool use, web search, and tool search without the limitations
|
||||
of the Chat Completions endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterator
|
||||
|
||||
import structlog
|
||||
|
||||
from turnstone.core.providers._openai_common import (
|
||||
RETRYABLE_ERROR_NAMES,
|
||||
apply_cache_retention,
|
||||
apply_temperature,
|
||||
apply_tool_search,
|
||||
extract_usage,
|
||||
format_citations,
|
||||
lookup_openai_capabilities,
|
||||
resolve_reasoning_effort,
|
||||
)
|
||||
from turnstone.core.providers._protocol import (
|
||||
CompletionResult,
|
||||
ModelCapabilities,
|
||||
StreamChunk,
|
||||
ToolCallDelta,
|
||||
)
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
def _convert_content_parts(parts: list[Any]) -> list[dict[str, Any]]:
|
||||
"""Convert Chat Completions content parts to Responses API format.
|
||||
|
||||
Handles text and image_url parts. The Responses API uses
|
||||
``input_image`` instead of ``image_url``.
|
||||
"""
|
||||
converted: list[dict[str, Any]] = []
|
||||
for part in parts:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
ptype = part.get("type", "")
|
||||
if ptype == "text":
|
||||
converted.append({"type": "input_text", "text": part.get("text", "")})
|
||||
elif ptype == "image_url":
|
||||
url_data = part.get("image_url", {})
|
||||
url = url_data.get("url", "") if isinstance(url_data, dict) else ""
|
||||
converted.append({"type": "input_image", "image_url": url})
|
||||
else:
|
||||
converted.append(part)
|
||||
return converted
|
||||
|
||||
|
||||
class OpenAIResponsesProvider:
|
||||
"""Provider for commercial OpenAI models via the Responses API.
|
||||
|
||||
Translates between turnstone's internal OpenAI Chat Completions-like
|
||||
message format and the Responses API input/output format.
|
||||
"""
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai"
|
||||
|
||||
def get_capabilities(self, model: str) -> ModelCapabilities:
|
||||
return lookup_openai_capabilities(model)
|
||||
|
||||
# -- message conversion --------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _convert_messages(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> tuple[str | None, list[dict[str, Any]]]:
|
||||
"""Convert Chat Completions messages to Responses API input items.
|
||||
|
||||
Returns ``(instructions, input_items)`` where *instructions* is the
|
||||
concatenated system/developer messages (or ``None``) and *input_items*
|
||||
is the Responses API ``input`` array.
|
||||
"""
|
||||
instructions_parts: list[str] = []
|
||||
items: list[dict[str, Any]] = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content")
|
||||
|
||||
if role in ("system", "developer"):
|
||||
if isinstance(content, str) and content:
|
||||
instructions_parts.append(content)
|
||||
elif isinstance(content, list):
|
||||
# Content parts — extract text
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
instructions_parts.append(part["text"])
|
||||
continue
|
||||
|
||||
if role == "user":
|
||||
item: dict[str, Any] = {"type": "message", "role": "user"}
|
||||
if isinstance(content, str):
|
||||
item["content"] = content
|
||||
elif isinstance(content, list):
|
||||
# Vision: content parts (text + image_url)
|
||||
item["content"] = _convert_content_parts(content)
|
||||
else:
|
||||
item["content"] = content or ""
|
||||
items.append(item)
|
||||
|
||||
elif role == "assistant":
|
||||
# With store=False, provider_blocks cannot be replayed as input
|
||||
# (output format != input format, and IDs aren't persisted).
|
||||
# Rebuild from the normalized content/tool_calls instead.
|
||||
|
||||
# Text content → assistant message (plain string for input)
|
||||
if content:
|
||||
items.append(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
}
|
||||
)
|
||||
|
||||
# Tool calls → function_call items
|
||||
for tc in msg.get("tool_calls") or []:
|
||||
func = tc.get("function", {})
|
||||
items.append(
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": tc.get("id", ""),
|
||||
"name": func.get("name", ""),
|
||||
"arguments": func.get("arguments", ""),
|
||||
}
|
||||
)
|
||||
|
||||
elif role == "tool":
|
||||
# Tool result → function_call_output
|
||||
output = content
|
||||
if isinstance(content, list):
|
||||
# Structured content (e.g. vision) — serialize to string
|
||||
output = json.dumps(content)
|
||||
items.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": msg.get("tool_call_id", ""),
|
||||
"output": output or "",
|
||||
}
|
||||
)
|
||||
|
||||
instructions = "\n\n".join(instructions_parts) if instructions_parts else None
|
||||
return instructions, items
|
||||
|
||||
# -- tool conversion -----------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _convert_tools(
|
||||
tools: list[dict[str, Any]] | None,
|
||||
caps: ModelCapabilities,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
"""Convert Chat Completions tool format to Responses API format.
|
||||
|
||||
Chat Completions: ``{"type": "function", "function": {"name", "description", "parameters"}}``
|
||||
Responses API: ``{"type": "function", "name", "description", "parameters", "strict": false}``
|
||||
|
||||
Also handles web_search injection for models that support it.
|
||||
"""
|
||||
if not tools:
|
||||
return None
|
||||
|
||||
converted: list[dict[str, Any]] = []
|
||||
has_web_search_func = False
|
||||
|
||||
for tool in tools:
|
||||
func = tool.get("function")
|
||||
if not func:
|
||||
converted.append(tool)
|
||||
continue
|
||||
|
||||
name = func.get("name", "")
|
||||
|
||||
# web_search function tool → native web_search_tool
|
||||
if name == "web_search" and caps.supports_web_search:
|
||||
has_web_search_func = True
|
||||
continue
|
||||
|
||||
item: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"name": name,
|
||||
"description": func.get("description", ""),
|
||||
"parameters": func.get("parameters", {}),
|
||||
"strict": False,
|
||||
}
|
||||
# Preserve defer_loading for tool search
|
||||
if tool.get("defer_loading"):
|
||||
item["defer_loading"] = True
|
||||
converted.append(item)
|
||||
|
||||
# Inject native web search tool
|
||||
if has_web_search_func or caps.supports_web_search:
|
||||
converted.append({"type": "web_search"})
|
||||
|
||||
# Responses API requires a tool_search tool when defer_loading is used
|
||||
if any(t.get("defer_loading") for t in converted):
|
||||
converted.append({"type": "tool_search"})
|
||||
|
||||
return converted if converted else None
|
||||
|
||||
# -- parameter building --------------------------------------------------
|
||||
|
||||
def _build_kwargs(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None,
|
||||
max_tokens: int,
|
||||
temperature: float,
|
||||
reasoning_effort: str,
|
||||
deferred_names: frozenset[str] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build the kwargs dict for ``client.responses.create/stream``."""
|
||||
caps = self.get_capabilities(model)
|
||||
|
||||
instructions, input_items = self._convert_messages(messages)
|
||||
tools = apply_tool_search(caps, tools, deferred_names)
|
||||
converted_tools = self._convert_tools(tools, caps)
|
||||
|
||||
# Ensure web search is always injected for search-capable models,
|
||||
# even when no function tools are registered (e.g. creative mode).
|
||||
if caps.supports_web_search:
|
||||
converted_tools = converted_tools or []
|
||||
if not any(t.get("type") == "web_search" for t in converted_tools):
|
||||
converted_tools.append({"type": "web_search"})
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"input": input_items,
|
||||
"max_output_tokens": max_tokens,
|
||||
"store": False,
|
||||
}
|
||||
|
||||
if instructions:
|
||||
kwargs["instructions"] = instructions
|
||||
|
||||
if converted_tools:
|
||||
kwargs["tools"] = converted_tools
|
||||
|
||||
apply_temperature(kwargs, caps, temperature, reasoning_effort)
|
||||
|
||||
# Reasoning effort → {"effort": value} dict (Responses API format)
|
||||
effort = resolve_reasoning_effort(caps, reasoning_effort)
|
||||
if effort:
|
||||
kwargs["reasoning"] = {"effort": effort}
|
||||
|
||||
apply_cache_retention(kwargs, model)
|
||||
return kwargs
|
||||
|
||||
# -- streaming -----------------------------------------------------------
|
||||
|
||||
def create_streaming(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
if extra_params:
|
||||
log.debug("openai.responses: extra_params ignored (not supported by Responses API)")
|
||||
kwargs = self._build_kwargs(
|
||||
model,
|
||||
messages,
|
||||
tools,
|
||||
max_tokens,
|
||||
temperature,
|
||||
reasoning_effort,
|
||||
deferred_names,
|
||||
)
|
||||
kwargs["stream"] = True
|
||||
|
||||
log.debug(
|
||||
"openai.responses.request",
|
||||
model=model,
|
||||
stream=True,
|
||||
max_tokens=max_tokens,
|
||||
input_items=len(kwargs.get("input", [])),
|
||||
tool_count=len(kwargs.get("tools", [])),
|
||||
)
|
||||
|
||||
stream = client.responses.create(**kwargs)
|
||||
if cancel_ref is not None:
|
||||
cancel_ref.append(stream)
|
||||
return self._iter_stream(stream)
|
||||
|
||||
def _iter_stream(self, stream: Any) -> Iterator[StreamChunk]:
|
||||
"""Convert Responses API stream events to StreamChunks."""
|
||||
first = True
|
||||
content_len = 0
|
||||
tool_call_count = 0
|
||||
last_finish: str | None = None
|
||||
completion_tokens: int | None = None
|
||||
# Track tool call indices by call_id for consistent ToolCallDelta.index
|
||||
tool_call_indices: dict[str, int] = {}
|
||||
# Collect output items for provider_blocks
|
||||
provider_blocks: list[dict[str, Any]] = []
|
||||
# Collect annotations across text parts
|
||||
annotations: list[Any] = []
|
||||
|
||||
for event in stream:
|
||||
event_type = getattr(event, "type", "")
|
||||
|
||||
# -- text content deltas --
|
||||
if event_type == "response.output_text.delta":
|
||||
delta_text = getattr(event, "delta", "")
|
||||
if delta_text:
|
||||
sc = StreamChunk(content_delta=delta_text)
|
||||
content_len += len(delta_text)
|
||||
if first:
|
||||
sc.is_first = True
|
||||
first = False
|
||||
yield sc
|
||||
continue
|
||||
|
||||
# -- reasoning deltas --
|
||||
if event_type in (
|
||||
"response.reasoning_text.delta",
|
||||
"response.reasoning_summary_text.delta",
|
||||
):
|
||||
delta_text = getattr(event, "delta", "")
|
||||
if delta_text:
|
||||
sc = StreamChunk(reasoning_delta=delta_text)
|
||||
if first:
|
||||
sc.is_first = True
|
||||
first = False
|
||||
yield sc
|
||||
continue
|
||||
|
||||
# -- new tool call (function_call output item added) --
|
||||
if event_type == "response.output_item.added":
|
||||
item = getattr(event, "item", None)
|
||||
if item and getattr(item, "type", "") == "function_call":
|
||||
call_id = getattr(item, "call_id", "")
|
||||
item_id = getattr(item, "id", "")
|
||||
name = getattr(item, "name", "")
|
||||
idx = len(tool_call_indices)
|
||||
# Index by item_id — argument deltas reference this, not call_id
|
||||
tool_call_indices[item_id] = idx
|
||||
sc = StreamChunk(
|
||||
tool_call_deltas=[ToolCallDelta(index=idx, id=call_id, name=name)]
|
||||
)
|
||||
tool_call_count += 1
|
||||
if first:
|
||||
sc.is_first = True
|
||||
first = False
|
||||
yield sc
|
||||
continue
|
||||
|
||||
# -- tool call argument deltas --
|
||||
if event_type == "response.function_call_arguments.delta":
|
||||
item_id = getattr(event, "item_id", "")
|
||||
delta_args = getattr(event, "delta", "")
|
||||
if delta_args:
|
||||
idx = tool_call_indices.get(item_id, 0)
|
||||
yield StreamChunk(
|
||||
tool_call_deltas=[ToolCallDelta(index=idx, arguments_delta=delta_args)]
|
||||
)
|
||||
continue
|
||||
|
||||
# -- web search status --
|
||||
if event_type == "response.web_search_call.searching":
|
||||
yield StreamChunk(info_delta="[Searching…]")
|
||||
continue
|
||||
if event_type == "response.web_search_call.completed":
|
||||
yield StreamChunk(info_delta="[Search complete]")
|
||||
continue
|
||||
|
||||
# -- output item done (capture for provider_blocks) --
|
||||
if event_type == "response.output_item.done":
|
||||
item = getattr(event, "item", None)
|
||||
if item:
|
||||
item_dict = item.model_dump() if hasattr(item, "model_dump") else {}
|
||||
if item_dict:
|
||||
provider_blocks.append(item_dict)
|
||||
# Collect annotations from completed text parts
|
||||
if getattr(item, "type", "") == "message":
|
||||
for content_part in getattr(item, "content", []):
|
||||
part_anns = getattr(content_part, "annotations", None)
|
||||
if part_anns:
|
||||
annotations.extend(part_anns)
|
||||
continue
|
||||
|
||||
# -- response completed --
|
||||
if event_type == "response.completed":
|
||||
response = getattr(event, "response", None)
|
||||
if response:
|
||||
status = getattr(response, "status", "completed")
|
||||
last_finish = "stop" if status == "completed" else "length"
|
||||
usage = extract_usage(getattr(response, "usage", None))
|
||||
if usage:
|
||||
completion_tokens = usage.completion_tokens
|
||||
sc = StreamChunk(
|
||||
finish_reason=last_finish,
|
||||
usage=usage,
|
||||
)
|
||||
if provider_blocks:
|
||||
sc.provider_blocks = provider_blocks
|
||||
yield sc
|
||||
continue
|
||||
|
||||
# -- error --
|
||||
if event_type == "response.failed":
|
||||
response = getattr(event, "response", None)
|
||||
error = getattr(response, "error", None) if response else None
|
||||
error_msg = getattr(error, "message", "Unknown error") if error else "Unknown error"
|
||||
raise RuntimeError(f"Responses API error: {error_msg}")
|
||||
|
||||
log.debug(
|
||||
"openai.responses.response",
|
||||
stream=True,
|
||||
finish_reason=last_finish,
|
||||
content_length=content_len,
|
||||
tool_call_count=tool_call_count,
|
||||
completion_tokens=completion_tokens,
|
||||
)
|
||||
|
||||
# Emit accumulated citations as a final info chunk
|
||||
if annotations:
|
||||
citation_text = format_citations("", annotations).strip()
|
||||
if citation_text:
|
||||
yield StreamChunk(info_delta=citation_text)
|
||||
|
||||
# -- non-streaming -------------------------------------------------------
|
||||
|
||||
def create_completion(
|
||||
self,
|
||||
*,
|
||||
client: Any,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.5,
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
) -> CompletionResult:
|
||||
if extra_params:
|
||||
log.debug("openai.responses: extra_params ignored (not supported by Responses API)")
|
||||
kwargs = self._build_kwargs(
|
||||
model,
|
||||
messages,
|
||||
tools,
|
||||
max_tokens,
|
||||
temperature,
|
||||
reasoning_effort,
|
||||
deferred_names,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
"openai.responses.request",
|
||||
model=model,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
input_items=len(kwargs.get("input", [])),
|
||||
tool_count=len(kwargs.get("tools", [])),
|
||||
)
|
||||
|
||||
response = client.responses.create(**kwargs)
|
||||
return self._parse_response(response)
|
||||
|
||||
def _parse_response(self, response: Any) -> CompletionResult:
|
||||
"""Convert a Responses API ``Response`` object to ``CompletionResult``."""
|
||||
content_parts: list[str] = []
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
provider_blocks: list[dict[str, Any]] = []
|
||||
all_annotations: list[Any] = []
|
||||
|
||||
for item in getattr(response, "output", []):
|
||||
item_type = getattr(item, "type", "")
|
||||
|
||||
if item_type == "message":
|
||||
for content_part in getattr(item, "content", []):
|
||||
part_type = getattr(content_part, "type", "")
|
||||
if part_type == "output_text":
|
||||
content_parts.append(getattr(content_part, "text", ""))
|
||||
anns = getattr(content_part, "annotations", None)
|
||||
if anns:
|
||||
all_annotations.extend(anns)
|
||||
elif part_type == "refusal":
|
||||
content_parts.append(f"[Refused: {getattr(content_part, 'refusal', '')}]")
|
||||
|
||||
elif item_type == "function_call":
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": getattr(item, "call_id", ""),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": getattr(item, "name", ""),
|
||||
"arguments": getattr(item, "arguments", ""),
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Capture all output items for provider_blocks (multi-turn)
|
||||
item_dict = item.model_dump() if hasattr(item, "model_dump") else {}
|
||||
if item_dict:
|
||||
provider_blocks.append(item_dict)
|
||||
|
||||
content = "".join(content_parts)
|
||||
if all_annotations:
|
||||
content = format_citations(content, all_annotations)
|
||||
|
||||
status = getattr(response, "status", "completed")
|
||||
finish_reason = "stop" if status == "completed" else "length"
|
||||
usage = extract_usage(getattr(response, "usage", None))
|
||||
|
||||
result = CompletionResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls if tool_calls else None,
|
||||
finish_reason=finish_reason,
|
||||
usage=usage,
|
||||
provider_blocks=provider_blocks,
|
||||
)
|
||||
log.debug(
|
||||
"openai.responses.response",
|
||||
stream=False,
|
||||
finish_reason=finish_reason,
|
||||
content_length=len(content),
|
||||
tool_call_count=len(tool_calls),
|
||||
completion_tokens=usage.completion_tokens if usage else None,
|
||||
)
|
||||
return result
|
||||
|
||||
# -- tool conversion (public interface) ----------------------------------
|
||||
|
||||
def convert_tools(
|
||||
self,
|
||||
tools: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
return tools # Conversion happens internally in _build_kwargs
|
||||
|
||||
# -- retryable errors ----------------------------------------------------
|
||||
|
||||
@property
|
||||
def retryable_error_names(self) -> frozenset[str]:
|
||||
return RETRYABLE_ERROR_NAMES
|
||||
@@ -236,14 +236,14 @@ def _math_exec_in_process(code: str, result_queue: multiprocessing.Queue[tuple[s
|
||||
):
|
||||
ns[name] = getattr(sympy, name)
|
||||
except ImportError:
|
||||
pass
|
||||
pass # optional dependency
|
||||
|
||||
try:
|
||||
import numpy as _np
|
||||
|
||||
ns["np"] = ns["numpy"] = _np
|
||||
except ImportError:
|
||||
pass
|
||||
pass # optional dependency
|
||||
|
||||
try:
|
||||
import scipy # type: ignore[import-untyped]
|
||||
@@ -260,7 +260,7 @@ def _math_exec_in_process(code: str, result_queue: multiprocessing.Queue[tuple[s
|
||||
ns["gamma"] = scipy.special.gamma
|
||||
ns["beta"] = scipy.special.beta
|
||||
except ImportError:
|
||||
pass
|
||||
pass # optional dependency
|
||||
|
||||
# Strip __builtins__ from all pre-imported modules so
|
||||
# module.__builtins__['__import__'] can't bypass _safe_import.
|
||||
|
||||
+341
-78
@@ -19,6 +19,7 @@ import mimetypes
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import tempfile
|
||||
@@ -161,6 +162,13 @@ _IMAGE_SIZE_CAP: int = 4 * 1024 * 1024
|
||||
# Upper bound on total skill content injected into system messages
|
||||
_MAX_SKILL_CONTENT: int = 32768
|
||||
|
||||
# Matches resource paths referenced in skill content (scripts/foo.py, etc.)
|
||||
_RESOURCE_PATH_RE = re.compile(
|
||||
r"(?<![/\w-])(?:scripts|references|assets)/[\w./-]+\."
|
||||
r"(?:json|yaml|yml|toml|cfg|ini|py|sh|js|ts|md|txt)"
|
||||
r"(?=[\s)\]}'\"`,;:\x60]|$)"
|
||||
)
|
||||
|
||||
|
||||
_TEMPLATE_VAR_RE = re.compile(r"\{\{(\w+)\}\}")
|
||||
|
||||
@@ -304,7 +312,7 @@ class ChatSession:
|
||||
self._provider: LLMProvider = (
|
||||
registry.get_provider(model_alias)
|
||||
if registry and model_alias
|
||||
else create_provider("openai")
|
||||
else create_provider("openai-compatible")
|
||||
)
|
||||
self._cached_capabilities: ModelCapabilities | None = None
|
||||
self.ui = ui
|
||||
@@ -417,6 +425,7 @@ class ChatSession:
|
||||
self._skill_name: str | None = skill
|
||||
self._skill_content: str | None = None
|
||||
self._skill_resources: dict[str, str] = {}
|
||||
self._skill_resources_dir: str | None = None
|
||||
self._load_skills()
|
||||
self._init_system_messages()
|
||||
self._save_config()
|
||||
@@ -583,6 +592,8 @@ class ChatSession:
|
||||
else:
|
||||
self._skill_content = None
|
||||
self._skill_resources = {}
|
||||
self._materialize_skill_resources()
|
||||
self._validate_skill_resources()
|
||||
|
||||
def set_skill(self, name: str | None) -> None:
|
||||
"""Set or clear the active skill."""
|
||||
@@ -613,6 +624,81 @@ class ChatSession:
|
||||
log.warning("skill_resources.load_failed", skill_id=skill_id, exc_info=True)
|
||||
return {}
|
||||
|
||||
def _cleanup_skill_resources(self) -> None:
|
||||
"""Remove materialized skill resources from disk."""
|
||||
d = self._skill_resources_dir
|
||||
if d is not None:
|
||||
shutil.rmtree(d, ignore_errors=True)
|
||||
self._skill_resources_dir = None
|
||||
|
||||
def _materialize_skill_resources(self) -> None:
|
||||
"""Write skill resources to a temp directory for subprocess access."""
|
||||
self._cleanup_skill_resources()
|
||||
if not self._skill_resources:
|
||||
return
|
||||
base = tempfile.mkdtemp(prefix=f"skill-{self._ws_id[:8]}-")
|
||||
written = 0
|
||||
for rel_path, content in self._skill_resources.items():
|
||||
normed = os.path.normpath(rel_path)
|
||||
if not normed or normed == "." or normed.startswith(("..", "/")):
|
||||
log.warning("skill_resources.bad_path", path=rel_path)
|
||||
continue
|
||||
if ".." in normed.split(os.sep):
|
||||
log.warning("skill_resources.bad_path", path=rel_path)
|
||||
continue
|
||||
full = os.path.join(base, normed)
|
||||
if not os.path.realpath(full).startswith(os.path.realpath(base)):
|
||||
log.warning("skill_resources.path_escape", path=rel_path)
|
||||
continue
|
||||
try:
|
||||
os.makedirs(os.path.dirname(full), exist_ok=True)
|
||||
with open(full, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
if normed.startswith("scripts/"):
|
||||
os.chmod(full, 0o755)
|
||||
written += 1
|
||||
except Exception:
|
||||
log.warning("skill_resources.write_failed", path=rel_path, exc_info=True)
|
||||
if written == 0:
|
||||
shutil.rmtree(base, ignore_errors=True)
|
||||
return
|
||||
self._skill_resources_dir = base
|
||||
log.info(
|
||||
"skill_resources.materialized",
|
||||
dir=base,
|
||||
count=written,
|
||||
)
|
||||
|
||||
def _skill_resource_env(self) -> dict[str, str]:
|
||||
"""Return extra env vars for bash when skill resources are materialized."""
|
||||
if not self._skill_resources_dir:
|
||||
return {}
|
||||
env: dict[str, str] = {"SKILL_RESOURCES_DIR": self._skill_resources_dir}
|
||||
scripts_dir = os.path.join(self._skill_resources_dir, "scripts")
|
||||
if os.path.isdir(scripts_dir):
|
||||
current_path = os.environ.get("PATH")
|
||||
if current_path:
|
||||
env["PATH"] = scripts_dir + os.pathsep + current_path
|
||||
else:
|
||||
env["PATH"] = scripts_dir
|
||||
return env
|
||||
|
||||
def _validate_skill_resources(self) -> None:
|
||||
"""Warn if skill content references resource paths not in skill_resources."""
|
||||
if not self._skill_content or not self._skill_name:
|
||||
return
|
||||
referenced = {os.path.normpath(p) for p in _RESOURCE_PATH_RE.findall(self._skill_content)}
|
||||
if not referenced:
|
||||
return
|
||||
available = {os.path.normpath(p) for p in self._skill_resources}
|
||||
missing = sorted(referenced - available)
|
||||
if missing:
|
||||
log.warning("skill_resources.missing", skill=self._skill_name, paths=missing)
|
||||
self.ui.on_info(
|
||||
f"Skill '{self._skill_name}' references {len(missing)} resource(s) "
|
||||
f"not bundled: {', '.join(missing)}"
|
||||
)
|
||||
|
||||
# -- MCP tool refresh ----------------------------------------------------
|
||||
|
||||
def _on_mcp_tools_changed(self) -> None:
|
||||
@@ -747,6 +833,7 @@ class ChatSession:
|
||||
self._mcp_prompt_cb = None
|
||||
if self._watch_runner:
|
||||
self._watch_runner.remove_dispatch_fn(self._ws_id)
|
||||
self._cleanup_skill_resources()
|
||||
|
||||
def _handle_mcp_refresh(self, arg: str) -> None:
|
||||
"""Handle ``/mcp refresh [server]``."""
|
||||
@@ -799,9 +886,34 @@ class ChatSession:
|
||||
self._tool_error_flags[call_id] = True
|
||||
self.ui.on_tool_result(call_id, name, output, is_error=is_error)
|
||||
|
||||
def _truncate_output(self, output: str) -> str:
|
||||
"""Truncate tool output to self.tool_truncation chars, keeping head + tail."""
|
||||
def _remaining_token_budget(self) -> int:
|
||||
"""Estimate how many tokens are available for new content.
|
||||
|
||||
Reserves a response budget (capped at 25% of context window, since
|
||||
``max_tokens`` is an upper bound, not guaranteed consumption) plus
|
||||
a 5% safety margin. Returns at least 0.
|
||||
"""
|
||||
used = self._system_tokens + sum(self._msg_tokens)
|
||||
response_reserve = min(self.max_tokens, self.context_window // 4)
|
||||
safety_margin = int(self.context_window * 0.05)
|
||||
return max(0, self.context_window - used - response_reserve - safety_margin)
|
||||
|
||||
def _truncate_output(self, output: str, remaining_budget_tokens: int | None = None) -> str:
|
||||
"""Truncate tool output, keeping head + tail.
|
||||
|
||||
The effective limit is the *minimum* of:
|
||||
- ``self.tool_truncation`` (fixed cap, defaults to 50% of context)
|
||||
- ``remaining_budget_tokens`` converted to chars (if provided)
|
||||
|
||||
This ensures a single tool result cannot overflow the context window
|
||||
even when the conversation is already partially full.
|
||||
"""
|
||||
limit = self.tool_truncation
|
||||
if remaining_budget_tokens is not None:
|
||||
budget_chars = int(remaining_budget_tokens * self._chars_per_token)
|
||||
limit = min(limit, budget_chars)
|
||||
if limit <= 0:
|
||||
return f"[Output truncated — {len(output)} chars exceeded context budget]"
|
||||
if len(output) <= limit:
|
||||
return output
|
||||
half = limit // 2
|
||||
@@ -837,10 +949,8 @@ class ChatSession:
|
||||
if asst_msg:
|
||||
snippet += f"\nAssistant: {asst_msg}"
|
||||
snippet += "\n\nTitle:"
|
||||
result = self._provider.create_completion(
|
||||
client=self.client,
|
||||
model=self.model,
|
||||
messages=[
|
||||
result = self._utility_completion(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
@@ -855,9 +965,6 @@ class ChatSession:
|
||||
{"role": "user", "content": snippet},
|
||||
],
|
||||
max_tokens=200,
|
||||
temperature=0.3,
|
||||
reasoning_effort="low",
|
||||
extra_params=self._provider_extra_params(reasoning_effort="low"),
|
||||
)
|
||||
raw = (result.content or "").strip()
|
||||
# Take first line, strip quotes
|
||||
@@ -893,6 +1000,13 @@ class ChatSession:
|
||||
self._msg_tokens = [
|
||||
max(1, int(self._msg_char_count(m) / self._chars_per_token)) for m in self.messages
|
||||
]
|
||||
log.info(
|
||||
"Resuming ws=%s: %d messages, provider=%s, model=%s",
|
||||
ws_id,
|
||||
len(messages),
|
||||
type(self._provider).__name__,
|
||||
self.model,
|
||||
)
|
||||
# Restore persisted config
|
||||
config = load_workstream_config(ws_id)
|
||||
if config:
|
||||
@@ -910,11 +1024,24 @@ class ChatSession:
|
||||
self.context_window = cfg.context_window
|
||||
if not self._manual_tool_truncation:
|
||||
self.tool_truncation = int(cfg.context_window * self._chars_per_token * 0.5)
|
||||
log.info(
|
||||
"Resume: resolved alias=%s → provider=%s, model=%s, ctx=%d",
|
||||
saved_alias,
|
||||
type(self._provider).__name__,
|
||||
model_name,
|
||||
cfg.context_window,
|
||||
)
|
||||
elif saved_model and saved_model != self.model:
|
||||
# No alias or alias no longer in registry — at least set the model name
|
||||
self.model = saved_model
|
||||
self._model_alias = None
|
||||
self._cached_capabilities = None
|
||||
log.warning(
|
||||
"Resume: alias %r not in registry, keeping default provider=%s for model=%s",
|
||||
saved_alias,
|
||||
type(self._provider).__name__,
|
||||
saved_model,
|
||||
)
|
||||
if "temperature" in config:
|
||||
self.temperature = float(config["temperature"])
|
||||
if "reasoning_effort" in config:
|
||||
@@ -969,29 +1096,38 @@ class ChatSession:
|
||||
self._chat_template_kwargs_base: dict[str, Any] = {
|
||||
"reasoning_effort": self.reasoning_effort,
|
||||
}
|
||||
self._chat_template_kwargs: dict[str, Any] = dict(self._chat_template_kwargs_base)
|
||||
|
||||
# -- Developer message --
|
||||
if self.creative_mode:
|
||||
dev_parts = [
|
||||
"# Instructions",
|
||||
"",
|
||||
"You are a creative writing partner. Use the analysis channel to "
|
||||
"think through structure, voice, and intent before drafting.",
|
||||
(
|
||||
"You are a creative writing partner. Use the analysis channel to "
|
||||
"think through structure, voice, and intent before drafting."
|
||||
),
|
||||
"",
|
||||
"Craft principles:",
|
||||
"- Ground scenes in concrete sensory detail — what is seen, heard, felt.",
|
||||
"- Vary rhythm. Short sentences hit hard. Longer ones carry the reader "
|
||||
"through texture and nuance, building toward something.",
|
||||
"- Dialogue should do at least two things: reveal character AND advance "
|
||||
"plot or tension. Cut anything that's just exchanging information.",
|
||||
"- Earn your abstractions. Don't say 'she felt sad' — show the thing "
|
||||
"that makes the reader feel it.",
|
||||
(
|
||||
"- Vary rhythm. Short sentences hit hard. Longer ones carry the reader "
|
||||
"through texture and nuance, building toward something."
|
||||
),
|
||||
(
|
||||
"- Dialogue should do at least two things: reveal character AND advance "
|
||||
"plot or tension. Cut anything that's just exchanging information."
|
||||
),
|
||||
(
|
||||
"- Earn your abstractions. Don't say 'she felt sad' — show the thing "
|
||||
"that makes the reader feel it."
|
||||
),
|
||||
"- Trust subtext. Leave room for the reader.",
|
||||
"",
|
||||
"Match the user's genre and tone. If they want literary fiction, write "
|
||||
"literary fiction. If they want pulp, write pulp with conviction. "
|
||||
"Never condescend to the form.",
|
||||
(
|
||||
"Match the user's genre and tone. If they want literary fiction, write "
|
||||
"literary fiction. If they want pulp, write pulp with conviction. "
|
||||
"Never condescend to the form."
|
||||
),
|
||||
]
|
||||
else:
|
||||
# Compose system message from modular components
|
||||
@@ -1003,7 +1139,7 @@ class ChatSession:
|
||||
if storage:
|
||||
db_policies = storage.list_prompt_policies()
|
||||
except Exception:
|
||||
pass
|
||||
log.debug("Failed to load prompt policies from storage", exc_info=True)
|
||||
now = datetime.now().astimezone()
|
||||
ctx = SessionContext(
|
||||
current_datetime=now.strftime("%Y-%m-%dT%H:%M"),
|
||||
@@ -1091,6 +1227,12 @@ class ChatSession:
|
||||
"Resource content omitted (total exceeds 8KB). "
|
||||
"Resource files are listed above by path and size."
|
||||
)
|
||||
if self._skill_resources_dir:
|
||||
lines.append(
|
||||
"\nResource files are materialized on disk. "
|
||||
"Scripts in scripts/ are on PATH and can be run by name. "
|
||||
"All files are under $SKILL_RESOURCES_DIR."
|
||||
)
|
||||
lines.append("</skill-resources>")
|
||||
dev_parts.append("\n".join(lines))
|
||||
# Skill catalog: disclose search-activated skills so the model
|
||||
@@ -1159,15 +1301,47 @@ class ChatSession:
|
||||
reasoning_effort: str | None = None,
|
||||
provider: LLMProvider | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Build provider-specific extra parameters."""
|
||||
"""Build provider-specific extra parameters.
|
||||
|
||||
``chat_template_kwargs`` is only meaningful for local model servers
|
||||
(``openai-compatible``). Commercial OpenAI rejects it as an unknown
|
||||
parameter, and handles ``reasoning_effort`` natively.
|
||||
"""
|
||||
prov = provider or self._provider
|
||||
if prov.provider_name == "openai":
|
||||
if prov.provider_name == "openai-compatible":
|
||||
kwargs = dict(self._chat_template_kwargs_base)
|
||||
if reasoning_effort:
|
||||
kwargs["reasoning_effort"] = reasoning_effort
|
||||
return {"chat_template_kwargs": kwargs}
|
||||
return None
|
||||
|
||||
def _utility_completion(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
*,
|
||||
max_tokens: int = 4096,
|
||||
temperature: float = 0.3,
|
||||
reasoning_effort: str = "low",
|
||||
) -> CompletionResult:
|
||||
"""Run a lightweight internal completion (title gen, compaction, extraction).
|
||||
|
||||
Threads ``reasoning_effort`` through both the direct keyword (for
|
||||
commercial providers) and ``extra_params`` (for local model servers)
|
||||
so callers don't need to duplicate it. ``max_tokens`` is clamped to
|
||||
the model's advertised output limit so small models don't error.
|
||||
"""
|
||||
caps = self._get_capabilities()
|
||||
clamped = min(max_tokens, caps.max_output_tokens) if caps.max_output_tokens else max_tokens
|
||||
return self._provider.create_completion(
|
||||
client=self.client,
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
max_tokens=clamped,
|
||||
temperature=temperature,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_params=self._provider_extra_params(reasoning_effort=reasoning_effort),
|
||||
)
|
||||
|
||||
# -- tool search helpers --------------------------------------------------
|
||||
|
||||
def _get_active_tools(self) -> list[dict[str, Any]] | None:
|
||||
@@ -1243,11 +1417,9 @@ class ChatSession:
|
||||
if self._health_monitor:
|
||||
self._health_monitor.record_success()
|
||||
return result
|
||||
except BaseException as primary_err:
|
||||
except Exception as primary_err:
|
||||
if self._health_monitor:
|
||||
self._health_monitor.record_failure()
|
||||
if isinstance(primary_err, (KeyboardInterrupt, SystemExit)):
|
||||
raise
|
||||
if not self._registry or not self._registry.fallback:
|
||||
raise
|
||||
# Try each fallback model. Fallbacks may use different backends;
|
||||
@@ -1275,6 +1447,21 @@ class ChatSession:
|
||||
) -> Iterator[StreamChunk]:
|
||||
"""Attempt a streaming API call with retries on transient errors."""
|
||||
prov = provider or self._provider
|
||||
raw_url = str(getattr(client, "base_url", getattr(client, "_base_url", "?")))
|
||||
safe_url = raw_url.split("?")[0] # strip query params (may contain keys)
|
||||
msg_count = len(msgs)
|
||||
role_counts: dict[str, int] = {}
|
||||
for m in msgs:
|
||||
r = m.get("role", "?")
|
||||
role_counts[r] = role_counts.get(r, 0) + 1
|
||||
log.debug(
|
||||
"API call: provider=%s model=%s base_url=%s msgs=%d roles=%s",
|
||||
type(prov).__name__,
|
||||
model,
|
||||
safe_url,
|
||||
msg_count,
|
||||
role_counts,
|
||||
)
|
||||
last_err: Exception | None = None
|
||||
for attempt in range(self._MAX_RETRIES + 1):
|
||||
self._check_cancelled()
|
||||
@@ -1294,6 +1481,29 @@ class ChatSession:
|
||||
)
|
||||
except Exception as e:
|
||||
ename = type(e).__name__
|
||||
cause_name = (
|
||||
type(e.__cause__).__name__
|
||||
if e.__cause__
|
||||
else (type(e.__context__).__name__ if e.__context__ else "None")
|
||||
)
|
||||
log.warning(
|
||||
"API error (attempt %d/%d): %s (cause=%s) "
|
||||
"provider=%s model=%s base_url=%s msgs=%d",
|
||||
attempt + 1,
|
||||
self._MAX_RETRIES + 1,
|
||||
ename,
|
||||
cause_name,
|
||||
type(prov).__name__,
|
||||
model,
|
||||
safe_url,
|
||||
msg_count,
|
||||
)
|
||||
log.debug(
|
||||
"API error details (attempt %d/%d)",
|
||||
attempt + 1,
|
||||
self._MAX_RETRIES + 1,
|
||||
exc_info=True,
|
||||
)
|
||||
if ename not in prov.retryable_error_names or attempt == self._MAX_RETRIES:
|
||||
raise
|
||||
last_err = e
|
||||
@@ -1395,7 +1605,44 @@ class ChatSession:
|
||||
self._emit_state("thinking")
|
||||
self.ui.on_thinking_start()
|
||||
try:
|
||||
stream = self._create_stream_with_retry(msgs)
|
||||
try:
|
||||
stream = self._create_stream_with_retry(msgs)
|
||||
except Exception as ctx_err:
|
||||
# Context overflow recovery: if the API rejects the
|
||||
# request due to exceeding the context window, compact
|
||||
# the conversation and retry once.
|
||||
err_text = str(ctx_err).lower()
|
||||
is_ctx_overflow = any(
|
||||
s in err_text
|
||||
for s in (
|
||||
"context length",
|
||||
"maximum context",
|
||||
"too many tokens",
|
||||
"prompt is too long",
|
||||
"input tokens",
|
||||
)
|
||||
)
|
||||
if not is_ctx_overflow:
|
||||
raise
|
||||
log.warning(
|
||||
"Context overflow detected (%s), compacting and retrying",
|
||||
type(ctx_err).__name__,
|
||||
)
|
||||
self.ui.on_info("\n[Context overflow — auto-compacting and retrying]")
|
||||
# Stop thinking indicator before compact (which has
|
||||
# its own thinking start/stop) to avoid nested spinners.
|
||||
self.ui.on_thinking_stop()
|
||||
try:
|
||||
self._compact_messages(auto=True)
|
||||
msgs = self._full_messages()
|
||||
self.ui.on_thinking_start()
|
||||
stream = self._create_stream_with_retry(msgs)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"Compact-and-retry failed, raising original error",
|
||||
exc_info=True,
|
||||
)
|
||||
raise ctx_err from None
|
||||
assistant_msg = self._stream_response(stream, my_generation)
|
||||
finally:
|
||||
# Only clear if this generation is still active —
|
||||
@@ -1561,6 +1808,12 @@ class ChatSession:
|
||||
tc_id, p["text"], _tc_names.get(tc_id, "")
|
||||
)
|
||||
|
||||
# Safety truncation: clamp output to remaining context budget
|
||||
# so a single large result cannot overflow the context window.
|
||||
if isinstance(output, str):
|
||||
budget = self._remaining_token_budget()
|
||||
output = self._truncate_output(output, remaining_budget_tokens=budget)
|
||||
|
||||
tool_msg: dict[str, Any] = {
|
||||
"role": "tool",
|
||||
"tool_call_id": tc_id,
|
||||
@@ -2184,7 +2437,8 @@ class ChatSession:
|
||||
"""Emit status info via the UI."""
|
||||
if not self._last_usage:
|
||||
return
|
||||
self.ui.on_status(self._last_usage, self.context_window, self.reasoning_effort)
|
||||
usage: dict[str, Any] = {**self._last_usage, "model": self.model}
|
||||
self.ui.on_status(usage, self.context_window, self.reasoning_effort)
|
||||
|
||||
# -- Conversation compaction ------------------------------------------------
|
||||
|
||||
@@ -2330,14 +2584,9 @@ class ChatSession:
|
||||
result: CompletionResult | None = None
|
||||
for attempt in range(self._MAX_RETRIES + 1):
|
||||
try:
|
||||
result = self._provider.create_completion(
|
||||
client=self.client,
|
||||
model=self.model,
|
||||
messages=summary_msgs,
|
||||
result = self._utility_completion(
|
||||
summary_msgs,
|
||||
max_tokens=summary_max_tokens,
|
||||
temperature=0.3,
|
||||
reasoning_effort="low",
|
||||
extra_params=self._provider_extra_params(reasoning_effort="low"),
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
@@ -2416,7 +2665,6 @@ class ChatSession:
|
||||
return None
|
||||
if self._judge is not None:
|
||||
return self._judge
|
||||
return None
|
||||
# Frozen config required for IntentJudge init (LLM client fields).
|
||||
# _judge_cfg already returns None when _judge_config is None, but
|
||||
# this guard makes the dependency explicit for type narrowing.
|
||||
@@ -2635,7 +2883,8 @@ class ChatSession:
|
||||
continue
|
||||
|
||||
cid, output = results[i]
|
||||
assert isinstance(output, str) # plan always returns text
|
||||
if not isinstance(output, str):
|
||||
raise TypeError(f"plan_agent must return str, got {type(output).__name__}")
|
||||
plan_path = f".plan-{self._ws_id}.md"
|
||||
|
||||
if not self.auto_approve:
|
||||
@@ -2688,7 +2937,10 @@ class ChatSession:
|
||||
with open(plan_path, "w") as f:
|
||||
f.write(output)
|
||||
except OSError:
|
||||
pass
|
||||
log.warning("Failed to write plan to %s", plan_path, exc_info=True)
|
||||
output += "\n\n---\nPlan could not be saved to disk."
|
||||
results[i] = (cid, output)
|
||||
continue
|
||||
|
||||
# Always include file path in the tool result so the
|
||||
# outer model knows where the plan lives on disk.
|
||||
@@ -2767,7 +3019,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": func_name,
|
||||
"header": f"\u2717 {func_name}: {exc}",
|
||||
"preview": f" {RED}{preview}{RESET}",
|
||||
"preview": f" {preview}",
|
||||
"needs_approval": False,
|
||||
"error": (
|
||||
f"JSON parse error for tool '{func_name}': {exc}\n"
|
||||
@@ -3157,9 +3409,10 @@ class ChatSession:
|
||||
"error": "Error: provide old_string/new_string or edits array, not both",
|
||||
}
|
||||
if has_batch:
|
||||
assert isinstance(raw_edits, list)
|
||||
# raw_edits is guaranteed to be a list by the has_batch check above
|
||||
batch_edits: list[Any] = raw_edits # type: ignore[assignment]
|
||||
edits: list[dict[str, Any]] = []
|
||||
for i, e in enumerate(raw_edits):
|
||||
for i, e in enumerate(batch_edits):
|
||||
if not isinstance(e, dict):
|
||||
return {
|
||||
"call_id": call_id,
|
||||
@@ -3359,7 +3612,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "man",
|
||||
"header": "\u2717 man: invalid page name",
|
||||
"preview": f" {RED}{page}{RESET}",
|
||||
"preview": f" {page}",
|
||||
"needs_approval": False,
|
||||
"error": f"Error: invalid page name {page!r}",
|
||||
}
|
||||
@@ -3405,7 +3658,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "web_fetch",
|
||||
"header": "\u2717 web_fetch: invalid url",
|
||||
"preview": f" {RED}{url}{RESET}",
|
||||
"preview": f" {url}",
|
||||
"needs_approval": False,
|
||||
"error": f"Error: URL must start with http:// or https:// (got {url!r})",
|
||||
}
|
||||
@@ -3416,12 +3669,12 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "web_fetch",
|
||||
"header": "\u2717 web_fetch: blocked (private network)",
|
||||
"preview": f" {RED}{url}{RESET}",
|
||||
"preview": f" {url}",
|
||||
"needs_approval": False,
|
||||
"error": f"Error: {ssrf_err}",
|
||||
}
|
||||
q_preview = question[:200] + ("..." if len(question) > 200 else "")
|
||||
preview = f" {DIM}{url}\n Q: {q_preview}{RESET}"
|
||||
preview = f" {url}\n Q: {q_preview}"
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"func_name": "web_fetch",
|
||||
@@ -3467,7 +3720,7 @@ class ChatSession:
|
||||
if topic not in ("general", "news", "finance"):
|
||||
topic = "general"
|
||||
q_preview = query[:200] + ("..." if len(query) > 200 else "")
|
||||
preview = f" {DIM}{q_preview}{RESET}"
|
||||
preview = f" {q_preview}"
|
||||
return {
|
||||
"call_id": call_id,
|
||||
"func_name": "web_search",
|
||||
@@ -3506,7 +3759,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "tool_search",
|
||||
"header": f"\u2699 tool_search: {query[:80]}",
|
||||
"preview": f" {DIM}{query}{RESET}",
|
||||
"preview": f" {query}",
|
||||
"needs_approval": False,
|
||||
"execute": self._exec_tool_search,
|
||||
"query": query,
|
||||
@@ -3540,7 +3793,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "task_agent",
|
||||
"header": "\u2699 task_agent (autonomous agent)",
|
||||
"preview": f" {DIM}{preview_text}{RESET}",
|
||||
"preview": f" {preview_text}",
|
||||
"needs_approval": True,
|
||||
"approval_label": "task_agent",
|
||||
"execute": self._exec_task,
|
||||
@@ -3564,7 +3817,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "plan_agent",
|
||||
"header": "\u2699 plan_agent (planning agent)",
|
||||
"preview": f" {DIM}{preview_text}{RESET}",
|
||||
"preview": f" {preview_text}",
|
||||
"needs_approval": True,
|
||||
"approval_label": "plan_agent",
|
||||
"execute": self._exec_plan,
|
||||
@@ -4045,7 +4298,7 @@ class ChatSession:
|
||||
if isinstance(parsed, list):
|
||||
return " ".join(str(t) for t in parsed)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
pass # falls back to raw string
|
||||
return raw
|
||||
|
||||
# Build corpus from name + description + tags + category
|
||||
@@ -4118,7 +4371,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": func_name,
|
||||
"header": f"\u2699 mcp:{display}",
|
||||
"preview": f"{DIM}{preview}{RESET}",
|
||||
"preview": preview,
|
||||
"needs_approval": True,
|
||||
"approval_label": func_name,
|
||||
"execute": self._exec_mcp_tool,
|
||||
@@ -4196,7 +4449,7 @@ class ChatSession:
|
||||
"call_id": call_id,
|
||||
"func_name": "read_resource",
|
||||
"header": "\u2699 read_resource",
|
||||
"preview": f"{DIM} uri: {uri}{RESET}",
|
||||
"preview": f" uri: {uri}",
|
||||
"needs_approval": True,
|
||||
"approval_label": f"mcp_resource__{self._normalize_resource_uri(uri)}",
|
||||
"execute": self._exec_read_resource,
|
||||
@@ -4338,7 +4591,7 @@ class ChatSession:
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
start_new_session=True,
|
||||
env=scrubbed_env(),
|
||||
env=scrubbed_env(extra=self._skill_resource_env()),
|
||||
)
|
||||
with self._procs_lock:
|
||||
self._active_procs.add(proc)
|
||||
@@ -4673,6 +4926,12 @@ class ChatSession:
|
||||
label_b = "(provided content)"
|
||||
lines_b = (content_b or "").splitlines(keepends=True)
|
||||
|
||||
# When content_b is a baseline, swap so diff reads as "what changed"
|
||||
# (--- old/baseline, +++ new/current file).
|
||||
if content_b is not None:
|
||||
lines_a, lines_b = lines_b, lines_a
|
||||
path_a, label_b = label_b, path_a
|
||||
|
||||
# Stream diff with early cutoff to avoid large allocations
|
||||
max_chars = self.tool_truncation or 262_144
|
||||
chunks: list[str] = []
|
||||
@@ -4730,12 +4989,10 @@ class ChatSession:
|
||||
tools = _without_tool(tools, "web_search")
|
||||
|
||||
# Build extra params for agent calls
|
||||
agent_extra: dict[str, Any] | None = None
|
||||
if agent_provider.provider_name == "openai":
|
||||
agent_kwargs = dict(self._chat_template_kwargs_base)
|
||||
if reasoning_effort:
|
||||
agent_kwargs["reasoning_effort"] = reasoning_effort
|
||||
agent_extra = {"chat_template_kwargs": agent_kwargs}
|
||||
agent_extra = self._provider_extra_params(
|
||||
reasoning_effort=reasoning_effort,
|
||||
provider=agent_provider,
|
||||
)
|
||||
|
||||
def _api_call(
|
||||
messages: list[dict[str, Any]],
|
||||
@@ -6015,26 +6272,30 @@ class ChatSession:
|
||||
return call_id, msg
|
||||
|
||||
if not text.strip():
|
||||
return call_id, "(empty response from URL)"
|
||||
msg = "Error: fetch returned empty response"
|
||||
self._report_tool_result(call_id, "web_fetch", msg, is_error=True)
|
||||
return call_id, msg
|
||||
|
||||
original_len = len(text)
|
||||
self.ui.on_info(f"fetched {original_len} chars, extracting...")
|
||||
|
||||
# Phase 2: truncate for summarization context
|
||||
max_content = 50_000
|
||||
# Phase 2: truncate for summarization context.
|
||||
# Reserve ~25% of the context window for the extraction prompt
|
||||
# overhead (system message, URL, question) and response tokens.
|
||||
# Convert token budget to chars using the calibrated ratio.
|
||||
max_content = int(self.context_window * self._chars_per_token * 0.75)
|
||||
max_content = min(max(max_content, 50_000), 500_000) # 50k–500k
|
||||
if len(text) > max_content:
|
||||
text = (
|
||||
text[: max_content // 2]
|
||||
+ f"\n\n... [{len(text) - max_content} chars omitted] ...\n\n"
|
||||
+ text[-(max_content // 2) :]
|
||||
)
|
||||
# Prefer the beginning — page content is usually top-heavy.
|
||||
text = text[:max_content] + f"\n\n... [{len(text) - max_content} chars truncated] ...\n"
|
||||
|
||||
# Phase 3: summarization API call
|
||||
# Phase 3: summarization API call.
|
||||
# Use a generous max_tokens so thinking models don't starve the
|
||||
# visible answer, and pass reasoning_effort="low" to avoid wasting
|
||||
# budget on deep reasoning for a simple extraction task.
|
||||
try:
|
||||
result = self._provider.create_completion(
|
||||
client=self.client,
|
||||
model=self.model,
|
||||
messages=[
|
||||
result = self._utility_completion(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
@@ -6054,11 +6315,12 @@ class ChatSession:
|
||||
),
|
||||
},
|
||||
],
|
||||
max_tokens=2000,
|
||||
max_tokens=8192,
|
||||
temperature=0.2,
|
||||
extra_params=self._provider_extra_params(),
|
||||
)
|
||||
answer = result.content or "(no answer)"
|
||||
answer = result.content or ""
|
||||
if not answer:
|
||||
answer = "Error: extraction returned no answer"
|
||||
except Exception as e:
|
||||
answer = f"Extraction failed (page was fetched but summarization errored): {e}"
|
||||
|
||||
@@ -6066,7 +6328,7 @@ class ChatSession:
|
||||
call_id,
|
||||
"web_fetch",
|
||||
answer,
|
||||
is_error=answer.startswith("Extraction failed"),
|
||||
is_error=answer.startswith(("Error:", "Extraction failed")),
|
||||
)
|
||||
|
||||
return call_id, answer
|
||||
@@ -6092,6 +6354,7 @@ class ChatSession:
|
||||
self._report_tool_result(call_id, "web_search", msg, is_error=True)
|
||||
return call_id, msg
|
||||
|
||||
output = self._truncate_output(output)
|
||||
self._report_tool_result(call_id, "web_search", output)
|
||||
return call_id, output
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user