mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-22 12:22:20 +00:00
Compare commits
380 Commits
feature/ze
...
deepseek-c
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
14785e4cde | ||
|
|
1dceffffa7 | ||
|
|
4b74a9cf81 | ||
|
|
3be2da6728 | ||
|
|
17d690e77e | ||
|
|
bf05c96dfc | ||
|
|
ed6e0c0189 | ||
|
|
412bfe479c | ||
|
|
861fda21b7 | ||
|
|
919dcf5535 | ||
|
|
cbc424e8e7 | ||
|
|
068fb3572f | ||
|
|
eaf74edbba | ||
|
|
28c4008892 | ||
|
|
a16bc1220c | ||
|
|
5d3e687876 | ||
|
|
68a7cd1dfb | ||
|
|
d96dc20e86 | ||
|
|
934afaa091 | ||
|
|
9ef01acac9 | ||
|
|
cbd38c15fe | ||
|
|
06c8a071a5 | ||
|
|
9e9c2bde57 | ||
|
|
222fd6ed45 | ||
|
|
4abd751f5f | ||
|
|
8f89171db1 | ||
|
|
5dc9d60bec | ||
|
|
feb76bc89d | ||
|
|
3b7f96560b | ||
|
|
23e206c93a | ||
|
|
3e0ca0daf9 | ||
|
|
d7bf1d6582 | ||
|
|
955090bd91 | ||
|
|
379c319e0d | ||
|
|
aaa80d47bc | ||
|
|
9633483fa4 | ||
|
|
9aa3408905 | ||
|
|
31ff1fd90c | ||
|
|
b9622689ea | ||
|
|
f9e8a3250d | ||
|
|
ade2d19be6 | ||
|
|
c4d0a1afba | ||
|
|
81817355cc | ||
|
|
d345d3b53f | ||
|
|
9fe13e9733 | ||
|
|
ee79a305ba | ||
|
|
ed2c8c9fe2 | ||
|
|
9161256d31 | ||
|
|
1a8f407142 | ||
|
|
eddc070628 | ||
|
|
f7bd250c97 | ||
|
|
f2ef63da62 | ||
|
|
29cfeaed9a | ||
|
|
e2b46dab9f | ||
|
|
632244e54f | ||
|
|
3251664513 | ||
|
|
ee16cd0495 | ||
|
|
f80d59182f | ||
|
|
70ef3c357c | ||
|
|
e4165f1dab | ||
|
|
41edae52b0 | ||
|
|
cb16da5543 | ||
|
|
d52b727bce | ||
|
|
6a2abd3439 | ||
|
|
1e4ed75179 | ||
|
|
efb5719679 | ||
|
|
3a1902b2b5 | ||
|
|
ba1978f830 | ||
|
|
2f51d62a46 | ||
|
|
1ca344ba7b | ||
|
|
dda9b848ba | ||
|
|
033d8c5a38 | ||
|
|
e73cdffdc5 | ||
|
|
e32b332eb5 | ||
|
|
89b84392a4 | ||
|
|
f37ff30598 | ||
|
|
547f45c185 | ||
|
|
a6e81a83cb | ||
|
|
e2143aa173 | ||
|
|
2eb257c2a6 | ||
|
|
490687bb71 | ||
|
|
966613847e | ||
|
|
d030d86f9a | ||
|
|
d9ab46b3bb | ||
|
|
6e596e3860 | ||
|
|
3c13be20cb | ||
|
|
9d5e14903b | ||
|
|
9e33d3b100 | ||
|
|
c88665f6d1 | ||
|
|
a8157b3e2d | ||
|
|
a07e6723d9 | ||
|
|
5bc7b741bc | ||
|
|
1db37cf084 | ||
|
|
6c5b103149 | ||
|
|
27ec052c17 | ||
|
|
00b813f6a8 | ||
|
|
91e5198a94 | ||
|
|
27f1cc3c42 | ||
|
|
a4d048f2b5 | ||
|
|
752d4f3803 | ||
|
|
164ed775c8 | ||
|
|
9d905758ec | ||
|
|
0d1cb66855 | ||
|
|
6e9932e0ac | ||
|
|
59773e972b | ||
|
|
7fd0cdf987 | ||
|
|
b1947d660a | ||
|
|
d45a9edcba | ||
|
|
4dbbb45240 | ||
|
|
a286b5efa0 | ||
|
|
26a0b04aef | ||
|
|
f7ccc25a7f | ||
|
|
8dd0ece501 | ||
|
|
10969e0719 | ||
|
|
b633455071 | ||
|
|
e280d1f8b1 | ||
|
|
81ad9f604b | ||
|
|
2345691176 | ||
|
|
a7615bc827 | ||
|
|
befdc5307e | ||
|
|
145777ffd0 | ||
|
|
37c2bea93d | ||
|
|
985e765285 | ||
|
|
a4b1330627 | ||
|
|
7b69284812 | ||
|
|
7fb1da7b58 | ||
|
|
38f9923469 | ||
|
|
5b986a3e15 | ||
|
|
56d9ff6b3f | ||
|
|
15b31e7c53 | ||
|
|
d0dcc3219d | ||
|
|
52db6fd308 | ||
|
|
ceca0e8efc | ||
|
|
309bc873e6 | ||
|
|
89109dc209 | ||
|
|
234ea19cad | ||
|
|
489b6eba14 | ||
|
|
66421cd17f | ||
|
|
cd7de3958c | ||
|
|
8c0ac499ef | ||
|
|
1d22155e05 | ||
|
|
ab2fb2afc5 | ||
|
|
884bae9fc4 | ||
|
|
92d581573d | ||
|
|
85a3d3adc0 | ||
|
|
2d7f03b2ed | ||
|
|
37bce70f76 | ||
|
|
038bc14757 | ||
|
|
287cbac5f9 | ||
|
|
28d91227af | ||
|
|
846d894a13 | ||
|
|
a2db3e2d57 | ||
|
|
e43ceb2e43 | ||
|
|
689a07f562 | ||
|
|
da487a850e | ||
|
|
fa6d3c76d0 | ||
|
|
ede1804d4b | ||
|
|
3392e8d4cb | ||
|
|
b5174d9753 | ||
|
|
1f2ff8a99c | ||
|
|
0c60644ba2 | ||
|
|
aca8d43a61 | ||
|
|
7fe4c1963b | ||
|
|
9a0919f149 | ||
|
|
d69ab913d4 | ||
|
|
42b8c332df | ||
|
|
aa682bf8ec | ||
|
|
c53e72e80a | ||
|
|
afd81aeca2 | ||
|
|
4e03145323 | ||
|
|
169686681f | ||
|
|
fa7d2804bb | ||
|
|
ccabc5d06d | ||
|
|
cb68227c88 | ||
|
|
c72fc7dd56 | ||
|
|
8c6d1f89dc | ||
|
|
1ab7e54bd7 | ||
|
|
42dceb0cd6 | ||
|
|
05115c3387 | ||
|
|
c0c8cafd00 | ||
|
|
16d6d66d17 | ||
|
|
3681fc8aab | ||
|
|
d95c09e00c | ||
|
|
3cdef5ec17 | ||
|
|
58e1620347 | ||
|
|
42b5a5e6c8 | ||
|
|
d84d249f2c | ||
|
|
8b6393f794 | ||
|
|
81146710ba | ||
|
|
25f39897d7 | ||
|
|
45bdfaee58 | ||
|
|
a1ae6e94e9 | ||
|
|
eebcc67c85 | ||
|
|
3f9e7f7728 | ||
|
|
a5b5549edd | ||
|
|
22eec0162c | ||
|
|
9f55da9bb8 | ||
|
|
16dce9ea81 | ||
|
|
963ee04619 | ||
|
|
b874b1f01c | ||
|
|
b6cca3d3a0 | ||
|
|
86ebc84f4c | ||
|
|
8a74c0543f | ||
|
|
3d16a9b988 | ||
|
|
40f98b99aa | ||
|
|
7fcae5b08d | ||
|
|
f50cb31749 | ||
|
|
85a1ea4b4c | ||
|
|
0f0f8c40bf | ||
|
|
a40d224ee7 | ||
|
|
a637acd8f4 | ||
|
|
58be0c7976 | ||
|
|
d934f3eead | ||
|
|
74b58e5fa3 | ||
|
|
c1497e0cfe | ||
|
|
30db582321 | ||
|
|
a1018776e9 | ||
|
|
5ef499e2ae | ||
|
|
c7c802c610 | ||
|
|
685368bb0a | ||
|
|
55e240d92a | ||
|
|
7da4ad3818 | ||
|
|
453337cb2c | ||
|
|
9e1934bfda | ||
|
|
31a9730cd8 | ||
|
|
d686e0e851 | ||
|
|
7708ed1c8b | ||
|
|
22ede7636a | ||
|
|
c9a5af0ffa | ||
|
|
072079b33d | ||
|
|
2abc28d295 | ||
|
|
7879b68cad | ||
|
|
9eff340c57 | ||
|
|
933ba105d8 | ||
|
|
205dc6c3b3 | ||
|
|
c1ada39278 | ||
|
|
200b32f1ff | ||
|
|
2b623a86c3 | ||
|
|
a68998e8ff | ||
|
|
905ae67015 | ||
|
|
6fde846be2 | ||
|
|
a833cf429e | ||
|
|
c5a205cf98 | ||
|
|
62185fbd38 | ||
|
|
236854bfe4 | ||
|
|
a7886c528f | ||
|
|
bb97e8dedb | ||
|
|
a7455a48d2 | ||
|
|
a090632804 | ||
|
|
d6b755ff81 | ||
|
|
a1c2a785a1 | ||
|
|
dc8ebe66e0 | ||
|
|
750193e99d | ||
|
|
4ee654331f | ||
|
|
b8fcf5bcc8 | ||
|
|
2ce8981f0c | ||
|
|
b76ff62f1c | ||
|
|
73ebfe8975 | ||
|
|
3dae03d731 | ||
|
|
a7b815b29f | ||
|
|
8a1f909083 | ||
|
|
a6c6d3e034 | ||
|
|
f2a0473fb3 | ||
|
|
33d2ef9ef1 | ||
|
|
f1c3b515fe | ||
|
|
74ddc48ff5 | ||
|
|
6db11ed88f | ||
|
|
4d302baeb5 | ||
|
|
063df51ced | ||
|
|
e97fa60b27 | ||
|
|
42b3840df6 | ||
|
|
a0b2b9466c | ||
|
|
9a553de011 | ||
|
|
901a6a9ba2 | ||
|
|
691927a996 | ||
|
|
8fcecf2c1f | ||
|
|
aed967dc44 | ||
|
|
ed8102d533 | ||
|
|
8fa8475cca | ||
|
|
4723b9db4d | ||
|
|
358ff25899 | ||
|
|
7267bb87b9 | ||
|
|
173f5fbcbd | ||
|
|
9006709f8d | ||
|
|
7d48b36be8 | ||
|
|
8b3fdaa545 | ||
|
|
b43d57df23 | ||
|
|
b9a5a9276e | ||
|
|
a79fdf7212 | ||
|
|
9560050946 | ||
|
|
dd88e9b172 | ||
|
|
e31b45fa9e | ||
|
|
1ed8b29d64 | ||
|
|
59f8d31719 | ||
|
|
93a368b1a2 | ||
|
|
d3dd346853 | ||
|
|
0198569a9a | ||
|
|
7e648cb5c2 | ||
|
|
deb75624f3 | ||
|
|
c5cb562165 | ||
|
|
8a89a38864 | ||
|
|
09e7f1f0bf | ||
|
|
d9d082ad5c | ||
|
|
9cd4ff5c21 | ||
|
|
11eb20a2d1 | ||
|
|
a63e81db06 | ||
|
|
8fc1b6484c | ||
|
|
9bc3feff62 | ||
|
|
cb22968ff3 | ||
|
|
e3bca39815 | ||
|
|
958f28fd82 | ||
|
|
c8c30d7cfd | ||
|
|
2c2124952f | ||
|
|
f56ba92ae8 | ||
|
|
1fa71ee806 | ||
|
|
97d164e181 | ||
|
|
9f3c915dec | ||
|
|
4ff6218140 | ||
|
|
67ea6aec43 | ||
|
|
41f346417a | ||
|
|
8a373276ce | ||
|
|
4857d741de | ||
|
|
eabe17cc26 | ||
|
|
8f298ae9c1 | ||
|
|
e5f3b7d755 | ||
|
|
3796e29cb4 | ||
|
|
812f50ea28 | ||
|
|
ce5fa136c2 | ||
|
|
4643aefb22 | ||
|
|
5aa902fd4a | ||
|
|
db144e0903 | ||
|
|
e20b20dbca | ||
|
|
aa700443d3 | ||
|
|
7e89b5bc6d | ||
|
|
8dbf036558 | ||
|
|
12f2cfecc9 | ||
|
|
a0f3378cf5 | ||
|
|
20b4c3a641 | ||
|
|
88f7ee734f | ||
|
|
326a0086b7 | ||
|
|
eb46651373 | ||
|
|
226188e22d | ||
|
|
89fc48eea2 | ||
|
|
f66d98e6b5 | ||
|
|
cfea4f9cd1 | ||
|
|
c74a877ab5 | ||
|
|
83a76e86fe | ||
|
|
66975ee271 | ||
|
|
5788d63892 | ||
|
|
d0c7cc6bd9 | ||
|
|
40096867d9 | ||
|
|
e002c0b66f | ||
|
|
608549d051 | ||
|
|
72f389c8d2 | ||
|
|
97daddb95a | ||
|
|
847f4b07a5 | ||
|
|
4bce04ada4 | ||
|
|
f1e2448620 | ||
|
|
28340a152c | ||
|
|
e28f6118e7 | ||
|
|
9fb6f54d12 | ||
|
|
3cb8d7b5dd | ||
|
|
ede076b881 | ||
|
|
449a0951f9 | ||
|
|
8c9ede2272 | ||
|
|
e794b09614 | ||
|
|
350714f23a | ||
|
|
9cc2e84f7c | ||
|
|
342a7f7f16 | ||
|
|
7694f20006 | ||
|
|
3ea0a267dd | ||
|
|
68e537f6fb | ||
|
|
814c39898d | ||
|
|
fdd4f12f5c | ||
|
|
6957d8c0d9 | ||
|
|
2c358276bf | ||
|
|
c5fd386c1e | ||
|
|
ce9834d7ec | ||
|
|
6ebe73f2f7 | ||
|
|
98aecb08f9 |
@@ -14,6 +14,7 @@ UPSTREAM_API_KEY=your-upstream-api-key
|
||||
# HTTP_URL=https://api.mynode.com
|
||||
# ONION_URL=http://mynode.onion (auto fetched from compose)
|
||||
# RELAYS="wss://relay.damus.io,wss://relay.nostr.band,wss://eden.nostr.land,wss://relay.routstr.com"
|
||||
# ENABLE_ANALYTICS_SHARING=true
|
||||
# CASHU_MINTS="https://mint.minibits.cash/Bitcoin,https://mint.cubabitcoin.org,https://ecashmint.otrta.me"
|
||||
# RECEIVE_LN_ADDRESS=
|
||||
|
||||
|
||||
12
.github/workflows/container.yml
vendored
12
.github/workflows/container.yml
vendored
@@ -14,6 +14,14 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v3
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Resolve git metadata
|
||||
id: gitmeta
|
||||
run: |
|
||||
echo "sha=$(git rev-parse --short=7 HEAD)" >> "$GITHUB_OUTPUT"
|
||||
echo "tag=$(git describe --tags --exact-match HEAD 2>/dev/null || true)" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v2
|
||||
@@ -29,7 +37,11 @@ jobs:
|
||||
uses: docker/build-push-action@v4
|
||||
with:
|
||||
context: .
|
||||
file: Dockerfile.full
|
||||
push: true
|
||||
build-args: |
|
||||
GIT_COMMIT=${{ steps.gitmeta.outputs.sha }}
|
||||
GIT_TAG=${{ steps.gitmeta.outputs.tag }}
|
||||
tags: |
|
||||
ghcr.io/routstr/proxy:latest
|
||||
ghcr.io/routstr/core:latest
|
||||
|
||||
2
.github/workflows/test.yml
vendored
2
.github/workflows/test.yml
vendored
@@ -67,7 +67,7 @@ jobs:
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "18"
|
||||
node-version: "20"
|
||||
cache: "pnpm"
|
||||
cache-dependency-path: ui/pnpm-lock.yaml
|
||||
|
||||
|
||||
29
Dockerfile
29
Dockerfile
@@ -1,26 +1,29 @@
|
||||
FROM ghcr.io/astral-sh/uv:python3.11-alpine
|
||||
FROM ghcr.io/astral-sh/uv:python3.11-bookworm-slim
|
||||
|
||||
# Install system dependencies required for secp256k1
|
||||
RUN apk add --no-cache \
|
||||
pkgconf \
|
||||
build-base \
|
||||
automake \
|
||||
autoconf \
|
||||
libtool \
|
||||
m4 \
|
||||
perl
|
||||
RUN apk add git
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
git \
|
||||
build-essential \
|
||||
pkg-config \
|
||||
libsecp256k1-dev \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY uv.lock pyproject.toml ./
|
||||
RUN mkdir -p /routstr
|
||||
|
||||
RUN uv add git+https://github.com/saschanaz/secp256k1-py.git#branch=upgrade060
|
||||
# RUN uv sync
|
||||
RUN uv sync --frozen --no-dev --no-install-project
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY . .
|
||||
|
||||
ARG GIT_COMMIT=""
|
||||
ARG GIT_TAG=""
|
||||
ENV GIT_COMMIT=${GIT_COMMIT}
|
||||
ENV GIT_TAG=${GIT_TAG}
|
||||
ENV PORT=8000
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ FROM node:23-alpine AS ui-builder
|
||||
WORKDIR /app/ui
|
||||
|
||||
# Install pnpm
|
||||
RUN corepack enable pnpm && corepack prepare pnpm@latest --activate
|
||||
RUN corepack enable pnpm && corepack prepare pnpm@10.15.0 --activate
|
||||
|
||||
# Copy UI source
|
||||
COPY ui/package.json ui/pnpm-lock.yaml* ./
|
||||
@@ -16,35 +16,38 @@ ENV NEXT_TELEMETRY_DISABLED=1
|
||||
RUN pnpm run build
|
||||
|
||||
# Stage 2: Build the Routstr Node
|
||||
FROM ghcr.io/astral-sh/uv:python3.11-alpine AS runner
|
||||
FROM ghcr.io/astral-sh/uv:python3.11-bookworm-slim AS runner
|
||||
|
||||
# Install system dependencies
|
||||
RUN apk add --no-cache \
|
||||
pkgconf \
|
||||
build-base \
|
||||
automake \
|
||||
autoconf \
|
||||
libtool \
|
||||
m4 \
|
||||
perl \
|
||||
git
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
git \
|
||||
build-essential \
|
||||
pkg-config \
|
||||
libsecp256k1-dev \
|
||||
autoconf \
|
||||
automake \
|
||||
libtool \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY uv.lock pyproject.toml ./
|
||||
|
||||
RUN uv sync --no-dev --no-install-project
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy the rest of the application (required for uv sync to find the package)
|
||||
COPY . .
|
||||
|
||||
# Install dependencies including the specific secp256k1 branch
|
||||
RUN uv add git+https://github.com/saschanaz/secp256k1-py.git#branch=upgrade060
|
||||
RUN uv sync --no-dev
|
||||
|
||||
# Copy the built UI from the ui-builder stage
|
||||
COPY --from=ui-builder /app/ui/out ./ui_out
|
||||
|
||||
ARG GIT_COMMIT=""
|
||||
ARG GIT_TAG=""
|
||||
ENV GIT_COMMIT=${GIT_COMMIT}
|
||||
ENV GIT_TAG=${GIT_TAG}
|
||||
ENV PORT=8000
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
# Run the application
|
||||
CMD ["/app/.venv/bin/fastapi", "run", "routstr", "--host", "0.0.0.0"]
|
||||
CMD ["/.venv/bin/fastapi", "run", "routstr", "--host", "0.0.0.0"]
|
||||
|
||||
6
Makefile
6
Makefile
@@ -98,7 +98,7 @@ docker-down:
|
||||
lint:
|
||||
@echo "🔍 Running linting checks..."
|
||||
$(RUFF) check .
|
||||
$(MYPY) routstr/ --ignore-missing-imports
|
||||
$(MYPY) .
|
||||
|
||||
format:
|
||||
@echo "✨ Formatting code..."
|
||||
@@ -107,7 +107,7 @@ format:
|
||||
|
||||
type-check:
|
||||
@echo "🔎 Running type checks..."
|
||||
$(MYPY) routstr/ --ignore-missing-imports
|
||||
$(MYPY) .
|
||||
|
||||
# Development setup
|
||||
dev-setup:
|
||||
@@ -234,7 +234,7 @@ ci-test:
|
||||
ci-lint:
|
||||
@echo "🤖 Running CI linting..."
|
||||
$(RUFF) check . --exit-non-zero-on-fix
|
||||
$(MYPY) routstr/ --ignore-missing-imports --no-error-summary
|
||||
$(MYPY) . --no-error-summary
|
||||
|
||||
# Debug helpers
|
||||
test-debug:
|
||||
|
||||
26
README.md
26
README.md
@@ -51,14 +51,26 @@ curl https://api.routstr.com/v1/chat/completions \
|
||||
|
||||
## Quick Start (Docker)
|
||||
|
||||
If you are a node runner, start a Routstr Core instance and configure upstream access in the dashboard.
|
||||
If you are a node runner, start a Routstr Core instance using Docker Compose:
|
||||
|
||||
```bash
|
||||
docker run -d \
|
||||
--name routstr-proxy \
|
||||
-p 8000:8000 \
|
||||
ghcr.io/routstr/proxy:latest
|
||||
```
|
||||
1. **Prepare your `.env`**:
|
||||
```bash
|
||||
ADMIN_PASSWORD=mysecretpassword
|
||||
NAME="My AI Node"
|
||||
DESCRIPTION="Fast access to models"
|
||||
NSEC=yournsec
|
||||
RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
```
|
||||
|
||||
2. **Start the services**:
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
3. **Configure**:
|
||||
Open [http://localhost:8000/admin/](http://localhost:8000/admin/) to connect your AI providers and set pricing.
|
||||
|
||||
For full instructions, see the **[Provider Quick Start Guide](https://docs.routstr.com/provider/quickstart/)**.
|
||||
|
||||
## Development
|
||||
|
||||
|
||||
@@ -8,8 +8,9 @@ services:
|
||||
args:
|
||||
# NEXT_PUBLIC_API_URL: ${NEXT_PUBLIC_API_URL:-http://127.0.0.1:8000}
|
||||
NEXT_PUBLIC_ADMIN_API_KEY: ${NEXT_PUBLIC_ADMIN_API_KEY:-}
|
||||
user: root
|
||||
volumes:
|
||||
- ./ui_out:/output
|
||||
- ./ui_out:/output:z
|
||||
command:
|
||||
["sh", "-c", "mkdir -p /output && cp -r /app/built/. /output/ && echo 'UI build copied to mounted volume' && ls -la /output/ && echo 'UI built and ready' && tail -f /dev/null"]
|
||||
|
||||
@@ -18,10 +19,10 @@ services:
|
||||
depends_on:
|
||||
- ui
|
||||
volumes:
|
||||
- .:/app
|
||||
- ./logs:/app/logs
|
||||
- .:/app:z
|
||||
- ./logs:/app/logs:z
|
||||
- tor-data:/var/lib/tor:ro
|
||||
- ./ui_out:/app/ui_out:ro
|
||||
- ./ui_out:/app/ui_out:ro,z
|
||||
env_file:
|
||||
- .env
|
||||
environment:
|
||||
|
||||
@@ -396,6 +396,8 @@ Authorization: Bearer sk-...
|
||||
}
|
||||
```
|
||||
|
||||
`balance` is the spendable balance used by request admission.
|
||||
|
||||
### Check Balance
|
||||
|
||||
Get current wallet balance.
|
||||
|
||||
@@ -155,7 +155,7 @@ The response includes your change in the same header:
|
||||
X-Cashu: cashuA7k2mNp4...
|
||||
```
|
||||
|
||||
This is fully stateless—no session, no `/v1/balance/refund` call needed. However, **streaming does not work with `X-Cashu`** because the refund can only be calculated after the full response is generated.
|
||||
This is fully stateless—no session, no `/v1/balance/refund` call needed. However, **streaming does not work with `X-Cashu`** because the refund can only be calculated after the full response is generated. If you lose the `X-Cashu` response header before claiming your change, you can reclaim the refund via `POST /v1/wallet/refund` by supplying the original payment token in the `x-cashu` header.
|
||||
|
||||
## Response Headers
|
||||
|
||||
|
||||
@@ -53,9 +53,11 @@ If your balance runs low, you don't need a new key. You can top up the existing
|
||||
|
||||
### Via Lightning
|
||||
|
||||
`POST /lightning/invoice` with `{"amount_sats": 1000, "purpose": "topup", "api_key": "sk-..."}`.
|
||||
`POST /lightning/invoice` with `Authorization: Bearer sk-...` header and body `{"amount_sats": 1000, "purpose": "topup"}`.
|
||||
*Once paid, the funds are added to your existing key.*
|
||||
|
||||
> Legacy: the endpoint is also exposed at `/v1/balance/lightning/invoice`, and accepts an `api_key` field in the body as a fallback for older clients. New integrations should use the RIP-08 path with the `Authorization` header.
|
||||
|
||||
### Via Cashu
|
||||
|
||||
`POST /v1/balance/topup` with `{"cashu_token": "..."}` and `Authorization: Bearer sk-...`.
|
||||
|
||||
@@ -78,9 +78,15 @@ Which mints to accept payments from:
|
||||
|
||||
Automatic profit withdrawal:
|
||||
|
||||
| Setting | Description |
|
||||
| --------------------- | ------------------------------- |
|
||||
| **Lightning Address** | Your LN address for withdrawals |
|
||||
| Setting | Description | Default |
|
||||
| ------------------------------------ | ---------------------------------------------------------------------------------------------------------- | ------- |
|
||||
| **Lightning Address** | Your LN address for withdrawals | — |
|
||||
| **Minimum Payout (sat)** | Min available balance (in sats) before profit is paid out. Applies to both `sat` and `msat` mints (auto-converted). | `210` |
|
||||
| **Payout Interval (seconds)** | How often the payout loop wakes up and checks balances | `900` |
|
||||
|
||||
All payout amounts must be positive. Set the minimums above your wallet's
|
||||
minimum-invoice constraint (typically 1 sat) and high enough to amortise
|
||||
routing fees.
|
||||
|
||||
### Security
|
||||
|
||||
@@ -97,6 +103,7 @@ Announce your node on the network:
|
||||
| **Npub** | Your Nostr public key |
|
||||
| **Nsec** | Your Nostr private key (for signing) |
|
||||
| **Relays** | Relays to publish announcements |
|
||||
| **Share Analytics** | Publish aggregate usage stats to Nostr |
|
||||
|
||||
See [Discovery](discovery.md) for details.
|
||||
|
||||
@@ -122,8 +129,11 @@ Use environment variables for:
|
||||
| `DESCRIPTION` | Node description | `A Routstr Node` |
|
||||
| `NPUB` | Nostr public key (bech32) | — |
|
||||
| `NSEC` | Nostr private key | — |
|
||||
| `ENABLE_ANALYTICS_SHARING` | Enable usage analytics sharing to Nostr | `true` |
|
||||
| `CASHU_MINTS` | Comma-separated mint URLs | `https://mint.minibits.cash/Bitcoin` |
|
||||
| `RECEIVE_LN_ADDRESS` | Lightning address for withdrawals | — |
|
||||
| `MIN_PAYOUT_SAT` | Min payout balance in sats (applies to all mints) | `210` |
|
||||
| `PAYOUT_INTERVAL_SECONDS` | Payout loop interval (seconds) | `900` |
|
||||
| `TOR_PROXY_URL` | SOCKS5 proxy for Tor | `socks5://127.0.0.1:9050` |
|
||||
| `CORS_ORIGINS` | Allowed CORS origins | `*` |
|
||||
| `RELAYS` | Nostr relays (comma-separated) | (default set) |
|
||||
|
||||
@@ -152,6 +152,7 @@ Manage which mints you accept payments from:
|
||||
|-------|-------------|
|
||||
| **Nsec** | Private key for signing announcements |
|
||||
| **Relays** | Where to publish your node advertisement |
|
||||
| **Share Analytics** | Toggle publishing aggregate usage stats to Nostr |
|
||||
|
||||
### Security
|
||||
|
||||
|
||||
@@ -2,20 +2,62 @@
|
||||
|
||||
Production deployment guide for Routstr Provider nodes.
|
||||
|
||||
## All-in-One Docker Image (Preferred)
|
||||
|
||||
The easiest way to deploy Routstr is using the all-in-one Docker image from Docker Hub, which includes both the FastAPI backend and the Next.js admin dashboard in a single container.
|
||||
|
||||
### Quick Start
|
||||
|
||||
```bash
|
||||
docker run -d \
|
||||
--name routstr \
|
||||
-p 8000:8000 \
|
||||
-v routstr-data:/app/data \
|
||||
-e DATABASE_URL="sqlite:////app/data/routstr.db" \
|
||||
9qeklajc/routstr:latest
|
||||
```
|
||||
|
||||
Access your node:
|
||||
- **API & Admin Dashboard**: http://localhost:8000
|
||||
|
||||
### Docker Compose Setup
|
||||
|
||||
Create `docker-compose.yml`:
|
||||
|
||||
```yaml
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
routstr:
|
||||
image: 9qeklajc/routstr:latest
|
||||
container_name: routstr
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- routstr-data:/app/data
|
||||
environment:
|
||||
DATABASE_URL: "sqlite:////app/data/routstr.db"
|
||||
ADMIN_KEY: "your-secure-admin-key"
|
||||
LOG_LEVEL: "info"
|
||||
|
||||
volumes:
|
||||
routstr-data:
|
||||
```
|
||||
|
||||
Start it:
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Docker Compose (Recommended)
|
||||
|
||||
For production, use Docker Compose with persistent storage and optional Tor support.
|
||||
|
||||
### Unified Setup (All-in-one)
|
||||
To build and run the node with the UI integrated in a single container using the multi-stage build:
|
||||
|
||||
```bash
|
||||
docker build -f Dockerfile.full -t routstr-full .
|
||||
docker run -d -p 8000:8000 --env-file .env routstr-full
|
||||
```
|
||||
|
||||
### Advanced Setup (Separated UI & Node)
|
||||
Use the included `compose.yml` for a more flexible setup that separates the UI build process from the node execution. This is useful for development or when you want to manage Tor as a separate service.
|
||||
Use the included `compose.yml` for a flexible setup that handles both the UI and the node execution. This is useful for development or when you want to manage Tor as a separate service.
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
@@ -184,20 +226,16 @@ docker compose up -d
|
||||
|
||||
## Building from Source
|
||||
|
||||
### Unified Image (UI + Node)
|
||||
The easiest way to build everything from source into a single production-ready image:
|
||||
### Using Docker Compose
|
||||
The easiest way to build everything from source:
|
||||
|
||||
```bash
|
||||
docker build -f Dockerfile.full -t routstr-full .
|
||||
docker compose build
|
||||
```
|
||||
|
||||
### Individual Components
|
||||
If you prefer building them separately or using Docker Compose:
|
||||
If you prefer building the node only (requires manual UI build first):
|
||||
|
||||
```bash
|
||||
# Build using compose
|
||||
docker compose build
|
||||
|
||||
# Or build the node only (requires manual UI build first)
|
||||
docker build -t routstr-node .
|
||||
```
|
||||
|
||||
@@ -35,6 +35,7 @@ ADMIN_PASSWORD=mysecretpassword
|
||||
# Node Identity
|
||||
NAME="My AI Node"
|
||||
DESCRIPTION="Fast access to models"
|
||||
NSEC=yournsec
|
||||
|
||||
# Lightning Payouts
|
||||
RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
@@ -43,32 +44,10 @@ RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
|
||||
## 2. Start the Node
|
||||
|
||||
You can run the pre-built image directly:
|
||||
The recommended way to run Routstr is using Docker Compose, which handles the node, the UI, and optional services like Tor.
|
||||
|
||||
```bash
|
||||
docker run -d \
|
||||
--name routstr \
|
||||
-p 8000:8000 \
|
||||
--env-file .env \
|
||||
-v routstr-data:/app/data \
|
||||
ghcr.io/routstr/proxy:latest
|
||||
```
|
||||
|
||||
*Note: The pre-built image does not contain the UI. For the all-in-one experience with the Admin Dashboard, use the Build from Source instructions below.*
|
||||
|
||||
### Build from Source (Recommended)
|
||||
|
||||
If you want to build the node and UI yourself from source, use the unified Dockerfile:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/routstr/routstr-core.git
|
||||
cd routstr-core
|
||||
# Edit your .env with ADMIN_PASSWORD and API keys
|
||||
cp .env.example .env
|
||||
nano .env
|
||||
|
||||
docker build -f Dockerfile.full -t routstr-local .
|
||||
docker run -d -p 8000:8000 --env-file .env --name routstr routstr-local
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
Verify it's running:
|
||||
@@ -77,6 +56,15 @@ Verify it's running:
|
||||
curl http://localhost:8000/v1/info
|
||||
```
|
||||
|
||||
### Build from Source (Optional)
|
||||
|
||||
If you've cloned the repository and want to build the images yourself:
|
||||
|
||||
```bash
|
||||
docker compose build
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Configure via Dashboard
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
|
||||
32
migrations/versions/02650cd6f028_add_routstr_fees_table.py
Normal file
32
migrations/versions/02650cd6f028_add_routstr_fees_table.py
Normal file
@@ -0,0 +1,32 @@
|
||||
"""add routstr_fees table
|
||||
|
||||
Revision ID: 02650cd6f028
|
||||
Revises: c3d4e5f6a7b8
|
||||
Create Date: 2026-04-24 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "02650cd6f028"
|
||||
down_revision = "c3d4e5f6a7b8"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"routstr_fees",
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("accumulated_msats", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("total_paid_msats", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("last_paid_at", sa.Integer(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
# Seed with a single row
|
||||
op.execute("INSERT INTO routstr_fees (id, accumulated_msats, total_paid_msats) VALUES (1, 0, 0)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("routstr_fees")
|
||||
@@ -0,0 +1,32 @@
|
||||
"""add provider_settings to upstream_providers
|
||||
|
||||
Revision ID: 614c0a740e68
|
||||
Revises: 06f81c0fc88d
|
||||
Create Date: 2026-02-13 22:36:53.608737
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "614c0a740e68"
|
||||
down_revision = "06f81c0fc88d"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Check if column exists before adding it
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [c["name"] for c in inspector.get_columns("upstream_providers")]
|
||||
|
||||
if "provider_settings" not in columns:
|
||||
op.add_column(
|
||||
"upstream_providers",
|
||||
sa.Column("provider_settings", sa.Text(), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("upstream_providers", "provider_settings")
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Add balance_limit, balance_limit_reset, validity_date to lightning_invoices
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: f1a2b3c4d5e6
|
||||
Create Date: 2026-06-03 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "a2b3c4d5e6f7"
|
||||
down_revision = "f1a2b3c4d5e6"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column("lightning_invoices", sa.Column("balance_limit", sa.Integer(), nullable=True))
|
||||
op.add_column("lightning_invoices", sa.Column("balance_limit_reset", sa.String(), nullable=True))
|
||||
op.add_column("lightning_invoices", sa.Column("validity_date", sa.Integer(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("lightning_invoices", "validity_date")
|
||||
op.drop_column("lightning_invoices", "balance_limit_reset")
|
||||
op.drop_column("lightning_invoices", "balance_limit")
|
||||
42
migrations/versions/a776ca70e5fe_add_cashu_refunds_table.py
Normal file
42
migrations/versions/a776ca70e5fe_add_cashu_refunds_table.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""add cashu_transactions table
|
||||
|
||||
Revision ID: a776ca70e5fe
|
||||
Revises: 614c0a740e68
|
||||
Create Date: 2026-03-11 22:00:01.554762
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a776ca70e5fe"
|
||||
down_revision = "614c0a740e68"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"cashu_transactions",
|
||||
sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("token", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("amount", sa.Integer(), nullable=False),
|
||||
sa.Column("unit", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("mint_url", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column(
|
||||
"type",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=False,
|
||||
server_default="out",
|
||||
),
|
||||
sa.Column("request_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column("created_at", sa.Integer(), nullable=False),
|
||||
sa.Column("collected", sa.Boolean(), nullable=False),
|
||||
sa.Column("swept", sa.Boolean(), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("cashu_transactions")
|
||||
@@ -0,0 +1,33 @@
|
||||
"""add forwarded_model_id to models
|
||||
|
||||
Revision ID: b1c2d3e4f5a6
|
||||
Revises: a776ca70e5fe
|
||||
Create Date: 2026-04-05 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b1c2d3e4f5a6"
|
||||
down_revision = "a776ca70e5fe"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"models",
|
||||
sa.Column(
|
||||
"forwarded_model_id",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
# Backfill: set forwarded_model_id = id for all existing rows
|
||||
op.execute("UPDATE models SET forwarded_model_id = id WHERE forwarded_model_id IS NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("models", "forwarded_model_id")
|
||||
@@ -0,0 +1,25 @@
|
||||
"""add reserved_at to api_keys
|
||||
|
||||
Revision ID: b5e7c9d1f3a2
|
||||
Revises: a2b3c4d5e6f7
|
||||
Create Date: 2026-06-12 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b5e7c9d1f3a2"
|
||||
down_revision = "a2b3c4d5e6f7"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# existing keys keep NULL
|
||||
# New reservations populate it via pay_for_request.
|
||||
op.add_column("api_keys", sa.Column("reserved_at", sa.Integer(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("api_keys", "reserved_at")
|
||||
@@ -0,0 +1,36 @@
|
||||
"""add source to cashu_transactions
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: b1c2d3e4f5a6
|
||||
Create Date: 2026-04-10 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "c3d4e5f6a7b8"
|
||||
down_revision = "b1c2d3e4f5a6"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [col["name"] for col in inspector.get_columns("cashu_transactions")]
|
||||
if "source" not in columns:
|
||||
op.add_column(
|
||||
"cashu_transactions",
|
||||
sa.Column(
|
||||
"source",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=False,
|
||||
server_default="x-cashu",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("cashu_transactions", "source")
|
||||
34
migrations/versions/cli_tokens_table.py
Normal file
34
migrations/versions/cli_tokens_table.py
Normal file
@@ -0,0 +1,34 @@
|
||||
"""add cli_tokens table
|
||||
|
||||
Revision ID: cli_tokens_001
|
||||
Revises: e8f9a0b1c2d3
|
||||
Create Date: 2026-04-25 00:00:00.000000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "cli_tokens_001"
|
||||
down_revision = "e8f9a0b1c2d3"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"cli_tokens",
|
||||
sa.Column("id", sa.String(), primary_key=True, nullable=False),
|
||||
sa.Column("token", sa.String(), nullable=False, unique=True),
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("created_at", sa.Integer(), nullable=False),
|
||||
sa.Column("last_used_at", sa.Integer(), nullable=True),
|
||||
sa.Column("expires_at", sa.Integer(), nullable=True),
|
||||
)
|
||||
op.create_index("ix_cli_tokens_token", "cli_tokens", ["token"], unique=True)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_cli_tokens_token", table_name="cli_tokens")
|
||||
op.drop_table("cli_tokens")
|
||||
@@ -0,0 +1,46 @@
|
||||
"""add api key link to cashu_transactions
|
||||
|
||||
Revision ID: d4e5f6a7b8c9
|
||||
Revises: c3d4e5f6a7b8
|
||||
Create Date: 2026-04-20 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "d4e5f6a7b8c9"
|
||||
down_revision = "c3d4e5f6a7b8"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [col["name"] for col in inspector.get_columns("cashu_transactions")]
|
||||
indexes = {index["name"] for index in inspector.get_indexes("cashu_transactions")}
|
||||
|
||||
if "api_key_hashed_key" not in columns:
|
||||
op.add_column(
|
||||
"cashu_transactions",
|
||||
sa.Column(
|
||||
"api_key_hashed_key",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
|
||||
if "ix_cashu_transactions_api_key_hashed_key" not in indexes:
|
||||
op.create_index(
|
||||
"ix_cashu_transactions_api_key_hashed_key",
|
||||
"cashu_transactions",
|
||||
["api_key_hashed_key"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_cashu_transactions_api_key_hashed_key", table_name="cashu_transactions")
|
||||
op.drop_column("cashu_transactions", "api_key_hashed_key")
|
||||
20
migrations/versions/e8f9a0b1c2d3_merge_heads.py
Normal file
20
migrations/versions/e8f9a0b1c2d3_merge_heads.py
Normal file
@@ -0,0 +1,20 @@
|
||||
"""merge heads: routstr_fees + api_key_to_cashu_transactions
|
||||
|
||||
Revision ID: e8f9a0b1c2d3
|
||||
Revises: 02650cd6f028, d4e5f6a7b8c9
|
||||
Create Date: 2026-04-24 00:00:00.000000
|
||||
"""
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "e8f9a0b1c2d3"
|
||||
down_revision = ("02650cd6f028", "d4e5f6a7b8c9")
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
@@ -0,0 +1,25 @@
|
||||
"""add created_at to api_keys
|
||||
|
||||
Revision ID: f1a2b3c4d5e6
|
||||
Revises: cli_tokens_001
|
||||
Create Date: 2026-06-01 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "f1a2b3c4d5e6"
|
||||
down_revision = "cli_tokens_001"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Nullable on purpose: existing keys keep NULL (unknown creation time) and
|
||||
# sort last; new keys get populated by the model's default_factory.
|
||||
op.add_column("api_keys", sa.Column("created_at", sa.Integer(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("api_keys", "created_at")
|
||||
355
plans/other-mints-balance-tracking-plan.md
Normal file
355
plans/other-mints-balance-tracking-plan.md
Normal file
@@ -0,0 +1,355 @@
|
||||
# Plan: Track unsupported incoming mints in `other_mints` and include them in balances
|
||||
|
||||
## Goal
|
||||
|
||||
When a Cashu token arrives from a mint that is **not** in `settings.cashu_mints`, we currently create/use a wallet for that mint and swap value into the primary mint. Any change/surplus left behind after `melt()` remains in the foreign `token_wallet`, but that balance is not surfaced by `fetch_all_balances()` because it only iterates over configured mints.
|
||||
|
||||
This plan adds persistent tracking for those foreign mints in a new database table called `other_mints`, and updates balance reporting to include them.
|
||||
|
||||
---
|
||||
|
||||
## Current behavior
|
||||
|
||||
### Incoming unsupported mint flow
|
||||
|
||||
In `routstr/wallet.py`:
|
||||
|
||||
- `recieve_token()` deserializes the token
|
||||
- if `token_obj.mint not in settings.cashu_mints`, it calls `swap_to_primary_mint(token_obj, wallet)`
|
||||
- `swap_to_primary_mint()` calls `token_wallet.melt(...)` to pay the primary mint invoice
|
||||
|
||||
### Important detail: change is retained, not discarded
|
||||
|
||||
The underlying Cashu wallet library keeps any melt change:
|
||||
|
||||
- `Wallet.melt()` constructs blank outputs for change
|
||||
- when the melt succeeds, returned change is reconstructed into proofs
|
||||
- those proofs are appended to `self.proofs` and stored in the wallet DB
|
||||
|
||||
So surplus from unsupported mints is **not discarded**, but it may become invisible operationally.
|
||||
|
||||
### Visibility problem
|
||||
|
||||
`fetch_all_balances()` currently only loops over:
|
||||
|
||||
- `settings.cashu_mints`
|
||||
- units `sat` and `msat`
|
||||
|
||||
This means balances left on unsupported mints are not shown in admin balance reporting.
|
||||
|
||||
---
|
||||
|
||||
## Proposed design
|
||||
|
||||
## 1. Add a new DB table: `other_mints`
|
||||
|
||||
Add a small table in `routstr/core/db.py` to persist unsupported mints we have seen in incoming tokens.
|
||||
|
||||
Suggested schema:
|
||||
|
||||
- `mint_url: str` primary key
|
||||
- `created_at: int`
|
||||
- `last_seen_at: int`
|
||||
|
||||
Minimal model:
|
||||
|
||||
```python
|
||||
class OtherMint(SQLModel, table=True):
|
||||
__tablename__ = "other_mints"
|
||||
|
||||
mint_url: str = Field(primary_key=True)
|
||||
created_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
last_seen_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
```
|
||||
|
||||
Why minimal:
|
||||
|
||||
- the only required function is mint discovery/tracking
|
||||
- unit handling can remain dynamic via existing balance queries over `sat` and `msat`
|
||||
|
||||
---
|
||||
|
||||
## 2. Add DB helpers for `other_mints`
|
||||
|
||||
In `routstr/core/db.py`, add helper functions:
|
||||
|
||||
### `register_other_mint(mint_url: str) -> None`
|
||||
|
||||
Behavior:
|
||||
|
||||
- if the mint is not present, insert it
|
||||
- if it already exists, update `last_seen_at`
|
||||
|
||||
### `list_other_mints(session) -> list[str]`
|
||||
|
||||
Behavior:
|
||||
|
||||
- return all tracked unsupported mint URLs
|
||||
|
||||
Optional later:
|
||||
|
||||
- `delete_other_mint(...)`
|
||||
- admin cleanup helpers
|
||||
|
||||
---
|
||||
|
||||
## 3. Register unsupported mints during token receipt
|
||||
|
||||
Update `recieve_token()` in `routstr/wallet.py`.
|
||||
|
||||
Current logic:
|
||||
|
||||
```python
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
```
|
||||
|
||||
Planned logic:
|
||||
|
||||
```python
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
await db.register_other_mint(token_obj.mint)
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
```
|
||||
|
||||
Why here:
|
||||
|
||||
- this is the earliest reliable point where we know the mint came in via an actual token
|
||||
- this is exactly the path that can leave foreign-mint change behind
|
||||
- it avoids needing to infer unsupported mints later from wallet internals
|
||||
|
||||
---
|
||||
|
||||
## 4. Update `fetch_all_balances()` to include `other_mints`
|
||||
|
||||
Current behavior only includes configured mints.
|
||||
|
||||
Planned behavior:
|
||||
|
||||
- load tracked unsupported mints from DB
|
||||
- combine them with `settings.cashu_mints`
|
||||
- dedupe while preserving order
|
||||
- fetch balances for all tracked mints across requested units
|
||||
|
||||
Conceptual flow:
|
||||
|
||||
```python
|
||||
tracked_mints = dedupe(settings.cashu_mints + other_mints_from_db)
|
||||
```
|
||||
|
||||
Then existing per-mint/per-unit balance logic can remain mostly unchanged.
|
||||
|
||||
This ensures that retained change on unsupported mints becomes visible in admin balance reporting.
|
||||
|
||||
---
|
||||
|
||||
## 5. Add a balance source marker
|
||||
|
||||
Extend `BalanceDetail` in `routstr/wallet.py` to identify whether a balance row comes from a configured mint or an `other_mints` entry.
|
||||
|
||||
Suggested field:
|
||||
|
||||
- `source: str` with values:
|
||||
- `"configured"`
|
||||
- `"other"`
|
||||
|
||||
Updated shape:
|
||||
|
||||
```python
|
||||
class BalanceDetail(TypedDict, total=False):
|
||||
mint_url: str
|
||||
unit: str
|
||||
source: str
|
||||
wallet_balance: int
|
||||
user_balance: int
|
||||
owner_balance: int
|
||||
error: str
|
||||
```
|
||||
|
||||
Why this helps:
|
||||
|
||||
- admin can distinguish normal configured wallet balances from foreign/unsupported balances
|
||||
- avoids confusion if unexpected mint URLs show up in the balances API/UI
|
||||
|
||||
---
|
||||
|
||||
## 6. Admin/API impact
|
||||
|
||||
Backend impact is minimal because `/admin/api/balances` already returns `fetch_all_balances()` output.
|
||||
|
||||
Effects:
|
||||
|
||||
- supported mints continue to show as before
|
||||
- tracked unsupported mints will also appear
|
||||
- UI can optionally display the new `source` field
|
||||
|
||||
No API contract break is expected if the frontend ignores unknown fields.
|
||||
|
||||
---
|
||||
|
||||
## 7. Payout behavior: do not change in phase 1
|
||||
|
||||
`periodic_payout()` currently only iterates over `settings.cashu_mints`.
|
||||
|
||||
Recommendation for this change:
|
||||
|
||||
- **do not** expand `periodic_payout()` to include `other_mints` yet
|
||||
- only improve visibility through balance reporting
|
||||
|
||||
Reason:
|
||||
|
||||
- automatic payout from unsupported/foreign mints may be operationally undesirable
|
||||
- visibility should come first, automation second
|
||||
|
||||
Possible future phase:
|
||||
|
||||
- add optional sweeping/payout support for `other_mints`
|
||||
- or provide an admin-triggered withdrawal/sweep flow
|
||||
|
||||
---
|
||||
|
||||
## 8. Logging improvements (optional)
|
||||
|
||||
Optional follow-up improvement in `swap_to_primary_mint()`:
|
||||
|
||||
- capture the return value from `token_wallet.melt(...)`
|
||||
- if feasible, log any reported change amount
|
||||
- otherwise, rely on wallet balance reporting to surface residual amounts
|
||||
|
||||
This is useful but not required for the first implementation.
|
||||
|
||||
---
|
||||
|
||||
## Files to change
|
||||
|
||||
### `routstr/core/db.py`
|
||||
|
||||
Add:
|
||||
|
||||
- `OtherMint` SQLModel
|
||||
- `register_other_mint()`
|
||||
- `list_other_mints()`
|
||||
|
||||
### `migrations/versions/<new_revision>_add_other_mints_table.py`
|
||||
|
||||
Create migration to add the `other_mints` table.
|
||||
|
||||
### `routstr/wallet.py`
|
||||
|
||||
Update:
|
||||
|
||||
- `recieve_token()` to register unsupported mints
|
||||
- `BalanceDetail` to include `source`
|
||||
- `fetch_all_balances()` to include both configured and tracked unsupported mints
|
||||
|
||||
### `routstr/core/admin.py`
|
||||
|
||||
Likely no backend changes required unless a dedicated `other_mints` API is desired.
|
||||
|
||||
---
|
||||
|
||||
## Behavior rules
|
||||
|
||||
### Register a mint when
|
||||
|
||||
- an incoming token is processed
|
||||
- the token mint is not in `settings.cashu_mints`
|
||||
|
||||
### Do not remove automatically when
|
||||
|
||||
- balance reaches zero
|
||||
|
||||
Reason:
|
||||
|
||||
- historical visibility is useful
|
||||
- avoids flapping entries in the admin balance list
|
||||
- mint may receive additional unsupported tokens later
|
||||
|
||||
Potential future enhancement:
|
||||
|
||||
- admin endpoint to prune zero-balance `other_mints`
|
||||
|
||||
---
|
||||
|
||||
## Edge cases
|
||||
|
||||
### A mint later becomes configured
|
||||
|
||||
If a mint in `other_mints` is later added to `settings.cashu_mints`:
|
||||
|
||||
- deduplication prevents duplicate balance rows
|
||||
- `source` should resolve to `configured`
|
||||
|
||||
### Unsupported mint with zero balance
|
||||
|
||||
A tracked unsupported mint may show zero balances.
|
||||
|
||||
Initial recommendation:
|
||||
|
||||
- allow it to appear
|
||||
- consider later filtering zero-balance `other` rows if the UI becomes noisy
|
||||
|
||||
### Units
|
||||
|
||||
Balance fetching can continue to query both `sat` and `msat` for each tracked mint.
|
||||
|
||||
If a mint has no proofs in one unit, current error/zero handling can continue to apply.
|
||||
|
||||
---
|
||||
|
||||
## Test plan
|
||||
|
||||
### DB tests
|
||||
|
||||
- registering a new unsupported mint inserts a row
|
||||
- registering the same mint again updates `last_seen_at` without duplication
|
||||
- listing other mints returns expected mint URLs
|
||||
|
||||
### Wallet tests
|
||||
|
||||
#### `recieve_token()`
|
||||
|
||||
- when mint is unsupported, `db.register_other_mint()` is called before swap
|
||||
- when mint is configured, `db.register_other_mint()` is not called
|
||||
|
||||
#### `fetch_all_balances()`
|
||||
|
||||
- includes configured mints
|
||||
- includes `other_mints` from DB
|
||||
- dedupes if a mint exists in both configured and other lists
|
||||
- sets `source` correctly
|
||||
|
||||
### Regression tests
|
||||
|
||||
- existing trusted mint balance reporting remains unchanged
|
||||
- `/admin/api/balances` continues to work
|
||||
|
||||
---
|
||||
|
||||
## Recommended implementation order
|
||||
|
||||
1. Add `OtherMint` model to `routstr/core/db.py`
|
||||
2. Add Alembic migration for `other_mints`
|
||||
3. Add `register_other_mint()` and `list_other_mints()` helpers
|
||||
4. Update `recieve_token()` to register unsupported mints
|
||||
5. Update `fetch_all_balances()` to union configured + tracked other mints
|
||||
6. Add `source` to `BalanceDetail`
|
||||
7. Add/adjust tests
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
This change solves an operational visibility problem:
|
||||
|
||||
- unsupported incoming mints can leave retained change in foreign wallets
|
||||
- those funds are currently preserved but not surfaced in balance reporting
|
||||
- introducing `other_mints` makes those mints discoverable and auditable
|
||||
- expanding `fetch_all_balances()` ensures their balances are visible in admin tooling
|
||||
|
||||
Recommended scope for the first pass:
|
||||
|
||||
- track unsupported mints in DB
|
||||
- include them in balance reporting
|
||||
- mark them as `source="other"`
|
||||
- do not yet change payout/sweeping behavior
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "routstr"
|
||||
version = "0.3.0"
|
||||
version = "0.4.3"
|
||||
description = "Payment proxy for your LLM endpoint using cashu and nostr."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -13,14 +13,14 @@ dependencies = [
|
||||
"greenlet>=3.2.1",
|
||||
"alembic>=1.13",
|
||||
"python-json-logger>=2.0.0",
|
||||
"cashu",
|
||||
"secp256k1",
|
||||
"cashu>=0.20",
|
||||
"marshmallow>=3.13,<4.0",
|
||||
"websockets>=12.0",
|
||||
"nostr>=0.0.2",
|
||||
"mdurl==0.1.2",
|
||||
"pillow>=10",
|
||||
"openai>=1.98.0",
|
||||
"litellm>=1.55.0",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
@@ -86,4 +86,3 @@ disallow_untyped_decorators = true
|
||||
|
||||
[tool.uv.sources]
|
||||
routstr = { workspace = true }
|
||||
secp256k1 = { git = "https://github.com/saschanaz/secp256k1-py", branch = "upgrade060" }
|
||||
|
||||
@@ -118,12 +118,32 @@ def create_model_mappings(
|
||||
|
||||
candidates: dict[str, list[tuple["Model", "BaseUpstreamProvider"]]] = {}
|
||||
unique_models: dict[str, "Model"] = {}
|
||||
seen_model_provider: set[tuple[str, str]] = set()
|
||||
|
||||
providers_by_db_id: dict[int, "BaseUpstreamProvider"] = {}
|
||||
for upstream in upstreams:
|
||||
db_id = getattr(upstream, "db_id", None)
|
||||
if isinstance(db_id, int):
|
||||
providers_by_db_id[db_id] = upstream
|
||||
|
||||
# Group upstreams by URL and keep only the one with the lowest fee for each URL
|
||||
upstreams_by_url: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
for upstream in upstreams:
|
||||
url = getattr(upstream, "base_url", "")
|
||||
if url not in upstreams_by_url:
|
||||
upstreams_by_url[url] = []
|
||||
upstreams_by_url[url].append(upstream)
|
||||
|
||||
filtered_upstreams: list["BaseUpstreamProvider"] = []
|
||||
for providers in upstreams_by_url.values():
|
||||
best_provider = min(providers, key=lambda p: p.provider_fee)
|
||||
filtered_upstreams.append(best_provider)
|
||||
|
||||
# Separate OpenRouter from other providers
|
||||
openrouter: "BaseUpstreamProvider" | None = None
|
||||
other_upstreams: list["BaseUpstreamProvider"] = []
|
||||
|
||||
for upstream in upstreams:
|
||||
for upstream in filtered_upstreams:
|
||||
base_url = getattr(upstream, "base_url", "")
|
||||
if base_url == "https://openrouter.ai/api/v1":
|
||||
openrouter = upstream
|
||||
@@ -134,6 +154,16 @@ def create_model_mappings(
|
||||
"""Get base model ID by removing provider prefix."""
|
||||
return model_id.split("/", 1)[1] if "/" in model_id else model_id
|
||||
|
||||
def get_provider_identity(upstream: "BaseUpstreamProvider") -> str:
|
||||
"""Get a stable provider identity used for deduplication."""
|
||||
db_id = getattr(upstream, "db_id", None)
|
||||
if isinstance(db_id, int):
|
||||
return f"db:{db_id}"
|
||||
|
||||
provider_type = str(getattr(upstream, "provider_type", "") or "").lower()
|
||||
base_url = str(getattr(upstream, "base_url", "") or "").lower()
|
||||
return f"{provider_type}|{base_url}"
|
||||
|
||||
def _add_candidate(
|
||||
alias: str, model: "Model", provider: "BaseUpstreamProvider"
|
||||
) -> None:
|
||||
@@ -148,6 +178,7 @@ def create_model_mappings(
|
||||
) -> None:
|
||||
"""Process all models from a given provider."""
|
||||
upstream_prefix = getattr(upstream, "upstream_name", None)
|
||||
provider_key = get_provider_identity(upstream)
|
||||
|
||||
for model in upstream.get_cached_models():
|
||||
if not model.enabled or model.id in disabled_model_ids:
|
||||
@@ -164,14 +195,15 @@ def create_model_mappings(
|
||||
|
||||
# Add to unique models
|
||||
base_id = get_base_model_id(model_to_use.id)
|
||||
if not is_openrouter or base_id not in unique_models:
|
||||
unique_key = model_to_use.forwarded_model_id or base_id
|
||||
if not is_openrouter or unique_key not in unique_models:
|
||||
unique_model = model_to_use.copy(
|
||||
update={
|
||||
"id": base_id,
|
||||
"upstream_provider_id": upstream.provider_type,
|
||||
}
|
||||
)
|
||||
unique_models[base_id] = unique_model
|
||||
unique_models[unique_key] = unique_model
|
||||
|
||||
# Get all aliases for this model
|
||||
aliases = resolve_model_alias(
|
||||
@@ -186,9 +218,14 @@ def create_model_mappings(
|
||||
if prefixed_id not in aliases:
|
||||
aliases.append(prefixed_id)
|
||||
|
||||
# Register forwarded_model_id as a routable alias
|
||||
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
|
||||
aliases.append(model_to_use.forwarded_model_id)
|
||||
|
||||
# Try to set each alias
|
||||
for alias in aliases:
|
||||
_add_candidate(alias, model_to_use, upstream)
|
||||
seen_model_provider.add((model_to_use.id.lower(), provider_key))
|
||||
|
||||
# Process non-OpenRouter providers first
|
||||
for upstream in other_upstreams:
|
||||
@@ -198,12 +235,115 @@ def create_model_mappings(
|
||||
if openrouter:
|
||||
process_provider_models(openrouter, is_openrouter=True)
|
||||
|
||||
# Include enabled DB overrides even when provider discovery misses models.
|
||||
# This is important for deployment-based providers like Azure.
|
||||
for model_id, override_data in overrides_by_id.items():
|
||||
if model_id in disabled_model_ids:
|
||||
continue
|
||||
override_row, provider_fee = override_data
|
||||
upstream_provider_id = getattr(override_row, "upstream_provider_id", None)
|
||||
if not isinstance(upstream_provider_id, int):
|
||||
continue
|
||||
|
||||
upstream_for_override = providers_by_db_id.get(upstream_provider_id)
|
||||
if upstream_for_override is None:
|
||||
continue
|
||||
|
||||
provider_key = get_provider_identity(upstream_for_override)
|
||||
dedupe_key = (model_id.lower(), provider_key)
|
||||
if dedupe_key in seen_model_provider:
|
||||
continue
|
||||
|
||||
try:
|
||||
model_to_use = _row_to_model(
|
||||
override_row, apply_provider_fee=True, provider_fee=provider_fee
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Skipping invalid model override while building model mappings",
|
||||
extra={
|
||||
"model_id": model_id,
|
||||
"upstream_provider_id": upstream_provider_id,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
},
|
||||
)
|
||||
continue
|
||||
if not model_to_use.enabled:
|
||||
continue
|
||||
|
||||
base_id = get_base_model_id(model_to_use.id)
|
||||
unique_key = model_to_use.forwarded_model_id or base_id
|
||||
is_openrouter = (
|
||||
getattr(upstream_for_override, "base_url", "")
|
||||
== "https://openrouter.ai/api/v1"
|
||||
)
|
||||
if not is_openrouter or unique_key not in unique_models:
|
||||
unique_model = model_to_use.copy(
|
||||
update={
|
||||
"id": base_id,
|
||||
"upstream_provider_id": upstream_for_override.provider_type,
|
||||
}
|
||||
)
|
||||
unique_models[unique_key] = unique_model
|
||||
|
||||
try:
|
||||
aliases = resolve_model_alias(
|
||||
model_to_use.id,
|
||||
model_to_use.canonical_slug,
|
||||
alias_ids=model_to_use.alias_ids,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Skipping model aliases for invalid override model",
|
||||
extra={
|
||||
"model_id": model_id,
|
||||
"upstream_provider_id": upstream_provider_id,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
},
|
||||
)
|
||||
continue
|
||||
|
||||
upstream_prefix = getattr(upstream_for_override, "upstream_name", None)
|
||||
if upstream_prefix and "/" not in model_to_use.id:
|
||||
prefixed_id = f"{upstream_prefix}/{model_to_use.id}"
|
||||
if prefixed_id not in aliases:
|
||||
aliases.append(prefixed_id)
|
||||
|
||||
# Register forwarded_model_id as a routable alias
|
||||
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
|
||||
aliases.append(model_to_use.forwarded_model_id)
|
||||
|
||||
for alias in aliases:
|
||||
_add_candidate(alias, model_to_use, upstream_for_override)
|
||||
seen_model_provider.add(dedupe_key)
|
||||
|
||||
# Sort candidates and build final maps
|
||||
model_instances: dict[str, "Model"] = {}
|
||||
provider_map: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
|
||||
def alias_priority(model: "Model", alias: str) -> int:
|
||||
"""Rank how strong the mapping of alias->model is."""
|
||||
"""Rank how strong the mapping of alias->model is.
|
||||
|
||||
forwarded_model_id is the most specific identifier (set per-provider
|
||||
instance), so a match there should beat a model_id match. This way,
|
||||
when multiple providers have the same model_id but different
|
||||
forwarded_model_ids, the one whose forwarded_model_id equals the
|
||||
requested alias wins.
|
||||
"""
|
||||
if (
|
||||
model.forwarded_model_id
|
||||
and model.forwarded_model_id.lower() == alias
|
||||
):
|
||||
return 5
|
||||
|
||||
if (
|
||||
model.id
|
||||
and model.id.lower() == alias
|
||||
):
|
||||
return 4
|
||||
|
||||
model_base = get_base_model_id(model.id)
|
||||
if model_base == alias:
|
||||
return 3
|
||||
|
||||
436
routstr/auth.py
436
routstr/auth.py
@@ -7,11 +7,12 @@ from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy import case
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlmodel import col, select, update
|
||||
|
||||
from .core import get_logger
|
||||
from .core.db import ApiKey, AsyncSession
|
||||
from .core.db import ApiKey, AsyncSession, accumulate_routstr_fee
|
||||
from .core.settings import settings
|
||||
from .payment.cost_calculation import (
|
||||
CostData,
|
||||
@@ -22,6 +23,13 @@ from .payment.cost_calculation import (
|
||||
from .wallet import credit_balance, deserialize_token_from_string
|
||||
|
||||
logger = get_logger(__name__)
|
||||
payments_logger = get_logger("routstr.payments")
|
||||
|
||||
# Routstr platform fee constants
|
||||
ROUTSTR_FEE_PERCENT: float = 2.1
|
||||
ROUTSTR_LN_ADDRESS: str = "npub130mznv74rxs032peqym6g3wqavh472623mt3z5w73xq9r6qqdufs7ql29s@npub.cash"
|
||||
ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS: int = 900
|
||||
ROUTSTR_FEE_DEFAULT_PAYOUT: int = 200
|
||||
|
||||
# TODO: implement prepaid api key (not like it was before)
|
||||
# PREPAID_API_KEY = os.environ.get("PREPAID_API_KEY", None)
|
||||
@@ -309,13 +317,13 @@ async def validate_bearer_key(
|
||||
extra={"key_hash": hashed_key[:8] + "..."},
|
||||
)
|
||||
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"AUTH: About to call credit_balance",
|
||||
extra={"token_preview": bearer_key[:50]},
|
||||
)
|
||||
try:
|
||||
msats = await credit_balance(bearer_key, new_key, session)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"AUTH: credit_balance returned successfully", extra={"msats": msats}
|
||||
)
|
||||
except Exception as credit_error:
|
||||
@@ -333,6 +341,11 @@ async def validate_bearer_key(
|
||||
"Token redemption returned zero or negative amount",
|
||||
extra={"msats": msats, "key_hash": hashed_key[:8] + "..."},
|
||||
)
|
||||
# Defense-in-depth: credit_balance now refuses to commit on a
|
||||
# zero/negative redemption, but if a row was nonetheless
|
||||
# persisted, drop it so we never leave an orphan zero-balance key.
|
||||
await session.delete(new_key)
|
||||
await session.commit()
|
||||
raise Exception("Token redemption failed")
|
||||
|
||||
await session.refresh(new_key)
|
||||
@@ -348,6 +361,8 @@ async def validate_bearer_key(
|
||||
)
|
||||
|
||||
return new_key
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Cashu token redemption failed",
|
||||
@@ -370,12 +385,12 @@ async def validate_bearer_key(
|
||||
},
|
||||
)
|
||||
|
||||
key_preview = bearer_key[:10] + "..." if len(bearer_key) > 10 else bearer_key
|
||||
logger.error(
|
||||
"Invalid API key format",
|
||||
f"Invalid API key format: preview={key_preview!r} length={len(bearer_key)} "
|
||||
f"(expected 'sk-...' or 'cashu...' token)",
|
||||
extra={
|
||||
"key_preview": bearer_key[:10] + "..."
|
||||
if len(bearer_key) > 10
|
||||
else bearer_key,
|
||||
"key_preview": key_preview,
|
||||
"key_length": len(bearer_key),
|
||||
},
|
||||
)
|
||||
@@ -384,7 +399,7 @@ async def validate_bearer_key(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Invalid API key",
|
||||
"message": "Invalid API key format. Expected an 'sk-...' API key or a 'cashu...' token.",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
@@ -519,31 +534,19 @@ async def pay_for_request(
|
||||
)
|
||||
|
||||
# Charge the base cost for the request atomically to avoid race conditions
|
||||
reserved_at_now = int(time.time())
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
reserved_at=reserved_at_now,
|
||||
total_requests=col(ApiKey.total_requests) + 1,
|
||||
)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also increment total_requests and reserved_balance on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_requests=col(ApiKey.total_requests) + 1,
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Concurrent request depleted balance",
|
||||
@@ -555,7 +558,6 @@ async def pay_for_request(
|
||||
},
|
||||
)
|
||||
|
||||
# Another concurrent request spent the balance first
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
@@ -567,6 +569,44 @@ async def pay_for_request(
|
||||
},
|
||||
)
|
||||
|
||||
# Also increment total_requests and reserved_balance on the child key if it's different.
|
||||
# The balance_limit guard is enforced atomically here — the Python pre-check above
|
||||
# is a fast-path rejection only and provides no concurrency guarantee.
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(
|
||||
(col(ApiKey.balance_limit).is_(None))
|
||||
| (
|
||||
col(ApiKey.total_spent)
|
||||
+ col(ApiKey.reserved_balance)
|
||||
+ cost_per_request
|
||||
<= col(ApiKey.balance_limit)
|
||||
)
|
||||
)
|
||||
.values(
|
||||
total_requests=col(ApiKey.total_requests) + 1,
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
reserved_at=reserved_at_now,
|
||||
)
|
||||
)
|
||||
child_result = await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
if child_result.rowcount == 0:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Balance limit exceeded: {key.balance_limit} mSats limit. {key.total_spent} already spent ({key.reserved_balance} reserved), {cost_per_request} required for this request.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "balance_limit_exceeded",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
await session.commit()
|
||||
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
@@ -582,20 +622,42 @@ async def pay_for_request(
|
||||
"total_requests": billing_key.total_requests,
|
||||
},
|
||||
)
|
||||
payments_logger.info(
|
||||
"RESERVE",
|
||||
extra={
|
||||
"event": "reserve",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_reserved": cost_per_request,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
},
|
||||
)
|
||||
|
||||
return cost_per_request
|
||||
|
||||
|
||||
async def revert_pay_for_request(
|
||||
key: ApiKey, session: AsyncSession, cost_per_request: int
|
||||
) -> None:
|
||||
) -> bool:
|
||||
"""Revert a previously reserved payment. Returns True if revert succeeded,
|
||||
False if the reservation was already released (prevents negative reserved_balance)."""
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
# Keep reserved_at while other reservations remain
|
||||
cleared_reserved_at = case(
|
||||
(col(ApiKey.reserved_balance) - cost_per_request > 0, col(ApiKey.reserved_at)),
|
||||
else_=None,
|
||||
)
|
||||
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
reserved_at=cleared_reserved_at,
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
)
|
||||
)
|
||||
@@ -607,17 +669,19 @@ async def revert_pay_for_request(
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
reserved_at=cleared_reserved_at,
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.error(
|
||||
"Failed to revert payment - insufficient reserved balance",
|
||||
logger.warning(
|
||||
"Revert skipped - reservation already released (no-op to prevent negative reserved_balance)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
@@ -625,28 +689,37 @@ async def revert_pay_for_request(
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"type": "payment_error",
|
||||
"code": "payment_error",
|
||||
}
|
||||
},
|
||||
)
|
||||
return False
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
payments_logger.info(
|
||||
"REVERT",
|
||||
extra={
|
||||
"event": "revert",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_reverted": cost_per_request,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
key: ApiKey, response_data: dict, session: AsyncSession, deducted_max_cost: int
|
||||
key: ApiKey,
|
||||
response_data: dict,
|
||||
session: AsyncSession,
|
||||
deducted_max_cost: int,
|
||||
) -> dict:
|
||||
"""
|
||||
Adjusts the payment based on token usage in the response.
|
||||
This is called after the initial payment and the upstream request is complete.
|
||||
Returns cost data to be included in the response.
|
||||
|
||||
The response's usage object is normalized with the default union parser in
|
||||
``calculate_cost``.
|
||||
"""
|
||||
billing_key = await get_billing_key(key, session)
|
||||
model = response_data.get("model", "unknown")
|
||||
@@ -669,17 +742,19 @@ async def adjust_payment_for_tokens(
|
||||
release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
|
||||
)
|
||||
)
|
||||
await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
result = await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also release on child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost
|
||||
@@ -688,14 +763,24 @@ async def adjust_payment_for_tokens(
|
||||
await session.exec(child_release_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
logger.warning(
|
||||
"Released reservation without charging (fallback)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
if result.rowcount == 0: # type: ignore[union-attr]
|
||||
logger.warning(
|
||||
"Release reservation skipped - already released (no-op to prevent negative reserved_balance)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Released reservation without charging (fallback)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to release reservation in fallback",
|
||||
@@ -706,7 +791,18 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost, session):
|
||||
async def _accumulate_fee(total_cost_msats: int) -> None:
|
||||
if total_cost_msats > 0 and ROUTSTR_FEE_PERCENT > 0:
|
||||
fee_msats = math.ceil(total_cost_msats * ROUTSTR_FEE_PERCENT / 100)
|
||||
try:
|
||||
await accumulate_routstr_fee(session, fee_msats)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to accumulate Routstr fee",
|
||||
extra={"error": str(e), "fee_msats": fee_msats},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost):
|
||||
case MaxCostData() as cost:
|
||||
logger.debug(
|
||||
"Using max cost data (no token adjustment)",
|
||||
@@ -718,11 +814,32 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
# Finalize by releasing reservation and charging max cost
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost before MaxCost finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": cost.total_msats,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
reserved_balance=safe_reserved,
|
||||
balance=col(ApiKey.balance) - cost.total_msats,
|
||||
total_spent=col(ApiKey.total_spent) + cost.total_msats,
|
||||
)
|
||||
@@ -731,13 +848,17 @@ async def adjust_payment_for_tokens(
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + cost.total_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=child_safe_reserved,
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
@@ -766,10 +887,30 @@ async def adjust_payment_for_tokens(
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost.total_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"new_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(cost.total_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": cost.total_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "max_cost",
|
||||
},
|
||||
)
|
||||
return cost.dict()
|
||||
|
||||
case CostData() as cost:
|
||||
@@ -789,6 +930,8 @@ async def adjust_payment_for_tokens(
|
||||
"cost_difference": cost_difference,
|
||||
"input_msats": cost.input_msats,
|
||||
"output_msats": cost.output_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -801,12 +944,32 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost on exact-cost finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": total_cost_msats,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
exact_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=exact_safe_reserved,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
)
|
||||
@@ -815,13 +978,17 @@ async def adjust_payment_for_tokens(
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_exact_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=child_exact_safe_reserved,
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
@@ -830,44 +997,56 @@ async def adjust_payment_for_tokens(
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
return cost.dict()
|
||||
|
||||
# this should never happen why do we handle this???
|
||||
if cost_difference > 0:
|
||||
# Need to charge more than reserved, finalize by releasing reservation and charging total
|
||||
logger.info(
|
||||
"Additional charge required for token usage",
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"additional_charge": cost_difference,
|
||||
"current_balance": billing_key.balance,
|
||||
"sufficient_balance": billing_key.balance >= cost_difference,
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "exact",
|
||||
},
|
||||
)
|
||||
return cost.dict()
|
||||
|
||||
# actual cost exceeded discounted reservation (due to tolerance_percentage)
|
||||
if cost_difference > 0:
|
||||
# Always release the reservation and charge min(actual_cost, balance).
|
||||
# Using a CASE expression makes this a single atomic UPDATE — no
|
||||
# multi-level fallback needed and balance can never go negative.
|
||||
chargeable = case(
|
||||
(col(ApiKey.balance) >= total_cost_msats, total_cost_msats),
|
||||
else_=col(ApiKey.balance),
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - chargeable,
|
||||
total_spent=col(ApiKey.total_spent) + chargeable,
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats),
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
@@ -875,11 +1054,10 @@ async def adjust_payment_for_tokens(
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
|
||||
cost.total_msats = total_cost_msats
|
||||
logger.info(
|
||||
"Finalized payment with additional charge",
|
||||
extra={
|
||||
@@ -890,9 +1068,29 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "overrun",
|
||||
},
|
||||
)
|
||||
else:
|
||||
# Guard fired: reservation was already released by a concurrent
|
||||
# finalization for this key. Nothing left to do.
|
||||
logger.warning(
|
||||
"Failed to finalize additional charge - releasing reservation",
|
||||
"Finalization skipped - reservation already released",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
@@ -900,7 +1098,6 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await release_reservation_only()
|
||||
else:
|
||||
# Refund some of the base cost
|
||||
refund = abs(cost_difference)
|
||||
@@ -915,12 +1112,33 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost on refund finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": total_cost_msats,
|
||||
"refund_amount": refund,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
refund_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=refund_safe_reserved,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
)
|
||||
@@ -929,13 +1147,17 @@ async def adjust_payment_for_tokens(
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_refund_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=child_refund_safe_reserved,
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
@@ -972,6 +1194,25 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"refunded": refund,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "refund",
|
||||
},
|
||||
)
|
||||
|
||||
return cost.dict()
|
||||
|
||||
@@ -997,18 +1238,8 @@ async def adjust_payment_for_tokens(
|
||||
}
|
||||
},
|
||||
)
|
||||
# Fallback: should not reach here, but release reservation just in case
|
||||
logger.error(
|
||||
"Unexpected fallback in adjust_payment_for_tokens - releasing reservation",
|
||||
extra={"key_hash": key.hashed_key[:8] + "...", "model": model},
|
||||
)
|
||||
await release_reservation_only()
|
||||
return {
|
||||
"base_msats": deducted_max_cost,
|
||||
"input_msats": 0,
|
||||
"output_msats": 0,
|
||||
"total_msats": deducted_max_cost,
|
||||
}
|
||||
# All calculate_cost variants are handled above.
|
||||
raise AssertionError("Unreachable: unhandled calculate_cost result")
|
||||
|
||||
|
||||
async def periodic_key_reset() -> None:
|
||||
@@ -1071,3 +1302,24 @@ async def periodic_key_reset() -> None:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in periodic_key_reset: {e}")
|
||||
|
||||
|
||||
STALE_RESERVATION_SWEEP_INTERVAL_SECONDS: int = 60
|
||||
|
||||
|
||||
async def periodic_stale_reservation_sweep() -> None:
|
||||
"""Background task that releases reservations leaked by client disconnects,
|
||||
crashes or abandoned streams.
|
||||
"""
|
||||
from .core.db import create_session, release_stale_reservations
|
||||
|
||||
while True:
|
||||
try:
|
||||
async with create_session() as session:
|
||||
await release_stale_reservations(
|
||||
session, settings.stale_reservation_timeout_seconds
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Error in periodic_stale_reservation_sweep")
|
||||
|
||||
await asyncio.sleep(STALE_RESERVATION_SWEEP_INTERVAL_SECONDS)
|
||||
|
||||
@@ -5,11 +5,18 @@ from time import monotonic
|
||||
from typing import Annotated, NoReturn
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel
|
||||
from sqlmodel import select
|
||||
from sqlmodel import col, or_, select, update
|
||||
|
||||
from .auth import get_billing_key, validate_bearer_key
|
||||
from .core.db import ApiKey, AsyncSession, get_session
|
||||
from .core.db import (
|
||||
ApiKey,
|
||||
AsyncSession,
|
||||
CashuTransaction,
|
||||
get_session,
|
||||
store_cashu_transaction,
|
||||
)
|
||||
from .core.logging import get_logger
|
||||
from .core.settings import settings
|
||||
from .lightning import lightning_router
|
||||
@@ -38,10 +45,9 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
billing_key = await get_billing_key(key, session)
|
||||
info = {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": billing_key.balance,
|
||||
"balance": billing_key.total_balance,
|
||||
"reserved": billing_key.reserved_balance,
|
||||
"is_child": key.parent_key_hash is not None,
|
||||
"parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
|
||||
"total_requests": key.total_requests,
|
||||
"total_spent": key.total_spent,
|
||||
"balance_limit": key.balance_limit,
|
||||
@@ -49,7 +55,9 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
"validity_date": key.validity_date,
|
||||
}
|
||||
|
||||
if not key.parent_key_hash:
|
||||
if key.parent_key_hash:
|
||||
info["parent_key_preview"] = key.parent_key_hash[:8] + "..."
|
||||
else:
|
||||
# Fetch child keys if this is a parent key
|
||||
statement = select(ApiKey).where(ApiKey.parent_key_hash == key.hashed_key)
|
||||
results = await session.exec(statement)
|
||||
@@ -154,9 +162,23 @@ async def topup_wallet_endpoint(
|
||||
raise HTTPException(status_code=400, detail="Token already spent")
|
||||
elif "invalid" in error_msg.lower() or "decode" in error_msg.lower():
|
||||
raise HTTPException(status_code=400, detail="Invalid token format")
|
||||
elif "insufficient" in error_msg.lower() or "melt fee" in error_msg.lower():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Token value is too small to cover swap fees. {error_msg}",
|
||||
)
|
||||
elif "failed to melt" in error_msg.lower():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Failed to swap foreign mint token. {error_msg}",
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Failed to redeem token")
|
||||
except Exception:
|
||||
raise HTTPException(status_code=400, detail=f"Failed to redeem token: {error_msg}")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"topup_wallet_endpoint: unhandled error",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
return {"msats": amount_msats}
|
||||
|
||||
@@ -190,23 +212,99 @@ async def _refund_cache_set(authorization: str, value: dict[str, str]) -> None:
|
||||
_refund_cache[key] = (expiry, value)
|
||||
|
||||
|
||||
@router.post("/refund")
|
||||
async def _lookup_key_no_create(
|
||||
bearer_value: str, session: AsyncSession
|
||||
) -> ApiKey | None:
|
||||
"""Look up an existing API key without creating one Used by the refund endpoint"""
|
||||
if bearer_value.startswith("sk-"):
|
||||
return await session.get(ApiKey, bearer_value[3:])
|
||||
if bearer_value.startswith("cashu"):
|
||||
hashed = hashlib.sha256(bearer_value.encode()).hexdigest()
|
||||
return await session.get(ApiKey, hashed)
|
||||
return None
|
||||
|
||||
|
||||
async def _restore_balance(
|
||||
session: AsyncSession, hashed_key: str, balance: int, reserved_balance: int, mint_url: str
|
||||
) -> None:
|
||||
"""Restore balance after a failed refund mint attempt."""
|
||||
restore_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == hashed_key)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) + balance,
|
||||
reserved_balance=col(ApiKey.reserved_balance) + reserved_balance,
|
||||
)
|
||||
)
|
||||
await session.exec(restore_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: balance restored after mint failure",
|
||||
extra={"hashed_key": hashed_key, "restored_balance": balance, "mint_url": mint_url},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/refund", response_model=None)
|
||||
async def refund_wallet_endpoint(
|
||||
authorization: Annotated[str, Header(...)],
|
||||
authorization: Annotated[str | None, Header()] = None,
|
||||
x_cashu: Annotated[str | None, Header()] = None,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, str]:
|
||||
if not authorization.startswith("Bearer "):
|
||||
) -> JSONResponse | dict[str, str]:
|
||||
if x_cashu:
|
||||
# Find the "in" transaction by the original payment token
|
||||
in_tx_result = await session.exec(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.token == x_cashu,
|
||||
CashuTransaction.type == "in",
|
||||
)
|
||||
)
|
||||
in_tx = in_tx_result.first()
|
||||
if in_tx is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
|
||||
# Use the request_id to find the associated "out" (refund) transaction
|
||||
if in_tx.request_id is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
|
||||
out_tx_result = await session.exec(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.request_id == in_tx.request_id,
|
||||
CashuTransaction.type == "out",
|
||||
)
|
||||
)
|
||||
out_tx = out_tx_result.first()
|
||||
if out_tx is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
if out_tx.swept:
|
||||
raise HTTPException(status_code=410, detail="Refund has been swept")
|
||||
|
||||
out_tx.collected = True
|
||||
session.add(out_tx)
|
||||
await session.commit()
|
||||
body: dict[str, str] = {"token": out_tx.token}
|
||||
if out_tx.unit == "sat":
|
||||
body["sats"] = str(out_tx.amount)
|
||||
else:
|
||||
body["msats"] = str(out_tx.amount)
|
||||
return JSONResponse(content=body, headers={"X-Cashu": out_tx.token})
|
||||
|
||||
if authorization is None or not authorization.startswith("Bearer "):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
|
||||
)
|
||||
|
||||
bearer_value: str = authorization[7:]
|
||||
key: ApiKey | None = await _lookup_key_no_create(bearer_value, session)
|
||||
if key is None:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Key not found. Deposit first via /v1/wallet/create before requesting a refund.",
|
||||
)
|
||||
|
||||
key: ApiKey = await validate_bearer_key(bearer_value, session)
|
||||
|
||||
if cached := await _refund_cache_get(bearer_value):
|
||||
return cached
|
||||
if key.total_balance <= 0:
|
||||
if cached := await _refund_cache_get(bearer_value):
|
||||
return cached
|
||||
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
@@ -214,6 +312,39 @@ async def refund_wallet_endpoint(
|
||||
detail="Cannot refund child key. Please refund the parent key instead.",
|
||||
)
|
||||
|
||||
if key.reserved_balance > 0:
|
||||
# Release the reservation if it is stale
|
||||
cutoff = int(time.time()) - settings.stale_reservation_timeout_seconds
|
||||
stale_release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) > 0)
|
||||
.where(
|
||||
or_(
|
||||
col(ApiKey.reserved_at).is_(None),
|
||||
col(ApiKey.reserved_at) < cutoff,
|
||||
)
|
||||
)
|
||||
.values(reserved_balance=0, reserved_at=None)
|
||||
)
|
||||
stale_result = await session.exec(stale_release_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
|
||||
if stale_result.rowcount == 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot refund key. There are ongoing requests for this api key.",
|
||||
)
|
||||
|
||||
await session.refresh(key)
|
||||
logger.warning(
|
||||
"refund_wallet_endpoint: released stale reservation before refund",
|
||||
extra={
|
||||
"hashed_key": key.hashed_key,
|
||||
"stale_timeout_seconds": settings.stale_reservation_timeout_seconds,
|
||||
},
|
||||
)
|
||||
|
||||
remaining_balance_msats: int = key.total_balance
|
||||
|
||||
if key.refund_currency == "sat":
|
||||
@@ -226,22 +357,51 @@ async def refund_wallet_endpoint(
|
||||
elif remaining_balance <= 0:
|
||||
raise HTTPException(status_code=400, detail="No balance to refund")
|
||||
|
||||
# Perform refund operation first, before modifying balance
|
||||
# Capture values before debit — the session may refresh key after commit
|
||||
pre_debit_balance = key.balance
|
||||
pre_debit_reserved = key.reserved_balance
|
||||
|
||||
# --- DEBIT FIRST: atomically zero the balance before minting tokens ---
|
||||
# This prevents the race where a concurrent topup/spend happens between
|
||||
# reading the balance and minting the refund token (double-spend).
|
||||
debit_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) == pre_debit_balance)
|
||||
.where(col(ApiKey.reserved_balance) == pre_debit_reserved)
|
||||
.values(balance=0, reserved_balance=0, reserved_at=None)
|
||||
)
|
||||
debit_result = await session.exec(debit_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
|
||||
if debit_result.rowcount == 0:
|
||||
# Balance changed between read and debit — another request is active
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Balance changed concurrently. Please retry the refund.",
|
||||
)
|
||||
|
||||
# --- MINT: balance is locked at zero, safe to create the refund token ---
|
||||
# Proofs from untrusted mints are swapped to primary_mint on receive.
|
||||
# Use primary_mint unless key.refund_mint_url is an explicitly trusted mint.
|
||||
effective_refund_mint = (
|
||||
key.refund_mint_url
|
||||
if key.refund_mint_url and key.refund_mint_url in settings.cashu_mints
|
||||
else settings.primary_mint
|
||||
)
|
||||
try:
|
||||
if key.refund_address:
|
||||
from .core.settings import settings as global_settings
|
||||
|
||||
await send_to_lnurl(
|
||||
remaining_balance,
|
||||
key.refund_currency or "sat",
|
||||
key.refund_mint_url or global_settings.primary_mint,
|
||||
effective_refund_mint,
|
||||
key.refund_address,
|
||||
)
|
||||
result = {"recipient": key.refund_address}
|
||||
else:
|
||||
refund_currency = key.refund_currency or "sat"
|
||||
token = await send_token(
|
||||
remaining_balance, refund_currency, key.refund_mint_url
|
||||
remaining_balance, refund_currency, effective_refund_mint
|
||||
)
|
||||
result = {"token": token}
|
||||
|
||||
@@ -250,32 +410,109 @@ async def refund_wallet_endpoint(
|
||||
else:
|
||||
result["msats"] = str(remaining_balance_msats)
|
||||
|
||||
if "token" in result:
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: cashu token issued",
|
||||
extra={
|
||||
"path": "/v1/wallet/refund",
|
||||
"token": result["token"],
|
||||
"amount": remaining_balance,
|
||||
"currency": key.refund_currency or "sat",
|
||||
},
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Re-raise HTTP exceptions (like 400 for balance too small)
|
||||
# Minting failed — restore the debited balance
|
||||
await _restore_balance(session, key.hashed_key, pre_debit_balance, pre_debit_reserved, key.refund_mint_url or "")
|
||||
raise
|
||||
except Exception as e:
|
||||
# If refund fails, don't modify the database
|
||||
# Minting failed — restore the debited balance
|
||||
await _restore_balance(session, key.hashed_key, pre_debit_balance, pre_debit_reserved, key.refund_mint_url or "")
|
||||
error_msg = str(e)
|
||||
logger.error(
|
||||
"refund_wallet_endpoint: mint/send failed",
|
||||
extra={
|
||||
"error": error_msg,
|
||||
"error_type": type(e).__name__,
|
||||
"hashed_key": key.hashed_key,
|
||||
"remaining_balance": remaining_balance,
|
||||
"refund_currency": key.refund_currency,
|
||||
"refund_mint_url": key.refund_mint_url,
|
||||
"has_refund_address": bool(key.refund_address),
|
||||
},
|
||||
)
|
||||
if (
|
||||
"mint" in error_msg.lower()
|
||||
or "connection" in error_msg.lower()
|
||||
or isinstance(e, Exception)
|
||||
and "ConnectError" in str(type(e))
|
||||
or "ConnectError" in str(type(e))
|
||||
):
|
||||
raise HTTPException(status_code=503, detail="Mint service unavailable")
|
||||
raise HTTPException(status_code=503, detail=f"Mint service unavailable: {error_msg}")
|
||||
else:
|
||||
raise HTTPException(status_code=500, detail="Refund failed")
|
||||
raise HTTPException(status_code=500, detail=f"Refund failed: {error_msg}")
|
||||
|
||||
await _refund_cache_set(bearer_value, result)
|
||||
|
||||
key.balance = 0
|
||||
key.reserved_balance = 0
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
if "token" in result:
|
||||
try:
|
||||
await store_cashu_transaction(
|
||||
token=result["token"],
|
||||
amount=remaining_balance,
|
||||
unit=key.refund_currency or "sat",
|
||||
mint_url=key.refund_mint_url,
|
||||
typ="out",
|
||||
collected=False,
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
except Exception:
|
||||
pass # store_cashu_transaction already logs
|
||||
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: refund successful",
|
||||
extra={
|
||||
"refunded_msats": remaining_balance_msats,
|
||||
"previous_reserved_balance": key.reserved_balance,
|
||||
},
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@router.get("/history")
|
||||
async def wallet_history(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, list[dict[str, str | int | bool | None]]]:
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot view child key history. Please use the parent key instead.",
|
||||
)
|
||||
|
||||
result = await session.exec(
|
||||
select(CashuTransaction)
|
||||
.where(CashuTransaction.api_key_hashed_key == key.hashed_key)
|
||||
.order_by(col(CashuTransaction.created_at).desc())
|
||||
)
|
||||
transactions = result.all()
|
||||
return {
|
||||
"transactions": [
|
||||
{
|
||||
"id": tx.id,
|
||||
"type": tx.type,
|
||||
"source": tx.source,
|
||||
"amount": tx.amount,
|
||||
"unit": tx.unit,
|
||||
"mint_url": tx.mint_url,
|
||||
"created_at": tx.created_at,
|
||||
"collected": tx.collected,
|
||||
"swept": tx.swept,
|
||||
}
|
||||
for tx in transactions
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/donate")
|
||||
async def donate(token: str, ref: str | None = None) -> str:
|
||||
try:
|
||||
@@ -326,10 +563,24 @@ async def create_child_key(
|
||||
detail=f"Insufficient balance to create {count} child keys. {total_cost} mSats required.",
|
||||
)
|
||||
|
||||
# Deduct cost from parent
|
||||
key.balance -= total_cost
|
||||
key.total_spent += total_cost
|
||||
session.add(key)
|
||||
# Deduct cost from parent atomically — guards against concurrent requests
|
||||
# that both pass the balance check above on stale in-memory state.
|
||||
deduct_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) - col(ApiKey.reserved_balance) >= total_cost)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) - total_cost,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost,
|
||||
)
|
||||
)
|
||||
result = await session.exec(deduct_stmt) # type: ignore[call-overload]
|
||||
|
||||
if result.rowcount == 0:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail=f"Insufficient balance to create {count} child keys. {total_cost} mSats required.",
|
||||
)
|
||||
|
||||
# Generate new keys
|
||||
import secrets
|
||||
@@ -354,6 +605,7 @@ async def create_child_key(
|
||||
new_keys.append("sk-" + new_key_hash)
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
|
||||
response_data = {
|
||||
"api_keys": new_keys,
|
||||
@@ -400,6 +652,7 @@ async def reset_child_key_spent(
|
||||
return {"success": True, "message": "Child key balance reset successfully."}
|
||||
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/{path:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE"],
|
||||
@@ -412,7 +665,7 @@ async def wallet_catch_all(path: str) -> NoReturn:
|
||||
)
|
||||
|
||||
|
||||
balance_router.include_router(lightning_router)
|
||||
balance_router.include_router(lightning_router, include_in_schema=False)
|
||||
balance_router.include_router(router)
|
||||
|
||||
deprecated_wallet_router = APIRouter(prefix="/v1/wallet", include_in_schema=False)
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
import asyncio
|
||||
import json
|
||||
import secrets
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, RootModel
|
||||
from pydantic.v1 import ValidationError as PydanticValidationError
|
||||
from sqlmodel import select
|
||||
|
||||
from ..payment.models import _row_to_model, list_models
|
||||
@@ -16,7 +18,15 @@ from ..wallet import (
|
||||
send_token,
|
||||
slow_filter_spend_proofs,
|
||||
)
|
||||
from .db import ApiKey, ModelRow, UpstreamProviderRow, create_session
|
||||
from .db import (
|
||||
ApiKey,
|
||||
CashuTransaction,
|
||||
CliToken,
|
||||
LightningInvoice,
|
||||
ModelRow,
|
||||
UpstreamProviderRow,
|
||||
create_session,
|
||||
)
|
||||
from .log_manager import log_manager
|
||||
from .logging import get_logger
|
||||
from .settings import SettingsService, settings
|
||||
@@ -27,40 +37,120 @@ admin_router = APIRouter(prefix="/admin", include_in_schema=False)
|
||||
|
||||
admin_sessions: dict[str, int] = {}
|
||||
ADMIN_SESSION_DURATION = 3600
|
||||
# Usage analytics remain queryable up to 12 months.
|
||||
MAX_USAGE_ANALYTICS_HOURS = 365 * 24
|
||||
|
||||
|
||||
def require_admin_api(request: Request) -> None:
|
||||
async def require_admin_api(request: Request) -> None:
|
||||
auth_header = request.headers.get("Authorization")
|
||||
if auth_header and auth_header.startswith("Bearer "):
|
||||
token = auth_header.split(" ", 1)[1]
|
||||
expiry = admin_sessions.get(token)
|
||||
if expiry and expiry > int(datetime.now(timezone.utc).timestamp()):
|
||||
if not auth_header or not auth_header.startswith("Bearer "):
|
||||
raise HTTPException(status_code=403, detail="Unauthorized")
|
||||
|
||||
token = auth_header.split(" ", 1)[1]
|
||||
now_ts = int(datetime.now(timezone.utc).timestamp())
|
||||
|
||||
# 1) Short-lived session token (in-memory)
|
||||
expiry = admin_sessions.get(token)
|
||||
if expiry and expiry > now_ts:
|
||||
return
|
||||
|
||||
# 2) Long-lived CLI token (DB-backed)
|
||||
async with create_session() as session:
|
||||
result = await session.exec(select(CliToken).where(CliToken.token == token))
|
||||
cli_token = result.first()
|
||||
if cli_token and (cli_token.expires_at is None or cli_token.expires_at > now_ts):
|
||||
cli_token.last_used_at = now_ts
|
||||
session.add(cli_token)
|
||||
await session.commit()
|
||||
return
|
||||
|
||||
raise HTTPException(status_code=403, detail="Unauthorized")
|
||||
|
||||
|
||||
@admin_router.get("/api/temporary-balances", dependencies=[Depends(require_admin_api)])
|
||||
async def get_temporary_balances_api(request: Request) -> list[dict[str, object]]:
|
||||
async def get_temporary_balances_api(
|
||||
request: Request,
|
||||
search: str | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> dict[str, object]:
|
||||
from sqlalchemy import case
|
||||
from sqlmodel import col, func
|
||||
|
||||
filters = []
|
||||
if search:
|
||||
pattern = f"%{search}%"
|
||||
filters.append(
|
||||
col(ApiKey.hashed_key).like(pattern)
|
||||
| col(ApiKey.refund_address).like(pattern)
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
result = await session.exec(select(ApiKey))
|
||||
base = select(ApiKey).where(*filters)
|
||||
|
||||
count_result = await session.exec(
|
||||
select(func.count()).select_from(base.subquery())
|
||||
)
|
||||
total = count_result.one()
|
||||
|
||||
# Aggregate totals across the whole (search-filtered) set, not just the
|
||||
# current page. Balance counts only parent (non-child) keys to avoid
|
||||
# double-counting, since child keys draw from their parent's balance.
|
||||
totals_result = await session.exec(
|
||||
select(
|
||||
func.coalesce(
|
||||
func.sum(
|
||||
case(
|
||||
(col(ApiKey.parent_key_hash).is_(None), ApiKey.balance),
|
||||
else_=0,
|
||||
)
|
||||
),
|
||||
0,
|
||||
),
|
||||
func.coalesce(func.sum(ApiKey.total_spent), 0),
|
||||
func.coalesce(func.sum(ApiKey.total_requests), 0),
|
||||
).where(*filters)
|
||||
)
|
||||
total_balance, total_spent, total_requests = totals_result.one()
|
||||
|
||||
# Latest created first; keys with no created_at (legacy rows) sort last.
|
||||
# Use an explicit CASE rather than relying on dialect NULL-ordering so
|
||||
# the behaviour is identical on SQLite and Postgres.
|
||||
stmt = (
|
||||
base.order_by(
|
||||
case((col(ApiKey.created_at).is_(None), 1), else_=0),
|
||||
col(ApiKey.created_at).desc(),
|
||||
)
|
||||
.offset(offset)
|
||||
.limit(limit)
|
||||
)
|
||||
result = await session.exec(stmt)
|
||||
api_keys = result.all()
|
||||
|
||||
return [
|
||||
{
|
||||
"hashed_key": key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"total_spent": key.total_spent,
|
||||
"total_requests": key.total_requests,
|
||||
"refund_address": key.refund_address,
|
||||
"key_expiry_time": key.key_expiry_time,
|
||||
"parent_key_hash": key.parent_key_hash,
|
||||
"balance_limit": key.balance_limit,
|
||||
"balance_limit_reset": key.balance_limit_reset,
|
||||
"validity_date": key.validity_date,
|
||||
}
|
||||
for key in api_keys
|
||||
]
|
||||
return {
|
||||
"balances": [
|
||||
{
|
||||
"hashed_key": key.hashed_key,
|
||||
"balance": key.balance,
|
||||
"total_spent": key.total_spent,
|
||||
"total_requests": key.total_requests,
|
||||
"refund_address": key.refund_address,
|
||||
"key_expiry_time": key.key_expiry_time,
|
||||
"parent_key_hash": key.parent_key_hash,
|
||||
"balance_limit": key.balance_limit,
|
||||
"balance_limit_reset": key.balance_limit_reset,
|
||||
"validity_date": key.validity_date,
|
||||
"created_at": key.created_at,
|
||||
}
|
||||
for key in api_keys
|
||||
],
|
||||
"total": total,
|
||||
"totals": {
|
||||
"total_balance": total_balance,
|
||||
"total_spent": total_spent,
|
||||
"total_requests": total_requests,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class ApiKeyUpdate(BaseModel):
|
||||
@@ -117,8 +207,8 @@ async def get_settings(request: Request) -> dict:
|
||||
return data
|
||||
|
||||
|
||||
class SettingsUpdate(BaseModel):
|
||||
__root__: dict[str, object]
|
||||
class SettingsUpdate(RootModel[dict[str, object]]):
|
||||
pass
|
||||
|
||||
|
||||
class PasswordUpdate(BaseModel):
|
||||
@@ -129,14 +219,19 @@ class PasswordUpdate(BaseModel):
|
||||
@admin_router.patch("/api/settings", dependencies=[Depends(require_admin_api)])
|
||||
async def update_settings(request: Request, update: SettingsUpdate) -> dict:
|
||||
# Remove sensitive fields from general settings update
|
||||
settings_data = update.__root__.copy()
|
||||
settings_data = update.root.copy()
|
||||
sensitive_fields = ["admin_password", "upstream_api_key", "nsec"]
|
||||
for field in sensitive_fields:
|
||||
if field in settings_data:
|
||||
del settings_data[field]
|
||||
|
||||
async with create_session() as session:
|
||||
new_settings = await SettingsService.update(settings_data, session)
|
||||
try:
|
||||
async with create_session() as session:
|
||||
new_settings = await SettingsService.update(settings_data, session)
|
||||
except PydanticValidationError as e:
|
||||
# Surface validation issues (e.g. non-positive payout amounts)
|
||||
# as a clean 400 instead of a 500.
|
||||
raise HTTPException(status_code=400, detail=e.errors()) from e
|
||||
data = new_settings.dict()
|
||||
if "upstream_api_key" in data:
|
||||
data["upstream_api_key"] = "[REDACTED]" if data["upstream_api_key"] else ""
|
||||
@@ -233,6 +328,73 @@ async def admin_logout(request: Request) -> dict[str, object]:
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
# ─── CLI Tokens (long-lived bearer tokens for CLI/agent use) ───
|
||||
|
||||
|
||||
class CliTokenCreate(BaseModel):
|
||||
name: str
|
||||
expires_in_days: int | None = None
|
||||
|
||||
|
||||
@admin_router.get("/api/cli-tokens", dependencies=[Depends(require_admin_api)])
|
||||
async def list_cli_tokens() -> list[dict[str, object]]:
|
||||
async with create_session() as session:
|
||||
result = await session.exec(select(CliToken))
|
||||
tokens = result.all()
|
||||
return [
|
||||
{
|
||||
"id": t.id,
|
||||
"name": t.name,
|
||||
"token_preview": f"{t.token[:8]}...{t.token[-4:]}",
|
||||
"created_at": t.created_at,
|
||||
"last_used_at": t.last_used_at,
|
||||
"expires_at": t.expires_at,
|
||||
}
|
||||
for t in tokens
|
||||
]
|
||||
|
||||
|
||||
@admin_router.post("/api/cli-tokens", dependencies=[Depends(require_admin_api)])
|
||||
async def create_cli_token(payload: CliTokenCreate) -> dict[str, object]:
|
||||
name = (payload.name or "").strip()
|
||||
if not name:
|
||||
raise HTTPException(status_code=400, detail="Name is required")
|
||||
|
||||
raw_token = secrets.token_urlsafe(32)
|
||||
expires_at: int | None = None
|
||||
if payload.expires_in_days is not None and payload.expires_in_days > 0:
|
||||
expires_at = int(datetime.now(timezone.utc).timestamp()) + (
|
||||
payload.expires_in_days * 86400
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
cli_token = CliToken(token=raw_token, name=name, expires_at=expires_at)
|
||||
session.add(cli_token)
|
||||
await session.commit()
|
||||
await session.refresh(cli_token)
|
||||
|
||||
return {
|
||||
"id": cli_token.id,
|
||||
"name": cli_token.name,
|
||||
"token": raw_token, # full token returned only on creation
|
||||
"created_at": cli_token.created_at,
|
||||
"expires_at": cli_token.expires_at,
|
||||
}
|
||||
|
||||
|
||||
@admin_router.delete(
|
||||
"/api/cli-tokens/{token_id}", dependencies=[Depends(require_admin_api)]
|
||||
)
|
||||
async def revoke_cli_token(token_id: str) -> dict[str, object]:
|
||||
async with create_session() as session:
|
||||
cli_token = await session.get(CliToken, token_id)
|
||||
if not cli_token:
|
||||
raise HTTPException(status_code=404, detail="Token not found")
|
||||
await session.delete(cli_token)
|
||||
await session.commit()
|
||||
return {"ok": True, "deleted_id": token_id}
|
||||
|
||||
|
||||
class WithdrawRequest(BaseModel):
|
||||
amount: int
|
||||
mint_url: str | None = None
|
||||
@@ -286,6 +448,7 @@ class ModelCreate(BaseModel):
|
||||
canonical_slug: str | None = None
|
||||
alias_ids: list[str] | None = None
|
||||
enabled: bool = True
|
||||
forwarded_model_id: str | None = None
|
||||
|
||||
|
||||
@admin_router.post(
|
||||
@@ -330,6 +493,7 @@ async def upsert_provider_model(
|
||||
json.dumps(payload.alias_ids) if payload.alias_ids else None
|
||||
)
|
||||
existing_row.enabled = payload.enabled
|
||||
existing_row.forwarded_model_id = payload.forwarded_model_id or payload.id
|
||||
|
||||
session.add(existing_row)
|
||||
await session.commit()
|
||||
@@ -362,6 +526,7 @@ async def upsert_provider_model(
|
||||
),
|
||||
upstream_provider_id=provider_id,
|
||||
enabled=payload.enabled,
|
||||
forwarded_model_id=payload.forwarded_model_id or payload.id,
|
||||
)
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
@@ -544,6 +709,7 @@ class UpstreamProviderCreate(BaseModel):
|
||||
api_version: str | None = None
|
||||
enabled: bool = True
|
||||
provider_fee: float = 1.01
|
||||
provider_settings: dict | None = None
|
||||
|
||||
|
||||
class UpstreamProviderUpdate(BaseModel):
|
||||
@@ -553,6 +719,7 @@ class UpstreamProviderUpdate(BaseModel):
|
||||
api_version: str | None = None
|
||||
enabled: bool | None = None
|
||||
provider_fee: float | None = None
|
||||
provider_settings: dict | None = None
|
||||
|
||||
|
||||
@admin_router.get("/api/upstream-providers", dependencies=[Depends(require_admin_api)])
|
||||
@@ -569,6 +736,9 @@ async def get_upstream_providers() -> list[dict[str, object]]:
|
||||
"api_version": p.api_version,
|
||||
"enabled": p.enabled,
|
||||
"provider_fee": p.provider_fee,
|
||||
"provider_settings": json.loads(p.provider_settings)
|
||||
if p.provider_settings
|
||||
else None,
|
||||
}
|
||||
for p in providers
|
||||
]
|
||||
@@ -598,6 +768,9 @@ async def create_upstream_provider(
|
||||
api_version=payload.api_version,
|
||||
enabled=payload.enabled,
|
||||
provider_fee=payload.provider_fee,
|
||||
provider_settings=json.dumps(payload.provider_settings)
|
||||
if payload.provider_settings
|
||||
else None,
|
||||
)
|
||||
session.add(provider)
|
||||
await session.commit()
|
||||
@@ -613,6 +786,7 @@ async def create_upstream_provider(
|
||||
"api_version": provider.api_version,
|
||||
"enabled": provider.enabled,
|
||||
"provider_fee": provider.provider_fee,
|
||||
"provider_settings": payload.provider_settings,
|
||||
}
|
||||
|
||||
|
||||
@@ -632,6 +806,9 @@ async def get_upstream_provider(provider_id: int) -> dict[str, object]:
|
||||
"api_version": provider.api_version,
|
||||
"enabled": provider.enabled,
|
||||
"provider_fee": provider.provider_fee,
|
||||
"provider_settings": json.loads(provider.provider_settings)
|
||||
if provider.provider_settings
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
@@ -658,6 +835,8 @@ async def update_upstream_provider(
|
||||
provider.enabled = payload.enabled
|
||||
if payload.provider_fee is not None:
|
||||
provider.provider_fee = payload.provider_fee
|
||||
if payload.provider_settings is not None:
|
||||
provider.provider_settings = json.dumps(payload.provider_settings)
|
||||
|
||||
session.add(provider)
|
||||
await session.commit()
|
||||
@@ -673,6 +852,9 @@ async def update_upstream_provider(
|
||||
"api_version": provider.api_version,
|
||||
"enabled": provider.enabled,
|
||||
"provider_fee": provider.provider_fee,
|
||||
"provider_settings": json.loads(provider.provider_settings)
|
||||
if provider.provider_settings
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
@@ -731,7 +913,7 @@ async def get_provider_models(provider_id: int) -> dict[str, object]:
|
||||
|
||||
db_model_ids = {model.id for model in db_models}
|
||||
filtered_remote_models = [
|
||||
m for m in upstream_models if m.name not in db_model_ids
|
||||
m for m in upstream_models if m.id not in db_model_ids
|
||||
]
|
||||
|
||||
return {
|
||||
@@ -794,6 +976,47 @@ class TopupRequest(BaseModel):
|
||||
amount: int
|
||||
|
||||
|
||||
class TopupTokenRequest(BaseModel):
|
||||
token: str
|
||||
|
||||
|
||||
@admin_router.post(
|
||||
"/api/upstream-providers/{provider_id}/topup-token",
|
||||
dependencies=[Depends(require_admin_api)],
|
||||
)
|
||||
async def topup_provider_with_token(
|
||||
provider_id: int, payload: TopupTokenRequest
|
||||
) -> dict:
|
||||
"""Redeem a Cashu token for an upstream provider."""
|
||||
async with create_session() as session:
|
||||
provider = await session.get(UpstreamProviderRow, provider_id)
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="Provider not found")
|
||||
|
||||
import httpx
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
clean_url = provider.base_url.rstrip("/")
|
||||
headers = {}
|
||||
if provider.api_key:
|
||||
headers["Authorization"] = f"Bearer {provider.api_key}"
|
||||
resp = await client.post(
|
||||
f"{clean_url}/v1/balance/topup",
|
||||
json={"cashu_token": payload.token},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
if resp.status_code == 200:
|
||||
return {"ok": True, "message": "Token redeemed successfully"}
|
||||
else:
|
||||
logger.error(f"Upstream token topup failed: {resp.text}")
|
||||
try:
|
||||
error_detail = resp.json()
|
||||
except Exception:
|
||||
error_detail = resp.text
|
||||
return {"ok": False, "message": f"Upstream error: {error_detail}"}
|
||||
|
||||
|
||||
@admin_router.post(
|
||||
"/api/upstream-providers/{provider_id}/topup",
|
||||
dependencies=[Depends(require_admin_api)],
|
||||
@@ -809,18 +1032,84 @@ async def initiate_provider_topup(
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="Provider not found")
|
||||
|
||||
upstream_instance = _instantiate_provider(provider)
|
||||
if not upstream_instance:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Could not instantiate provider"
|
||||
)
|
||||
|
||||
try:
|
||||
logger.info(
|
||||
f"Initiating top-up for provider {provider_id}",
|
||||
extra={"amount": payload.amount},
|
||||
)
|
||||
|
||||
# For Routstr providers, we might be doing a Lightning top-up or a direct token transfer
|
||||
if provider.provider_type == "routstr":
|
||||
# UI sends sats for Routstr topup
|
||||
import httpx
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
clean_url = provider.base_url.rstrip("/")
|
||||
request_json = {
|
||||
"amount_sats": int(payload.amount),
|
||||
"purpose": "topup",
|
||||
"api_key": provider.api_key,
|
||||
}
|
||||
headers = (
|
||||
{"Authorization": f"Bearer {provider.api_key}"}
|
||||
if provider.api_key
|
||||
else {}
|
||||
)
|
||||
|
||||
last_status_code = 500
|
||||
last_error_detail: object = "Failed to create top-up invoice"
|
||||
|
||||
# Some upstream Routstr nodes fail the first invoice request after warm-up
|
||||
# and succeed immediately on retry. Retry once here so the UI stays single-click.
|
||||
for attempt in range(2):
|
||||
resp = await client.post(
|
||||
f"{clean_url}/v1/balance/lightning/invoice",
|
||||
json=request_json,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
return {
|
||||
"ok": True,
|
||||
"topup_data": {
|
||||
"payment_request": data.get("bolt11"),
|
||||
"invoice_id": data.get("invoice_id"),
|
||||
"status": "pending",
|
||||
},
|
||||
}
|
||||
|
||||
logger.error(
|
||||
f"Upstream topup request failed: {resp.text}",
|
||||
extra={
|
||||
"provider_id": provider_id,
|
||||
"attempt": attempt + 1,
|
||||
"status_code": resp.status_code,
|
||||
},
|
||||
)
|
||||
try:
|
||||
last_error_detail = resp.json()
|
||||
except Exception:
|
||||
last_error_detail = resp.text
|
||||
last_status_code = resp.status_code
|
||||
|
||||
if resp.status_code < 500 or attempt == 1:
|
||||
break
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
raise HTTPException(
|
||||
status_code=last_status_code, detail=last_error_detail
|
||||
)
|
||||
|
||||
upstream_instance = _instantiate_provider(provider)
|
||||
if not upstream_instance:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Could not instantiate provider"
|
||||
)
|
||||
|
||||
topup_data = await upstream_instance.initiate_topup(payload.amount)
|
||||
|
||||
logger.info(
|
||||
"Top-up initiated successfully",
|
||||
extra={
|
||||
@@ -871,6 +1160,25 @@ async def check_topup_status(provider_id: int, invoice_id: str) -> dict[str, obj
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="Provider not found")
|
||||
|
||||
# For Routstr providers, proxy the status check
|
||||
if provider.provider_type == "routstr":
|
||||
import httpx
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
clean_url = provider.base_url.rstrip("/")
|
||||
resp = await client.get(
|
||||
f"{clean_url}/v1/balance/lightning/invoice/{invoice_id}/status",
|
||||
headers={"Authorization": f"Bearer {provider.api_key}"}
|
||||
if provider.api_key
|
||||
else {},
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
status_data = resp.json()
|
||||
return {"ok": True, "paid": status_data.get("status") == "paid"}
|
||||
else:
|
||||
logger.error(f"Upstream status check failed: {resp.text}")
|
||||
return {"ok": False, "paid": False}
|
||||
|
||||
upstream_instance = _instantiate_provider(provider)
|
||||
if not upstream_instance:
|
||||
raise HTTPException(
|
||||
@@ -898,7 +1206,7 @@ async def check_topup_status(provider_id: int, invoice_id: str) -> dict[str, obj
|
||||
dependencies=[Depends(require_admin_api)],
|
||||
)
|
||||
async def get_provider_balance(provider_id: int) -> dict[str, object]:
|
||||
"""Get the current account balance for the upstream provider."""
|
||||
"""Get the current balance for an upstream provider account."""
|
||||
from ..upstream.helpers import _instantiate_provider
|
||||
|
||||
async with create_session() as session:
|
||||
@@ -906,6 +1214,61 @@ async def get_provider_balance(provider_id: int) -> dict[str, object]:
|
||||
if not provider:
|
||||
raise HTTPException(status_code=404, detail="Provider not found")
|
||||
|
||||
# For Routstr providers, proxy the balance check
|
||||
if provider.provider_type == "routstr":
|
||||
import httpx
|
||||
|
||||
clean_url = provider.base_url.rstrip("/")
|
||||
headers = {}
|
||||
if provider.api_key:
|
||||
headers["Authorization"] = f"Bearer {provider.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
try:
|
||||
resp = await client.get(
|
||||
f"{clean_url}/v1/balance/info",
|
||||
headers=headers,
|
||||
)
|
||||
except httpx.TimeoutException as exc:
|
||||
logger.error(
|
||||
"Timed out fetching Routstr provider balance",
|
||||
extra={
|
||||
"provider_id": provider_id,
|
||||
"base_url": clean_url,
|
||||
"upstream_url": f"{clean_url}/v1/balance/info",
|
||||
"error": str(exc),
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=504,
|
||||
detail="Timed out contacting upstream Routstr provider",
|
||||
) from exc
|
||||
except httpx.RequestError as exc:
|
||||
logger.error(
|
||||
"Failed to fetch Routstr provider balance",
|
||||
extra={
|
||||
"provider_id": provider_id,
|
||||
"base_url": clean_url,
|
||||
"upstream_url": f"{clean_url}/v1/balance/info",
|
||||
"error": str(exc),
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail="Failed to contact upstream Routstr provider",
|
||||
) from exc
|
||||
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
# Return balance in sats
|
||||
balance = data.get("balance", 0)
|
||||
if isinstance(balance, (int, float)):
|
||||
return {"ok": True, "balance_data": balance // 1000}
|
||||
return {"ok": True, "balance_data": balance}
|
||||
else:
|
||||
logger.error(f"Failed to fetch Routstr balance: {resp.text}")
|
||||
return {"ok": False, "balance_data": None}
|
||||
|
||||
upstream_instance = _instantiate_provider(provider)
|
||||
if not upstream_instance:
|
||||
raise HTTPException(
|
||||
@@ -943,18 +1306,55 @@ async def get_usage_metrics(
|
||||
default=15, ge=1, le=1440, description="Time interval in minutes"
|
||||
),
|
||||
hours: int = Query(
|
||||
default=24, ge=1, le=168, description="Hours of history to analyze"
|
||||
default=24,
|
||||
ge=1,
|
||||
le=MAX_USAGE_ANALYTICS_HOURS,
|
||||
description="Hours of history to analyze",
|
||||
),
|
||||
) -> dict:
|
||||
"""Get usage metrics aggregated by time interval."""
|
||||
return log_manager.get_usage_metrics(interval=interval, hours=hours)
|
||||
|
||||
|
||||
@admin_router.get("/api/usage/dashboard", dependencies=[Depends(require_admin_api)])
|
||||
async def get_usage_dashboard(
|
||||
request: Request,
|
||||
interval: int = Query(
|
||||
default=15, ge=1, le=1440, description="Time interval in minutes"
|
||||
),
|
||||
hours: int = Query(
|
||||
default=24,
|
||||
ge=1,
|
||||
le=MAX_USAGE_ANALYTICS_HOURS,
|
||||
description="Hours of history to analyze",
|
||||
),
|
||||
error_limit: int = Query(
|
||||
default=100, ge=1, le=1000, description="Maximum number of errors to return"
|
||||
),
|
||||
model_limit: int = Query(
|
||||
default=20, ge=1, le=100, description="Maximum number of models to return"
|
||||
),
|
||||
) -> dict:
|
||||
"""
|
||||
Get all dashboard analytics in one request.
|
||||
This runs one combined aggregation pass and avoids repeated scans.
|
||||
"""
|
||||
return log_manager.get_usage_dashboard(
|
||||
interval=interval,
|
||||
hours=hours,
|
||||
error_limit=error_limit,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
|
||||
@admin_router.get("/api/usage/summary", dependencies=[Depends(require_admin_api)])
|
||||
async def get_usage_summary(
|
||||
request: Request,
|
||||
hours: int = Query(
|
||||
default=24, ge=1, le=168, description="Hours of history to analyze"
|
||||
default=24,
|
||||
ge=1,
|
||||
le=MAX_USAGE_ANALYTICS_HOURS,
|
||||
description="Hours of history to analyze",
|
||||
),
|
||||
) -> dict:
|
||||
"""Get summary statistics for the specified time period."""
|
||||
@@ -965,7 +1365,10 @@ async def get_usage_summary(
|
||||
async def get_error_details(
|
||||
request: Request,
|
||||
hours: int = Query(
|
||||
default=24, ge=1, le=168, description="Hours of history to analyze"
|
||||
default=24,
|
||||
ge=1,
|
||||
le=MAX_USAGE_ANALYTICS_HOURS,
|
||||
description="Hours of history to analyze",
|
||||
),
|
||||
limit: int = Query(
|
||||
default=100, ge=1, le=1000, description="Maximum number of errors to return"
|
||||
@@ -981,7 +1384,10 @@ async def get_error_details(
|
||||
async def get_revenue_by_model(
|
||||
request: Request,
|
||||
hours: int = Query(
|
||||
default=24, ge=1, le=168, description="Hours of history to analyze"
|
||||
default=24,
|
||||
ge=1,
|
||||
le=MAX_USAGE_ANALYTICS_HOURS,
|
||||
description="Hours of history to analyze",
|
||||
),
|
||||
limit: int = Query(
|
||||
default=20, ge=1, le=100, description="Maximum number of models to return"
|
||||
@@ -1075,3 +1481,175 @@ async def get_log_dates_api(request: Request) -> dict[str, object]:
|
||||
continue
|
||||
|
||||
return {"dates": dates}
|
||||
|
||||
|
||||
@admin_router.get("/api/transactions", dependencies=[Depends(require_admin_api)])
|
||||
async def get_transactions_api(
|
||||
type: str | None = None,
|
||||
status: str | None = None,
|
||||
search: str | None = None,
|
||||
source: str | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> dict:
|
||||
async with create_session() as session:
|
||||
from sqlmodel import col, func
|
||||
|
||||
base = select(CashuTransaction)
|
||||
if type:
|
||||
base = base.where(CashuTransaction.type == type)
|
||||
if source:
|
||||
if source == "x-cashu":
|
||||
base = base.where(
|
||||
(CashuTransaction.source == "x-cashu")
|
||||
| (CashuTransaction.source == None) # noqa: E711
|
||||
)
|
||||
else:
|
||||
base = base.where(CashuTransaction.source == source)
|
||||
if status:
|
||||
if status == "collected":
|
||||
base = base.where(CashuTransaction.collected == True) # noqa: E712
|
||||
elif status == "swept":
|
||||
base = base.where(CashuTransaction.swept == True) # noqa: E712
|
||||
elif status == "pending":
|
||||
base = base.where(
|
||||
CashuTransaction.collected == False, # noqa: E712
|
||||
CashuTransaction.swept == False, # noqa: E712
|
||||
)
|
||||
|
||||
if search:
|
||||
search_pattern = f"%{search}%"
|
||||
base = base.where(
|
||||
(col(CashuTransaction.id).like(search_pattern))
|
||||
| (col(CashuTransaction.token).like(search_pattern))
|
||||
| (col(CashuTransaction.request_id).like(search_pattern))
|
||||
| (col(CashuTransaction.api_key_hashed_key).like(search_pattern))
|
||||
)
|
||||
|
||||
count_result = await session.exec(
|
||||
select(func.count()).select_from(base.subquery())
|
||||
)
|
||||
total = count_result.one()
|
||||
|
||||
stmt = base.order_by(col(CashuTransaction.created_at).desc()).offset(offset).limit(limit)
|
||||
results = await session.exec(stmt)
|
||||
transactions = results.all()
|
||||
|
||||
return {
|
||||
"transactions": [tx.dict() for tx in transactions],
|
||||
"total": total,
|
||||
}
|
||||
|
||||
|
||||
@admin_router.get(
|
||||
"/api/lightning-invoices", dependencies=[Depends(require_admin_api)]
|
||||
)
|
||||
async def get_lightning_invoices_api(
|
||||
status: str | None = None,
|
||||
purpose: str | None = None,
|
||||
search: str | None = None,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
) -> dict:
|
||||
async with create_session() as session:
|
||||
from sqlmodel import col, func
|
||||
|
||||
base = select(LightningInvoice)
|
||||
if status:
|
||||
base = base.where(LightningInvoice.status == status)
|
||||
if purpose:
|
||||
base = base.where(LightningInvoice.purpose == purpose)
|
||||
if search:
|
||||
pattern = f"%{search}%"
|
||||
base = base.where(
|
||||
(col(LightningInvoice.id).like(pattern))
|
||||
| (col(LightningInvoice.bolt11).like(pattern))
|
||||
| (col(LightningInvoice.payment_hash).like(pattern))
|
||||
| (col(LightningInvoice.api_key_hash).like(pattern))
|
||||
)
|
||||
|
||||
count_result = await session.exec(
|
||||
select(func.count()).select_from(base.subquery())
|
||||
)
|
||||
total = count_result.one()
|
||||
|
||||
stmt = (
|
||||
base.order_by(col(LightningInvoice.created_at).desc())
|
||||
.offset(offset)
|
||||
.limit(limit)
|
||||
)
|
||||
results = await session.exec(stmt)
|
||||
invoices = results.all()
|
||||
|
||||
return {
|
||||
"invoices": [inv.dict() for inv in invoices],
|
||||
"total": total,
|
||||
}
|
||||
|
||||
|
||||
@admin_router.post(
|
||||
"/api/upstream-providers/{provider_id}/routstr/refund",
|
||||
dependencies=[Depends(require_admin_api)],
|
||||
)
|
||||
async def refund_routstr_provider_balance(provider_id: int) -> dict[str, object]:
|
||||
"""Refund balance from an upstream Routstr provider back to the local wallet."""
|
||||
from ..upstream.helpers import _instantiate_provider
|
||||
from ..upstream.routstr import RoutstrUpstreamProvider
|
||||
|
||||
async with create_session() as session:
|
||||
provider_row = await session.get(UpstreamProviderRow, provider_id)
|
||||
if not provider_row:
|
||||
raise HTTPException(status_code=404, detail="Provider not found")
|
||||
|
||||
if provider_row.provider_type != "routstr":
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Refund only supported for Routstr providers"
|
||||
)
|
||||
|
||||
provider = _instantiate_provider(provider_row)
|
||||
if not isinstance(provider, RoutstrUpstreamProvider):
|
||||
raise HTTPException(status_code=400, detail="Invalid provider instance")
|
||||
|
||||
try:
|
||||
# Request refund from upstream
|
||||
data = await provider.refund_balance()
|
||||
if "error" in data:
|
||||
# If the upstream returned an OpenAI-style error (like the model unknown error)
|
||||
# it means the request likely didn't even reach the refund endpoint handler
|
||||
# but was intercepted by the proxy layer.
|
||||
error_info = data.get("error", {})
|
||||
message = (
|
||||
error_info.get("message")
|
||||
if isinstance(error_info, dict)
|
||||
else str(error_info)
|
||||
)
|
||||
return {
|
||||
"ok": False,
|
||||
"message": f"Upstream refund failed: {message}",
|
||||
}
|
||||
|
||||
token = data.get("token")
|
||||
if not token:
|
||||
return {"ok": False, "message": "Upstream did not return a token"}
|
||||
|
||||
# Receive token into local wallet
|
||||
from ..wallet import recieve_token
|
||||
|
||||
try:
|
||||
# Use current wallet to receive
|
||||
await recieve_token(token)
|
||||
return {
|
||||
"ok": True,
|
||||
"message": "Successfully received refund from upstream provider",
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to receive refund token: {e}")
|
||||
return {
|
||||
"ok": False,
|
||||
"message": f"Failed to receive refund token: {str(e)}",
|
||||
"token": token,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Refund failed for provider {provider_id}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@@ -2,14 +2,17 @@ import os
|
||||
import pathlib
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from alembic.util.exc import CommandError
|
||||
from sqlalchemy import UniqueConstraint
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio.engine import create_async_engine
|
||||
from sqlmodel import Field, Relationship, SQLModel, func, select, update
|
||||
from sqlmodel import Field, Relationship, SQLModel, col, func, select, update
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from .logging import get_logger
|
||||
@@ -30,6 +33,14 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
reserved_balance: int = Field(
|
||||
default=0, description="Reserved balance in millisatoshis (msats)"
|
||||
)
|
||||
reserved_at: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Unix timestamp of the most recent balance reservation. Used to "
|
||||
"detect and release stale reservations (e.g. after client "
|
||||
"disconnects). NULL when no reservation has been made yet."
|
||||
),
|
||||
)
|
||||
refund_address: str | None = Field(
|
||||
default=None,
|
||||
description="Lightning address to refund remaining balance after key expires",
|
||||
@@ -42,6 +53,14 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
default=0, description="Total spent in millisatoshis (msats)"
|
||||
)
|
||||
total_requests: int = Field(default=0)
|
||||
created_at: int | None = Field(
|
||||
default_factory=lambda: int(time.time()),
|
||||
nullable=True,
|
||||
description=(
|
||||
"Unix timestamp when the key was created. Nullable: keys created "
|
||||
"before this column existed have no value and sort last."
|
||||
),
|
||||
)
|
||||
refund_mint_url: str | None = Field(
|
||||
default=None,
|
||||
description="URL of the mint used to create the cashu-token",
|
||||
@@ -76,11 +95,34 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
|
||||
|
||||
async def reset_all_reserved_balances(session: AsyncSession) -> None:
|
||||
logger.info("Resetting all reserved balances to 0")
|
||||
stmt = update(ApiKey).values(reserved_balance=0)
|
||||
stmt = update(ApiKey).values(reserved_balance=0, reserved_at=None)
|
||||
await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
logger.info("Reserved balances reset successfully")
|
||||
logger.info("Reset reserved balances on startup")
|
||||
|
||||
|
||||
async def release_stale_reservations(
|
||||
session: AsyncSession, max_age_seconds: int
|
||||
) -> int:
|
||||
"""Release reservations whose last reserve is older than max_age_seconds.
|
||||
"""
|
||||
cutoff = int(time.time()) - max_age_seconds
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.reserved_balance) > 0)
|
||||
.where(col(ApiKey.reserved_at).is_not(None))
|
||||
.where(col(ApiKey.reserved_at) < cutoff)
|
||||
.values(reserved_balance=0, reserved_at=None)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
released = int(result.rowcount or 0)
|
||||
if released:
|
||||
logger.warning(
|
||||
"Released stale balance reservations",
|
||||
extra={"released_keys": released, "max_age_seconds": max_age_seconds},
|
||||
)
|
||||
return released
|
||||
|
||||
|
||||
class ModelRow(SQLModel, table=True): # type: ignore
|
||||
@@ -103,6 +145,10 @@ class ModelRow(SQLModel, table=True): # type: ignore
|
||||
default=None, description="JSON array of model alias IDs"
|
||||
)
|
||||
enabled: bool = Field(default=True, description="Whether this model is enabled")
|
||||
forwarded_model_id: str | None = Field(
|
||||
default=None,
|
||||
description="Model ID to use when forwarding requests to upstream provider. Defaults to id if not set.",
|
||||
)
|
||||
upstream_provider: "UpstreamProviderRow" = Relationship(back_populates="models")
|
||||
|
||||
|
||||
@@ -126,6 +172,85 @@ class LightningInvoice(SQLModel, table=True): # type: ignore
|
||||
)
|
||||
expires_at: int = Field(description="Unix timestamp when invoice expires")
|
||||
paid_at: int | None = Field(default=None, description="Unix timestamp when paid")
|
||||
balance_limit: int | None = Field(
|
||||
default=None,
|
||||
description="Max spendable msats for the created key",
|
||||
)
|
||||
balance_limit_reset: str | None = Field(
|
||||
default=None,
|
||||
description="Reset policy for balance limit (daily, weekly, monthly)",
|
||||
)
|
||||
validity_date: int | None = Field(
|
||||
default=None,
|
||||
description="Unix timestamp after which the created key expires",
|
||||
)
|
||||
|
||||
|
||||
class CashuTransaction(SQLModel, table=True): # type: ignore
|
||||
__tablename__ = "cashu_transactions"
|
||||
|
||||
id: str = Field(
|
||||
primary_key=True,
|
||||
default_factory=lambda: uuid.uuid4().hex,
|
||||
description="Unique transaction identifier",
|
||||
)
|
||||
token: str = Field(description="Serialized Cashu token")
|
||||
amount: int = Field(description="Amount in the token's unit")
|
||||
unit: str = Field(description="Token unit (sat or msat)")
|
||||
mint_url: str | None = Field(default=None, description="Mint URL for the token")
|
||||
type: str = Field(default="out", description="Transaction type: in or out")
|
||||
request_id: str | None = Field(default=None, description="Associated request ID")
|
||||
created_at: int = Field(
|
||||
default_factory=lambda: int(time.time()),
|
||||
description="Unix timestamp",
|
||||
)
|
||||
collected: bool = Field(default=False)
|
||||
swept: bool = Field(default=False)
|
||||
source: str = Field(
|
||||
default="x-cashu",
|
||||
description="Payment source: x-cashu or apikey",
|
||||
)
|
||||
api_key_hashed_key: str | None = Field(
|
||||
default=None,
|
||||
foreign_key="api_keys.hashed_key",
|
||||
index=True,
|
||||
description="Associated API key hash for wallet history",
|
||||
)
|
||||
|
||||
|
||||
async def store_cashu_transaction(
|
||||
token: str,
|
||||
amount: int,
|
||||
unit: str,
|
||||
mint_url: str | None = None,
|
||||
typ: str = "out",
|
||||
request_id: str | None = None,
|
||||
collected: bool = False,
|
||||
created_at: int | None = None,
|
||||
source: str = "x-cashu",
|
||||
api_key_hashed_key: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
async with create_session() as session:
|
||||
tx = CashuTransaction(
|
||||
token=token,
|
||||
amount=amount,
|
||||
unit=unit,
|
||||
mint_url=mint_url,
|
||||
type=typ,
|
||||
request_id=request_id,
|
||||
collected=collected,
|
||||
created_at=created_at or int(time.time()),
|
||||
source=source,
|
||||
api_key_hashed_key=api_key_hashed_key,
|
||||
)
|
||||
session.add(tx)
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to store cashu transaction: {e} (type={typ})",
|
||||
extra={"error": str(e), "type": typ},
|
||||
)
|
||||
|
||||
|
||||
class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
@@ -148,12 +273,75 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
provider_fee: float = Field(
|
||||
default=1.01, description="Provider fee multiplier (default 1%)"
|
||||
)
|
||||
provider_settings: str | None = Field(
|
||||
default=None, description="JSON string for provider-specific settings"
|
||||
)
|
||||
models: list["ModelRow"] = Relationship(
|
||||
back_populates="upstream_provider",
|
||||
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
|
||||
)
|
||||
|
||||
|
||||
class RoutstrFee(SQLModel, table=True): # type: ignore
|
||||
__tablename__ = "routstr_fees"
|
||||
id: int = Field(default=1, primary_key=True)
|
||||
accumulated_msats: int = Field(default=0)
|
||||
total_paid_msats: int = Field(default=0)
|
||||
last_paid_at: int | None = Field(default=None)
|
||||
|
||||
|
||||
class CliToken(SQLModel, table=True): # type: ignore
|
||||
"""Long-lived authorization token for CLI/agent use against admin endpoints."""
|
||||
|
||||
__tablename__ = "cli_tokens"
|
||||
id: str = Field(
|
||||
primary_key=True, default_factory=lambda: uuid.uuid4().hex
|
||||
)
|
||||
token: str = Field(unique=True, index=True, description="Bearer token value")
|
||||
name: str = Field(description="Human-readable label for this token")
|
||||
created_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
last_used_at: int | None = Field(default=None)
|
||||
expires_at: int | None = Field(
|
||||
default=None, description="Optional expiry unix timestamp; null = never expires"
|
||||
)
|
||||
|
||||
|
||||
async def accumulate_routstr_fee(session: AsyncSession, amount_msats: int) -> None:
|
||||
stmt = (
|
||||
update(RoutstrFee)
|
||||
.where(col(RoutstrFee.id) == 1)
|
||||
.values(accumulated_msats=RoutstrFee.accumulated_msats + amount_msats)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
if result.rowcount == 0:
|
||||
session.add(RoutstrFee(id=1, accumulated_msats=amount_msats))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_routstr_fee(session: AsyncSession) -> RoutstrFee:
|
||||
fee = await session.get(RoutstrFee, 1)
|
||||
if fee is None:
|
||||
fee = RoutstrFee(id=1, accumulated_msats=0, total_paid_msats=0)
|
||||
session.add(fee)
|
||||
await session.commit()
|
||||
await session.refresh(fee)
|
||||
return fee
|
||||
|
||||
|
||||
async def reset_routstr_fee(session: AsyncSession, paid_msats: int) -> None:
|
||||
stmt = (
|
||||
update(RoutstrFee)
|
||||
.where(col(RoutstrFee.id) == 1)
|
||||
.values(
|
||||
accumulated_msats=RoutstrFee.accumulated_msats - paid_msats,
|
||||
total_paid_msats=RoutstrFee.total_paid_msats + paid_msats,
|
||||
last_paid_at=int(time.time()),
|
||||
)
|
||||
)
|
||||
await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def balances_for_mint_and_unit(
|
||||
db_session: AsyncSession, mint_url: str, unit: str
|
||||
) -> int:
|
||||
@@ -224,6 +412,17 @@ def fix_cashu_migrations() -> None:
|
||||
logger.warning(f"Could not check/fix Cashu database {db_file}: {e}")
|
||||
|
||||
|
||||
def _clear_alembic_version() -> None:
|
||||
"""Clear the alembic_version table so stamp/upgrade can proceed."""
|
||||
sync_url = DATABASE_URL.replace("+aiosqlite", "")
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
eng = create_engine(sync_url)
|
||||
with eng.begin() as conn:
|
||||
conn.execute(text("DELETE FROM alembic_version"))
|
||||
eng.dispose()
|
||||
|
||||
|
||||
def run_migrations() -> None:
|
||||
"""Run Alembic migrations programmatically."""
|
||||
try:
|
||||
@@ -245,8 +444,30 @@ def run_migrations() -> None:
|
||||
# Set the database URL in the config
|
||||
alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)
|
||||
|
||||
# Run migrations to the latest revision
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
try:
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
except CommandError as e:
|
||||
if "Can't locate revision" in str(e):
|
||||
logger.warning(
|
||||
"Database stamped with unknown revision (likely from another branch). "
|
||||
"Re-stamping to current head.",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
_clear_alembic_version()
|
||||
command.stamp(alembic_cfg, "head")
|
||||
else:
|
||||
raise
|
||||
except OperationalError as e:
|
||||
if "duplicate column name" in str(e).lower():
|
||||
logger.warning(
|
||||
"Migration hit a column that already exists (likely added via "
|
||||
"create_all on another branch). Stamping to current head.",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
_clear_alembic_version()
|
||||
command.stamp(alembic_cfg, "head")
|
||||
else:
|
||||
raise
|
||||
|
||||
logger.info("Database migrations completed successfully")
|
||||
|
||||
|
||||
@@ -22,16 +22,20 @@ async def http_exception_handler(request: Request, exc: Exception) -> JSONRespon
|
||||
# Get status code and detail - works for both FastAPI and Starlette HTTPException
|
||||
status_code = getattr(exc, "status_code", 500)
|
||||
detail = getattr(exc, "detail", str(exc))
|
||||
path = request.url.path
|
||||
|
||||
logger.warning(
|
||||
"HTTP exception",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"status_code": status_code,
|
||||
"detail": detail,
|
||||
"path": request.url.path,
|
||||
},
|
||||
)
|
||||
# 4xx is client behaviour; the uvicorn access log already records it.
|
||||
# Only 5xx warrants a server-side warning/error log here.
|
||||
if status_code >= 500:
|
||||
logger.error(
|
||||
f"HTTP {status_code} on {path}: {detail}",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"status_code": status_code,
|
||||
"detail": detail,
|
||||
"path": path,
|
||||
},
|
||||
)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,50 +3,61 @@ Logging configuration for Routstr.
|
||||
|
||||
CRITICAL LOG MESSAGES FOR USAGE STATISTICS:
|
||||
===========================================
|
||||
The following log messages are parsed by the usage tracking system (routstr/core/admin.py).
|
||||
The following log messages are parsed by the usage tracking system
|
||||
(routstr/core/usage_analytics_store.py and routstr/core/log_manager.py).
|
||||
DO NOT modify or remove these messages without updating the usage tracking logic:
|
||||
|
||||
1. "Received proxy request" (INFO) - routstr/proxy.py
|
||||
- Used to count total incoming requests
|
||||
- Includes model information in context
|
||||
|
||||
2. "Payment adjustment completed for streaming" (INFO) - routstr/upstream/base.py
|
||||
"Payment adjustment completed for non-streaming" (INFO) - routstr/upstream/base.py
|
||||
2. "Calculated token-based cost" (INFO) - routstr/auth.py
|
||||
- Used to track successful completions and revenue
|
||||
- The 'cost_data.total_msats' field is extracted for revenue calculation
|
||||
- Must include 'cost_data' in extra dict
|
||||
- The 'token_cost', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
|
||||
|
||||
3. "Payment processed successfully" (INFO) - routstr/auth.py
|
||||
3. "Max cost payment finalized" (INFO) - routstr/auth.py
|
||||
- Used as the successful completion fallback when token usage is unavailable
|
||||
- The 'charged_amount', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
|
||||
|
||||
4. "Payment processed successfully" (INFO) - routstr/auth.py
|
||||
- Used to count successful payment processing events
|
||||
- Tracks payment-related metrics
|
||||
|
||||
4. "Upstream request failed, revert payment" (WARNING) - routstr/proxy.py
|
||||
5. "Upstream request failed, revert payment" (WARNING) - routstr/proxy.py
|
||||
- Used to track failed requests and refunds
|
||||
- The 'max_cost_for_model' field is extracted for refund calculation
|
||||
- Must include 'max_cost_for_model' in extra dict
|
||||
|
||||
5. Any ERROR level logs with "upstream" in the message
|
||||
6. Any ERROR level logs with "upstream" in the message
|
||||
- Used to count upstream provider errors
|
||||
- Helps identify service reliability issues
|
||||
|
||||
If you need to modify these messages, ensure you also update the parsing logic in:
|
||||
- routstr/core/admin.py:_aggregate_metrics_by_time()
|
||||
- routstr/core/admin.py:_get_summary_stats()
|
||||
- routstr/core/admin.py:get_revenue_by_model()
|
||||
- routstr/core/usage_analytics_store.py
|
||||
- routstr/core/log_manager.py
|
||||
"""
|
||||
|
||||
import logging.config
|
||||
import logging.handlers
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import tomllib
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pythonjsonlogger import jsonlogger
|
||||
from rich.console import Console
|
||||
from rich.logging import RichHandler
|
||||
|
||||
# Only use RichHandler when stdout is a real TTY. In non-TTY contexts
|
||||
# (docker logs, pipes, CI) Rich pads every line to width and wraps long
|
||||
# records, producing visually-empty trailing whitespace and split records.
|
||||
# A plain StreamHandler avoids both problems.
|
||||
_stdout_is_tty = sys.stdout.isatty()
|
||||
_console = Console(soft_wrap=True) if _stdout_is_tty else None
|
||||
|
||||
# Define custom TRACE level
|
||||
TRACE_LEVEL = 5
|
||||
logging.addLevelName(TRACE_LEVEL, "TRACE")
|
||||
@@ -259,6 +270,26 @@ def setup_logging() -> None:
|
||||
if console_enabled:
|
||||
handlers.append("console")
|
||||
|
||||
if _stdout_is_tty:
|
||||
console_handler: dict[str, Any] = {
|
||||
"()": RichHandler,
|
||||
"level": log_level,
|
||||
"show_time": False,
|
||||
"show_path": False,
|
||||
"rich_tracebacks": True,
|
||||
"markup": True,
|
||||
"console": _console,
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
}
|
||||
else:
|
||||
console_handler = {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": log_level,
|
||||
"formatter": "plain",
|
||||
"stream": "ext://sys.stdout",
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
}
|
||||
|
||||
LOGGING_CONFIG = {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
@@ -268,6 +299,10 @@ def setup_logging() -> None:
|
||||
"format": "%(asctime)s %(name)s %(levelname)s %(message)s %(pathname)s %(lineno)d %(version)s %(request_id)s",
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
"plain": {
|
||||
"format": "%(asctime)s %(levelname)-7s %(name)s %(message)s",
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
},
|
||||
"filters": {
|
||||
"version_filter": {"()": VersionFilter},
|
||||
@@ -275,15 +310,7 @@ def setup_logging() -> None:
|
||||
"security_filter": {"()": SecurityFilter},
|
||||
},
|
||||
"handlers": {
|
||||
"console": {
|
||||
"()": RichHandler,
|
||||
"level": log_level,
|
||||
"show_time": False,
|
||||
"show_path": False,
|
||||
"rich_tracebacks": True,
|
||||
"markup": True,
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
},
|
||||
"console": console_handler,
|
||||
"file": {
|
||||
"()": DailyRotatingFileHandler,
|
||||
"level": log_level,
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
@@ -9,32 +8,38 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, RedirectResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
from starlette.types import Scope
|
||||
|
||||
from ..auth import periodic_key_reset
|
||||
from ..auth import periodic_key_reset, periodic_stale_reservation_sweep
|
||||
from ..balance import balance_router, deprecated_wallet_router
|
||||
from ..nostr import announce_provider, providers_cache_refresher
|
||||
from ..lightning import lightning_router, periodic_invoice_watcher
|
||||
from ..nostr import (
|
||||
announce_provider,
|
||||
providers_cache_refresher,
|
||||
publish_usage_analytics,
|
||||
)
|
||||
from ..nostr.discovery import providers_router
|
||||
from ..payment.models import models_router, update_sats_pricing
|
||||
from ..payment.price import update_prices_periodically
|
||||
from ..proxy import initialize_upstreams, proxy_router, refresh_model_maps_periodically
|
||||
from ..wallet import periodic_payout
|
||||
from ..upstream.auto_topup import periodic_auto_topup
|
||||
from ..upstream.litellm_routing import configure_litellm
|
||||
from ..wallet import periodic_payout, periodic_refund_sweep, periodic_routstr_fee_payout
|
||||
from .admin import admin_router
|
||||
from .db import create_session, init_db, run_migrations
|
||||
from .exceptions import general_exception_handler, http_exception_handler
|
||||
from .logging import get_logger, setup_logging
|
||||
from .middleware import LoggingMiddleware
|
||||
from .not_found import _NOT_FOUND_HTML, not_found_catch_all # noqa: F401
|
||||
from .settings import SettingsService
|
||||
from .settings import settings as global_settings
|
||||
from .version import __version__
|
||||
|
||||
# Initialize logging first
|
||||
setup_logging()
|
||||
logger = get_logger(__name__)
|
||||
|
||||
if os.getenv("VERSION_SUFFIX") is not None:
|
||||
__version__ = f"0.3.0-{os.getenv('VERSION_SUFFIX')}"
|
||||
else:
|
||||
__version__ = "0.3.0"
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
@@ -44,12 +49,22 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
pricing_task = None
|
||||
payout_task = None
|
||||
nip91_task = None
|
||||
analytics_task = None
|
||||
providers_task = None
|
||||
models_refresh_task = None
|
||||
model_maps_refresh_task = None
|
||||
key_reset_task = None
|
||||
stale_reservation_task = None
|
||||
auto_topup_task = None
|
||||
refund_sweep_task = None
|
||||
routstr_fee_task = None
|
||||
invoice_watcher_task = None
|
||||
|
||||
try:
|
||||
# Apply litellm-wide settings (drop_params, chat-completions URL,
|
||||
# debug logging) before any upstream provider dispatches a request.
|
||||
configure_litellm()
|
||||
|
||||
# Run database migrations on startup
|
||||
run_migrations()
|
||||
|
||||
@@ -94,16 +109,27 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
btc_price_task = asyncio.create_task(update_prices_periodically())
|
||||
pricing_task = asyncio.create_task(update_sats_pricing())
|
||||
if global_settings.models_refresh_interval_seconds > 0:
|
||||
# Pass the accessor (not its current value) so the loop sees providers
|
||||
# added/changed via reinitialize_upstreams() instead of staying pinned
|
||||
# to the startup snapshot.
|
||||
models_refresh_task = asyncio.create_task(
|
||||
refresh_upstreams_models_periodically(get_upstreams())
|
||||
refresh_upstreams_models_periodically(get_upstreams)
|
||||
)
|
||||
model_maps_refresh_task = asyncio.create_task(refresh_model_maps_periodically())
|
||||
payout_task = asyncio.create_task(periodic_payout())
|
||||
if global_settings.nsec:
|
||||
nip91_task = asyncio.create_task(announce_provider())
|
||||
analytics_task = asyncio.create_task(publish_usage_analytics())
|
||||
if global_settings.providers_refresh_interval_seconds > 0:
|
||||
providers_task = asyncio.create_task(providers_cache_refresher())
|
||||
key_reset_task = asyncio.create_task(periodic_key_reset())
|
||||
stale_reservation_task = asyncio.create_task(
|
||||
periodic_stale_reservation_sweep()
|
||||
)
|
||||
auto_topup_task = asyncio.create_task(periodic_auto_topup())
|
||||
refund_sweep_task = asyncio.create_task(periodic_refund_sweep())
|
||||
routstr_fee_task = asyncio.create_task(periodic_routstr_fee_payout())
|
||||
invoice_watcher_task = asyncio.create_task(periodic_invoice_watcher())
|
||||
|
||||
yield
|
||||
|
||||
@@ -127,6 +153,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
payout_task.cancel()
|
||||
if nip91_task is not None:
|
||||
nip91_task.cancel()
|
||||
if analytics_task is not None:
|
||||
analytics_task.cancel()
|
||||
if providers_task is not None:
|
||||
providers_task.cancel()
|
||||
if models_refresh_task is not None:
|
||||
@@ -135,6 +163,16 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
model_maps_refresh_task.cancel()
|
||||
if key_reset_task is not None:
|
||||
key_reset_task.cancel()
|
||||
if stale_reservation_task is not None:
|
||||
stale_reservation_task.cancel()
|
||||
if auto_topup_task is not None:
|
||||
auto_topup_task.cancel()
|
||||
if refund_sweep_task is not None:
|
||||
refund_sweep_task.cancel()
|
||||
if routstr_fee_task is not None:
|
||||
routstr_fee_task.cancel()
|
||||
if invoice_watcher_task is not None:
|
||||
invoice_watcher_task.cancel()
|
||||
|
||||
try:
|
||||
tasks_to_wait = []
|
||||
@@ -146,6 +184,8 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
tasks_to_wait.append(payout_task)
|
||||
if nip91_task is not None:
|
||||
tasks_to_wait.append(nip91_task)
|
||||
if analytics_task is not None:
|
||||
tasks_to_wait.append(analytics_task)
|
||||
if providers_task is not None:
|
||||
tasks_to_wait.append(providers_task)
|
||||
if models_refresh_task is not None:
|
||||
@@ -154,6 +194,16 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
tasks_to_wait.append(model_maps_refresh_task)
|
||||
if key_reset_task is not None:
|
||||
tasks_to_wait.append(key_reset_task)
|
||||
if stale_reservation_task is not None:
|
||||
tasks_to_wait.append(stale_reservation_task)
|
||||
if auto_topup_task is not None:
|
||||
tasks_to_wait.append(auto_topup_task)
|
||||
if refund_sweep_task is not None:
|
||||
tasks_to_wait.append(refund_sweep_task)
|
||||
if routstr_fee_task is not None:
|
||||
tasks_to_wait.append(routstr_fee_task)
|
||||
if invoice_watcher_task is not None:
|
||||
tasks_to_wait.append(invoice_watcher_task)
|
||||
|
||||
if tasks_to_wait:
|
||||
await asyncio.gather(*tasks_to_wait, return_exceptions=True)
|
||||
@@ -165,6 +215,23 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
)
|
||||
|
||||
|
||||
class _ImmutableStaticFiles(StaticFiles):
|
||||
"""Static files with long Cache-Control for content-hashed Next.js assets.
|
||||
|
||||
Files under `/_next/static/` are emitted with content hashes in their
|
||||
filenames and never mutate, so we serve them with a one-year immutable
|
||||
cache header so browsers and CDNs stop revalidating on every reload.
|
||||
"""
|
||||
|
||||
async def get_response(self, path: str, scope: Scope) -> StarletteResponse:
|
||||
response = await super().get_response(path, scope)
|
||||
if response.status_code == 200:
|
||||
response.headers["Cache-Control"] = (
|
||||
"public, max-age=31536000, immutable"
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
app = FastAPI(version=__version__, lifespan=lifespan)
|
||||
|
||||
|
||||
@@ -174,7 +241,7 @@ app.add_middleware(
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
expose_headers=["x-routstr-request-id"],
|
||||
expose_headers=["x-routstr-request-id", "x-cashu"],
|
||||
)
|
||||
|
||||
# Add logging middleware
|
||||
@@ -211,7 +278,7 @@ if UI_DIST_PATH.exists() and UI_DIST_PATH.is_dir():
|
||||
|
||||
app.mount(
|
||||
"/_next",
|
||||
StaticFiles(directory=UI_DIST_PATH / "_next", check_dir=True),
|
||||
_ImmutableStaticFiles(directory=UI_DIST_PATH / "_next", check_dir=True),
|
||||
name="next-static",
|
||||
)
|
||||
|
||||
@@ -219,100 +286,70 @@ if UI_DIST_PATH.exists() and UI_DIST_PATH.is_dir():
|
||||
async def serve_root_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
# Add explicit route for /index.txt to redirect to /
|
||||
# Serve the App Router RSC payload for the home page.
|
||||
@app.get("/index.txt", include_in_schema=False)
|
||||
async def redirect_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/")
|
||||
async def serve_root_rsc() -> FileResponse:
|
||||
return FileResponse(
|
||||
UI_DIST_PATH / "index.txt", media_type="text/x-component"
|
||||
)
|
||||
|
||||
# Next.js is built with `trailingSlash: true`, so all UI page URLs end
|
||||
# with a slash (e.g. `/login/`). The proxy router catches `/{path:path}`
|
||||
# before FastAPI's `redirect_slashes` logic can normalize the URL, so we
|
||||
# must register both the with-slash and without-slash variants here.
|
||||
UI_PAGES = (
|
||||
"dashboard",
|
||||
"login",
|
||||
"model",
|
||||
"providers",
|
||||
"settings",
|
||||
"transactions",
|
||||
"balances",
|
||||
"logs",
|
||||
"usage",
|
||||
"unauthorized",
|
||||
)
|
||||
|
||||
def _register_ui_page(name: str) -> None:
|
||||
page_dir = UI_DIST_PATH / name
|
||||
index_html = page_dir / "index.html"
|
||||
index_txt = page_dir / "index.txt"
|
||||
|
||||
async def serve_page() -> FileResponse:
|
||||
return FileResponse(index_html)
|
||||
|
||||
async def serve_page_rsc() -> FileResponse:
|
||||
return FileResponse(index_txt, media_type="text/x-component")
|
||||
|
||||
app.add_api_route(
|
||||
f"/{name}",
|
||||
serve_page,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_ui",
|
||||
)
|
||||
app.add_api_route(
|
||||
f"/{name}/",
|
||||
serve_page,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_ui_slash",
|
||||
)
|
||||
app.add_api_route(
|
||||
f"/{name}/index.txt",
|
||||
serve_page_rsc,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_rsc",
|
||||
)
|
||||
|
||||
for _page in UI_PAGES:
|
||||
_register_ui_page(_page)
|
||||
|
||||
@app.get("/admin")
|
||||
async def admin_redirect() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
@app.get("/dashboard", include_in_schema=False)
|
||||
async def serve_dashboard_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
@app.get("/login", include_in_schema=False)
|
||||
async def serve_login_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "login" / "index.html")
|
||||
|
||||
# Add explicit route for /login/index.txt to redirect to /login
|
||||
@app.get("/login/index.txt", include_in_schema=False)
|
||||
async def redirect_login_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/login")
|
||||
|
||||
@app.get("/model", include_in_schema=False)
|
||||
async def serve_models_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "model" / "index.html")
|
||||
|
||||
# Add explicit route for /model/index.txt to redirect to /model
|
||||
@app.get("/model/index.txt", include_in_schema=False)
|
||||
async def redirect_model_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/model")
|
||||
|
||||
@app.get("/providers", include_in_schema=False)
|
||||
async def serve_providers_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "providers" / "index.html")
|
||||
|
||||
# Add explicit route for /providers/index.txt to redirect to /providers
|
||||
@app.get("/providers/index.txt", include_in_schema=False)
|
||||
async def redirect_providers_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/providers")
|
||||
|
||||
@app.get("/settings", include_in_schema=False)
|
||||
async def serve_settings_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "settings" / "index.html")
|
||||
|
||||
# Add explicit route for /settings/index.txt to redirect to /settings
|
||||
@app.get("/settings/index.txt", include_in_schema=False)
|
||||
async def redirect_settings_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/settings")
|
||||
|
||||
@app.get("/transactions", include_in_schema=False)
|
||||
async def serve_transactions_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "transactions" / "index.html")
|
||||
|
||||
# Add explicit route for /transactions/index.txt to redirect to /transactions
|
||||
@app.get("/transactions/index.txt", include_in_schema=False)
|
||||
async def redirect_transactions_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/transactions")
|
||||
|
||||
@app.get("/balances", include_in_schema=False)
|
||||
async def serve_balances_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "balances" / "index.html")
|
||||
|
||||
# Add explicit route for /balances/index.txt to redirect to /balances
|
||||
@app.get("/balances/index.txt", include_in_schema=False)
|
||||
async def redirect_balances_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/balances")
|
||||
|
||||
@app.get("/logs", include_in_schema=False)
|
||||
async def serve_logs_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "logs" / "index.html")
|
||||
|
||||
# Add explicit route for /logs/index.txt to redirect to /logs
|
||||
@app.get("/logs/index.txt", include_in_schema=False)
|
||||
async def redirect_logs_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/logs")
|
||||
|
||||
@app.get("/usage", include_in_schema=False)
|
||||
async def serve_usage_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "usage" / "index.html")
|
||||
|
||||
# Add explicit route for /usage/index.txt to redirect to /usage
|
||||
@app.get("/usage/index.txt", include_in_schema=False)
|
||||
async def redirect_usage_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/usage")
|
||||
|
||||
@app.get("/unauthorized", include_in_schema=False)
|
||||
async def serve_unauthorized_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "unauthorized" / "index.html")
|
||||
|
||||
# Add explicit route for /unauthorized/index.txt to redirect to /unauthorized
|
||||
@app.get("/unauthorized/index.txt", include_in_schema=False)
|
||||
async def redirect_unauthorized_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/unauthorized")
|
||||
|
||||
@app.get("/favicon.ico", include_in_schema=False)
|
||||
async def serve_favicon() -> FileResponse:
|
||||
icon_path = UI_DIST_PATH / "icon.ico"
|
||||
@@ -324,9 +361,6 @@ if UI_DIST_PATH.exists() and UI_DIST_PATH.is_dir():
|
||||
async def serve_icon() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "icon.ico")
|
||||
|
||||
app.mount(
|
||||
"/static", StaticFiles(directory=UI_DIST_PATH, check_dir=True), name="ui-static"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"UI dist directory not found at {UI_DIST_PATH}, skipping static file serving"
|
||||
@@ -346,6 +380,7 @@ else:
|
||||
app.include_router(models_router)
|
||||
app.include_router(admin_router)
|
||||
app.include_router(balance_router)
|
||||
app.include_router(lightning_router)
|
||||
app.include_router(deprecated_wallet_router)
|
||||
app.include_router(providers_router)
|
||||
app.include_router(proxy_router)
|
||||
|
||||
@@ -14,8 +14,54 @@ logger = get_logger(__name__)
|
||||
request_id_context: ContextVar[str | None] = ContextVar("request_id")
|
||||
|
||||
|
||||
# Methods that are never logged: HEAD requests are health probes from
|
||||
# monitoring/load balancers, OPTIONS are CORS preflights — both are framework
|
||||
# chatter, not user-meaningful events.
|
||||
_SKIP_LOG_METHODS: frozenset[str] = frozenset({"HEAD", "OPTIONS"})
|
||||
|
||||
# Path prefixes to skip. Includes Next.js static chunks and the admin
|
||||
# dashboard's internal polling API (/admin/api/*) which the UI hits on a timer
|
||||
# to refresh balances, logs, providers, etc. — high volume, low diagnostic
|
||||
# value. Mutating admin actions are recorded separately in the audit log.
|
||||
_SKIP_LOG_PREFIXES: tuple[str, ...] = (
|
||||
"/_next/",
|
||||
"/admin/api/",
|
||||
)
|
||||
|
||||
# Exact paths to skip. RSC payload prefetches (`*/index.txt`) fire automatically
|
||||
# as the user hovers near `<Link>`s, and `/v1/wallet/info` is polled by the UI.
|
||||
_SKIP_LOG_EXACT: frozenset[str] = frozenset(
|
||||
{
|
||||
"/favicon.ico",
|
||||
"/icon.ico",
|
||||
"/v1/wallet/info",
|
||||
"/index.txt",
|
||||
"/login/index.txt",
|
||||
"/model/index.txt",
|
||||
"/providers/index.txt",
|
||||
"/settings/index.txt",
|
||||
"/transactions/index.txt",
|
||||
"/balances/index.txt",
|
||||
"/logs/index.txt",
|
||||
"/usage/index.txt",
|
||||
"/unauthorized/index.txt",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _should_log(method: str, path: str) -> bool:
|
||||
if method in _SKIP_LOG_METHODS:
|
||||
return False
|
||||
if path in _SKIP_LOG_EXACT:
|
||||
return False
|
||||
return not any(path.startswith(prefix) for prefix in _SKIP_LOG_PREFIXES)
|
||||
|
||||
|
||||
class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
"""Middleware to log detailed request and response information."""
|
||||
"""Middleware to log proxy interactions and page navigation.
|
||||
|
||||
Skips logging for static assets and Next.js chunks to avoid noise.
|
||||
"""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
||||
# Generate request ID
|
||||
@@ -25,62 +71,20 @@ class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
# Set request ID in context for logging
|
||||
token = request_id_context.set(request_id)
|
||||
|
||||
path = request.url.path
|
||||
should_log = _should_log(request.method, path)
|
||||
|
||||
# Start timing
|
||||
start_time = time.time()
|
||||
|
||||
# Log request details
|
||||
request_body = None
|
||||
if request.method in ["POST", "PUT", "PATCH"]:
|
||||
try:
|
||||
# Only read body for non-streaming requests
|
||||
if hasattr(request, "_body"):
|
||||
request_body = await request.body()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Extract request info
|
||||
client_host = None
|
||||
if request.client:
|
||||
client_host = request.client.host
|
||||
|
||||
# Log incoming request
|
||||
logger.info(
|
||||
"Incoming request",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"query_params": dict(request.query_params),
|
||||
"client_host": client_host,
|
||||
"headers": {
|
||||
k: v
|
||||
for k, v in request.headers.items()
|
||||
if k.lower()
|
||||
not in [
|
||||
"authorization",
|
||||
"x-cashu",
|
||||
"cookie",
|
||||
"cf-connecting-ip",
|
||||
"cf-ipcountry",
|
||||
"x-forwarded-for",
|
||||
"x-real-ip",
|
||||
]
|
||||
},
|
||||
"body_size": len(request_body) if request_body else 0,
|
||||
},
|
||||
)
|
||||
|
||||
# Log at TRACE level for full body (security filter will redact sensitive data)
|
||||
if request_body and hasattr(logger, "exception"):
|
||||
logger.exception(
|
||||
"Request body",
|
||||
if should_log:
|
||||
logger.info(
|
||||
"Incoming request",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"body": request_body.decode("utf-8", errors="ignore")[
|
||||
:1000
|
||||
], # Limit size
|
||||
"path": path,
|
||||
"query_params": dict(request.query_params),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -88,39 +92,33 @@ class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
try:
|
||||
response = await call_next(request)
|
||||
|
||||
# Calculate duration
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Log response
|
||||
logger.info(
|
||||
"Request completed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"status_code": response.status_code,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
"client_host": client_host,
|
||||
},
|
||||
)
|
||||
if should_log:
|
||||
duration = time.time() - start_time
|
||||
logger.info(
|
||||
"Request completed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": path,
|
||||
"status_code": response.status_code,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
},
|
||||
)
|
||||
if hasattr(response, "headers"):
|
||||
response.headers["x-routstr-request-id"] = request_id
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
# Calculate duration
|
||||
# Always log failures, even for skipped paths, so we don't lose errors.
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Log error
|
||||
logger.error(
|
||||
"Request failed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"path": path,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
"client_host": client_host,
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
},
|
||||
|
||||
55
routstr/core/not_found.py
Normal file
55
routstr/core/not_found.py
Normal file
@@ -0,0 +1,55 @@
|
||||
"""Shared 404 handler used by the proxy catch-all and tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, Response
|
||||
|
||||
_NOT_FOUND_HTML_FILE = Path(__file__).parent.parent.parent / "ui_out" / "404.html"
|
||||
|
||||
|
||||
def _read_not_found_html() -> str | None:
|
||||
try:
|
||||
return _NOT_FOUND_HTML_FILE.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
_NOT_FOUND_HTML: str | None = _read_not_found_html()
|
||||
|
||||
|
||||
def build_not_found_response(request: Request, path: str) -> Response:
|
||||
"""Return a 404 response.
|
||||
|
||||
HTML 404 page only for GET requests from browsers (Accept: text/html).
|
||||
All POST requests and API clients receive a JSON 404.
|
||||
"""
|
||||
accept = request.headers.get("accept", "").lower()
|
||||
prefers_html = (
|
||||
request.method == "GET"
|
||||
and "text/html" in accept
|
||||
and "application/json" not in accept
|
||||
)
|
||||
request_id = getattr(request.state, "request_id", "unknown")
|
||||
|
||||
if prefers_html and _NOT_FOUND_HTML is not None:
|
||||
return HTMLResponse(content=_NOT_FOUND_HTML, status_code=404)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=404,
|
||||
content={
|
||||
"error": {
|
||||
"message": f"Path '/{path}' not found",
|
||||
"type": "not_found",
|
||||
"code": 404,
|
||||
},
|
||||
"request_id": request_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def not_found_catch_all(request: Request, path: str) -> Response:
|
||||
"""ASGI handler form of :func:`build_not_found_response`."""
|
||||
return build_not_found_response(request, path)
|
||||
@@ -41,6 +41,15 @@ class Settings(BaseSettings):
|
||||
primary_mint: str = Field(default="", env="PRIMARY_MINT_URL")
|
||||
primary_mint_unit: str = Field(default="sat", env="PRIMARY_MINT_UNIT")
|
||||
|
||||
# Lightning payout configuration
|
||||
# Minimum available balance (in satoshis) before profit is paid out over
|
||||
# Lightning
|
||||
min_payout_sat: int = Field(default=210, gt=0, env="MIN_PAYOUT_SAT")
|
||||
# Interval (seconds) between periodic payout attempts. Must be positive.
|
||||
payout_interval_seconds: int = Field(
|
||||
default=900, gt=0, env="PAYOUT_INTERVAL_SECONDS"
|
||||
)
|
||||
|
||||
# Pricing
|
||||
# Default behavior: derive pricing from MODELS
|
||||
# If fixed_pricing is True -> use fixed_cost_per_request and ignore tokens
|
||||
@@ -52,12 +61,18 @@ class Settings(BaseSettings):
|
||||
exchange_fee: float = Field(default=1.005, env="EXCHANGE_FEE")
|
||||
upstream_provider_fee: float = Field(default=1.05, env="UPSTREAM_PROVIDER_FEE")
|
||||
tolerance_percentage: float = Field(default=1.0, env="TOLERANCE_PERCENTAGE")
|
||||
child_key_cost: int = Field(default=1000, env="CHILD_KEY_COST")
|
||||
child_key_cost: int = Field(default=0, env="CHILD_KEY_COST")
|
||||
# Minimum per-request charge in millisatoshis when model pricing is free/zero
|
||||
min_request_msat: int = Field(default=1, env="MIN_REQUEST_MSAT")
|
||||
reset_reserved_balance_on_startup: bool = Field(
|
||||
default=True, env="RESET_RESERVED_BALANCE_ON_STARTUP"
|
||||
) # deactivate in horizontal scaling setups
|
||||
# Reservations older than this are considered leaked (client disconnect,
|
||||
# crash, abandoned stream) and released by the background sweeper and the
|
||||
# refund endpoint.
|
||||
stale_reservation_timeout_seconds: int = Field(
|
||||
default=300, env="STALE_RESERVATION_TIMEOUT_SECONDS"
|
||||
)
|
||||
|
||||
# Network
|
||||
cors_origins: list[str] = Field(default_factory=lambda: ["*"], env="CORS_ORIGINS")
|
||||
@@ -74,6 +89,7 @@ class Settings(BaseSettings):
|
||||
enable_pricing_refresh: bool = Field(default=True, env="ENABLE_PRICING_REFRESH")
|
||||
enable_models_refresh: bool = Field(default=True, env="ENABLE_MODELS_REFRESH")
|
||||
refund_cache_ttl_seconds: int = Field(default=3600, env="REFUND_CACHE_TTL_SECONDS")
|
||||
refund_sweep_ttl_seconds: int = Field(default=604800, env="REFUND_SWEEP_TTL_SECONDS")
|
||||
|
||||
# Logging
|
||||
log_level: str = Field(default="INFO", env="LOG_LEVEL")
|
||||
@@ -92,6 +108,20 @@ class Settings(BaseSettings):
|
||||
|
||||
# Discovery
|
||||
relays: list[str] = Field(default_factory=list, env="RELAYS")
|
||||
enable_analytics_sharing: bool = Field(
|
||||
default=True, env="ENABLE_ANALYTICS_SHARING"
|
||||
)
|
||||
|
||||
def _normalize_settings_data(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Discard unknown keys from persisted settings."""
|
||||
normalized: dict[str, Any] = {}
|
||||
known_fields = Settings.__fields__
|
||||
|
||||
for key, value in data.items():
|
||||
if key in known_fields:
|
||||
normalized[key] = value
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def _compute_primary_mint(cashu_mints: list[str]) -> str:
|
||||
@@ -231,16 +261,21 @@ class SettingsService:
|
||||
|
||||
db_id, db_data, _updated_at = row
|
||||
try:
|
||||
db_json = (
|
||||
db_json_raw = (
|
||||
json.loads(db_data) if isinstance(db_data, str) else dict(db_data)
|
||||
)
|
||||
if not isinstance(db_json_raw, dict):
|
||||
db_json_raw = {}
|
||||
except Exception:
|
||||
db_json = {}
|
||||
db_json_raw = {}
|
||||
db_json = _normalize_settings_data(db_json_raw)
|
||||
|
||||
valid_fields = set(env_resolved.dict().keys())
|
||||
merged_dict: dict[str, Any] = dict(env_resolved.dict())
|
||||
merged_dict.update(
|
||||
{k: v for k, v in db_json.items() if v not in (None, "", [], {})}
|
||||
{k: v for k, v in db_json.items() if v not in (None, "", [], {}) and k in valid_fields}
|
||||
)
|
||||
merged_dict = Settings(**merged_dict).dict()
|
||||
|
||||
# Ensure primary_mint is consistent with cashu_mints if not explicitly set
|
||||
if not merged_dict.get("primary_mint"):
|
||||
@@ -248,7 +283,7 @@ class SettingsService:
|
||||
merged_dict.get("cashu_mints", [])
|
||||
)
|
||||
|
||||
if any(k not in db_json for k in merged_dict.keys()):
|
||||
if db_json_raw != merged_dict:
|
||||
await db_session.exec( # type: ignore
|
||||
text(
|
||||
"UPDATE settings SET data = :data, updated_at = :updated_at WHERE id = 1"
|
||||
@@ -271,7 +306,7 @@ class SettingsService:
|
||||
) -> Settings:
|
||||
async with cls._lock:
|
||||
current = cls.get()
|
||||
candidate_dict = {**current.dict(), **partial}
|
||||
candidate_dict = {**current.dict(), **_normalize_settings_data(partial)}
|
||||
candidate = Settings(**candidate_dict)
|
||||
from sqlmodel import text
|
||||
|
||||
@@ -305,8 +340,10 @@ class SettingsService:
|
||||
raise RuntimeError("Settings row missing")
|
||||
(data_str,) = row
|
||||
data = json.loads(data_str) if isinstance(data_str, str) else dict(data_str)
|
||||
valid_fields = set(settings.dict().keys())
|
||||
# Update in-place
|
||||
for k, v in data.items():
|
||||
setattr(settings, k, v)
|
||||
if k in valid_fields:
|
||||
setattr(settings, k, v)
|
||||
cls._current = settings
|
||||
return settings
|
||||
|
||||
1380
routstr/core/usage_analytics_store.py
Normal file
1380
routstr/core/usage_analytics_store.py
Normal file
File diff suppressed because it is too large
Load Diff
81
routstr/core/version.py
Normal file
81
routstr/core/version.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""Application version resolution.
|
||||
|
||||
Priority order:
|
||||
|
||||
1. ``VERSION_SUFFIX`` env var (manual override; preserves prior behaviour).
|
||||
2. Bare base version when HEAD is on the matching release tag (detected via
|
||||
``GIT_TAG`` env or ``git describe --tags --exact-match HEAD``).
|
||||
3. ``GIT_COMMIT`` env var (build-time injection) -> ``<base>+g<sha>``.
|
||||
4. Local ``.git`` lookup (source checkouts) -> ``<base>+g<sha>``.
|
||||
5. Fallback: bare base version.
|
||||
|
||||
The ``+g<sha>`` form is PEP 440 local-version syntax so the result remains a
|
||||
valid package version.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
BASE_VERSION = "0.4.3"
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_GIT_TIMEOUT_SECONDS = 2.0
|
||||
|
||||
|
||||
def _run_git(*args: str) -> str | None:
|
||||
try:
|
||||
result = subprocess.run( # noqa: S603 - fixed argv, no shell
|
||||
["git", *args],
|
||||
cwd=_REPO_ROOT,
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=_GIT_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (FileNotFoundError, subprocess.SubprocessError, OSError):
|
||||
return None
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
return result.stdout.strip() or None
|
||||
|
||||
|
||||
def _git_short_sha() -> str | None:
|
||||
sha = os.getenv("GIT_COMMIT", "").strip()
|
||||
if sha:
|
||||
return sha[:7]
|
||||
return _run_git("rev-parse", "--short=7", "HEAD")
|
||||
|
||||
|
||||
def _on_tagged_release() -> bool:
|
||||
tag = os.getenv("GIT_TAG", "").strip()
|
||||
if tag:
|
||||
return tag.lstrip("v") == BASE_VERSION
|
||||
described = _run_git("describe", "--tags", "--exact-match", "HEAD")
|
||||
if described and described.lstrip("v") == BASE_VERSION:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_version() -> str:
|
||||
suffix = os.getenv("VERSION_SUFFIX")
|
||||
if suffix is not None:
|
||||
return f"{BASE_VERSION}-{suffix}"
|
||||
|
||||
if _on_tagged_release():
|
||||
return BASE_VERSION
|
||||
|
||||
sha = _git_short_sha()
|
||||
if not sha:
|
||||
return BASE_VERSION
|
||||
|
||||
return f"{BASE_VERSION}+g{sha}"
|
||||
|
||||
|
||||
__version__ = get_version()
|
||||
|
||||
__all__ = ["BASE_VERSION", "__version__", "get_version"]
|
||||
@@ -1,13 +1,14 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import secrets
|
||||
import time
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
from sqlmodel import select
|
||||
from sqlmodel import col, select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from .core.db import ApiKey, LightningInvoice, get_session
|
||||
from .core.db import ApiKey, LightningInvoice, create_session, get_session
|
||||
from .core.logging import get_logger
|
||||
from .core.settings import settings
|
||||
from .wallet import get_wallet
|
||||
@@ -19,15 +20,29 @@ lightning_router = APIRouter(prefix="/lightning")
|
||||
|
||||
class InvoiceCreateRequest(BaseModel):
|
||||
amount_sats: int = Field(gt=0, le=1_000_000, description="Amount in satoshis")
|
||||
purpose: str = Field(description="create or topup", pattern="^(create|topup)$")
|
||||
purpose: str = Field(
|
||||
default="create",
|
||||
description="create or topup",
|
||||
pattern="^(create|topup)$",
|
||||
)
|
||||
api_key: str | None = Field(
|
||||
default=None, description="Required for topup operations"
|
||||
default=None,
|
||||
description="Deprecated: legacy field for topup. Prefer Authorization header.",
|
||||
)
|
||||
balance_limit: int | None = Field(default=None)
|
||||
balance_limit_reset: str | None = Field(default=None)
|
||||
validity_date: int | None = Field(default=None)
|
||||
|
||||
|
||||
def _extract_bearer_api_key(authorization: str | None) -> str | None:
|
||||
if not authorization:
|
||||
return None
|
||||
token = authorization.strip()
|
||||
if token.lower().startswith("bearer "):
|
||||
token = token[7:].strip()
|
||||
return token or None
|
||||
|
||||
|
||||
class InvoiceCreateResponse(BaseModel):
|
||||
invoice_id: str
|
||||
bolt11: str
|
||||
@@ -64,18 +79,21 @@ def generate_invoice_id() -> str:
|
||||
@lightning_router.post("/invoice", response_model=InvoiceCreateResponse)
|
||||
async def create_invoice(
|
||||
request: InvoiceCreateRequest,
|
||||
authorization: str | None = Header(default=None),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> InvoiceCreateResponse:
|
||||
if request.purpose == "topup" and not request.api_key:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="api_key is required for topup operations"
|
||||
)
|
||||
api_key_token = _extract_bearer_api_key(authorization) or request.api_key
|
||||
|
||||
if request.purpose == "topup" and request.api_key:
|
||||
if not request.api_key.startswith("sk-"):
|
||||
if request.purpose == "topup":
|
||||
if not api_key_token:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Authorization bearer api key is required for topup",
|
||||
)
|
||||
if not api_key_token.startswith("sk-"):
|
||||
raise HTTPException(status_code=400, detail="Invalid API key format")
|
||||
|
||||
api_key = await session.get(ApiKey, request.api_key[3:])
|
||||
api_key = await session.get(ApiKey, api_key_token[3:])
|
||||
if not api_key:
|
||||
raise HTTPException(status_code=404, detail="API key not found")
|
||||
|
||||
@@ -95,7 +113,7 @@ async def create_invoice(
|
||||
description=description,
|
||||
payment_hash=payment_hash,
|
||||
status="pending",
|
||||
api_key_hash=request.api_key[3:] if request.api_key else None,
|
||||
api_key_hash=api_key_token[3:] if api_key_token else None,
|
||||
purpose=request.purpose,
|
||||
balance_limit=request.balance_limit,
|
||||
balance_limit_reset=request.balance_limit_reset,
|
||||
@@ -142,13 +160,13 @@ async def get_invoice_status(
|
||||
if not invoice:
|
||||
raise HTTPException(status_code=404, detail="Invoice not found")
|
||||
|
||||
if invoice.status == "pending":
|
||||
await check_invoice_payment(invoice, session)
|
||||
|
||||
if invoice.status == "pending" and int(time.time()) > invoice.expires_at:
|
||||
invoice.status = "expired"
|
||||
await session.commit()
|
||||
|
||||
if invoice.status == "pending":
|
||||
await check_invoice_payment(invoice, session)
|
||||
|
||||
api_key = None
|
||||
if invoice.status == "paid" and invoice.purpose == "create":
|
||||
if invoice.api_key_hash:
|
||||
@@ -251,6 +269,9 @@ async def create_api_key_from_invoice(
|
||||
balance=invoice.amount_sats * 1000, # Convert to msats
|
||||
refund_currency="sat",
|
||||
refund_mint_url=settings.primary_mint,
|
||||
balance_limit=invoice.balance_limit,
|
||||
balance_limit_reset=invoice.balance_limit_reset,
|
||||
validity_date=invoice.validity_date,
|
||||
)
|
||||
|
||||
session.add(api_key)
|
||||
@@ -274,3 +295,41 @@ async def topup_api_key_from_invoice(
|
||||
|
||||
api_key.balance += invoice.amount_sats * 1000 # Convert to msats
|
||||
await session.flush()
|
||||
|
||||
|
||||
INVOICE_WATCH_INTERVAL_SECONDS = 5
|
||||
INVOICE_WATCH_BATCH_LIMIT = 100
|
||||
|
||||
|
||||
async def periodic_invoice_watcher() -> None:
|
||||
"""Background task: detect paid Lightning invoices and credit balances.
|
||||
|
||||
Removes the need for clients to poll the status endpoint after paying.
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
async with create_session() as session:
|
||||
now = int(time.time())
|
||||
result = await session.exec(
|
||||
select(LightningInvoice)
|
||||
.where(
|
||||
LightningInvoice.status == "pending",
|
||||
col(LightningInvoice.expires_at) > now,
|
||||
)
|
||||
.limit(INVOICE_WATCH_BATCH_LIMIT)
|
||||
)
|
||||
pending = result.all()
|
||||
for invoice in pending:
|
||||
try:
|
||||
await check_invoice_payment(invoice, session)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Invoice watcher failed for invoice",
|
||||
extra={"invoice_id": invoice.id, "error": str(e)},
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Invoice watcher loop error: {e}")
|
||||
|
||||
await asyncio.sleep(INVOICE_WATCH_INTERVAL_SECONDS)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from .analytics import publish_usage_analytics
|
||||
from .discovery import providers_cache_refresher
|
||||
from .listing import announce_provider
|
||||
|
||||
__all__ = ["providers_cache_refresher", "announce_provider"]
|
||||
__all__ = ["providers_cache_refresher", "announce_provider", "publish_usage_analytics"]
|
||||
|
||||
419
routstr/nostr/analytics.py
Normal file
419
routstr/nostr/analytics.py
Normal file
@@ -0,0 +1,419 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Nostr usage analytics publisher.
|
||||
Publishes a single replaceable analytics snapshot for each provider.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from nostr.event import Event
|
||||
from nostr.key import PrivateKey
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.log_manager import log_manager
|
||||
from ..core.settings import settings
|
||||
from .listing import nsec_to_keypair, publish_to_relay
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
ANALYTICS_KIND = 38422
|
||||
ANALYTICS_SCHEMA = "routstr.analytics.snapshot.v1"
|
||||
DEFAULT_RELAYS = [
|
||||
"wss://relay.nostr.band",
|
||||
"wss://relay.damus.io",
|
||||
"wss://relay.routstr.com",
|
||||
"wss://nos.lol",
|
||||
]
|
||||
PUBLISH_INTERVAL_SECONDS = 15 * 60
|
||||
DISABLED_POLL_SECONDS = 60
|
||||
DASHBOARD_WINDOW_HOURS = 24
|
||||
DASHBOARD_INTERVAL_MINUTES = 60
|
||||
MODEL_LIMIT = 20
|
||||
WINDOW_DEFINITIONS: tuple[tuple[str, int, int], ...] = (
|
||||
("24h", 24, 60),
|
||||
("7d", 7 * 24, 6 * 60),
|
||||
("30d", 30 * 24, 24 * 60),
|
||||
("3m", 90 * 24, 24 * 60),
|
||||
("1y", 365 * 24, 7 * 24 * 60),
|
||||
)
|
||||
|
||||
|
||||
def _event_to_dict(ev: Event) -> dict[str, Any]:
|
||||
return {
|
||||
"id": ev.id,
|
||||
"pubkey": ev.public_key,
|
||||
"created_at": ev.created_at,
|
||||
"kind": int(ev.kind) if not isinstance(ev.kind, int) else ev.kind,
|
||||
"tags": ev.tags,
|
||||
"content": ev.content,
|
||||
"sig": ev.signature,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_provider_id(public_key_hex: str) -> str:
|
||||
explicit_provider_id = (settings.provider_id or "").strip()
|
||||
if explicit_provider_id:
|
||||
return explicit_provider_id
|
||||
return public_key_hex[:12]
|
||||
|
||||
|
||||
def _resolve_endpoint_urls() -> list[str]:
|
||||
urls: list[str] = []
|
||||
http_url = (settings.http_url or "").strip()
|
||||
onion_url = (settings.onion_url or "").strip()
|
||||
|
||||
if http_url and http_url != "http://localhost:8000":
|
||||
urls.append(http_url)
|
||||
|
||||
if onion_url:
|
||||
if onion_url.endswith(".onion") and not (
|
||||
onion_url.startswith("http://") or onion_url.startswith("https://")
|
||||
):
|
||||
onion_url = f"http://{onion_url}"
|
||||
urls.append(onion_url)
|
||||
|
||||
return urls
|
||||
|
||||
|
||||
def _resolve_relays() -> list[str]:
|
||||
configured = [url.strip() for url in settings.relays if url.strip()]
|
||||
return configured if configured else list(DEFAULT_RELAYS)
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
if isinstance(value, bool):
|
||||
return int(value)
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return int(float(value))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _to_float(value: Any) -> float:
|
||||
if isinstance(value, bool):
|
||||
return float(int(value))
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return 0.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def _aggregate_top_model_usage(
|
||||
model_usage_mix: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||||
top_models_raw = model_usage_mix.get("top_models", [])
|
||||
mix_metrics_raw = model_usage_mix.get("metrics", [])
|
||||
|
||||
top_models = [model for model in top_models_raw if isinstance(model, str)]
|
||||
metrics = [row for row in mix_metrics_raw if isinstance(row, dict)]
|
||||
|
||||
model_totals: dict[str, dict[str, float | int]] = {
|
||||
model: {
|
||||
"successful_requests": 0,
|
||||
"revenue_msats": 0.0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
for model in top_models
|
||||
}
|
||||
others = {
|
||||
"successful_requests": 0,
|
||||
"revenue_msats": 0.0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
|
||||
for metric in metrics:
|
||||
model_counts = metric.get("model_counts", {})
|
||||
model_revenue = metric.get("model_revenue_msats", {})
|
||||
model_tokens = metric.get("model_tokens", {})
|
||||
|
||||
if isinstance(model_counts, dict):
|
||||
for model, count in model_counts.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["successful_requests"] += _to_int(count)
|
||||
|
||||
if isinstance(model_revenue, dict):
|
||||
for model, amount in model_revenue.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["revenue_msats"] += _to_float(amount)
|
||||
|
||||
if isinstance(model_tokens, dict):
|
||||
for model, token_count in model_tokens.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["total_tokens"] += _to_int(token_count)
|
||||
|
||||
others["successful_requests"] += _to_int(metric.get("others", 0))
|
||||
others["revenue_msats"] += _to_float(metric.get("others_revenue_msats", 0.0))
|
||||
others["total_tokens"] += _to_int(metric.get("others_tokens", 0))
|
||||
|
||||
model_rows = [
|
||||
{
|
||||
"model": model,
|
||||
"successful_requests": int(values["successful_requests"]),
|
||||
"revenue_msats": float(values["revenue_msats"]),
|
||||
"total_tokens": int(values["total_tokens"]),
|
||||
}
|
||||
for model, values in model_totals.items()
|
||||
]
|
||||
model_rows.sort(
|
||||
key=lambda row: _to_int(row.get("successful_requests", 0)),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
return model_rows, others
|
||||
|
||||
|
||||
def _build_summary_payload(summary: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"total_requests": _to_int(summary.get("total_requests", 0)),
|
||||
"successful_chat_completions": _to_int(
|
||||
summary.get("successful_chat_completions", 0)
|
||||
),
|
||||
"failed_requests": _to_int(summary.get("failed_requests", 0)),
|
||||
"success_rate": _to_float(summary.get("success_rate", 0.0)),
|
||||
"unique_models_count": _to_int(summary.get("unique_models_count", 0)),
|
||||
"input_tokens": _to_int(summary.get("input_tokens", 0)),
|
||||
"output_tokens": _to_int(summary.get("output_tokens", 0)),
|
||||
"total_tokens": _to_int(summary.get("total_tokens", 0)),
|
||||
"revenue_msats": _to_float(summary.get("revenue_msats", 0.0)),
|
||||
"refunds_msats": _to_float(summary.get("refunds_msats", 0.0)),
|
||||
"net_revenue_msats": _to_float(summary.get("net_revenue_msats", 0.0)),
|
||||
"revenue_sats": _to_float(summary.get("revenue_sats", 0.0)),
|
||||
"refunds_sats": _to_float(summary.get("refunds_sats", 0.0)),
|
||||
"net_revenue_sats": _to_float(summary.get("net_revenue_sats", 0.0)),
|
||||
}
|
||||
|
||||
|
||||
def _build_window_payload(
|
||||
*,
|
||||
hours: int,
|
||||
interval_minutes: int,
|
||||
model_limit: int,
|
||||
) -> dict[str, Any]:
|
||||
dashboard = log_manager.get_usage_dashboard(
|
||||
interval=interval_minutes,
|
||||
hours=hours,
|
||||
error_limit=1,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
summary = dashboard.get("summary", {})
|
||||
model_usage_mix = dashboard.get("model_usage_mix", {})
|
||||
|
||||
summary_payload = _build_summary_payload(summary if isinstance(summary, dict) else {})
|
||||
usage_mix_payload = model_usage_mix if isinstance(model_usage_mix, dict) else {}
|
||||
top_model_usage, others_usage = _aggregate_top_model_usage(usage_mix_payload)
|
||||
|
||||
return {
|
||||
"window_hours": hours,
|
||||
"interval_minutes": interval_minutes,
|
||||
"summary": summary_payload,
|
||||
"model_usage_mix": usage_mix_payload,
|
||||
"top_model_usage": top_model_usage,
|
||||
"others_usage": others_usage,
|
||||
}
|
||||
|
||||
|
||||
def build_stats_snapshot_payload(
|
||||
provider_id: str,
|
||||
*,
|
||||
public_key_hex: str,
|
||||
generated_at: int,
|
||||
window_hours: int = DASHBOARD_WINDOW_HOURS,
|
||||
interval_minutes: int = DASHBOARD_INTERVAL_MINUTES,
|
||||
model_limit: int = MODEL_LIMIT,
|
||||
) -> dict[str, Any]:
|
||||
_ = (window_hours, interval_minutes)
|
||||
windows: dict[str, dict[str, Any]] = {}
|
||||
for key, hours, window_interval_minutes in WINDOW_DEFINITIONS:
|
||||
windows[key] = _build_window_payload(
|
||||
hours=hours,
|
||||
interval_minutes=window_interval_minutes,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
primary_window = windows.get("24h", {})
|
||||
summary_payload = (
|
||||
primary_window.get("summary", {})
|
||||
if isinstance(primary_window.get("summary", {}), dict)
|
||||
else {}
|
||||
)
|
||||
usage_mix_payload = (
|
||||
primary_window.get("model_usage_mix", {})
|
||||
if isinstance(primary_window.get("model_usage_mix", {}), dict)
|
||||
else {}
|
||||
)
|
||||
top_model_usage = (
|
||||
primary_window.get("top_model_usage", [])
|
||||
if isinstance(primary_window.get("top_model_usage", []), list)
|
||||
else []
|
||||
)
|
||||
others_usage = (
|
||||
primary_window.get("others_usage", {})
|
||||
if isinstance(primary_window.get("others_usage", {}), dict)
|
||||
else {}
|
||||
)
|
||||
|
||||
return {
|
||||
"schema": ANALYTICS_SCHEMA,
|
||||
"generated_at": generated_at,
|
||||
"provider_id": provider_id,
|
||||
"pubkey": public_key_hex,
|
||||
"npub": settings.npub or "",
|
||||
"endpoint_urls": _resolve_endpoint_urls(),
|
||||
"window_hours": DASHBOARD_WINDOW_HOURS,
|
||||
"interval_minutes": DASHBOARD_INTERVAL_MINUTES,
|
||||
"summary": summary_payload,
|
||||
"model_usage_mix": usage_mix_payload,
|
||||
"top_model_usage": top_model_usage,
|
||||
"others_usage": others_usage,
|
||||
"windows": windows,
|
||||
}
|
||||
|
||||
|
||||
def create_stats_snapshot_event(
|
||||
private_key_hex: str,
|
||||
provider_id: str,
|
||||
payload_json: str,
|
||||
*,
|
||||
d_tag: str,
|
||||
) -> dict[str, Any]:
|
||||
private_key = PrivateKey(bytes.fromhex(private_key_hex))
|
||||
tags = [
|
||||
["d", d_tag],
|
||||
["provider", provider_id],
|
||||
["schema", ANALYTICS_SCHEMA],
|
||||
]
|
||||
|
||||
event = Event(
|
||||
public_key=private_key.public_key.hex(),
|
||||
content=payload_json,
|
||||
kind=ANALYTICS_KIND,
|
||||
tags=tags,
|
||||
)
|
||||
private_key.sign_event(event)
|
||||
return _event_to_dict(event)
|
||||
|
||||
|
||||
def _fingerprint_payload(payload: dict[str, Any]) -> str:
|
||||
normalized = dict(payload)
|
||||
# Ignore generated timestamp for semantic dedupe.
|
||||
normalized.pop("generated_at", None)
|
||||
payload_json = json.dumps(normalized, separators=(",", ":"), sort_keys=True)
|
||||
return hashlib.sha256(payload_json.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
async def publish_usage_analytics() -> None:
|
||||
last_payload_hash: str | None = None
|
||||
|
||||
parsed_nsec: str | None = None
|
||||
private_key_hex: str | None = None
|
||||
public_key_hex: str | None = None
|
||||
provider_id: str | None = None
|
||||
warned_missing_nsec = False
|
||||
|
||||
logger.info("Usage analytics sharing task started")
|
||||
|
||||
while True:
|
||||
try:
|
||||
if not settings.enable_analytics_sharing:
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
nsec = (settings.nsec or "").strip()
|
||||
if not nsec:
|
||||
if not warned_missing_nsec:
|
||||
logger.info("NSEC is not configured; skipping analytics sharing to Nostr")
|
||||
warned_missing_nsec = True
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
warned_missing_nsec = False
|
||||
if nsec != parsed_nsec or private_key_hex is None or public_key_hex is None:
|
||||
keypair = nsec_to_keypair(nsec)
|
||||
if not keypair:
|
||||
logger.error("Invalid NSEC; analytics sharing is paused")
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
private_key_hex, public_key_hex = keypair
|
||||
parsed_nsec = nsec
|
||||
provider_id = _resolve_provider_id(public_key_hex)
|
||||
last_payload_hash = None
|
||||
|
||||
if private_key_hex is None or public_key_hex is None:
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
relay_urls = _resolve_relays()
|
||||
if not relay_urls:
|
||||
logger.warning("No Nostr relays configured; analytics sharing skipped")
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
resolved_provider_id = provider_id or _resolve_provider_id(public_key_hex)
|
||||
now_ts = int(time.time())
|
||||
payload = build_stats_snapshot_payload(
|
||||
resolved_provider_id,
|
||||
public_key_hex=public_key_hex,
|
||||
generated_at=now_ts,
|
||||
)
|
||||
|
||||
payload_hash = _fingerprint_payload(payload)
|
||||
if last_payload_hash == payload_hash:
|
||||
await asyncio.sleep(PUBLISH_INTERVAL_SECONDS)
|
||||
continue
|
||||
|
||||
payload_json = json.dumps(payload, separators=(",", ":"), sort_keys=True)
|
||||
d_tag = f"{resolved_provider_id}:stats"
|
||||
event = create_stats_snapshot_event(
|
||||
private_key_hex,
|
||||
resolved_provider_id,
|
||||
payload_json,
|
||||
d_tag=d_tag,
|
||||
)
|
||||
|
||||
success_count = 0
|
||||
for relay_url in relay_urls:
|
||||
if await publish_to_relay(relay_url, event):
|
||||
success_count += 1
|
||||
|
||||
if success_count > 0:
|
||||
last_payload_hash = payload_hash
|
||||
|
||||
logger.info(
|
||||
"Published analytics snapshot (success=%s/%s provider=%s)",
|
||||
success_count,
|
||||
len(relay_urls),
|
||||
resolved_provider_id,
|
||||
extra={
|
||||
"relay_success_count": success_count,
|
||||
"relay_total": len(relay_urls),
|
||||
"provider_id": resolved_provider_id,
|
||||
},
|
||||
)
|
||||
await asyncio.sleep(PUBLISH_INTERVAL_SECONDS)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Usage analytics sharing task cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Usage analytics sharing error",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
@@ -26,7 +26,7 @@ logger = get_logger(__name__)
|
||||
|
||||
def get_app_version() -> str | None:
|
||||
try:
|
||||
from ..core.main import __version__ as imported_version
|
||||
from ..core.version import __version__ as imported_version
|
||||
|
||||
return imported_version
|
||||
except Exception:
|
||||
|
||||
@@ -3,9 +3,17 @@ import math
|
||||
from pydantic.v1 import BaseModel
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.db import AsyncSession
|
||||
from ..core.settings import settings
|
||||
from .price import sats_usd_price
|
||||
from .usage import normalize_usage, parse_token_count
|
||||
|
||||
__all__ = [
|
||||
"CostData",
|
||||
"CostDataError",
|
||||
"MaxCostData",
|
||||
"calculate_cost",
|
||||
"parse_token_count",
|
||||
]
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -16,6 +24,12 @@ class CostData(BaseModel):
|
||||
output_msats: int
|
||||
total_msats: int
|
||||
total_usd: float = 0.0
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_msats: int = 0
|
||||
cache_creation_msats: int = 0
|
||||
|
||||
|
||||
class MaxCostData(CostData):
|
||||
@@ -27,11 +41,11 @@ class CostDataError(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
async def calculate_cost( # todo: can be sync
|
||||
response_data: dict, max_cost: int, session: AsyncSession
|
||||
async def calculate_cost(
|
||||
response_data: dict,
|
||||
max_cost: int,
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
"""
|
||||
Calculate the cost of an API request based on token usage.
|
||||
"""Calculate the cost of an API request based on token usage.
|
||||
|
||||
Args:
|
||||
response_data: Response data containing usage information
|
||||
@@ -39,6 +53,9 @@ async def calculate_cost( # todo: can be sync
|
||||
|
||||
Returns:
|
||||
Cost data or error information
|
||||
|
||||
The response's usage object is normalized with the default union parser;
|
||||
this function holds no vendor-dialect knowledge of its own.
|
||||
"""
|
||||
logger.debug(
|
||||
"Starting cost calculation",
|
||||
@@ -49,12 +66,21 @@ async def calculate_cost( # todo: can be sync
|
||||
},
|
||||
)
|
||||
|
||||
if "usage" not in response_data or response_data["usage"] is None:
|
||||
usage = normalize_usage(response_data.get("usage"))
|
||||
|
||||
if usage is None:
|
||||
logger.warning(
|
||||
"No usage data in response, using base cost only",
|
||||
"No usage data in response — billing at MaxCostData with zero "
|
||||
"tokens. Dashboard will show this request as `(0+0)`. Most "
|
||||
"common cause: upstream stream did not include a final usage "
|
||||
"chunk (OpenAI-compat backends require "
|
||||
"`stream_options.include_usage=true`).",
|
||||
extra={
|
||||
"max_cost_msats": max_cost,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"response_keys": sorted(response_data.keys())
|
||||
if isinstance(response_data, dict)
|
||||
else None,
|
||||
},
|
||||
)
|
||||
return MaxCostData(
|
||||
@@ -63,47 +89,58 @@ async def calculate_cost( # todo: can be sync
|
||||
output_msats=0,
|
||||
total_msats=0,
|
||||
total_usd=0.0,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
usage_data = response_data["usage"]
|
||||
usage_data = response_data.get("usage") or {}
|
||||
if not isinstance(usage_data, dict):
|
||||
usage_data = {}
|
||||
|
||||
usd_cost = 0.0
|
||||
|
||||
# Prioritize cost_details.upstream_inference_cost
|
||||
if "cost_details" in usage_data:
|
||||
usd_cost = float(
|
||||
usage_data["cost_details"].get("upstream_inference_cost", 0) or 0
|
||||
)
|
||||
|
||||
# Fallback to cost field if upstream_inference_cost is 0
|
||||
if usd_cost == 0 and "cost" in usage_data:
|
||||
try:
|
||||
usd_cost = float(usage_data.get("cost", 0) or 0)
|
||||
except Exception:
|
||||
pass
|
||||
input_tokens = usage.input_tokens
|
||||
output_tokens = usage.output_tokens
|
||||
cache_read_tokens = usage.cache_read_tokens
|
||||
cache_creation_tokens = usage.cache_write_tokens
|
||||
|
||||
# Try USD cost first
|
||||
usd_cost = _resolve_usd_cost(usage_data, response_data)
|
||||
if usd_cost > 0:
|
||||
try:
|
||||
sats_per_usd = 1.0 / sats_usd_price()
|
||||
cost_in_sats = usd_cost * sats_per_usd
|
||||
cost_in_msats = math.ceil(cost_in_sats * 1000)
|
||||
|
||||
logger.info(
|
||||
"Using cost from usage data/details",
|
||||
if input_tokens == 0 and output_tokens == 0:
|
||||
logger.warning(
|
||||
"Upstream reported a USD cost but no token counts — "
|
||||
"billing the USD-derived cost while the dashboard will "
|
||||
"show this request as `(0+0)` tokens. Check that the "
|
||||
"upstream actually emits `usage.input_tokens` and "
|
||||
"`usage.output_tokens` (OpenAI-compat streams require "
|
||||
"`stream_options.include_usage=true`).",
|
||||
extra={
|
||||
"usd_cost": usd_cost,
|
||||
"cost_in_sats": cost_in_sats,
|
||||
"cost_in_msats": cost_in_msats,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"usd_cost": usd_cost,
|
||||
"usage_keys": sorted(usage_data.keys())
|
||||
if isinstance(usage_data, dict)
|
||||
else None,
|
||||
},
|
||||
)
|
||||
|
||||
return CostData(
|
||||
base_msats=-1,
|
||||
input_msats=-1, # Cost field doesn't break down by token type
|
||||
output_msats=-1,
|
||||
total_msats=cost_in_msats,
|
||||
total_usd=usd_cost,
|
||||
try:
|
||||
input_usd = _coerce_usd(
|
||||
usage_data.get("cost_details", {}).get("input_cost", 0)
|
||||
)
|
||||
output_usd = _coerce_usd(
|
||||
usage_data.get("cost_details", {}).get("output_cost", 0)
|
||||
)
|
||||
return _calculate_from_usd_cost(
|
||||
usd_cost,
|
||||
input_usd,
|
||||
output_usd,
|
||||
input_tokens,
|
||||
cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
output_tokens,
|
||||
response_data,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
@@ -114,69 +151,35 @@ async def calculate_cost( # todo: can be sync
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
# Fall through to token-based calculation
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS: float = (
|
||||
float(settings.fixed_per_1k_input_tokens) * 1000.0
|
||||
)
|
||||
MSATS_PER_1K_OUTPUT_TOKENS: float = (
|
||||
float(settings.fixed_per_1k_output_tokens) * 1000.0
|
||||
)
|
||||
# Fall back to token-based pricing
|
||||
try:
|
||||
pricing_rates = _get_pricing_rates(response_data)
|
||||
except ValueError as e:
|
||||
return CostDataError(message=str(e), code="pricing_error")
|
||||
|
||||
if not settings.fixed_pricing:
|
||||
response_model = response_data.get("model", "")
|
||||
logger.debug(
|
||||
"Using model-based pricing",
|
||||
extra={"model": response_model},
|
||||
)
|
||||
if pricing_rates is None:
|
||||
input_rate = float(settings.fixed_per_1k_input_tokens) * 1000.0
|
||||
output_rate = float(settings.fixed_per_1k_output_tokens) * 1000.0
|
||||
cache_read_rate = input_rate
|
||||
cache_creation_rate = input_rate
|
||||
else:
|
||||
input_rate, output_rate, cache_read_rate, cache_creation_rate = pricing_rates
|
||||
|
||||
from ..proxy import get_model_instance
|
||||
|
||||
model_obj = get_model_instance(response_model)
|
||||
|
||||
if not model_obj:
|
||||
logger.error(
|
||||
"Invalid model in response",
|
||||
extra={"response_model": response_model},
|
||||
)
|
||||
return CostDataError(
|
||||
message=f"Invalid model in response: {response_model}",
|
||||
code="model_not_found",
|
||||
)
|
||||
|
||||
if not model_obj.sats_pricing:
|
||||
logger.error(
|
||||
"Model pricing not defined",
|
||||
extra={"model": response_model, "model_id": response_model},
|
||||
)
|
||||
return CostDataError(
|
||||
message="Model pricing not defined", code="pricing_not_found"
|
||||
)
|
||||
|
||||
try:
|
||||
mspp = float(model_obj.sats_pricing.prompt)
|
||||
mspc = float(model_obj.sats_pricing.completion)
|
||||
except Exception:
|
||||
return CostDataError(message="Invalid pricing data", code="pricing_invalid")
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS = mspp * 1_000_000.0
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = mspc * 1_000_000.0
|
||||
|
||||
logger.info(
|
||||
"Applied model-specific pricing",
|
||||
extra={
|
||||
"model": response_model,
|
||||
"input_price_msats_per_1k": MSATS_PER_1K_INPUT_TOKENS,
|
||||
"output_price_msats_per_1k": MSATS_PER_1K_OUTPUT_TOKENS,
|
||||
},
|
||||
)
|
||||
|
||||
if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS):
|
||||
if not (input_rate and output_rate):
|
||||
logger.warning(
|
||||
"No token pricing configured, using base cost",
|
||||
"No token pricing configured — billing at flat MaxCostData. "
|
||||
"Token counts %s in the upstream response but cannot be "
|
||||
"priced; the request will appear in dashboards with the "
|
||||
"raw counts and a fixed max-cost charge.",
|
||||
"are present"
|
||||
if (input_tokens > 0 or output_tokens > 0)
|
||||
else "are zero",
|
||||
extra={
|
||||
"base_cost_msats": max_cost,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
},
|
||||
)
|
||||
return MaxCostData(
|
||||
@@ -184,35 +187,218 @@ async def calculate_cost( # todo: can be sync
|
||||
input_msats=0,
|
||||
output_msats=0,
|
||||
total_msats=max_cost,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
input_tokens = usage_data.get("prompt_tokens", 0)
|
||||
output_tokens = usage_data.get("completion_tokens", 0)
|
||||
|
||||
# added for response api
|
||||
input_tokens = (
|
||||
input_tokens if input_tokens != 0 else usage_data.get("input_tokens", 0)
|
||||
)
|
||||
output_tokens = (
|
||||
output_tokens if output_tokens != 0 else usage_data.get("output_tokens", 0)
|
||||
return _calculate_from_tokens(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
input_rate,
|
||||
output_rate,
|
||||
cache_read_rate,
|
||||
cache_creation_rate,
|
||||
response_data,
|
||||
)
|
||||
|
||||
# added for response api
|
||||
input_tokens = (
|
||||
input_tokens
|
||||
if input_tokens != 0
|
||||
else response_data.get("usage", {}).get("input_tokens", 0)
|
||||
)
|
||||
output_tokens = (
|
||||
output_tokens
|
||||
if output_tokens != 0
|
||||
else response_data.get("usage", {}).get("output_tokens", 0)
|
||||
|
||||
# ============================================================================
|
||||
# Helper Functions (ordered by call sequence in calculate_cost)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _coerce_usd(value: object) -> float:
|
||||
"""Coerce a value to USD float, handling various formats safely."""
|
||||
if value is None or isinstance(value, bool):
|
||||
return 0.0
|
||||
if not isinstance(value, (int, float, str)):
|
||||
return 0.0
|
||||
try:
|
||||
return max(0.0, float(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float:
|
||||
"""Resolve USD cost with clear priority order.
|
||||
|
||||
Priority: cost_details.total_cost → total_cost → cost (in both usage and response).
|
||||
"""
|
||||
cost_details = usage_data.get("cost_details")
|
||||
if isinstance(cost_details, dict):
|
||||
cost = _coerce_usd(cost_details.get("total_cost"))
|
||||
if cost > 0:
|
||||
return cost
|
||||
|
||||
for source in [usage_data, response_data]:
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
for field in ("total_cost", "cost"):
|
||||
cost = _coerce_usd(source.get(field))
|
||||
if cost > 0:
|
||||
return cost
|
||||
|
||||
return 0.0
|
||||
|
||||
|
||||
def _get_pricing_rates(
|
||||
response_data: dict,
|
||||
) -> tuple[float, float, float, float] | None:
|
||||
"""Get model-based pricing rates or None if using fixed pricing.
|
||||
|
||||
Returns: (input_rate, output_rate, cache_read_rate, cache_write_rate)
|
||||
"""
|
||||
if settings.fixed_pricing:
|
||||
return None
|
||||
|
||||
from ..proxy import get_model_instance
|
||||
|
||||
response_model = response_data.get("model", "")
|
||||
model_obj = get_model_instance(response_model)
|
||||
|
||||
if not model_obj:
|
||||
logger.error("Invalid model in response", extra={"response_model": response_model})
|
||||
raise ValueError(f"Invalid model: {response_model}")
|
||||
|
||||
if not model_obj.sats_pricing:
|
||||
logger.error(
|
||||
"Model pricing not defined",
|
||||
extra={"model": response_model, "model_id": response_model},
|
||||
)
|
||||
raise ValueError("Model pricing not defined")
|
||||
|
||||
try:
|
||||
mspp = float(model_obj.sats_pricing.prompt)
|
||||
mspc = float(model_obj.sats_pricing.completion)
|
||||
mscr = float(model_obj.sats_pricing.input_cache_read or 0)
|
||||
mscw = float(model_obj.sats_pricing.input_cache_write or 0)
|
||||
|
||||
mspp_1k = mspp * 1_000_000.0
|
||||
mspc_1k = mspc * 1_000_000.0
|
||||
mscr_1k = mscr * 1_000_000.0 if mscr > 0 else mspp_1k
|
||||
mscw_1k = mscw * 1_000_000.0 if mscw > 0 else mspp_1k
|
||||
|
||||
logger.info(
|
||||
"Applied model-specific pricing",
|
||||
extra={
|
||||
"model": response_model,
|
||||
"input_price_msats_per_1k": mspp_1k,
|
||||
"output_price_msats_per_1k": mspc_1k,
|
||||
"cache_read_price_msats_per_1k": mscr_1k,
|
||||
"cache_write_price_msats_per_1k": mscw_1k,
|
||||
},
|
||||
)
|
||||
return mspp_1k, mspc_1k, mscr_1k, mscw_1k
|
||||
except Exception as e:
|
||||
logger.error("Invalid pricing data", extra={"error": str(e)})
|
||||
raise ValueError("Invalid pricing data") from e
|
||||
|
||||
|
||||
def _resolve_provider_fee(model_id: str) -> float:
|
||||
"""Resolve the provider fee multiplier for the given model id.
|
||||
|
||||
Falls back to 1.0 (no markup) when the provider cannot be resolved so
|
||||
the USD cost path never silently double-applies or omits the fee.
|
||||
"""
|
||||
from ..proxy import get_provider_for_model
|
||||
|
||||
if not model_id:
|
||||
return 1.0
|
||||
providers = get_provider_for_model(model_id)
|
||||
if not providers:
|
||||
return 1.0
|
||||
return float(providers[0].provider_fee)
|
||||
|
||||
|
||||
def _calculate_from_usd_cost(
|
||||
usd_cost: float,
|
||||
input_usd: float,
|
||||
output_usd: float,
|
||||
input_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_creation_tokens: int,
|
||||
output_tokens: int,
|
||||
response_data: dict,
|
||||
) -> CostData:
|
||||
"""Calculate cost from USD figures, deriving input/output split from tokens."""
|
||||
provider_fee = _resolve_provider_fee(response_data.get("model", ""))
|
||||
usd_cost = usd_cost * provider_fee
|
||||
input_usd = input_usd * provider_fee
|
||||
output_usd = output_usd * provider_fee
|
||||
sats_per_usd = 1.0 / sats_usd_price()
|
||||
cost_in_sats = usd_cost * sats_per_usd
|
||||
cost_in_msats = math.ceil(cost_in_sats * 1000)
|
||||
|
||||
if input_usd > 0 or output_usd > 0:
|
||||
input_msats = int((input_usd * sats_per_usd) * 1000)
|
||||
output_msats = int((output_usd * sats_per_usd) * 1000)
|
||||
else:
|
||||
effective_input_tokens = (
|
||||
input_tokens + cache_read_tokens + cache_creation_tokens
|
||||
)
|
||||
total_tokens = effective_input_tokens + output_tokens
|
||||
input_msats = (
|
||||
int(cost_in_msats * effective_input_tokens / total_tokens)
|
||||
if total_tokens > 0
|
||||
else 0
|
||||
)
|
||||
output_msats = cost_in_msats - input_msats
|
||||
|
||||
logger.info(
|
||||
"Using cost from usage data/details",
|
||||
extra={
|
||||
"usd_cost": usd_cost,
|
||||
"cost_in_sats": cost_in_sats,
|
||||
"cost_in_msats": cost_in_msats,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
|
||||
input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3)
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=input_msats,
|
||||
output_msats=output_msats,
|
||||
total_msats=cost_in_msats,
|
||||
total_usd=usd_cost,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3)
|
||||
token_based_cost = math.ceil(input_msats + output_msats)
|
||||
|
||||
def _calculate_from_tokens(
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_creation_tokens: int,
|
||||
input_rate: float,
|
||||
output_rate: float,
|
||||
cache_read_rate: float,
|
||||
cache_creation_rate: float,
|
||||
response_data: dict,
|
||||
) -> CostData:
|
||||
"""Calculate cost from token counts using pricing rates."""
|
||||
calc_input_msats = round(input_tokens / 1000 * input_rate, 3)
|
||||
calc_output_msats = round(output_tokens / 1000 * output_rate, 3)
|
||||
calc_cache_read_msats = round(cache_read_tokens / 1000 * cache_read_rate, 3)
|
||||
calc_cache_write_msats = round(
|
||||
cache_creation_tokens / 1000 * cache_creation_rate, 3
|
||||
)
|
||||
token_based_cost = math.ceil(
|
||||
calc_input_msats
|
||||
+ calc_output_msats
|
||||
+ calc_cache_read_msats
|
||||
+ calc_cache_write_msats
|
||||
)
|
||||
total_usd = (token_based_cost / 1000.0) * sats_usd_price()
|
||||
|
||||
logger.info(
|
||||
@@ -220,18 +406,30 @@ async def calculate_cost( # todo: can be sync
|
||||
extra={
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"input_cost_msats": input_msats,
|
||||
"output_cost_msats": output_msats,
|
||||
"cache_read_input_tokens": cache_read_tokens,
|
||||
"cache_creation_input_tokens": cache_creation_tokens,
|
||||
"input_cost_msats": calc_input_msats,
|
||||
"output_cost_msats": calc_output_msats,
|
||||
"cache_read_cost_msats": calc_cache_read_msats,
|
||||
"cache_creation_cost_msats": calc_cache_write_msats,
|
||||
"total_cost_msats": token_based_cost,
|
||||
"total_usd": total_usd,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
|
||||
visible_input_msats = int(calc_input_msats + calc_cache_read_msats)
|
||||
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=int(input_msats),
|
||||
output_msats=int(output_msats),
|
||||
input_msats=visible_input_msats,
|
||||
output_msats=int(calc_output_msats),
|
||||
total_msats=token_based_cost,
|
||||
total_usd=total_usd,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=int(calc_cache_read_msats),
|
||||
cache_creation_msats=int(calc_cache_write_msats),
|
||||
)
|
||||
|
||||
@@ -29,15 +29,13 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
|
||||
},
|
||||
)
|
||||
elif auth := headers.get("authorization", None):
|
||||
cashu_token = auth.split(" ")[1] if len(auth.split(" ")) > 1 else ""
|
||||
logger.debug(
|
||||
"Using Authorization header token",
|
||||
"Skipping preflight token balance check for Authorization header",
|
||||
extra={
|
||||
"token_preview": cashu_token[:20] + "..."
|
||||
if len(cashu_token) > 20
|
||||
else cashu_token
|
||||
"auth_preview": auth[:20] + "..." if len(auth) > 20 else auth,
|
||||
},
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.error("No authentication token provided")
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
@@ -75,7 +73,7 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
|
||||
|
||||
if max_cost_for_model > amount_msat:
|
||||
raise HTTPException(
|
||||
status_code=413,
|
||||
status_code=402,
|
||||
detail={
|
||||
"reason": "Insufficient balance",
|
||||
"amount_required_msat": max_cost_for_model,
|
||||
|
||||
@@ -3,11 +3,12 @@ import json
|
||||
import random
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel as V2BaseModel
|
||||
from pydantic.v1 import BaseModel
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from ..core.db import ModelRow, get_session
|
||||
from ..core.db import ModelRow, UpstreamProviderRow, get_session
|
||||
from ..core.logging import get_logger
|
||||
from ..core.settings import settings
|
||||
from .price import sats_usd_price
|
||||
@@ -16,6 +17,24 @@ logger = get_logger(__name__)
|
||||
|
||||
models_router = APIRouter()
|
||||
|
||||
_MODEL_TEST_ENDPOINT_PATHS = {
|
||||
"chat-completions": "chat/completions",
|
||||
"completions": "completions",
|
||||
"embeddings": "embeddings",
|
||||
"responses": "responses",
|
||||
}
|
||||
|
||||
# Cap the caller-supplied test payload to avoid forwarding oversized bodies
|
||||
# upstream on the operator's credentials.
|
||||
_MODEL_TEST_MAX_REQUEST_BYTES = 64 * 1024
|
||||
|
||||
|
||||
async def _require_admin_api(request: Request) -> None:
|
||||
"""Require admin auth without creating an import-time cycle with core.admin."""
|
||||
from ..core.admin import require_admin_api
|
||||
|
||||
await require_admin_api(request)
|
||||
|
||||
|
||||
class Architecture(BaseModel):
|
||||
modality: str
|
||||
@@ -60,11 +79,54 @@ class Model(BaseModel):
|
||||
upstream_provider_id: int | str | None = None
|
||||
canonical_slug: str | None = None
|
||||
alias_ids: list[str] | None = None
|
||||
forwarded_model_id: str | None = None
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(self.id)
|
||||
|
||||
|
||||
def backfill_cache_pricing(model_id: str, pricing: Pricing) -> Pricing:
|
||||
"""Fill missing cache rates from litellm's bundled cost map.
|
||||
|
||||
The OpenRouter model feed omits ``input_cache_read``/``input_cache_write``
|
||||
for many models (most DeepSeek entries, openai/gpt-4o, ...). Without a
|
||||
cache rate, billing falls back to the full input rate, which overcharges
|
||||
cache reads (DeepSeek hits are 10x cheaper) and undercharges Anthropic
|
||||
cache writes (1.25x). litellm ships per-model USD rates keyed by the exact
|
||||
OpenRouter id (deepseek/deepseek-chat) or by the bare model name
|
||||
(gpt-4o, claude-sonnet-4-5), so both spellings are tried.
|
||||
|
||||
Rates already present (e.g. provided by OpenRouter) are authoritative and
|
||||
never overwritten. Unknown models are returned unchanged.
|
||||
"""
|
||||
needs_read = (pricing.input_cache_read or 0.0) <= 0.0
|
||||
needs_write = (pricing.input_cache_write or 0.0) <= 0.0
|
||||
if not (needs_read or needs_write):
|
||||
return pricing
|
||||
|
||||
import litellm
|
||||
|
||||
info: dict | None = None
|
||||
for key in (model_id, model_id.split("/", 1)[-1]):
|
||||
candidate = litellm.model_cost.get(key)
|
||||
if isinstance(candidate, dict):
|
||||
info = candidate
|
||||
break
|
||||
if info is None:
|
||||
return pricing
|
||||
|
||||
updated = Pricing.parse_obj(pricing.dict())
|
||||
if needs_read:
|
||||
read_rate = info.get("cache_read_input_token_cost")
|
||||
if isinstance(read_rate, (int, float)) and read_rate > 0:
|
||||
updated.input_cache_read = float(read_rate)
|
||||
if needs_write:
|
||||
write_rate = info.get("cache_creation_input_token_cost")
|
||||
if isinstance(write_rate, (int, float)) and write_rate > 0:
|
||||
updated.input_cache_write = float(write_rate)
|
||||
return updated
|
||||
|
||||
|
||||
def _has_valid_pricing(model: dict) -> bool:
|
||||
"""Check if model has valid pricing (not free, no negative values)."""
|
||||
pricing = model.get("pricing", {})
|
||||
@@ -177,6 +239,7 @@ def _row_to_model(
|
||||
upstream_provider_id=row.upstream_provider_id,
|
||||
canonical_slug=getattr(row, "canonical_slug", None),
|
||||
alias_ids=json.loads(row.alias_ids) if row.alias_ids else None,
|
||||
forwarded_model_id=getattr(row, "forwarded_model_id", None) or row.id,
|
||||
)
|
||||
|
||||
if apply_provider_fee:
|
||||
@@ -329,6 +392,7 @@ def _update_model_sats_pricing(model: Model, sats_to_usd: float) -> Model:
|
||||
upstream_provider_id=model.upstream_provider_id,
|
||||
canonical_slug=model.canonical_slug,
|
||||
alias_ids=model.alias_ids,
|
||||
forwarded_model_id=model.forwarded_model_id,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -347,6 +411,9 @@ async def _update_sats_pricing_once() -> None:
|
||||
from ..proxy import get_upstreams, refresh_model_maps
|
||||
|
||||
upstreams = get_upstreams()
|
||||
if not upstreams:
|
||||
return
|
||||
|
||||
sats_to_usd = sats_usd_price()
|
||||
|
||||
updated_count = 0
|
||||
@@ -356,11 +423,14 @@ async def _update_sats_pricing_once() -> None:
|
||||
for m in upstream.get_cached_models()
|
||||
]
|
||||
upstream._models_cache = updated_models
|
||||
upstream._models_by_id = {m.id: m for m in updated_models}
|
||||
upstream._models_by_id = {m.forwarded_model_id or m.id: m for m in updated_models}
|
||||
updated_count += len(updated_models)
|
||||
|
||||
if updated_count > 0:
|
||||
logger.info("Updated sats pricing", extra={"models_updated": updated_count})
|
||||
logger.info(
|
||||
f"Updated sats pricing for {updated_count} models",
|
||||
extra={"models_updated": updated_count},
|
||||
)
|
||||
await refresh_model_maps()
|
||||
|
||||
|
||||
@@ -402,11 +472,110 @@ async def update_sats_pricing() -> None:
|
||||
logger.error(f"Error updating sats pricing: {e}")
|
||||
|
||||
|
||||
class ModelTestRequest(V2BaseModel):
|
||||
model_id: str
|
||||
endpoint_type: str
|
||||
request_data: dict
|
||||
|
||||
|
||||
@models_router.post(
|
||||
"/api/models/test", dependencies=[Depends(_require_admin_api)]
|
||||
)
|
||||
async def test_model(
|
||||
payload: ModelTestRequest,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Test a model by sending a request through its configured upstream provider."""
|
||||
from sqlmodel import select
|
||||
|
||||
result = await session.execute(
|
||||
select(ModelRow).where(ModelRow.id == payload.model_id)
|
||||
)
|
||||
model_row = result.scalars().first()
|
||||
|
||||
if not model_row:
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Model '{payload.model_id}' not found in database",
|
||||
"status_code": 404,
|
||||
}
|
||||
|
||||
provider = await session.get(UpstreamProviderRow, model_row.upstream_provider_id)
|
||||
if not provider:
|
||||
return {
|
||||
"success": False,
|
||||
"error": "Upstream provider not found",
|
||||
"status_code": 404,
|
||||
}
|
||||
|
||||
endpoint_path = _MODEL_TEST_ENDPOINT_PATHS.get(payload.endpoint_type)
|
||||
if endpoint_path is None:
|
||||
raise HTTPException(status_code=400, detail="Unsupported endpoint_type")
|
||||
|
||||
actual_model_id = model_row.forwarded_model_id or model_row.id
|
||||
request_data = dict(payload.request_data)
|
||||
request_data["model"] = actual_model_id
|
||||
|
||||
try:
|
||||
request_size = len(json.dumps(request_data).encode("utf-8"))
|
||||
except (TypeError, ValueError):
|
||||
raise HTTPException(status_code=400, detail="Invalid request_data")
|
||||
if request_size > _MODEL_TEST_MAX_REQUEST_BYTES:
|
||||
raise HTTPException(status_code=413, detail="request_data too large")
|
||||
|
||||
base_url = provider.base_url.rstrip("/")
|
||||
url = f"{base_url}/{endpoint_path}"
|
||||
|
||||
logger.info(
|
||||
"admin model test",
|
||||
extra={
|
||||
"model_id": payload.model_id,
|
||||
"forwarded_model_id": actual_model_id,
|
||||
"endpoint_type": payload.endpoint_type,
|
||||
"upstream_provider_id": model_row.upstream_provider_id,
|
||||
"request_bytes": request_size,
|
||||
},
|
||||
)
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {provider.api_key}",
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, json=request_data, headers=headers)
|
||||
try:
|
||||
response_data = response.json()
|
||||
except Exception:
|
||||
response_data = {"raw": response.text}
|
||||
|
||||
return {
|
||||
"success": response.status_code < 400,
|
||||
"data": response_data,
|
||||
"status_code": response.status_code,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"status_code": 500,
|
||||
}
|
||||
|
||||
|
||||
@models_router.get("/v1/models")
|
||||
@models_router.get("/models", include_in_schema=False)
|
||||
@models_router.get("/v1/models/", include_in_schema=False)
|
||||
@models_router.get("/models")
|
||||
@models_router.get("/models/", include_in_schema=False)
|
||||
async def models(session: AsyncSession = Depends(get_session)) -> dict:
|
||||
"""Get all available models from all providers with database overrides applied."""
|
||||
from ..proxy import get_unique_models
|
||||
|
||||
items = get_unique_models()
|
||||
return {"data": items}
|
||||
data = []
|
||||
for model in items:
|
||||
m = model.dict()
|
||||
if model.forwarded_model_id:
|
||||
m["id"] = model.forwarded_model_id
|
||||
data.append(m)
|
||||
return {"data": data}
|
||||
|
||||
131
routstr/payment/usage.py
Normal file
131
routstr/payment/usage.py
Normal file
@@ -0,0 +1,131 @@
|
||||
"""Vendor-agnostic normalization of upstream usage objects.
|
||||
|
||||
Upstream providers report token usage in vendor dialects that differ in field
|
||||
names and in whether cached tokens are included in the input count:
|
||||
|
||||
* OpenAI / Azure / xAI / Groq / Moonshot / Qwen / Gemini-compat: cache reads in
|
||||
``prompt_tokens_details.cached_tokens``, included in ``prompt_tokens``.
|
||||
* OpenRouter: same as OpenAI plus cache *writes* in
|
||||
``prompt_tokens_details.cache_write_tokens``, also included in
|
||||
``prompt_tokens``.
|
||||
* litellm-normalized: same nesting, but names the write field
|
||||
``prompt_tokens_details.cache_creation_tokens`` (and additionally mirrors the
|
||||
Anthropic top-level fields), with ``prompt_tokens`` as the grand total.
|
||||
* Anthropic native: ``cache_read_input_tokens`` / ``cache_creation_input_tokens``
|
||||
top-level, additive to (not included in) ``input_tokens``.
|
||||
* DeepSeek: ``prompt_cache_hit_tokens`` / ``prompt_cache_miss_tokens``, with
|
||||
``prompt_tokens = hit + miss``.
|
||||
|
||||
What decides whether cached tokens must be subtracted out of the input count is
|
||||
**which prompt field the vendor uses**, not which cache field appears:
|
||||
|
||||
* ``prompt_tokens`` present -> cached + cache-write tokens are *included* in it
|
||||
(OpenAI family, DeepSeek, OpenRouter, litellm); subtract both so
|
||||
``input_tokens`` holds only the regular-rate portion.
|
||||
* only ``input_tokens`` (Anthropic native) -> cached tokens are *additive*;
|
||||
leave ``input_tokens`` untouched.
|
||||
|
||||
``normalize_usage`` maps all of them onto one canonical ``NormalizedUsage``
|
||||
shape so billing code needs no vendor knowledge. The known dialects' field
|
||||
names do not collide, so a single union parser is safe; a vendor whose fields
|
||||
would genuinely conflict needs a dedicated branch here.
|
||||
"""
|
||||
|
||||
from pydantic.v1 import BaseModel
|
||||
|
||||
|
||||
class NormalizedUsage(BaseModel):
|
||||
"""Canonical token usage: input_tokens never includes cached tokens."""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
cache_write_tokens: int = 0
|
||||
|
||||
|
||||
def parse_token_count(value: object) -> int:
|
||||
"""Parse a token count from various formats (int, float, str, bool)."""
|
||||
if isinstance(value, bool):
|
||||
return 0
|
||||
if isinstance(value, int):
|
||||
return max(0, value)
|
||||
if isinstance(value, float):
|
||||
return max(0, int(value))
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return max(0, int(float(value)))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _first_token_count(usage_data: dict, *fields: str) -> int:
|
||||
"""Return the first positive token count among the given fields."""
|
||||
for field in fields:
|
||||
value = parse_token_count(usage_data.get(field, 0))
|
||||
if value > 0:
|
||||
return value
|
||||
return 0
|
||||
|
||||
|
||||
def _extract_cache_tokens(usage_data: dict) -> tuple[int, int]:
|
||||
"""Pull (cache_read, cache_write) across all known dialects.
|
||||
|
||||
Precedence (highest first), independent for reads and writes:
|
||||
|
||||
* Anthropic top-level: ``cache_read_input_tokens`` /
|
||||
``cache_creation_input_tokens``.
|
||||
* Nested ``prompt_tokens_details``: ``cached_tokens`` for reads;
|
||||
``cache_creation_tokens`` (litellm) or ``cache_write_tokens``
|
||||
(OpenRouter) for writes.
|
||||
* DeepSeek: ``prompt_cache_hit_tokens`` for reads (no write concept).
|
||||
"""
|
||||
cache_read = parse_token_count(usage_data.get("cache_read_input_tokens", 0))
|
||||
cache_write = parse_token_count(usage_data.get("cache_creation_input_tokens", 0))
|
||||
|
||||
prompt_details = usage_data.get("prompt_tokens_details")
|
||||
if isinstance(prompt_details, dict):
|
||||
if not cache_read:
|
||||
cache_read = parse_token_count(prompt_details.get("cached_tokens", 0))
|
||||
if not cache_write:
|
||||
cache_write = _first_token_count(
|
||||
prompt_details, "cache_creation_tokens", "cache_write_tokens"
|
||||
)
|
||||
|
||||
if not cache_read:
|
||||
# DeepSeek: prompt_tokens = prompt_cache_hit_tokens + prompt_cache_miss_tokens
|
||||
cache_read = parse_token_count(usage_data.get("prompt_cache_hit_tokens", 0))
|
||||
|
||||
return cache_read, cache_write
|
||||
|
||||
|
||||
def normalize_usage(usage_data: object) -> NormalizedUsage | None:
|
||||
"""Map a vendor usage dict onto the canonical shape, or None if absent.
|
||||
|
||||
Cached reads and writes are subtracted from the input count exactly once,
|
||||
only for dialects that report a ``prompt_tokens`` grand total that already
|
||||
includes them (OpenAI family, DeepSeek, OpenRouter, litellm). Anthropic
|
||||
native reports them additively under ``input_tokens`` and is left untouched.
|
||||
"""
|
||||
if not isinstance(usage_data, dict):
|
||||
return None
|
||||
|
||||
output_tokens = _first_token_count(
|
||||
usage_data, "completion_tokens", "output_tokens"
|
||||
)
|
||||
cache_read, cache_write = _extract_cache_tokens(usage_data)
|
||||
|
||||
# ``prompt_tokens`` is the inclusive grand total; ``input_tokens`` (Anthropic
|
||||
# native) excludes cached tokens. The field chosen decides whether to subtract.
|
||||
if "prompt_tokens" in usage_data:
|
||||
input_tokens = parse_token_count(usage_data.get("prompt_tokens", 0))
|
||||
input_tokens = max(0, input_tokens - cache_read - cache_write)
|
||||
else:
|
||||
input_tokens = parse_token_count(usage_data.get("input_tokens", 0))
|
||||
|
||||
return NormalizedUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read,
|
||||
cache_write_tokens=cache_write,
|
||||
)
|
||||
269
routstr/proxy.py
269
routstr/proxy.py
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
@@ -17,6 +18,7 @@ from .core.db import (
|
||||
get_session,
|
||||
)
|
||||
from .core.exceptions import UpstreamError
|
||||
from .core.not_found import build_not_found_response
|
||||
from .core.settings import settings
|
||||
from .payment.helpers import (
|
||||
calculate_discounted_max_cost,
|
||||
@@ -27,6 +29,7 @@ from .payment.helpers import (
|
||||
from .payment.models import Model
|
||||
from .upstream import BaseUpstreamProvider
|
||||
from .upstream.helpers import init_upstreams
|
||||
from .upstream.request_correction import correct_request, extract_error_message
|
||||
|
||||
logger = get_logger(__name__)
|
||||
proxy_router = APIRouter()
|
||||
@@ -69,7 +72,25 @@ def get_upstreams() -> list[BaseUpstreamProvider]:
|
||||
|
||||
def get_model_instance(model_id: str) -> Model | None:
|
||||
"""Get Model instance by ID from global cache."""
|
||||
return _model_instances.get(model_id.lower())
|
||||
if not model_id:
|
||||
return None
|
||||
|
||||
model_id_lower = model_id.lower()
|
||||
# Try exact match first
|
||||
if model := _model_instances.get(model_id_lower):
|
||||
return model
|
||||
|
||||
# Try stripping common version suffixes (e.g., -20251222)
|
||||
# This handles cases where upstream returns a specific version
|
||||
# but we only track the base model name.
|
||||
import re
|
||||
|
||||
base_model_id = re.sub(r"-\d{8}$", "", model_id_lower)
|
||||
if base_model_id != model_id_lower:
|
||||
if model := _model_instances.get(base_model_id):
|
||||
return model
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None:
|
||||
@@ -132,27 +153,79 @@ async def refresh_model_maps_periodically() -> None:
|
||||
)
|
||||
|
||||
|
||||
_API_PATH_PREFIXES = (
|
||||
"v1/",
|
||||
"responses",
|
||||
"chat/",
|
||||
"completions",
|
||||
"models",
|
||||
"embeddings",
|
||||
"audio/",
|
||||
"images/",
|
||||
"moderations",
|
||||
"providers",
|
||||
"tee/",
|
||||
)
|
||||
|
||||
|
||||
@proxy_router.api_route("/{path:path}", methods=["GET", "POST"], response_model=None)
|
||||
async def proxy(
|
||||
request: Request, path: str, session: AsyncSession = Depends(get_session)
|
||||
) -> Response | StreamingResponse:
|
||||
headers = dict(request.headers)
|
||||
# GET requests must hit a known API prefix; otherwise return a 404 (HTML
|
||||
# for browsers, JSON for API clients). POST requests are always forwarded
|
||||
# so that OpenAI-style endpoints work with or without the `v1/` prefix
|
||||
# (e.g. `/chat/completions` as well as `/v1/chat/completions`).
|
||||
if request.method == "GET" and not path.startswith(_API_PATH_PREFIXES):
|
||||
return build_not_found_response(request, path)
|
||||
|
||||
if "x-cashu" not in headers and "authorization" not in headers.keys():
|
||||
return create_error_response(
|
||||
"unauthorized", "Unauthorized", 401, request=request
|
||||
)
|
||||
headers = dict(request.headers)
|
||||
|
||||
is_responses_api = path.startswith("v1/responses") or path.startswith("responses")
|
||||
request_body = await request.body()
|
||||
request_body_dict = parse_request_body_json(request_body, path)
|
||||
|
||||
# /tee/* GET requests (e.g. attestation) don't map to models — just
|
||||
# forward to all enabled upstreams without model/cost/auth lookups.
|
||||
if request.method == "GET" and path.startswith("tee/"):
|
||||
all_upstreams = _upstreams
|
||||
last_error_response = None
|
||||
for i, upstream in enumerate(all_upstreams):
|
||||
try:
|
||||
headers = upstream.prepare_headers(dict(request.headers))
|
||||
response = await upstream.forward_get_request(request, path, headers)
|
||||
if response.status_code in [502, 429] and i < len(all_upstreams) - 1:
|
||||
logger.warning(
|
||||
"Upstream %s returned %s for tee GET %s, trying next",
|
||||
upstream.provider_type,
|
||||
response.status_code,
|
||||
path,
|
||||
)
|
||||
continue
|
||||
return response
|
||||
except UpstreamError as e:
|
||||
logger.warning(
|
||||
"Upstream %s failed for tee GET %s: %s",
|
||||
upstream.provider_type,
|
||||
path,
|
||||
e,
|
||||
)
|
||||
if i == len(all_upstreams) - 1:
|
||||
last_error_response = create_error_response(
|
||||
"upstream_error", str(e), 502, request=request
|
||||
)
|
||||
continue
|
||||
return last_error_response or create_error_response(
|
||||
"upstream_error", "All upstreams failed", 502, request=request
|
||||
)
|
||||
|
||||
if is_responses_api:
|
||||
model_id = extract_model_from_responses_request(request_body_dict)
|
||||
else:
|
||||
model_id = request_body_dict.get("model", "unknown")
|
||||
|
||||
model_obj = get_model_instance(model_id)
|
||||
|
||||
if not model_obj:
|
||||
return create_error_response(
|
||||
"invalid_model", f"Model '{model_id}' not found", 400, request=request
|
||||
@@ -196,7 +269,15 @@ async def proxy(
|
||||
)
|
||||
except UpstreamError as e:
|
||||
logger.warning(
|
||||
f"Upstream {upstream.provider_type} failed (x-cashu): {e}"
|
||||
"Upstream %s failed (x-cashu) for model=%s: %s",
|
||||
upstream.provider_type,
|
||||
model_id,
|
||||
e,
|
||||
extra={
|
||||
"provider": upstream.provider_type,
|
||||
"model": model_id,
|
||||
"status_code": e.status_code,
|
||||
},
|
||||
)
|
||||
if i == len(upstreams) - 1:
|
||||
last_error = e
|
||||
@@ -211,7 +292,7 @@ async def proxy(
|
||||
|
||||
elif auth := headers.get("authorization", None):
|
||||
key = await get_bearer_token_key(
|
||||
headers, path, session, auth, max_cost_for_model
|
||||
headers, path, session, auth, max_cost_for_model, model_id
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -273,46 +354,86 @@ async def proxy(
|
||||
if request_body_dict:
|
||||
await pay_for_request(key, max_cost_for_model, session)
|
||||
|
||||
# Tracks request params already removed in response to upstream rejections,
|
||||
# shared across providers so a stripped param stays stripped on failover and
|
||||
# the reactive retry can never loop unboundedly.
|
||||
already_stripped: set[str] = set()
|
||||
|
||||
for i, upstream in enumerate(upstreams):
|
||||
headers = upstream.prepare_headers(dict(request.headers))
|
||||
|
||||
try:
|
||||
try:
|
||||
if is_responses_api:
|
||||
response = await upstream.forward_responses_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
while True:
|
||||
try:
|
||||
if is_responses_api:
|
||||
response = await upstream.forward_responses_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
else:
|
||||
response = await upstream.forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
except UpstreamError:
|
||||
# Let the outer UpstreamError handler manage retry/revert
|
||||
raise
|
||||
except Exception as e:
|
||||
# Unexpected error (not an upstream failure) — revert and propagate
|
||||
logger.error(
|
||||
"Unexpected error in upstream request, reverting payment",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"max_cost_for_model": max_cost_for_model,
|
||||
},
|
||||
)
|
||||
else:
|
||||
response = await upstream.forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
await revert_pay_for_request(key, session, max_cost_for_model)
|
||||
raise
|
||||
|
||||
# Reactive recovery: some models reject one specific request
|
||||
# param (e.g. newer Anthropic models deprecating `temperature`).
|
||||
# When the upstream 400s naming such a param, strip it from the
|
||||
# body and retry the SAME upstream. ``already_stripped`` bounds
|
||||
# this to one retry per distinct param so it always terminates.
|
||||
if response.status_code == 400:
|
||||
correction = correct_request(
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
extract_error_message(response),
|
||||
already_stripped,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Upstream request failed, ensuring payment is reverted",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"max_cost_for_model": max_cost_for_model,
|
||||
},
|
||||
)
|
||||
await revert_pay_for_request(key, session, max_cost_for_model)
|
||||
raise
|
||||
if correction is not None:
|
||||
request_body, bad_param = correction.body, correction.label
|
||||
already_stripped.add(bad_param)
|
||||
logger.warning(
|
||||
"Upstream %s rejected param '%s' for model=%s; "
|
||||
"stripping and retrying same upstream",
|
||||
upstream.provider_type,
|
||||
bad_param,
|
||||
model_id,
|
||||
extra={
|
||||
"provider": upstream.provider_type,
|
||||
"model": model_id,
|
||||
"stripped_param": bad_param,
|
||||
"path": path,
|
||||
},
|
||||
)
|
||||
continue
|
||||
break
|
||||
|
||||
if response.status_code != 200:
|
||||
# Check if we should retry (502 Upstream Error or 429 Rate Limit)
|
||||
@@ -337,10 +458,14 @@ async def proxy(
|
||||
)
|
||||
|
||||
logger.warning(
|
||||
f"Upstream {upstream.provider_type} returned {response.status_code}, trying next provider",
|
||||
"Upstream %s returned %s for model=%s, trying next provider",
|
||||
upstream.provider_type,
|
||||
response.status_code,
|
||||
model_id,
|
||||
extra={
|
||||
"status_code": response.status_code,
|
||||
"upstream": upstream.provider_type,
|
||||
"provider": upstream.provider_type,
|
||||
"model": model_id,
|
||||
},
|
||||
)
|
||||
continue
|
||||
@@ -348,26 +473,53 @@ async def proxy(
|
||||
# 4xx error (user error), or other non-retryable error, or last provider failed
|
||||
await revert_pay_for_request(key, session, max_cost_for_model)
|
||||
logger.warning(
|
||||
"Upstream request failed, revert payment",
|
||||
"Upstream request failed, revert payment "
|
||||
"(provider=%s model=%s status=%s path=%s)",
|
||||
upstream.provider_type,
|
||||
model_id,
|
||||
response.status_code,
|
||||
path,
|
||||
extra={
|
||||
"status_code": response.status_code,
|
||||
"path": path,
|
||||
"provider": upstream.provider_type,
|
||||
"model": model_id,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"key_balance": key.balance,
|
||||
"max_cost_for_model": max_cost_for_model,
|
||||
"upstream_headers": response.headers
|
||||
if hasattr(response, "headers")
|
||||
else None,
|
||||
},
|
||||
)
|
||||
return response
|
||||
|
||||
return response
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.warning(
|
||||
"Client disconnected mid-request, reverting reservation",
|
||||
extra={
|
||||
"path": path,
|
||||
"model": model_id,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"max_cost_for_model": max_cost_for_model,
|
||||
},
|
||||
)
|
||||
await asyncio.shield(
|
||||
revert_pay_for_request(key, session, max_cost_for_model)
|
||||
)
|
||||
raise
|
||||
|
||||
except UpstreamError as e:
|
||||
logger.warning(
|
||||
f"Upstream {upstream.provider_type} failed: {e}",
|
||||
extra={"retry": i < len(upstreams) - 1},
|
||||
"Upstream %s failed for model=%s: %s",
|
||||
upstream.provider_type,
|
||||
model_id,
|
||||
e,
|
||||
extra={
|
||||
"provider": upstream.provider_type,
|
||||
"model": model_id,
|
||||
"status_code": e.status_code,
|
||||
"retry": i < len(upstreams) - 1,
|
||||
},
|
||||
)
|
||||
|
||||
# If this was the last provider
|
||||
@@ -387,10 +539,16 @@ async def proxy(
|
||||
|
||||
|
||||
async def get_bearer_token_key(
|
||||
headers: dict, path: str, session: AsyncSession, auth: str, min_cost: int = 0
|
||||
headers: dict,
|
||||
path: str,
|
||||
session: AsyncSession,
|
||||
auth: str,
|
||||
min_cost: int = 0,
|
||||
model_id: str = "unknown",
|
||||
) -> ApiKey:
|
||||
"""Handle bearer token authentication proxy requests."""
|
||||
bearer_key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else ""
|
||||
parts = auth.split()
|
||||
bearer_key = parts[1] if len(parts) > 1 and parts[0].lower() == "bearer" else ""
|
||||
refund_address = headers.get("Refund-LNURL", None)
|
||||
key_expiry_time = headers.get("Key-Expiry-Time", None)
|
||||
|
||||
@@ -454,15 +612,16 @@ async def get_bearer_token_key(
|
||||
)
|
||||
return key
|
||||
except Exception as e:
|
||||
key_preview = bearer_key[:20] + "..." if len(bearer_key) > 20 else bearer_key
|
||||
logger.error(
|
||||
"Bearer token validation failed",
|
||||
f"Bearer token validation failed: {type(e).__name__}: {e} path={path} model={model_id!r} min_cost={min_cost} key={key_preview!r}",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"bearer_key_preview": bearer_key[:20] + "..."
|
||||
if len(bearer_key) > 20
|
||||
else bearer_key,
|
||||
"model_id": model_id,
|
||||
"min_cost_msat": min_cost,
|
||||
"bearer_key_preview": key_preview,
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
@@ -10,6 +10,7 @@ from .openai import OpenAIUpstreamProvider
|
||||
from .openrouter import OpenRouterUpstreamProvider
|
||||
from .perplexity import PerplexityUpstreamProvider
|
||||
from .ppqai import PPQAIUpstreamProvider
|
||||
from .routstr import RoutstrUpstreamProvider
|
||||
from .xai import XAIUpstreamProvider
|
||||
|
||||
upstream_provider_classes: list[type[BaseUpstreamProvider]] = [
|
||||
@@ -24,6 +25,7 @@ upstream_provider_classes: list[type[BaseUpstreamProvider]] = [
|
||||
OpenRouterUpstreamProvider,
|
||||
PerplexityUpstreamProvider,
|
||||
PPQAIUpstreamProvider,
|
||||
RoutstrUpstreamProvider,
|
||||
XAIUpstreamProvider,
|
||||
]
|
||||
"""List of all upstream classes"""
|
||||
|
||||
@@ -13,6 +13,8 @@ class AnthropicUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "anthropic"
|
||||
default_base_url = "https://api.anthropic.com/v1"
|
||||
platform_url = "https://console.anthropic.com/settings/keys"
|
||||
supports_anthropic_messages = True
|
||||
litellm_provider_prefix = "anthropic/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
157
routstr/upstream/auto_topup.py
Normal file
157
routstr/upstream/auto_topup.py
Normal file
@@ -0,0 +1,157 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from sqlmodel import select
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.db import UpstreamProviderRow, create_session
|
||||
from ..wallet import send_token
|
||||
from .routstr import RoutstrUpstreamProvider
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Check every 60 seconds
|
||||
AUTO_TOPUP_INTERVAL_SECONDS = 60
|
||||
|
||||
|
||||
async def periodic_auto_topup() -> None:
|
||||
"""Background task that monitors Routstr provider balances and auto-tops up when below threshold.
|
||||
|
||||
For each Routstr provider with auto_topup enabled in provider_settings:
|
||||
1. Checks the upstream balance via get_balance()
|
||||
2. If balance < topup_threshold, creates a cashu token from the configured mint
|
||||
3. Sends the token to the upstream provider via topup()
|
||||
"""
|
||||
# Wait for initial startup to complete
|
||||
await asyncio.sleep(30)
|
||||
logger.info("Auto top-up worker started")
|
||||
|
||||
while True:
|
||||
try:
|
||||
await _run_auto_topup_cycle()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Auto top-up cycle failed",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
|
||||
await asyncio.sleep(AUTO_TOPUP_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
async def _run_auto_topup_cycle() -> None:
|
||||
"""Single cycle: check all eligible providers and top up if needed."""
|
||||
async with create_session() as session:
|
||||
query = select(UpstreamProviderRow).where(
|
||||
UpstreamProviderRow.provider_type == "routstr",
|
||||
UpstreamProviderRow.enabled == True, # noqa: E712
|
||||
)
|
||||
result = await session.exec(query)
|
||||
providers = result.all()
|
||||
|
||||
for row in providers:
|
||||
try:
|
||||
await _check_and_topup(row)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Auto top-up failed for provider",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"base_url": row.base_url,
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _check_and_topup(row: UpstreamProviderRow) -> None:
|
||||
"""Check a single provider's balance and top up if below threshold."""
|
||||
# Parse provider settings
|
||||
settings: dict = {}
|
||||
if row.provider_settings:
|
||||
try:
|
||||
settings = json.loads(row.provider_settings)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return
|
||||
|
||||
if not settings.get("auto_topup"):
|
||||
return
|
||||
|
||||
threshold = settings.get("topup_threshold")
|
||||
amount = settings.get("topup_amount_limit")
|
||||
mint_url = settings.get("topup_mint_url")
|
||||
|
||||
if not threshold or not amount or not mint_url:
|
||||
logger.warning(
|
||||
"Auto top-up enabled but missing configuration",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"has_threshold": bool(threshold),
|
||||
"has_amount": bool(amount),
|
||||
"has_mint": bool(mint_url),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
if not row.api_key:
|
||||
return
|
||||
|
||||
# Instantiate provider and check balance
|
||||
provider = RoutstrUpstreamProvider.from_db_row(row)
|
||||
balance = await provider.get_balance()
|
||||
|
||||
if balance is None:
|
||||
logger.warning(
|
||||
"Could not fetch balance for auto top-up",
|
||||
extra={"provider_id": row.id, "base_url": row.base_url},
|
||||
)
|
||||
return
|
||||
|
||||
if balance >= threshold * 1000:
|
||||
return
|
||||
|
||||
# Balance is below threshold - create token and top up
|
||||
logger.info(
|
||||
"Auto top-up triggered",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"balance": balance,
|
||||
"threshold": threshold,
|
||||
"topup_amount": amount,
|
||||
"mint_url": mint_url,
|
||||
},
|
||||
)
|
||||
|
||||
print(amount, mint_url)
|
||||
try:
|
||||
token = await send_token(amount, "sat", mint_url)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to create cashu token for auto top-up",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"amount": amount,
|
||||
"mint_url": mint_url,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
result = await provider.topup(token)
|
||||
|
||||
if "error" in result:
|
||||
logger.error(
|
||||
"Auto top-up upstream call failed",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"error": result["error"],
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Auto top-up completed successfully",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"amount": amount,
|
||||
"new_balance_approx": balance + amount,
|
||||
},
|
||||
)
|
||||
@@ -4,6 +4,7 @@ from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
|
||||
class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
@@ -12,6 +13,7 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "azure"
|
||||
default_base_url = None
|
||||
platform_url = "https://portal.azure.com/"
|
||||
litellm_provider_prefix = "azure/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -58,19 +60,52 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
"platform_url": cls.platform_url,
|
||||
}
|
||||
|
||||
def prepare_headers(self, request_headers: dict) -> dict:
|
||||
"""Prepare headers for Azure OpenAI, adding api-key."""
|
||||
headers = super().prepare_headers(request_headers)
|
||||
if self.api_key:
|
||||
headers["api-key"] = self.api_key
|
||||
headers.pop("Authorization", None)
|
||||
headers.pop("authorization", None)
|
||||
return headers
|
||||
|
||||
def prepare_params(
|
||||
self, path: str, query_params: Mapping[str, str] | None
|
||||
) -> Mapping[str, str]:
|
||||
"""Prepare query parameters for Azure OpenAI, adding API version.
|
||||
|
||||
Args:
|
||||
path: Request path
|
||||
query_params: Original query parameters from the client
|
||||
|
||||
Returns:
|
||||
Query parameters dict with Azure API version added for chat completions
|
||||
"""
|
||||
"""Prepare query parameters for Azure OpenAI, adding API version."""
|
||||
params = dict(query_params or {})
|
||||
if path.endswith("chat/completions"):
|
||||
params["api-version"] = self.api_version
|
||||
version = (self.api_version or "").replace("\ufeff", "").strip()
|
||||
if not version or version.lower() == "v1":
|
||||
version = "2024-02-15-preview"
|
||||
params["api-version"] = version
|
||||
return params
|
||||
|
||||
def normalize_request_path(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Build Azure deployment-specific request path."""
|
||||
clean_path = super().normalize_request_path(path, model_obj).lstrip("/")
|
||||
if model_obj is None:
|
||||
return clean_path
|
||||
|
||||
deployment_id = getattr(
|
||||
model_obj, "canonical_slug", None
|
||||
) or self.transform_model_name(model_obj.id)
|
||||
deployment_id = deployment_id.split("/")[-1]
|
||||
return f"openai/deployments/{deployment_id}/{clean_path}"
|
||||
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Use endpoint root, stripping accidental /openai/v1 suffix if present."""
|
||||
base_url = self.base_url.rstrip("/")
|
||||
marker = "/openai/v1"
|
||||
if marker in base_url:
|
||||
base_url = base_url.split(marker, 1)[0].rstrip("/")
|
||||
return base_url
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
"""Extract deployment name from model ID."""
|
||||
if "/" in model_id:
|
||||
return model_id.split("/")[-1]
|
||||
return model_id
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
157
routstr/upstream/cache_breakpoints.py
Normal file
157
routstr/upstream/cache_breakpoints.py
Normal file
@@ -0,0 +1,157 @@
|
||||
"""Inject explicit prompt-cache breakpoints into OpenAI-shaped requests.
|
||||
|
||||
Some upstreams cache *explicitly*: the request must carry
|
||||
``cache_control: {"type": "ephemeral"}`` markers on the content blocks that
|
||||
should be cached. Two model families use this identical wire format:
|
||||
|
||||
* **Anthropic Claude** — direct or via OpenRouter's ``anthropic/*`` models.
|
||||
* **Alibaba's explicit-cache models on OpenRouter** — ``qwen/qwen3-max``,
|
||||
``qwen/qwen-plus``, ``qwen/qwen3.6-plus``, ``qwen/qwen3-coder-plus``,
|
||||
``qwen/qwen3-coder-flash`` and ``deepseek/deepseek-v3.2`` — which OpenRouter
|
||||
documents as using "the same syntax as Anthropic explicit caching".
|
||||
|
||||
Every other provider routstr proxies (OpenAI, Azure, xAI/Grok, Groq, Moonshot,
|
||||
default DeepSeek, Gemini implicit, Fireworks) caches *automatically* and needs
|
||||
no markers — they are left untouched.
|
||||
|
||||
A client that doesn't know it is talking to one of these models *through*
|
||||
routstr (e.g. an OpenAI-compatible coding agent pointed at a routstr URL) never
|
||||
emits the markers — it only adds them when it recognises the provider as
|
||||
OpenRouter. So caching silently never engages over routstr even though the same
|
||||
client caches fine talking to OpenRouter directly.
|
||||
|
||||
This module restores caching by stamping the standard breakpoints onto the
|
||||
forwarded body — the system prompt, the last tool, and the last conversation
|
||||
message (the format allows up to four; we use three, matching the common
|
||||
agent convention) — but only when the client supplied none of its own, so
|
||||
explicit client control always wins. The caller is responsible for only
|
||||
applying this toward an upstream that accepts the markers (OpenRouter /
|
||||
Anthropic), so they never leak to a provider that would reject them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
# The single ephemeral marker stamped onto each chosen breakpoint. A 5-minute
|
||||
# TTL (the default for ``ephemeral``) — deliberately not the 1h tier, which
|
||||
# carries a higher cache-write premium and should stay opt-in.
|
||||
EPHEMERAL_CACHE_CONTROL: dict[str, str] = {"type": "ephemeral"}
|
||||
|
||||
# Alibaba's explicit-cache models on OpenRouter. Matched as substrings of the
|
||||
# model id (any spelling routstr carries). Snapshot endpoints that OpenRouter
|
||||
# documents as *not* supporting explicit caching (e.g. ``qwen3.5-plus-02-15``)
|
||||
# are different families and deliberately absent from this list.
|
||||
_ALIBABA_EXPLICIT_CACHE_SLUGS: tuple[str, ...] = (
|
||||
"qwen3-max",
|
||||
"qwen-plus",
|
||||
"qwen3.6-plus",
|
||||
"qwen3-coder-plus",
|
||||
"qwen3-coder-flash",
|
||||
"deepseek-v3.2",
|
||||
)
|
||||
|
||||
|
||||
def is_explicit_cache_model(model_id: str | None, *fallbacks: str | None) -> bool:
|
||||
"""True when the target model uses the explicit ``cache_control`` dialect.
|
||||
|
||||
Covers the Claude family (broadly — every Claude model supports it) and
|
||||
Alibaba's documented explicit-cache models, across the id spellings routstr
|
||||
carries: the OpenRouter id (``anthropic/claude-...``, ``qwen/qwen3-max``),
|
||||
the bare upstream id, and any forwarded/canonical alias.
|
||||
"""
|
||||
for candidate in (model_id, *fallbacks):
|
||||
if not candidate:
|
||||
continue
|
||||
lowered = candidate.lower()
|
||||
if "claude" in lowered or "anthropic/" in lowered:
|
||||
return True
|
||||
if any(slug in lowered for slug in _ALIBABA_EXPLICIT_CACHE_SLUGS):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _has_cache_control(obj: Any) -> bool:
|
||||
"""Recursively detect any client-supplied ``cache_control`` marker."""
|
||||
if isinstance(obj, dict):
|
||||
if "cache_control" in obj:
|
||||
return True
|
||||
return any(_has_cache_control(v) for v in obj.values())
|
||||
if isinstance(obj, list):
|
||||
return any(_has_cache_control(v) for v in obj)
|
||||
return False
|
||||
|
||||
|
||||
def body_has_cache_control(data: dict) -> bool:
|
||||
"""True when the request already carries cache_control on messages/tools."""
|
||||
return _has_cache_control(data.get("messages")) or _has_cache_control(
|
||||
data.get("tools")
|
||||
)
|
||||
|
||||
|
||||
def _stamp_text_content(message: dict) -> bool:
|
||||
"""Add the ephemeral marker to a message's last text block.
|
||||
|
||||
A string content is promoted to the array form Anthropic requires for
|
||||
cache markers; an existing array gets the marker on its last text part.
|
||||
Returns True when a marker was placed.
|
||||
"""
|
||||
content = message.get("content")
|
||||
if isinstance(content, str):
|
||||
if not content:
|
||||
return False
|
||||
message["content"] = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": content,
|
||||
"cache_control": dict(EPHEMERAL_CACHE_CONTROL),
|
||||
}
|
||||
]
|
||||
return True
|
||||
if isinstance(content, list):
|
||||
for part in reversed(content):
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
part["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _stamp_system_prompt(messages: list) -> None:
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and message.get("role") in (
|
||||
"system",
|
||||
"developer",
|
||||
):
|
||||
_stamp_text_content(message)
|
||||
return
|
||||
|
||||
|
||||
def _stamp_last_tool(tools: Any) -> None:
|
||||
if isinstance(tools, list) and tools and isinstance(tools[-1], dict):
|
||||
tools[-1]["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL)
|
||||
|
||||
|
||||
def _stamp_last_conversation_message(messages: list) -> None:
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") in ("user", "assistant"):
|
||||
if _stamp_text_content(message):
|
||||
return
|
||||
|
||||
|
||||
def inject_anthropic_cache_breakpoints(data: dict) -> bool:
|
||||
"""Stamp ephemeral cache breakpoints onto an OpenAI-shaped chat body.
|
||||
|
||||
Mutates ``data`` in place (the established convention in
|
||||
``prepare_request_body``) and returns True when anything changed. No-ops
|
||||
when the body isn't chat-shaped or the client already set cache_control.
|
||||
"""
|
||||
messages = data.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return False
|
||||
if body_has_cache_control(data):
|
||||
return False
|
||||
|
||||
_stamp_system_prompt(messages)
|
||||
_stamp_last_tool(data.get("tools"))
|
||||
_stamp_last_conversation_message(messages)
|
||||
return True
|
||||
@@ -26,14 +26,14 @@ class GeminiClient(BaseAPIClient):
|
||||
max_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
from openai import NOT_GIVEN
|
||||
from openai import omit
|
||||
|
||||
response = await self.client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages, # type: ignore
|
||||
temperature=temperature if temperature is not None else NOT_GIVEN,
|
||||
max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
|
||||
top_p=kwargs.get("top_p", NOT_GIVEN),
|
||||
temperature=temperature if temperature is not None else omit,
|
||||
max_tokens=max_tokens if max_tokens is not None else omit,
|
||||
top_p=kwargs.get("top_p", omit),
|
||||
)
|
||||
return response.model_dump()
|
||||
|
||||
@@ -45,7 +45,7 @@ class GeminiClient(BaseAPIClient):
|
||||
max_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncGenerator[dict[str, Any], None]:
|
||||
from openai import NOT_GIVEN
|
||||
from openai import omit
|
||||
|
||||
usage_callback = kwargs.get("usage_callback")
|
||||
completion_callback = kwargs.get("completion_callback")
|
||||
@@ -55,9 +55,9 @@ class GeminiClient(BaseAPIClient):
|
||||
messages=messages, # type: ignore
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
temperature=temperature if temperature is not None else NOT_GIVEN,
|
||||
max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
|
||||
top_p=kwargs.get("top_p", NOT_GIVEN),
|
||||
temperature=temperature if temperature is not None else omit,
|
||||
max_tokens=max_tokens if max_tokens is not None else omit,
|
||||
top_p=kwargs.get("top_p", omit),
|
||||
)
|
||||
|
||||
final_usage = None
|
||||
|
||||
108
routstr/upstream/count_tokens.py
Normal file
108
routstr/upstream/count_tokens.py
Normal file
@@ -0,0 +1,108 @@
|
||||
"""Local handling of Anthropic ``/v1/messages/count_tokens`` for upstreams
|
||||
that do not natively expose the endpoint.
|
||||
|
||||
Most non-Anthropic upstreams (OpenAI-compat, Gemini OpenAI-compat,
|
||||
OpenRouter chat-completions, generic providers) return 400/404 when asked
|
||||
to ``POST /messages/count_tokens``. Claude Code and other Anthropic SDK
|
||||
clients call this endpoint before each turn to size context windows and
|
||||
trigger compaction, so a failure breaks the whole chat.
|
||||
|
||||
We answer locally. ``litellm.token_counter`` understands the Anthropic
|
||||
message shape and the per-model tokenizers, so we prefer it. If it raises
|
||||
(unknown model, encoding lookup failure, ...), we fall back to the
|
||||
project's own ``estimate_tokens`` heuristic, which is always defined and
|
||||
never raises.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
from fastapi.responses import Response
|
||||
|
||||
from ..core import get_logger
|
||||
from ..payment.helpers import estimate_tokens
|
||||
from ..payment.models import Model
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _parse_request_body(request_body: bytes | None) -> dict[str, Any]:
|
||||
if not request_body:
|
||||
return {}
|
||||
try:
|
||||
parsed = json.loads(request_body)
|
||||
except (ValueError, TypeError):
|
||||
return {}
|
||||
return parsed if isinstance(parsed, dict) else {}
|
||||
|
||||
|
||||
def _count_with_litellm(model: str, body: dict[str, Any]) -> int:
|
||||
messages = body.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
messages = []
|
||||
|
||||
system = body.get("system")
|
||||
if isinstance(system, str) and system:
|
||||
messages = [{"role": "system", "content": system}, *messages]
|
||||
elif isinstance(system, list):
|
||||
text = "".join(
|
||||
block.get("text", "")
|
||||
for block in system
|
||||
if isinstance(block, dict) and block.get("type") == "text"
|
||||
)
|
||||
if text:
|
||||
messages = [{"role": "system", "content": text}, *messages]
|
||||
|
||||
tools = body.get("tools") if isinstance(body.get("tools"), list) else None
|
||||
|
||||
return int(
|
||||
litellm.token_counter(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def count_tokens_locally(
|
||||
request_body: bytes | None,
|
||||
model_obj: Model | None,
|
||||
) -> Response:
|
||||
"""Return an Anthropic-compatible count_tokens response without
|
||||
touching the upstream. Always returns 200; never raises."""
|
||||
body = _parse_request_body(request_body)
|
||||
|
||||
model_name = ""
|
||||
if model_obj is not None:
|
||||
model_name = model_obj.forwarded_model_id or model_obj.id or ""
|
||||
if not model_name:
|
||||
body_model = body.get("model")
|
||||
if isinstance(body_model, str):
|
||||
model_name = body_model
|
||||
|
||||
input_tokens: int
|
||||
try:
|
||||
input_tokens = _count_with_litellm(model_name, body)
|
||||
except Exception as exc:
|
||||
messages = body.get("messages")
|
||||
fallback_messages = messages if isinstance(messages, list) else []
|
||||
input_tokens = estimate_tokens(fallback_messages)
|
||||
logger.debug(
|
||||
"litellm token_counter failed; using local estimator",
|
||||
extra={
|
||||
"model": model_name,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
"estimated_tokens": input_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
payload = {"input_tokens": max(0, int(input_tokens))}
|
||||
return Response(
|
||||
content=json.dumps(payload).encode(),
|
||||
status_code=200,
|
||||
media_type="application/json",
|
||||
)
|
||||
@@ -12,6 +12,7 @@ class FireworksUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "fireworks"
|
||||
default_base_url = "https://api.fireworks.ai/inference/v1"
|
||||
platform_url = "https://app.fireworks.ai/settings/users/api-keys"
|
||||
litellm_provider_prefix = "fireworks_ai/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -1,17 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from . import gemini_messages
|
||||
from .base import BaseUpstreamProvider
|
||||
from .clients.gemini import GeminiClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
from ..core.logging import get_logger
|
||||
@@ -20,9 +16,19 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Gemini provider — proxies through Gemini's OpenAI-compat surface.
|
||||
|
||||
The chat-completions, embeddings, and models paths all flow through
|
||||
:meth:`BaseUpstreamProvider.forward_request`; we only override
|
||||
``get_request_base_url`` to redirect to ``{base}/openai/...`` and
|
||||
``_dispatch_anthropic_messages`` to inject thought-signatures on the
|
||||
/v1/messages path (see :mod:`gemini_messages` for that rationale).
|
||||
"""
|
||||
|
||||
provider_type = "gemini"
|
||||
default_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
platform_url = "https://aistudio.google.com/app/apikey"
|
||||
litellm_provider_prefix = "gemini/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -39,7 +45,7 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
|
||||
@property
|
||||
def client(self) -> GeminiClient:
|
||||
"""Get or create the Gemini API client."""
|
||||
"""Get or create the Gemini API client (used for the models listing)."""
|
||||
if self._client is None:
|
||||
self._client = GeminiClient(api_key=self.api_key)
|
||||
return self._client
|
||||
@@ -65,240 +71,66 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
}
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
return model_id.removeprefix("gemini/")
|
||||
"""Reduce a routstr model id to the bare upstream Gemini name.
|
||||
|
||||
async def forward_request(
|
||||
Gemini's OpenAI-compat surface expects the literal model id
|
||||
(e.g. ``gemini-3.1-flash-lite-preview``) — no ``gemini/`` provider
|
||||
prefix and no ``google/`` vendor sub-prefix. Take the last path
|
||||
segment so we tolerate any of:
|
||||
|
||||
``gemini-2.0-flash``
|
||||
``gemini/gemini-2.0-flash``
|
||||
``gemini/google/gemini-3.1-flash-lite-preview``
|
||||
"""
|
||||
return model_id.rsplit("/", 1)[-1]
|
||||
|
||||
@property
|
||||
def compat_base_url(self) -> str:
|
||||
"""Gemini's OpenAI-compat surface, regardless of what's stored.
|
||||
|
||||
Stored ``base_url`` may be ``.../v1beta`` (the native Gemini API
|
||||
root) or ``.../v1beta/openai`` (already pointed at the compat
|
||||
surface). Normalize to the latter.
|
||||
"""
|
||||
return self.base_url.rstrip("/").removesuffix("/openai") + "/openai"
|
||||
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Route every proxied request to the OpenAI-compat surface.
|
||||
|
||||
Required because the stored ``base_url`` typically points at the
|
||||
native Gemini API (``/v1beta``), but :meth:`forward_request`
|
||||
forwards OpenAI-shaped paths (``/chat/completions``,
|
||||
``/embeddings``, ``/models``) which only exist under the
|
||||
``/openai`` subtree.
|
||||
"""
|
||||
return self.compat_base_url
|
||||
|
||||
async def _dispatch_anthropic_messages(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
# Remove provider prefix from model ID for Gemini API
|
||||
if "/" in model_obj.id:
|
||||
model_obj.id = model_obj.id.split("/", 1)[1]
|
||||
model_obj: "Model",
|
||||
*,
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Dispatch /v1/messages via the gemini-specific httpx path.
|
||||
|
||||
if not path.startswith("chat/completions"):
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
|
||||
if not request_body:
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
|
||||
try:
|
||||
openai_data = json.loads(request_body)
|
||||
messages = openai_data.get("messages", [])
|
||||
temperature = openai_data.get("temperature")
|
||||
max_tokens = openai_data.get("max_tokens")
|
||||
top_p = openai_data.get("top_p")
|
||||
is_streaming = openai_data.get("stream", False)
|
||||
|
||||
logger.info(
|
||||
"Processing Gemini request with client abstraction",
|
||||
extra={
|
||||
"model": model_obj.id,
|
||||
"is_streaming": is_streaming,
|
||||
"message_count": len(messages),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
if is_streaming:
|
||||
final_usage_data: dict | None = None
|
||||
|
||||
def usage_callback(usage_data: dict[str, Any]) -> None:
|
||||
"""Callback to capture usage data during streaming"""
|
||||
nonlocal final_usage_data
|
||||
final_usage_data = usage_data
|
||||
|
||||
async def completion_callback(
|
||||
model: str, usage_data: dict[str, Any] | None
|
||||
) -> None:
|
||||
"""Callback to handle payment when streaming completes"""
|
||||
nonlocal final_usage_data
|
||||
if usage_data:
|
||||
final_usage_data = usage_data
|
||||
|
||||
payment_data = {
|
||||
"model": model,
|
||||
"usage": final_usage_data,
|
||||
}
|
||||
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
from ..core.db import create_session
|
||||
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if fresh_key:
|
||||
try:
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
payment_data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Gemini streaming payment finalized",
|
||||
extra={
|
||||
"cost_data": cost_data,
|
||||
"usage_data": final_usage_data,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing Gemini streaming payment",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
response_generator = self.client.generate_content_stream(
|
||||
model=model_obj.id,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
top_p=top_p,
|
||||
usage_callback=usage_callback,
|
||||
completion_callback=completion_callback,
|
||||
)
|
||||
|
||||
async def stream_with_cost() -> AsyncGenerator[bytes, None]:
|
||||
payment_finalized = False
|
||||
|
||||
async def finalize_payment() -> None:
|
||||
nonlocal payment_finalized
|
||||
if payment_finalized:
|
||||
return
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
from ..core.db import create_session
|
||||
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(
|
||||
key.__class__, key.hashed_key
|
||||
)
|
||||
if fresh_key:
|
||||
try:
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
{
|
||||
"model": model_obj.id,
|
||||
"usage": final_usage_data,
|
||||
},
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
payment_finalized = True
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing Gemini streaming payment in fallback",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
async for chunk in response_generator:
|
||||
sse_data = f"data: {json.dumps(chunk)}\n\n"
|
||||
yield sse_data.encode()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in Gemini streaming response",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
if not payment_finalized:
|
||||
await finalize_payment()
|
||||
|
||||
return StreamingResponse(
|
||||
stream_with_cost(),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
|
||||
)
|
||||
|
||||
else:
|
||||
openai_format_response = await self.client.generate_content(
|
||||
model=model_obj.id,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
top_p=top_p,
|
||||
)
|
||||
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, openai_format_response, session, max_cost_for_model
|
||||
)
|
||||
openai_format_response["cost"] = cost_data
|
||||
|
||||
logger.info(
|
||||
"Gemini non-streaming payment completed",
|
||||
extra={
|
||||
"cost_data": cost_data,
|
||||
"model": model_obj.id,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=json.dumps(openai_format_response),
|
||||
media_type="application/json",
|
||||
headers={"Cache-Control": "no-cache"},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in Gemini forward_request",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
See :mod:`routstr.upstream.gemini_messages` for the full rationale
|
||||
(thought-signature injection, why litellm + the openai SDK can't
|
||||
carry the required ``extra_content`` field).
|
||||
"""
|
||||
return await gemini_messages.dispatch_gemini_messages(
|
||||
request_body=request_body,
|
||||
model_obj=model_obj,
|
||||
base_url=self.compat_base_url,
|
||||
api_key=self.api_key,
|
||||
transform_model_name=self.transform_model_name,
|
||||
log_extra=log_extra,
|
||||
)
|
||||
|
||||
async def _fetch_provider_models(self) -> dict:
|
||||
"""Fetch models from Gemini API."""
|
||||
"""Fetch models from Gemini API via the OpenAI-compat client."""
|
||||
try:
|
||||
models_data = await self.client.list_models()
|
||||
|
||||
|
||||
480
routstr/upstream/gemini_messages.py
Normal file
480
routstr/upstream/gemini_messages.py
Normal file
@@ -0,0 +1,480 @@
|
||||
"""Custom /v1/messages dispatcher for Gemini's OpenAI-compat endpoint.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
|
||||
Gemini 2.5 / 3 thinking models reject inbound ``functionCall`` parts that
|
||||
lack a ``thought_signature`` field once any prior turn in the conversation
|
||||
contains a function call. Anthropic-Messages clients (Claude Code etc.)
|
||||
have no concept of thought signatures, so multi-turn tool conversations
|
||||
fail with::
|
||||
|
||||
Function call is missing a thought_signature in functionCall parts.
|
||||
|
||||
Google's published escape hatch (https://ai.google.dev/gemini-api/docs/
|
||||
thought-signatures, FAQ #1) is the dummy signature
|
||||
``"skip_thought_signature_validator"`` placed at
|
||||
``tool_calls[i].extra_content.google.thought_signature`` for every tool
|
||||
call in the request. The hatch is documented specifically for
|
||||
"transferring a trace from a different model that does not include thought
|
||||
signatures" — exactly our case.
|
||||
|
||||
Why we can't reach the wire via litellm
|
||||
---------------------------------------
|
||||
|
||||
``litellm.anthropic.messages.acreate`` flows through the openai SDK, whose
|
||||
pydantic ``ChatCompletionMessageToolCall`` model silently drops unknown
|
||||
fields like ``extra_content``. Litellm has no openai-compat translator
|
||||
that emits ``extra_content.google.thought_signature``. So we bypass both
|
||||
litellm and the openai SDK at the transport layer.
|
||||
|
||||
Pipeline
|
||||
--------
|
||||
|
||||
1. Translate Anthropic body → OpenAI body via litellm's
|
||||
``AnthropicAdapter`` (the same translator
|
||||
``litellm.anthropic.messages.acreate`` uses internally).
|
||||
2. Inject ``extra_content.google.thought_signature`` on every
|
||||
``tool_calls[]`` entry.
|
||||
3. Set ``reasoning_effort="none"`` to disable Gemini's thinking pass.
|
||||
4. POST directly to ``{base_url}/chat/completions`` with ``stream=true``
|
||||
via ``httpx`` (preserves arbitrary fields verbatim).
|
||||
5. Translate OpenAI streaming chunks → Anthropic SSE events.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from typing import Any, Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.exceptions import UpstreamError
|
||||
from ..payment.models import Model
|
||||
from .messages_dispatch import (
|
||||
ANTHROPIC_ONLY_FIELDS,
|
||||
aggregate_anthropic_events_to_message,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
DUMMY_THOUGHT_SIGNATURE = "skip_thought_signature_validator"
|
||||
|
||||
# Mapping: OpenAI finish_reason → Anthropic stop_reason
|
||||
_FINISH_TO_STOP = {
|
||||
"stop": "end_turn",
|
||||
"length": "max_tokens",
|
||||
"tool_calls": "tool_use",
|
||||
"function_call": "tool_use",
|
||||
"content_filter": "refusal",
|
||||
}
|
||||
|
||||
|
||||
def inject_thought_signatures(messages: list[dict]) -> None:
|
||||
"""Add ``extra_content.google.thought_signature`` to every tool_call.
|
||||
|
||||
Mutates ``messages`` in place. Idempotent: existing signatures are not
|
||||
overwritten.
|
||||
"""
|
||||
for msg in messages:
|
||||
tool_calls = msg.get("tool_calls")
|
||||
if not isinstance(tool_calls, list):
|
||||
continue
|
||||
for tc in tool_calls:
|
||||
if not isinstance(tc, dict):
|
||||
continue
|
||||
extra = tc.get("extra_content")
|
||||
if not isinstance(extra, dict):
|
||||
extra = tc["extra_content"] = {}
|
||||
google_cfg = extra.get("google")
|
||||
if not isinstance(google_cfg, dict):
|
||||
google_cfg = extra["google"] = {}
|
||||
google_cfg.setdefault("thought_signature", DUMMY_THOUGHT_SIGNATURE)
|
||||
|
||||
|
||||
def _translate_anthropic_to_openai(body: dict, model: str) -> dict:
|
||||
"""Use litellm's translator to convert an Anthropic /messages body to
|
||||
OpenAI /chat/completions kwargs.
|
||||
|
||||
Imported lazily because the litellm internal path is heavy and not
|
||||
needed for any other code path in routstr.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( # noqa: E501
|
||||
AnthropicAdapter,
|
||||
)
|
||||
|
||||
kwargs = {"model": model, **body}
|
||||
translated = AnthropicAdapter().translate_completion_input_params(kwargs)
|
||||
if translated is None:
|
||||
raise UpstreamError(
|
||||
"Failed to translate Anthropic body to OpenAI format",
|
||||
status_code=500,
|
||||
)
|
||||
return dict(translated)
|
||||
|
||||
|
||||
def _sse_event(event_type: str, payload: dict) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
async def _openai_chunks_to_anthropic_events(
|
||||
line_iter: AsyncIterator[str], requested_model: str | None
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
"""Translate an OpenAI chat-completions SSE byte stream into the
|
||||
Anthropic-Messages SSE event sequence.
|
||||
|
||||
Maintains per-chunk state across:
|
||||
|
||||
* one optional text content block (lazy-opened on first text delta)
|
||||
* any number of tool_use blocks indexed by openai's ``delta.tool_calls[].index``
|
||||
* final ``stop_reason`` / ``usage`` carried out via ``message_delta`` /
|
||||
``message_stop``
|
||||
"""
|
||||
msg_id = f"msg_{uuid.uuid4().hex[:24]}"
|
||||
started = False
|
||||
text_block_idx: int | None = None
|
||||
tool_block_indices: dict[int, int] = {}
|
||||
next_block_idx = 0
|
||||
final_finish_reason: str | None = None
|
||||
final_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
|
||||
|
||||
def open_text_block() -> bytes:
|
||||
nonlocal text_block_idx, next_block_idx
|
||||
text_block_idx = next_block_idx
|
||||
next_block_idx += 1
|
||||
return _sse_event(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": text_block_idx,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
)
|
||||
|
||||
def open_tool_block(delta_idx: int, tc: dict) -> bytes:
|
||||
nonlocal next_block_idx
|
||||
block_idx = next_block_idx
|
||||
next_block_idx += 1
|
||||
tool_block_indices[delta_idx] = block_idx
|
||||
fn = tc.get("function") or {}
|
||||
return _sse_event(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": block_idx,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id") or f"toolu_{uuid.uuid4().hex[:24]}",
|
||||
"name": fn.get("name") or "",
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def close_block(idx: int) -> bytes:
|
||||
return _sse_event(
|
||||
"content_block_stop",
|
||||
{"type": "content_block_stop", "index": idx},
|
||||
)
|
||||
|
||||
async for raw_line in line_iter:
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
payload = line[5:].lstrip()
|
||||
if not payload or payload == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
chunk = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if not isinstance(chunk, dict):
|
||||
continue
|
||||
|
||||
if not started:
|
||||
started = True
|
||||
yield _sse_event(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": chunk.get("id") or msg_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": requested_model or chunk.get("model") or "",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
usage = chunk.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
in_tok = usage.get("prompt_tokens") or usage.get("input_tokens") or 0
|
||||
out_tok = usage.get("completion_tokens") or usage.get("output_tokens") or 0
|
||||
if in_tok:
|
||||
final_usage["input_tokens"] = int(in_tok)
|
||||
if out_tok:
|
||||
final_usage["output_tokens"] = int(out_tok)
|
||||
|
||||
choices = chunk.get("choices") or []
|
||||
if not choices:
|
||||
continue
|
||||
choice = choices[0] if isinstance(choices[0], dict) else {}
|
||||
delta = choice.get("delta") or {}
|
||||
if not isinstance(delta, dict):
|
||||
delta = {}
|
||||
|
||||
# Text delta
|
||||
text = delta.get("content")
|
||||
if isinstance(text, str) and text:
|
||||
if text_block_idx is None:
|
||||
yield open_text_block()
|
||||
yield _sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": text_block_idx,
|
||||
"delta": {"type": "text_delta", "text": text},
|
||||
},
|
||||
)
|
||||
|
||||
# Tool call deltas
|
||||
tool_calls_delta = delta.get("tool_calls")
|
||||
if isinstance(tool_calls_delta, list):
|
||||
for tc in tool_calls_delta:
|
||||
if not isinstance(tc, dict):
|
||||
continue
|
||||
d_idx = int(tc.get("index") or 0)
|
||||
if d_idx not in tool_block_indices:
|
||||
yield open_tool_block(d_idx, tc)
|
||||
block_idx = tool_block_indices[d_idx]
|
||||
fn = tc.get("function") or {}
|
||||
args = fn.get("arguments")
|
||||
if isinstance(args, str) and args:
|
||||
yield _sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": block_idx,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": args,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
finish = choice.get("finish_reason")
|
||||
if finish:
|
||||
final_finish_reason = finish
|
||||
|
||||
# Close any open content blocks
|
||||
if text_block_idx is not None:
|
||||
yield close_block(text_block_idx)
|
||||
for block_idx in tool_block_indices.values():
|
||||
yield close_block(block_idx)
|
||||
|
||||
# message_delta with stop_reason and usage
|
||||
stop_reason = _FINISH_TO_STOP.get(final_finish_reason or "", "end_turn")
|
||||
yield _sse_event(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": final_usage,
|
||||
},
|
||||
)
|
||||
yield _sse_event("message_stop", {"type": "message_stop"})
|
||||
|
||||
|
||||
async def _post_and_stream(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
payload: dict,
|
||||
log_extra: dict[str, Any] | None,
|
||||
) -> tuple[httpx.AsyncClient, httpx.Response]:
|
||||
"""POST to upstream chat-completions and return (client, response) for
|
||||
streaming. Caller is responsible for closing both."""
|
||||
url = f"{base_url.rstrip('/')}/chat/completions"
|
||||
client = httpx.AsyncClient(timeout=httpx.Timeout(120.0, read=120.0))
|
||||
try:
|
||||
request = client.build_request(
|
||||
"POST",
|
||||
url,
|
||||
json=payload,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
)
|
||||
response = await client.send(request, stream=True)
|
||||
except Exception as exc:
|
||||
await client.aclose()
|
||||
logger.error(
|
||||
"Gemini messages dispatch HTTP error",
|
||||
extra={"error": str(exc), "url": url, **(log_extra or {})},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to reach Gemini upstream: {exc}", status_code=502
|
||||
) from exc
|
||||
|
||||
if response.status_code >= 400:
|
||||
try:
|
||||
body_bytes = await response.aread()
|
||||
finally:
|
||||
await response.aclose()
|
||||
await client.aclose()
|
||||
body_text = body_bytes.decode("utf-8", errors="replace")
|
||||
logger.error(
|
||||
"Gemini messages dispatch upstream error",
|
||||
extra={
|
||||
"status_code": response.status_code,
|
||||
"body": body_text[:1000],
|
||||
"url": url,
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Upstream error via gemini compat: {body_text}",
|
||||
status_code=response.status_code,
|
||||
)
|
||||
|
||||
return client, response
|
||||
|
||||
|
||||
async def dispatch_gemini_messages(
|
||||
*,
|
||||
request_body: bytes | None,
|
||||
model_obj: Model,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
transform_model_name: Callable[[str], str],
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Dispatch a /v1/messages request to Gemini's OpenAI-compat endpoint
|
||||
with thought-signature injection.
|
||||
|
||||
Returns ``(client_stream, result, requested_model)`` where ``result``
|
||||
is either an ``AsyncIterator[bytes]`` of Anthropic-format SSE events
|
||||
(for streaming clients) or an Anthropic Message dict (after the caller
|
||||
aggregates).
|
||||
"""
|
||||
if not request_body:
|
||||
raise UpstreamError(
|
||||
"Missing request body for /v1/messages", status_code=400
|
||||
)
|
||||
|
||||
try:
|
||||
body: dict = json.loads(request_body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise UpstreamError(
|
||||
f"Invalid JSON in /v1/messages body: {exc}", status_code=400
|
||||
) from exc
|
||||
|
||||
body.pop("model", None)
|
||||
client_stream = bool(body.pop("stream", False))
|
||||
|
||||
# Anthropic-Messages-only fields that don't translate to OpenAI
|
||||
# Chat Completions. litellm's translator passes through unknown
|
||||
# top-level fields verbatim and Gemini's compat surface 400s on
|
||||
# unknown names like ``context_management`` / ``output_config``.
|
||||
dropped: dict[str, Any] = {}
|
||||
for field in ANTHROPIC_ONLY_FIELDS:
|
||||
if field in body:
|
||||
dropped[field] = body.pop(field)
|
||||
if dropped:
|
||||
logger.debug(
|
||||
"Dropped anthropic-only fields before gemini compat dispatch",
|
||||
extra={"dropped_keys": sorted(dropped.keys())},
|
||||
)
|
||||
|
||||
requested_model = (
|
||||
(model_obj.forwarded_model_id or model_obj.id) if model_obj else None
|
||||
)
|
||||
upstream_model = transform_model_name(model_obj.id)
|
||||
|
||||
openai_kwargs = _translate_anthropic_to_openai(body, upstream_model)
|
||||
|
||||
messages = openai_kwargs.get("messages") or []
|
||||
if isinstance(messages, list):
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
# Disable Gemini's thinking pass; the dummy signature already lifts
|
||||
# validation, but skipping thinking entirely avoids degraded model
|
||||
# output and keeps tool-calling deterministic.
|
||||
openai_kwargs.setdefault("reasoning_effort", "none")
|
||||
openai_kwargs["stream"] = True
|
||||
openai_kwargs["model"] = upstream_model
|
||||
# OpenAI-compat backends (including Gemini's) only emit a final
|
||||
# ``usage`` chunk when the request opts in via this flag. Without it
|
||||
# the cost-calculation pipeline can't read real token counts and
|
||||
# falls back to MaxCostData billing.
|
||||
existing_stream_options = openai_kwargs.get("stream_options")
|
||||
merged_stream_options = (
|
||||
dict(existing_stream_options)
|
||||
if isinstance(existing_stream_options, dict)
|
||||
else {}
|
||||
)
|
||||
merged_stream_options.setdefault("include_usage", True)
|
||||
openai_kwargs["stream_options"] = merged_stream_options
|
||||
|
||||
logger.info(
|
||||
"Dispatching /v1/messages via gemini compat (httpx)",
|
||||
extra={
|
||||
"model": upstream_model,
|
||||
"client_stream": client_stream,
|
||||
"messages_with_tool_calls": sum(
|
||||
1 for m in messages if isinstance(m, dict) and m.get("tool_calls")
|
||||
),
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
|
||||
http_client, response = await _post_and_stream(
|
||||
base_url, api_key, openai_kwargs, log_extra
|
||||
)
|
||||
|
||||
async def line_iter() -> AsyncGenerator[str, None]:
|
||||
try:
|
||||
async for line in response.aiter_lines():
|
||||
yield line
|
||||
finally:
|
||||
await response.aclose()
|
||||
await http_client.aclose()
|
||||
|
||||
anthropic_event_iter = _openai_chunks_to_anthropic_events(
|
||||
line_iter(), requested_model
|
||||
)
|
||||
|
||||
if not client_stream:
|
||||
# Aggregate the Anthropic SSE byte stream into a single Message dict
|
||||
# so the rest of the pipeline (cost calc, metadata injection,
|
||||
# response building) can treat it identically to a non-streaming
|
||||
# litellm response.
|
||||
try:
|
||||
aggregated = await aggregate_anthropic_events_to_message(
|
||||
anthropic_event_iter
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to aggregate Gemini compat events into message",
|
||||
extra={"error": str(exc), **(log_extra or {})},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to aggregate upstream stream: {exc}",
|
||||
status_code=502,
|
||||
) from exc
|
||||
return client_stream, aggregated, requested_model
|
||||
|
||||
return client_stream, anthropic_event_iter, requested_model
|
||||
@@ -12,6 +12,7 @@ class GroqUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "groq"
|
||||
default_base_url = "https://api.groq.com/openai/v1"
|
||||
platform_url = "https://console.groq.com/keys"
|
||||
litellm_provider_prefix = "groq/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.settings import Settings
|
||||
@@ -122,12 +122,16 @@ async def get_all_models_with_overrides(
|
||||
|
||||
|
||||
async def refresh_upstreams_models_periodically(
|
||||
upstreams: list[BaseUpstreamProvider],
|
||||
upstreams_provider: (
|
||||
Callable[[], list[BaseUpstreamProvider]] | list[BaseUpstreamProvider]
|
||||
),
|
||||
) -> None:
|
||||
"""Background task to periodically refresh models cache for all providers.
|
||||
|
||||
Args:
|
||||
upstreams: List of upstream provider instances
|
||||
upstreams_provider: Either a callable returning the live upstream list
|
||||
(preferred — picks up providers added/changed via reinitialize_upstreams),
|
||||
or a static list (legacy, will go stale after reinitialize_upstreams).
|
||||
"""
|
||||
import asyncio
|
||||
import random
|
||||
@@ -139,9 +143,14 @@ async def refresh_upstreams_models_periodically(
|
||||
logger.info("Provider models refresh disabled (interval <= 0)")
|
||||
return
|
||||
|
||||
def _resolve_upstreams() -> list[BaseUpstreamProvider]:
|
||||
if callable(upstreams_provider):
|
||||
return upstreams_provider()
|
||||
return upstreams_provider
|
||||
|
||||
while True:
|
||||
try:
|
||||
for upstream in upstreams:
|
||||
for upstream in _resolve_upstreams():
|
||||
try:
|
||||
await upstream.refresh_models_cache()
|
||||
except Exception as e:
|
||||
@@ -188,13 +197,14 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
existing_providers = result.all()
|
||||
|
||||
if not existing_providers:
|
||||
logger.info(
|
||||
"No upstream providers found in database, seeding from settings"
|
||||
)
|
||||
await _seed_providers_from_settings(session, settings)
|
||||
await session.commit()
|
||||
result = await session.exec(select(UpstreamProviderRow))
|
||||
existing_providers = result.all()
|
||||
if existing_providers:
|
||||
logger.info(
|
||||
f"Seeded {len(existing_providers)} upstream providers from settings"
|
||||
)
|
||||
|
||||
async def _init_single_provider(
|
||||
provider_row: UpstreamProviderRow,
|
||||
@@ -205,6 +215,9 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
|
||||
provider = _instantiate_provider(provider_row)
|
||||
if provider:
|
||||
# Keep provider DB id on runtime instance so model mapping can
|
||||
# bind DB overrides to the correct upstream.
|
||||
setattr(provider, "db_id", provider_row.id)
|
||||
await provider.refresh_models_cache()
|
||||
logger.debug(
|
||||
f"Initialized {provider_row.provider_type} provider",
|
||||
|
||||
162
routstr/upstream/litellm_routing.py
Normal file
162
routstr/upstream/litellm_routing.py
Normal file
@@ -0,0 +1,162 @@
|
||||
"""Map an upstream `base_url` to the correct litellm provider prefix.
|
||||
|
||||
Used by `BaseUpstreamProvider.get_litellm_provider_prefix` so that custom /
|
||||
generic provider rows (which inherit the base class default) get routed to
|
||||
the right litellm backend instead of falling back to `openai/`.
|
||||
|
||||
The table is compiled from litellm 1.74's
|
||||
`litellm/litellm_core_utils/get_llm_provider_logic.py` (the
|
||||
`openai_compatible_endpoints` table) plus the providers documented at
|
||||
https://docs.litellm.ai/docs/providers. Substring match is used so that
|
||||
URLs with paths, ports, regional subdomains, etc. all resolve correctly.
|
||||
|
||||
Order matters: more specific needles must appear before more generic ones
|
||||
(e.g. ``openai.azure.com`` before ``api.openai.com``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import litellm
|
||||
|
||||
DEFAULT_PREFIX = "openai/"
|
||||
|
||||
LITELLM_HOST_PREFIX_MAP: tuple[tuple[str, str], ...] = (
|
||||
# Azure must win over api.openai.com because the host ends with
|
||||
# `openai.azure.com` and we don't want it picked up as plain OpenAI.
|
||||
("openai.azure.com", "azure/"),
|
||||
# Google
|
||||
("generativelanguage.googleapis.com", "gemini/"),
|
||||
("aiplatform.googleapis.com", "vertex_ai/"),
|
||||
# First-class providers with native litellm prefixes
|
||||
("api.openai.com", "openai/"),
|
||||
("api.anthropic.com", "anthropic/"),
|
||||
("api.groq.com", "groq/"),
|
||||
("api.fireworks.ai", "fireworks_ai/"),
|
||||
("api.x.ai", "xai/"),
|
||||
("api.perplexity.ai", "perplexity/"),
|
||||
("openrouter.ai", "openrouter/"),
|
||||
("api.deepseek.com", "deepseek/"),
|
||||
("api.together.xyz", "together_ai/"),
|
||||
("codestral.mistral.ai", "codestral/"),
|
||||
("api.mistral.ai", "mistral/"),
|
||||
("api.cohere.com", "cohere_chat/"),
|
||||
("api.cohere.ai", "cohere_chat/"),
|
||||
("api.deepinfra.com", "deepinfra/"),
|
||||
("api.endpoints.anyscale.com", "anyscale/"),
|
||||
("api.cerebras.ai", "cerebras/"),
|
||||
("inference.baseten.co", "baseten/"),
|
||||
("api.sambanova.ai", "sambanova/"),
|
||||
("api.ai21.com", "ai21_chat/"),
|
||||
("api.friendli.ai", "friendliai/"),
|
||||
("api.galadriel.com", "galadriel/"),
|
||||
("api.llama.com", "meta_llama/"),
|
||||
("api.featherless.ai", "featherless_ai/"),
|
||||
("inference.api.nscale.com", "nscale/"),
|
||||
("dashscope-intl.aliyuncs.com", "dashscope/"),
|
||||
("api.moonshot.ai", "moonshot/"),
|
||||
("api.moonshot.cn", "moonshot/"),
|
||||
("api.minimax.io", "minimax/"),
|
||||
("api.minimaxi.com", "minimax/"),
|
||||
("platform.publicai.co", "publicai/"),
|
||||
("api.synthetic.new", "synthetic/"),
|
||||
("api.stima.tech", "apertis/"),
|
||||
("nano-gpt.com", "nano-gpt/"),
|
||||
("api.poe.com", "poe/"),
|
||||
("llm.chutes.ai", "chutes/"),
|
||||
("api.v0.dev", "v0/"),
|
||||
("api.lambda.ai", "lambda_ai/"),
|
||||
("api.hyperbolic.xyz", "hyperbolic/"),
|
||||
("ai-gateway.vercel.sh", "vercel_ai_gateway/"),
|
||||
("api.inference.wandb.ai", "wandb/"),
|
||||
("integrate.api.nvidia.com", "nvidia_nim/"),
|
||||
("api.studio.nebius.com", "nebius/"),
|
||||
("api.novita.ai", "novita/"),
|
||||
("ark.cn-beijing.volces.com", "volcengine/"),
|
||||
("api.voyageai.com", "voyage/"),
|
||||
("api.jina.ai", "jina_ai/"),
|
||||
("api.aimlapi.com", "aiml/"),
|
||||
("api.snowflakecomputing.com", "snowflake/"),
|
||||
("databricks.com", "databricks/"),
|
||||
("huggingface.co", "huggingface/"),
|
||||
)
|
||||
|
||||
# Substrings that indicate an Ollama deployment regardless of port/scheme.
|
||||
OLLAMA_HOST_HINTS: tuple[str, ...] = (
|
||||
"localhost:11434",
|
||||
"127.0.0.1:11434",
|
||||
"ollama",
|
||||
)
|
||||
|
||||
|
||||
def detect_litellm_prefix(
|
||||
base_url: str | None, default: str = DEFAULT_PREFIX
|
||||
) -> str:
|
||||
"""Return the litellm provider prefix (`"<provider>/"`) for `base_url`.
|
||||
|
||||
Falls back to `default` when the host doesn't match any known provider.
|
||||
The default is `openai/` because every unmatched OpenAI-compatible
|
||||
server is, by definition, an OpenAI-compatible server.
|
||||
"""
|
||||
if not base_url:
|
||||
return default
|
||||
|
||||
parsed = urlsplit(base_url)
|
||||
host = parsed.netloc.lower() or base_url.lower()
|
||||
|
||||
for needle, prefix in LITELLM_HOST_PREFIX_MAP:
|
||||
if needle in host:
|
||||
return prefix
|
||||
|
||||
if any(hint in host for hint in OLLAMA_HOST_HINTS):
|
||||
return "ollama_chat/"
|
||||
|
||||
return default
|
||||
|
||||
|
||||
_configured = False
|
||||
|
||||
|
||||
def configure_litellm() -> None:
|
||||
"""Apply litellm global settings used by the messages-dispatch path.
|
||||
|
||||
Idempotent: safe to call from both app startup and module-level
|
||||
initializers without side effects on the second invocation.
|
||||
|
||||
Settings applied:
|
||||
|
||||
* ``LITELLM_DEBUG=1`` enables litellm's verbose debug logger.
|
||||
* Forces the Anthropic-messages adapter to call OpenAI Chat Completions
|
||||
(POST ``/chat/completions``) instead of the Responses API (POST
|
||||
``/responses``) for ``openai/``-prefixed providers. OpenAI-compatible
|
||||
upstreams like Google's generativelanguage compat endpoint expose
|
||||
``/chat/completions`` but not ``/responses``, which would 404. Set
|
||||
``LITELLM_USE_RESPONSES_API_FOR_ANTHROPIC_MESSAGES=1`` to opt out.
|
||||
* Silently drops Anthropic-Messages-only parameters (``thinking``,
|
||||
``cache_control``, ``context_management``, ...) when translating to
|
||||
providers that don't accept them, instead of raising
|
||||
``UnsupportedParamsError``. Set ``LITELLM_STRICT_PARAMS=1`` to opt
|
||||
out.
|
||||
"""
|
||||
global _configured
|
||||
if _configured:
|
||||
return
|
||||
|
||||
if os.getenv("LITELLM_DEBUG") == "1":
|
||||
try:
|
||||
litellm._turn_on_debug() # type: ignore[no-untyped-call]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if os.getenv("LITELLM_USE_RESPONSES_API_FOR_ANTHROPIC_MESSAGES") != "1":
|
||||
try:
|
||||
litellm.use_chat_completions_url_for_anthropic_messages = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if os.getenv("LITELLM_STRICT_PARAMS") != "1":
|
||||
litellm.drop_params = True
|
||||
|
||||
_configured = True
|
||||
561
routstr/upstream/messages_dispatch.py
Normal file
561
routstr/upstream/messages_dispatch.py
Normal file
@@ -0,0 +1,561 @@
|
||||
"""Pure helpers for translating ``/v1/messages`` to upstream chat completions
|
||||
via litellm.
|
||||
|
||||
This module owns the litellm/Anthropic-Messages translation layer:
|
||||
|
||||
* SSE parsing (``parse_sse_blocks``, ``events_from_chunk``)
|
||||
* Payload coercion (``coerce_litellm_payload``)
|
||||
* Stream aggregation (``aggregate_anthropic_events_to_message``) — drains
|
||||
a streamed Anthropic event sequence into a single Message dict
|
||||
* Per-event annotation for streaming (``annotate_event``,
|
||||
``stream_annotated_events``) — handles the model-rewrite + token-tally
|
||||
bookkeeping shared by the bearer-key and x-cashu streaming paths
|
||||
* The dispatch entry point (``dispatch_anthropic_messages``)
|
||||
* Refund math (``compute_refund``)
|
||||
|
||||
Nothing in here touches ``BaseUpstreamProvider``; the thin instance methods
|
||||
on the provider class forward to these functions and only retain logic that
|
||||
genuinely needs ``self`` (cost adjustment, metadata injection, refund
|
||||
sending).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from typing import Any, Callable, NamedTuple, cast
|
||||
|
||||
import litellm
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.exceptions import UpstreamError
|
||||
from ..payment.models import Model
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Anthropic-Messages-only fields that don't translate to OpenAI
|
||||
# Chat Completions. ``litellm.drop_params`` only filters *known*
|
||||
# unsupported params; these newer/extension fields get passed through
|
||||
# verbatim and the upstream rejects them with a 400. Pop them here so the
|
||||
# request reaches the upstream cleanly.
|
||||
ANTHROPIC_ONLY_FIELDS: tuple[str, ...] = (
|
||||
"thinking",
|
||||
"cache_control",
|
||||
"context_management",
|
||||
"output_config",
|
||||
"mcp_servers",
|
||||
"service_tier",
|
||||
"anthropic_version",
|
||||
"anthropic_beta",
|
||||
)
|
||||
|
||||
|
||||
def coerce_litellm_payload(payload: object) -> dict:
|
||||
"""Convert a litellm event into a plain dict.
|
||||
|
||||
Non-streaming responses come back as Anthropic-shaped pydantic models
|
||||
or dicts. Streaming may yield raw bytes/str (SSE-encoded); those go
|
||||
through ``events_from_chunk`` instead, not here.
|
||||
"""
|
||||
if isinstance(payload, dict):
|
||||
return dict(payload)
|
||||
if hasattr(payload, "model_dump"):
|
||||
return cast(dict, payload.model_dump())
|
||||
raise TypeError(f"Cannot coerce {type(payload).__name__} to dict")
|
||||
|
||||
|
||||
def parse_sse_blocks(buffer: bytes) -> tuple[list[dict], bytes]:
|
||||
"""Parse complete SSE event blocks out of a byte buffer.
|
||||
|
||||
Returns (events, remaining_buffer). Events are JSON objects parsed from
|
||||
one or more ``data:`` lines per block. Comments, blank lines, and
|
||||
``[DONE]`` sentinels are ignored. A trailing partial block is preserved
|
||||
in remaining_buffer.
|
||||
"""
|
||||
events: list[dict] = []
|
||||
while True:
|
||||
sep = buffer.find(b"\n\n")
|
||||
if sep < 0:
|
||||
sep_rn = buffer.find(b"\r\n\r\n")
|
||||
if sep_rn < 0:
|
||||
break
|
||||
block = buffer[:sep_rn]
|
||||
buffer = buffer[sep_rn + 4 :]
|
||||
else:
|
||||
block = buffer[:sep]
|
||||
buffer = buffer[sep + 2 :]
|
||||
|
||||
data_lines: list[str] = []
|
||||
for raw_line in block.replace(b"\r\n", b"\n").split(b"\n"):
|
||||
line = raw_line.decode("utf-8", errors="replace")
|
||||
if line.startswith(":"):
|
||||
continue
|
||||
if line.startswith("data:"):
|
||||
data_lines.append(line[5:].lstrip())
|
||||
if not data_lines:
|
||||
continue
|
||||
payload = "\n".join(data_lines).strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(obj, dict):
|
||||
events.append(obj)
|
||||
return events, buffer
|
||||
|
||||
|
||||
def events_from_chunk(
|
||||
chunk: object, sse_buffer: bytes
|
||||
) -> tuple[list[dict], bytes]:
|
||||
"""Normalize a stream chunk into one or more event dicts.
|
||||
|
||||
``litellm.anthropic.messages.acreate(stream=True)`` yields raw SSE
|
||||
bytes in practice; some adapters yield strings or typed events. Handle
|
||||
all three.
|
||||
"""
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
sse_buffer += bytes(chunk)
|
||||
events, sse_buffer = parse_sse_blocks(sse_buffer)
|
||||
return events, sse_buffer
|
||||
if isinstance(chunk, str):
|
||||
sse_buffer += chunk.encode("utf-8")
|
||||
events, sse_buffer = parse_sse_blocks(sse_buffer)
|
||||
return events, sse_buffer
|
||||
return [coerce_litellm_payload(chunk)], sse_buffer
|
||||
|
||||
|
||||
async def aggregate_anthropic_events_to_message(
|
||||
iterator: AsyncIterator[Any],
|
||||
) -> dict:
|
||||
"""Drain an Anthropic-Messages event iterator into a single Message dict.
|
||||
|
||||
Produces the shape ``litellm.anthropic.messages.acreate(stream=False)``
|
||||
would have returned. Used to transparently stream from upstream while
|
||||
still returning a non-streaming response to the client. Lets us
|
||||
sidestep upstream quirks (e.g. Fireworks rejects ``max_tokens > 4096``
|
||||
unless ``stream=true``) without leaking that into client-visible
|
||||
behavior.
|
||||
"""
|
||||
sse_buffer = b""
|
||||
message: dict = {}
|
||||
blocks: list[dict] = []
|
||||
partial_json: dict[int, str] = {}
|
||||
final_stop_reason: str | None = None
|
||||
final_stop_sequence: str | None = None
|
||||
final_usage: dict[str, Any] = {}
|
||||
final_model: str | None = None
|
||||
|
||||
async for chunk in iterator:
|
||||
events, sse_buffer = events_from_chunk(chunk, sse_buffer)
|
||||
for event in events:
|
||||
etype = event.get("type")
|
||||
if etype == "message_start":
|
||||
raw = event.get("message") or {}
|
||||
if isinstance(raw, dict):
|
||||
message = dict(raw)
|
||||
existing = message.get("content")
|
||||
blocks = list(existing) if isinstance(existing, list) else []
|
||||
usage = message.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
final_usage = dict(usage)
|
||||
if isinstance(message.get("model"), str):
|
||||
final_model = message["model"]
|
||||
elif etype == "content_block_start":
|
||||
idx = int(event.get("index") or 0)
|
||||
cb = event.get("content_block") or {}
|
||||
cb_dict = dict(cb) if isinstance(cb, dict) else {}
|
||||
while len(blocks) <= idx:
|
||||
blocks.append({})
|
||||
blocks[idx] = cb_dict
|
||||
elif etype == "content_block_delta":
|
||||
idx = int(event.get("index") or 0)
|
||||
if idx >= len(blocks):
|
||||
continue
|
||||
delta = event.get("delta") or {}
|
||||
if not isinstance(delta, dict):
|
||||
continue
|
||||
dtype = delta.get("type")
|
||||
block = blocks[idx]
|
||||
if dtype == "text_delta":
|
||||
block["text"] = (block.get("text") or "") + (
|
||||
delta.get("text") or ""
|
||||
)
|
||||
elif dtype == "input_json_delta":
|
||||
partial_json[idx] = partial_json.get(idx, "") + (
|
||||
delta.get("partial_json") or ""
|
||||
)
|
||||
elif dtype == "thinking_delta":
|
||||
block["thinking"] = (block.get("thinking") or "") + (
|
||||
delta.get("thinking") or ""
|
||||
)
|
||||
elif dtype == "signature_delta":
|
||||
block["signature"] = (block.get("signature") or "") + (
|
||||
delta.get("signature") or ""
|
||||
)
|
||||
elif etype == "content_block_stop":
|
||||
idx = int(event.get("index") or 0)
|
||||
raw_json = partial_json.pop(idx, None)
|
||||
if raw_json is not None and idx < len(blocks):
|
||||
try:
|
||||
blocks[idx]["input"] = (
|
||||
json.loads(raw_json) if raw_json else {}
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
blocks[idx]["input"] = raw_json
|
||||
elif etype == "message_delta":
|
||||
delta = event.get("delta") or {}
|
||||
if isinstance(delta, dict):
|
||||
if "stop_reason" in delta:
|
||||
final_stop_reason = delta.get("stop_reason")
|
||||
if "stop_sequence" in delta:
|
||||
final_stop_sequence = delta.get("stop_sequence")
|
||||
usage = event.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
final_usage.update(usage)
|
||||
# message_stop: nothing to merge
|
||||
|
||||
if not message:
|
||||
# Upstream returned no message_start; expose what we can so the
|
||||
# client at least sees the assembled content.
|
||||
message = {
|
||||
"id": "",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
}
|
||||
|
||||
message["content"] = blocks
|
||||
if final_model and not message.get("model"):
|
||||
message["model"] = final_model
|
||||
if final_stop_reason is not None:
|
||||
message["stop_reason"] = final_stop_reason
|
||||
if final_stop_sequence is not None:
|
||||
message["stop_sequence"] = final_stop_sequence
|
||||
if final_usage:
|
||||
existing_usage = message.get("usage")
|
||||
merged = dict(existing_usage) if isinstance(existing_usage, dict) else {}
|
||||
merged.update(final_usage)
|
||||
message["usage"] = merged
|
||||
return message
|
||||
|
||||
|
||||
class AnnotatedEvent(NamedTuple):
|
||||
"""One Anthropic SSE event after model-rewrite + token-tally bookkeeping.
|
||||
|
||||
``sse_bytes`` is the wire-ready ``event:`` / ``data:`` block; the two
|
||||
streaming paths in ``BaseUpstreamProvider`` consume ``sse_bytes`` plus
|
||||
the tallies and only differ in whether they stream live or buffer
|
||||
first.
|
||||
|
||||
``cache_read_input_tokens`` and ``cache_creation_input_tokens`` are
|
||||
surfaced separately so the cost path can price them against the cache
|
||||
rate rather than fold them silently into the regular input bucket.
|
||||
|
||||
``total_cost`` / ``input_cost`` / ``output_cost`` carry any
|
||||
USD cost figures the upstream attached to this event (from
|
||||
``usage.cost``, ``usage.total_cost``, or ``usage.cost_details``) so the
|
||||
streaming paths can re-embed them in the rebuilt ``usage`` dict and let
|
||||
``calculate_cost`` convert directly USD→sats instead of falling back to
|
||||
token-based math.
|
||||
"""
|
||||
|
||||
event: dict
|
||||
sse_bytes: bytes
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
total_cost: float
|
||||
input_cost: float
|
||||
output_cost: float
|
||||
model: str | None
|
||||
|
||||
|
||||
def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent:
|
||||
"""Rewrite ``model`` fields and extract per-event token / model info.
|
||||
|
||||
Mutates ``event`` in place when ``requested_model`` is set so the
|
||||
upstream's true model name doesn't leak to the client.
|
||||
"""
|
||||
if requested_model:
|
||||
msg = event.get("message")
|
||||
if isinstance(msg, dict) and "model" in msg:
|
||||
msg["model"] = requested_model
|
||||
if "model" in event:
|
||||
event["model"] = requested_model
|
||||
|
||||
in_tokens = 0
|
||||
out_tokens = 0
|
||||
cache_read_tokens = 0
|
||||
cache_create_tokens = 0
|
||||
total_cost = 0.0
|
||||
input_cost = 0.0
|
||||
output_cost = 0.0
|
||||
model: str | None = None
|
||||
|
||||
def _coerce_float(value: object) -> float:
|
||||
if value is None or isinstance(value, bool):
|
||||
return 0.0
|
||||
if not isinstance(value, (int, float, str)):
|
||||
return 0.0
|
||||
try:
|
||||
return max(0.0, float(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
def _accumulate(usage: dict) -> None:
|
||||
nonlocal in_tokens, out_tokens, cache_read_tokens, cache_create_tokens
|
||||
nonlocal total_cost, input_cost, output_cost
|
||||
in_tokens += int(usage.get("input_tokens") or 0)
|
||||
out_tokens += int(usage.get("output_tokens") or 0)
|
||||
cache_read_tokens += int(usage.get("cache_read_input_tokens") or 0)
|
||||
cache_create_tokens += int(usage.get("cache_creation_input_tokens") or 0)
|
||||
total_cost += _coerce_float(usage.get("total_cost"))
|
||||
input_cost += _coerce_float(usage.get("input_cost"))
|
||||
output_cost += _coerce_float(usage.get("output_cost"))
|
||||
|
||||
msg_for_meta = event.get("message")
|
||||
if isinstance(msg_for_meta, dict):
|
||||
if msg_for_meta.get("model"):
|
||||
model = str(msg_for_meta["model"])
|
||||
usage = msg_for_meta.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
_accumulate(usage)
|
||||
|
||||
if isinstance(event.get("usage"), dict):
|
||||
_accumulate(event["usage"])
|
||||
|
||||
# Some upstreams (notably OpenRouter-style proxies) attach cost fields
|
||||
# directly at the event root rather than inside ``usage``.
|
||||
for field in ("total_cost", "cost"):
|
||||
total_cost = max(total_cost, _coerce_float(event.get(field)))
|
||||
input_cost = max(input_cost, _coerce_float(event.get("input_cost")))
|
||||
output_cost = max(output_cost, _coerce_float(event.get("output_cost")))
|
||||
root_cost_details = event.get("cost_details")
|
||||
if isinstance(root_cost_details, dict):
|
||||
total_cost = max(
|
||||
total_cost,
|
||||
_coerce_float(root_cost_details.get("total_cost")),
|
||||
)
|
||||
input_cost = max(
|
||||
input_cost,
|
||||
_coerce_float(root_cost_details.get("input_cost")),
|
||||
)
|
||||
output_cost = max(
|
||||
output_cost,
|
||||
_coerce_float(root_cost_details.get("output_cost")),
|
||||
)
|
||||
|
||||
event_type = str(event.get("type") or "")
|
||||
payload = json.dumps(event)
|
||||
if event_type:
|
||||
sse_bytes = f"event: {event_type}\ndata: {payload}\n\n".encode()
|
||||
else:
|
||||
sse_bytes = f"data: {payload}\n\n".encode()
|
||||
|
||||
return AnnotatedEvent(
|
||||
event,
|
||||
sse_bytes,
|
||||
in_tokens,
|
||||
out_tokens,
|
||||
cache_read_tokens,
|
||||
cache_create_tokens,
|
||||
total_cost,
|
||||
input_cost,
|
||||
output_cost,
|
||||
model,
|
||||
)
|
||||
|
||||
|
||||
async def stream_annotated_events(
|
||||
iterator: AsyncIterator[Any],
|
||||
requested_model: str | None,
|
||||
) -> AsyncGenerator[AnnotatedEvent, None]:
|
||||
"""Yield annotated, SSE-serialized events from a litellm stream.
|
||||
|
||||
Both streaming paths in ``BaseUpstreamProvider`` consume this; the only
|
||||
divergence between them — yield-as-you-go vs buffer-then-replay — stays
|
||||
in the caller.
|
||||
"""
|
||||
sse_buffer = b""
|
||||
async for chunk in iterator:
|
||||
events, sse_buffer = events_from_chunk(chunk, sse_buffer)
|
||||
for event in events:
|
||||
yield annotate_event(event, requested_model)
|
||||
|
||||
|
||||
def embed_usd_costs(
|
||||
usage: dict,
|
||||
total_cost: float,
|
||||
input_cost: float,
|
||||
output_cost: float,
|
||||
) -> None:
|
||||
"""Mutate ``usage`` so ``calculate_cost`` will pick up the USD totals.
|
||||
|
||||
Mirrors the upstream shape: when any USD figure is present, attach
|
||||
``cost`` (used by the simple-fallback branch in ``calculate_cost``) and
|
||||
a ``cost_details`` block (used by the preferred branch — also gives the
|
||||
input/output USD split when we have one).
|
||||
"""
|
||||
if total_cost <= 0 and input_cost <= 0 and output_cost <= 0:
|
||||
return
|
||||
|
||||
cost_details: dict[str, float] = {}
|
||||
effective_total = total_cost
|
||||
if effective_total <= 0 and (input_cost > 0 or output_cost > 0):
|
||||
effective_total = input_cost + output_cost
|
||||
|
||||
if effective_total > 0:
|
||||
cost_details["total_cost"] = effective_total
|
||||
usage["cost"] = effective_total
|
||||
if input_cost > 0:
|
||||
cost_details["input_cost"] = input_cost
|
||||
if output_cost > 0:
|
||||
cost_details["output_cost"] = output_cost
|
||||
if cost_details:
|
||||
usage["cost_details"] = cost_details
|
||||
|
||||
|
||||
def compute_refund(amount: int, unit: str, cost_msats: int) -> int:
|
||||
if unit == "msat":
|
||||
return amount - cost_msats
|
||||
if unit == "sat":
|
||||
return amount - (cost_msats + 999) // 1000
|
||||
raise ValueError(f"Invalid unit: {unit}")
|
||||
|
||||
|
||||
async def dispatch_anthropic_messages(
|
||||
*,
|
||||
request_body: bytes | None,
|
||||
model_obj: Model,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
provider_prefix: str,
|
||||
transform_model_name: Callable[[str], str],
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Call ``litellm.anthropic.messages.acreate`` and return
|
||||
``(client_stream, result, requested_model)``.
|
||||
|
||||
Shared by the bearer-key and x-cashu paths. Raises :class:`UpstreamError`
|
||||
on bad input or upstream failure.
|
||||
"""
|
||||
if not request_body:
|
||||
raise UpstreamError(
|
||||
"Missing request body for /v1/messages", status_code=400
|
||||
)
|
||||
|
||||
try:
|
||||
body: dict = json.loads(request_body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise UpstreamError(
|
||||
f"Invalid JSON in /v1/messages body: {exc}", status_code=400
|
||||
) from exc
|
||||
|
||||
body.pop("model", None)
|
||||
# `stream` here is what the **client** asked for. Upstream is always
|
||||
# streamed (see `upstream_stream` below); when the client asked for a
|
||||
# non-streaming response we drain and aggregate the events into a
|
||||
# single Anthropic Message dict before returning. This sidesteps
|
||||
# provider-specific non-streaming caps (e.g. Fireworks rejects
|
||||
# `max_tokens > 4096` unless `stream=true`).
|
||||
client_stream = bool(body.pop("stream", False))
|
||||
upstream_stream = True
|
||||
|
||||
dropped: dict[str, Any] = {}
|
||||
for field in ANTHROPIC_ONLY_FIELDS:
|
||||
if field in body:
|
||||
dropped[field] = body.pop(field)
|
||||
if dropped:
|
||||
logger.debug(
|
||||
"Dropped anthropic-only fields before litellm dispatch",
|
||||
extra={"dropped_keys": sorted(dropped.keys())},
|
||||
)
|
||||
|
||||
# Convention: `model.id` is the canonical upstream model name;
|
||||
# `forwarded_model_id` is the public alias the internal API exposes
|
||||
# and echoes back to the client.
|
||||
requested_model = (
|
||||
(model_obj.forwarded_model_id or model_obj.id) if model_obj else None
|
||||
)
|
||||
upstream_model = transform_model_name(model_obj.id)
|
||||
litellm_model = f"{provider_prefix}{upstream_model}"
|
||||
|
||||
kwargs: dict = {
|
||||
"model": litellm_model,
|
||||
"api_base": base_url,
|
||||
"api_key": api_key,
|
||||
"stream": upstream_stream,
|
||||
**body,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Dispatching /v1/messages via litellm",
|
||||
extra={
|
||||
"model": litellm_model,
|
||||
"resolved_provider": provider_prefix.rstrip("/"),
|
||||
"client_stream": client_stream,
|
||||
"upstream_stream": upstream_stream,
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
result = await litellm.anthropic.messages.acreate(**kwargs)
|
||||
except Exception as exc:
|
||||
exc_message = getattr(exc, "message", None) or str(exc) or repr(exc)
|
||||
exc_status = getattr(exc, "status_code", None)
|
||||
exc_response = getattr(exc, "response", None)
|
||||
response_text = None
|
||||
if exc_response is not None:
|
||||
try:
|
||||
response_text = getattr(exc_response, "text", str(exc_response))
|
||||
except Exception:
|
||||
response_text = "<unreadable>"
|
||||
logger.error(
|
||||
"litellm dispatch failed",
|
||||
extra={
|
||||
"error": exc_message,
|
||||
"error_type": type(exc).__name__,
|
||||
"status_code": exc_status,
|
||||
"llm_provider": getattr(exc, "llm_provider", None),
|
||||
"body": getattr(exc, "body", None),
|
||||
"response_text": response_text,
|
||||
"model": litellm_model,
|
||||
"api_base": base_url,
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Upstream error via litellm: {exc_message}",
|
||||
status_code=exc_status if isinstance(exc_status, int) else 502,
|
||||
) from exc
|
||||
|
||||
if not client_stream and hasattr(result, "__aiter__"):
|
||||
# Client asked for a non-streaming response but we always stream
|
||||
# from upstream — drain the events into a single Anthropic Message
|
||||
# dict so the rest of the pipeline can treat it as if upstream had
|
||||
# returned non-streaming. Some litellm adapters return a
|
||||
# non-streaming dict even when ``stream=True``; in that case,
|
||||
# leave the result as-is.
|
||||
try:
|
||||
aggregated: Any = await aggregate_anthropic_events_to_message(
|
||||
cast(AsyncIterator[Any], result)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to aggregate streamed events into message",
|
||||
extra={
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
"model": litellm_model,
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to aggregate upstream stream: {exc}",
|
||||
status_code=502,
|
||||
) from exc
|
||||
return client_stream, aggregated, requested_model
|
||||
|
||||
return client_stream, result, requested_model
|
||||
@@ -3,13 +3,11 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
from ..core.logging import get_logger
|
||||
@@ -23,6 +21,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "ollama"
|
||||
default_base_url = "http://localhost:11434"
|
||||
platform_url = None
|
||||
litellm_provider_prefix = "ollama_chat/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -67,38 +66,11 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Strip 'ollama/' prefix for Ollama API compatibility."""
|
||||
return model_id.removeprefix("ollama/")
|
||||
|
||||
async def forward_request(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Override to use OpenAI-compatible endpoint for proxy requests."""
|
||||
if path.startswith("v1/"):
|
||||
path = path.replace("v1/", "")
|
||||
|
||||
original_base_url = self.base_url
|
||||
self.base_url = f"{self.base_url}/v1"
|
||||
|
||||
try:
|
||||
result = await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
return result
|
||||
finally:
|
||||
self.base_url = original_base_url
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: Model | None = None
|
||||
) -> str:
|
||||
"""Route proxy traffic through Ollama's OpenAI-compatible /v1 endpoint."""
|
||||
return f"{self.base_url.rstrip('/')}/v1"
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
"""Fetch models from Ollama API using /api/tags endpoint."""
|
||||
@@ -213,7 +185,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
except Exception:
|
||||
self._models_cache = models_with_fees
|
||||
|
||||
self._models_by_id = {m.id: m for m in self._models_cache}
|
||||
self._models_by_id = {m.forwarded_model_id or m.id: m for m in self._models_cache}
|
||||
logger.info(
|
||||
f"Refreshed models cache for {self.base_url}",
|
||||
extra={"model_count": len(models)},
|
||||
|
||||
@@ -15,6 +15,35 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "openrouter"
|
||||
default_base_url = "https://openrouter.ai/api/v1"
|
||||
platform_url = "https://openrouter.ai/settings/keys"
|
||||
supports_anthropic_messages = True
|
||||
litellm_provider_prefix = "openrouter/"
|
||||
|
||||
def _apply_provider_field(self, response_json: object) -> None:
|
||||
"""Stamp the ``provider`` field for OpenRouter responses.
|
||||
|
||||
OpenRouter is a router, not the real serving provider, so a bare
|
||||
``"openrouter"`` value carries no useful information. Rules:
|
||||
|
||||
- Real upstream sub-provider (e.g. ``"GMICloud"``) -> ``"openrouter:GMICloud"``.
|
||||
- Missing sub-provider, or one that merely echoes ``"openrouter"`` ->
|
||||
``"unknown"``.
|
||||
- Idempotent: re-stamping never produces ``"openrouter:openrouter:..."``;
|
||||
the ``openrouter:`` prefix appears at most once.
|
||||
"""
|
||||
if not isinstance(response_json, dict):
|
||||
return
|
||||
provider_type = (self.provider_type or "").strip()
|
||||
existing = response_json.get("provider")
|
||||
sub = existing.strip() if isinstance(existing, str) else ""
|
||||
# Strip any already-applied "openrouter:" prefixes (idempotency).
|
||||
prefix = f"{provider_type}:"
|
||||
while sub.lower().startswith(prefix.lower()):
|
||||
sub = sub[len(prefix) :].strip()
|
||||
# No real sub-provider, or it just echoes our own router name.
|
||||
if not sub or sub.lower() == provider_type.lower():
|
||||
response_json["provider"] = "unknown"
|
||||
return
|
||||
response_json["provider"] = f"{provider_type}:{sub}"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.06):
|
||||
"""Initialize OpenRouter provider with API key.
|
||||
|
||||
@@ -13,6 +13,7 @@ class PerplexityUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "perplexity"
|
||||
default_base_url = "https://api.perplexity.ai/"
|
||||
platform_url = "https://www.perplexity.ai/account/api/keys"
|
||||
litellm_provider_prefix = "perplexity/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
from pydantic.v1 import BaseModel, Field
|
||||
|
||||
from ..core.logging import get_logger
|
||||
from ..payment.models import Architecture, Model, Pricing, async_fetch_openrouter_models
|
||||
@@ -16,18 +16,20 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class PPQAIModelPricing(BaseModel):
|
||||
ui: dict[str, float]
|
||||
api: dict[str, float]
|
||||
ui: Optional[dict[str, float]] = None
|
||||
api: Optional[dict[str, float]] = None
|
||||
input_per_1M_tokens: Optional[float] = Field(None, alias="input_per_1M_tokens")
|
||||
output_per_1M_tokens: Optional[float] = Field(None, alias="output_per_1M_tokens")
|
||||
|
||||
|
||||
class PPQAIModel(BaseModel):
|
||||
id: str
|
||||
provider: str
|
||||
provider: Optional[str] = None
|
||||
name: str
|
||||
created_at: int
|
||||
context_length: int
|
||||
pricing: PPQAIModelPricing
|
||||
popular: bool
|
||||
popular: bool = False
|
||||
|
||||
|
||||
class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
@@ -134,31 +136,54 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
if or_model:
|
||||
if input_price := ppqai_model.pricing.api.get(
|
||||
"input_per_1M"
|
||||
):
|
||||
input_price = None
|
||||
if ppqai_model.pricing.api:
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M"
|
||||
)
|
||||
elif ppqai_model.pricing.input_per_1M_tokens:
|
||||
input_price = ppqai_model.pricing.input_per_1M_tokens
|
||||
|
||||
if input_price is not None:
|
||||
or_model.pricing.prompt = input_price / 1_000_000
|
||||
if output_price := ppqai_model.pricing.api.get(
|
||||
"output_per_1M"
|
||||
):
|
||||
|
||||
output_price = None
|
||||
if ppqai_model.pricing.api:
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M"
|
||||
)
|
||||
elif ppqai_model.pricing.output_per_1M_tokens:
|
||||
output_price = ppqai_model.pricing.output_per_1M_tokens
|
||||
|
||||
if output_price is not None:
|
||||
or_model.pricing.completion = output_price / 1_000_000
|
||||
|
||||
if cl := ppqai_model.context_length:
|
||||
or_model.context_length = cl
|
||||
models.append(or_model)
|
||||
else:
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M", 0.0
|
||||
)
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M", 0.0
|
||||
)
|
||||
input_price = 0.0
|
||||
if ppqai_model.pricing.api:
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M", 0.0
|
||||
)
|
||||
elif ppqai_model.pricing.input_per_1M_tokens:
|
||||
input_price = ppqai_model.pricing.input_per_1M_tokens
|
||||
|
||||
output_price = 0.0
|
||||
if ppqai_model.pricing.api:
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M", 0.0
|
||||
)
|
||||
elif ppqai_model.pricing.output_per_1M_tokens:
|
||||
output_price = ppqai_model.pricing.output_per_1M_tokens
|
||||
|
||||
models.append(
|
||||
Model(
|
||||
id=ppqai_model.id,
|
||||
name=ppqai_model.name,
|
||||
created=ppqai_model.created_at // 1000,
|
||||
description=f"{ppqai_model.provider} model",
|
||||
description=f"{ppqai_model.provider or 'PPQ.AI'} model",
|
||||
context_length=ppqai_model.context_length,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
@@ -283,7 +308,6 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
print(f"Payload: {payload}", "sending to", url)
|
||||
response = await client.post(url, headers=headers, json=payload)
|
||||
response.raise_for_status()
|
||||
invoice_data = response.json()
|
||||
|
||||
142
routstr/upstream/request_correction.py
Normal file
142
routstr/upstream/request_correction.py
Normal file
@@ -0,0 +1,142 @@
|
||||
"""Reactive request-correction layer.
|
||||
|
||||
When an upstream rejects a request with a recoverable 4xx error, this layer
|
||||
tries to *fix* the request body and let the caller retry the same upstream
|
||||
instead of failing outright. It is provider-agnostic: correctors key off the
|
||||
upstream's own error wording, so the same recovery works across every provider.
|
||||
|
||||
The layer is a small pipeline of :data:`Corrector` callables. Each corrector
|
||||
inspects the parsed request body and the upstream error message and either
|
||||
returns a corrected body (plus a short label identifying the fix) or declines
|
||||
by returning ``None``. Adding a new reactive fix means writing one corrector
|
||||
and adding it to :data:`DEFAULT_CORRECTORS` — no changes to the proxy loop.
|
||||
|
||||
All corrections are immutable: a corrector never mutates the body it is given,
|
||||
it returns a new ``dict``. The proxy threads an ``applied`` set of fix labels
|
||||
through retries so each distinct fix is applied at most once, guaranteeing the
|
||||
retry loop always terminates.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastapi.responses import Response
|
||||
|
||||
from ..core import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# Matches upstream error text that names a single rejected request parameter,
|
||||
# e.g. "`temperature` is deprecated for this model." or
|
||||
# "parameter 'top_p' is not supported". Keys off the upstream's own wording so
|
||||
# a 400 about an unsupported sampling/option field can be recovered by stripping
|
||||
# that field and retrying the same upstream.
|
||||
_UNSUPPORTED_PARAM_RE = re.compile(
|
||||
r"[`'\"]?(?P<param>[a-zA-Z_][a-zA-Z0-9_]*)[`'\"]?\s+is\s+"
|
||||
r"(?:deprecated|not\s+supported|unsupported|no\s+longer\s+supported)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
# A corrector inspects the parsed request body and the upstream error message
|
||||
# and returns ``(new_body_dict, label)`` for a fix it can apply, or ``None`` to
|
||||
# decline. ``label`` identifies the fix so it is applied at most once per request.
|
||||
Corrector = Callable[[dict, str], "tuple[dict, str] | None"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Correction:
|
||||
"""A successful request correction ready to retry.
|
||||
|
||||
``body`` is the corrected JSON body (encoded), ``label`` identifies the fix
|
||||
that was applied (e.g. the stripped param name) so the caller can guard
|
||||
against applying the same fix twice.
|
||||
"""
|
||||
|
||||
body: bytes
|
||||
label: str
|
||||
|
||||
|
||||
def extract_error_message(response: Response) -> str:
|
||||
"""Best-effort extraction of an error message string from a proxy Response."""
|
||||
body_bytes = getattr(response, "body", None)
|
||||
if not body_bytes:
|
||||
return ""
|
||||
try:
|
||||
data = json.loads(body_bytes)
|
||||
except Exception:
|
||||
return body_bytes.decode("utf-8", errors="ignore")[:500]
|
||||
if isinstance(data, dict):
|
||||
err = data.get("error")
|
||||
if isinstance(err, dict):
|
||||
msg = err.get("message") or err.get("detail")
|
||||
if isinstance(msg, str):
|
||||
return msg
|
||||
elif isinstance(err, str):
|
||||
return err
|
||||
if isinstance(data.get("message"), str):
|
||||
return data["message"]
|
||||
return ""
|
||||
|
||||
|
||||
def strip_unsupported_param(
|
||||
body: dict, error_message: str
|
||||
) -> tuple[dict, str] | None:
|
||||
"""Drop a top-level param the upstream named as unsupported/deprecated.
|
||||
|
||||
Returns ``(new_body, param)`` (a new dict, original untouched) when the
|
||||
error names a top-level param present in the body, otherwise ``None``.
|
||||
"""
|
||||
match = _UNSUPPORTED_PARAM_RE.search(error_message)
|
||||
if not match:
|
||||
return None
|
||||
param = match.group("param")
|
||||
if param not in body:
|
||||
return None
|
||||
new_body = {k: v for k, v in body.items() if k != param}
|
||||
return new_body, param
|
||||
|
||||
|
||||
# Ordered pipeline of correctors tried on each recoverable rejection.
|
||||
DEFAULT_CORRECTORS: tuple[Corrector, ...] = (strip_unsupported_param,)
|
||||
|
||||
|
||||
def correct_request(
|
||||
request_body: bytes,
|
||||
error_message: str,
|
||||
applied: set[str],
|
||||
correctors: Sequence[Corrector] = DEFAULT_CORRECTORS,
|
||||
) -> Correction | None:
|
||||
"""Try to correct a rejected request body so it can be retried.
|
||||
|
||||
Runs each corrector in order against the parsed body and ``error_message``.
|
||||
The first corrector that proposes a fix whose ``label`` is not already in
|
||||
``applied`` wins; its result is returned as a :class:`Correction`. Returns
|
||||
``None`` when nothing parses, nothing matches, or every proposed fix was
|
||||
already applied — the caller then treats the response as a normal failure.
|
||||
|
||||
``applied`` is read-only here; the caller records the returned ``label`` to
|
||||
bound retries and guarantee forward progress.
|
||||
"""
|
||||
if not request_body or not error_message:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(request_body)
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
for corrector in correctors:
|
||||
result = corrector(data, error_message)
|
||||
if result is None:
|
||||
continue
|
||||
new_body, label = result
|
||||
if label in applied:
|
||||
continue
|
||||
return Correction(body=json.dumps(new_body).encode(), label=label)
|
||||
return None
|
||||
180
routstr/upstream/routstr.py
Normal file
180
routstr/upstream/routstr.py
Normal file
@@ -0,0 +1,180 @@
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
|
||||
from ..core import get_logger
|
||||
from ..payment.models import Model
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class RoutstrUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Upstream provider for communicating with another Routstr instance."""
|
||||
|
||||
provider_type = "routstr"
|
||||
default_base_url = None
|
||||
platform_url = None
|
||||
# Upstream Routstr nodes serve `/v1/messages` natively, so forward the
|
||||
# request as-is instead of round-tripping through litellm's
|
||||
# Anthropic→OpenAI translator.
|
||||
supports_anthropic_messages = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
provider_fee: float = 1.01,
|
||||
provider_settings: dict | None = None,
|
||||
):
|
||||
"""Initialize Routstr provider.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the upstream Routstr instance
|
||||
api_key: API key for the upstream Routstr instance
|
||||
provider_fee: Provider fee multiplier
|
||||
provider_settings: Provider-specific settings (auto-topup, etc.)
|
||||
"""
|
||||
# Ensure base_url doesn't end with /v1 as BaseUpstreamProvider appends it if needed
|
||||
# but Routstr paths are usually absolute from base.
|
||||
super().__init__(
|
||||
base_url=base_url.rstrip("/"),
|
||||
api_key=api_key,
|
||||
provider_fee=provider_fee,
|
||||
)
|
||||
self.settings = provider_settings or {}
|
||||
|
||||
def normalize_request_path(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Preserve the ``v1/`` prefix when forwarding to an upstream Routstr.
|
||||
"""
|
||||
return path.lstrip("/")
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "RoutstrUpstreamProvider":
|
||||
import json
|
||||
|
||||
settings = {}
|
||||
if provider_row.provider_settings:
|
||||
try:
|
||||
settings = json.loads(provider_row.provider_settings)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return cls(
|
||||
base_url=provider_row.base_url,
|
||||
api_key=provider_row.api_key,
|
||||
provider_fee=provider_row.provider_fee,
|
||||
provider_settings=settings,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_provider_metadata(cls) -> dict[str, object]:
|
||||
return {
|
||||
"id": cls.provider_type,
|
||||
"name": "Routstr Node",
|
||||
"default_base_url": "",
|
||||
"fixed_base_url": False,
|
||||
"platform_url": cls.platform_url,
|
||||
"can_create_account": False,
|
||||
"can_topup": True,
|
||||
"can_show_balance": True,
|
||||
}
|
||||
|
||||
async def get_balance(self) -> float | None:
|
||||
"""Fetch balance from the upstream Routstr node.
|
||||
|
||||
Returns:
|
||||
Balance in satoshis, or None if failed
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/info"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get(url, headers=headers, timeout=10.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
# Routstr balance info usually contains 'balance' in msats or sats
|
||||
# Check for msats and convert to sats
|
||||
if "balance_msats" in data:
|
||||
return float(data["balance_msats"]) / 1000.0
|
||||
return float(data.get("balance", 0))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to fetch balance from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return None
|
||||
|
||||
async def topup(self, cashu_token: str) -> dict[str, Any]:
|
||||
"""Top up balance on the upstream Routstr node.
|
||||
|
||||
Args:
|
||||
cashu_token: Cashu token to deposit
|
||||
|
||||
Returns:
|
||||
Dict containing top-up result
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/topup"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
payload = {"cashu_token": cashu_token}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.post(
|
||||
url, headers=headers, json=payload, timeout=30.0
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to topup upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return {"error": str(e)}
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
"""Fetch models from the upstream Routstr node."""
|
||||
url = f"{self.base_url}/v1/models"
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get(url, headers={}, timeout=15.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
models = data.get("data", [])
|
||||
return [Model(**m) for m in models]
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to fetch models from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return []
|
||||
|
||||
async def refund_balance(self) -> dict[str, Any]:
|
||||
"""Request a refund from the upstream Routstr node.
|
||||
|
||||
Returns:
|
||||
Dict containing refund result and token
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/refund"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.post(url, headers=headers, timeout=30.0)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to request refund from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return {"error": str(e)}
|
||||
@@ -13,6 +13,7 @@ class XAIUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "x-ai"
|
||||
default_base_url = "https://api.x.ai/v1"
|
||||
platform_url = "https://console.x.ai/"
|
||||
litellm_provider_prefix = "xai/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -1,16 +1,32 @@
|
||||
import asyncio
|
||||
import math
|
||||
import time
|
||||
import typing
|
||||
from typing import TypedDict
|
||||
|
||||
from cashu.core.base import Proof, Token
|
||||
from cashu.core.mint_info import MintInfo as _CashuMintInfo
|
||||
from cashu.wallet.helpers import deserialize_token_from_string
|
||||
from cashu.wallet.wallet import Wallet
|
||||
from sqlmodel import col, update
|
||||
from pydantic_core import PydanticUndefined
|
||||
from sqlmodel import col, select, update
|
||||
|
||||
from .core import db, get_logger
|
||||
from .core.db import store_cashu_transaction
|
||||
from .core.settings import settings
|
||||
from .payment.lnurl import raw_send_to_lnurl
|
||||
|
||||
# cashu still declares Optional[X] without explicit defaults on MintInfo.
|
||||
# Under pydantic v2 those are required, but real mints omit many of them.
|
||||
# Default Optional fields to None at import time so balance fetches don't 422.
|
||||
for _name, _field in _CashuMintInfo.model_fields.items():
|
||||
_annot = _field.annotation
|
||||
_is_optional = typing.get_origin(_annot) is typing.Union and type(
|
||||
None
|
||||
) in typing.get_args(_annot)
|
||||
if _is_optional and _field.default is PydanticUndefined:
|
||||
_field.default = None
|
||||
_CashuMintInfo.model_rebuild(force=True)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -32,24 +48,66 @@ async def recieve_token(
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
|
||||
await wallet.load_mint(keyset_id=token_obj.keysets[0])
|
||||
|
||||
wallet.verify_proofs_dleq(token_obj.proofs)
|
||||
await wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
|
||||
return token_obj.amount, token_obj.unit, token_obj.mint
|
||||
|
||||
|
||||
async def send(amount: int, unit: str, mint_url: str | None = None) -> tuple[int, str]:
|
||||
"""Internal send function - returns amount and serialized token"""
|
||||
wallet: Wallet = await get_wallet(mint_url or settings.primary_mint, unit)
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
wallet, mint_url or settings.primary_mint, unit
|
||||
effective_mint_url = mint_url or settings.primary_mint
|
||||
wallet: Wallet = await get_wallet(effective_mint_url, unit)
|
||||
proofs = get_proofs_per_mint_and_unit(wallet, effective_mint_url, unit)
|
||||
proofs_for_mint = sum(p.amount for p in proofs)
|
||||
|
||||
# Fallback: proofs from untrusted source mints are swapped to primary_mint
|
||||
# during receive, so the user's preferred refund_mint_url may have no proofs
|
||||
# even though the global wallet has the balance.
|
||||
if proofs_for_mint < amount and effective_mint_url != settings.primary_mint:
|
||||
logger.info(
|
||||
f"send: insufficient proofs at {effective_mint_url} "
|
||||
f"(have {proofs_for_mint}, need {amount}), falling back to primary_mint={settings.primary_mint}"
|
||||
)
|
||||
effective_mint_url = settings.primary_mint
|
||||
wallet = await get_wallet(effective_mint_url, unit)
|
||||
proofs = get_proofs_per_mint_and_unit(wallet, effective_mint_url, unit)
|
||||
proofs_for_mint = sum(p.amount for p in proofs)
|
||||
|
||||
all_mint_urls = list({k.mint_url for k in wallet.keysets.values()})
|
||||
proof_summary = {
|
||||
f"{k.mint_url}/{k.unit.name}": sum(p.amount for p in wallet.proofs if p.id == k.id)
|
||||
for k in wallet.keysets.values()
|
||||
}
|
||||
# Show ALL proofs in DB by keyset_id, regardless of whether the loaded wallet
|
||||
# knows about that keyset. This reveals proofs orphaned under stale keysets.
|
||||
raw_proofs_by_keyset: dict[str, int] = {}
|
||||
for p in wallet.proofs:
|
||||
raw_proofs_by_keyset[p.id] = raw_proofs_by_keyset.get(p.id, 0) + p.amount
|
||||
logger.info(
|
||||
f"send: proof inventory | mint={effective_mint_url} unit={unit} amount={amount} "
|
||||
f"primary_mint={settings.primary_mint} proofs_for_mint={proofs_for_mint} "
|
||||
f"all_mints={all_mint_urls} by_keyset={proof_summary} "
|
||||
f"raw_proofs_by_keyset_id={raw_proofs_by_keyset} "
|
||||
f"total_wallet_proofs={sum(p.amount for p in wallet.proofs)}"
|
||||
)
|
||||
|
||||
# Reserve proofs only after serialization succeeds — if serialize_proofs or
|
||||
# swap_to_send fails mid-way, proofs stay unreserved so dashboard balance
|
||||
# doesn't go negative.
|
||||
send_proofs, _ = await wallet.select_to_send(
|
||||
proofs, amount, set_reserved=True, include_fees=False
|
||||
)
|
||||
token = await wallet.serialize_proofs(
|
||||
send_proofs, include_dleq=False, legacy=False, memo=None
|
||||
proofs, amount, set_reserved=False, include_fees=False
|
||||
)
|
||||
try:
|
||||
token = await wallet.serialize_proofs(
|
||||
send_proofs, include_dleq=False, legacy=False, memo=None
|
||||
)
|
||||
except Exception:
|
||||
await wallet.set_reserved_for_send(send_proofs, reserved=False)
|
||||
raise
|
||||
await wallet.set_reserved_for_send(send_proofs, reserved=True)
|
||||
return amount, token
|
||||
|
||||
|
||||
@@ -58,15 +116,89 @@ async def send_token(amount: int, unit: str, mint_url: str | None = None) -> str
|
||||
return token
|
||||
|
||||
|
||||
async def _calculate_swap_amount(
|
||||
amount_msat: int,
|
||||
token_unit: str,
|
||||
token_mint_url: str,
|
||||
token_wallet: Wallet,
|
||||
primary_wallet: Wallet,
|
||||
proofs: list,
|
||||
) -> int:
|
||||
"""
|
||||
Calculate the amount to mint on the primary mint after accounting for
|
||||
melt fees and NUT-02 input fees on the foreign mint.
|
||||
"""
|
||||
if settings.primary_mint_unit == "sat":
|
||||
receive_amount = amount_msat // 1000
|
||||
else:
|
||||
receive_amount = amount_msat
|
||||
|
||||
if token_mint_url == settings.primary_mint:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: skipping fee estimation (same mint)",
|
||||
extra={"minted_amount": receive_amount},
|
||||
)
|
||||
return int(receive_amount)
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: estimating fees",
|
||||
extra={
|
||||
"dummy_amount": receive_amount,
|
||||
"unit": settings.primary_mint_unit,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
dummy_mint_quote = await primary_wallet.request_mint(receive_amount)
|
||||
dummy_melt_quote = await token_wallet.melt_quote(dummy_mint_quote.request)
|
||||
|
||||
fee_reserve = dummy_melt_quote.fee_reserve
|
||||
input_fees = token_wallet.get_fees_for_proofs(proofs)
|
||||
if token_unit == "sat":
|
||||
fee_msat = (fee_reserve + input_fees) * 1000
|
||||
else:
|
||||
fee_msat = fee_reserve + input_fees
|
||||
|
||||
amount_msat_after_fee = amount_msat - fee_msat
|
||||
|
||||
if settings.primary_mint_unit == "sat":
|
||||
minted_amount = int(amount_msat_after_fee // 1000)
|
||||
else:
|
||||
minted_amount = int(amount_msat_after_fee)
|
||||
|
||||
if minted_amount <= 0:
|
||||
raise ValueError(f"Fees ({fee_reserve + input_fees} {token_unit}) exceed token amount")
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: fee estimation result",
|
||||
extra={
|
||||
"token_amount_sat": amount_msat // 1000,
|
||||
"estimated_fee_sat": fee_msat // 1000,
|
||||
"input_fees": input_fees,
|
||||
"minted_amount": minted_amount,
|
||||
"minted_unit": settings.primary_mint_unit,
|
||||
},
|
||||
)
|
||||
return minted_amount
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: fee estimation failed",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
raise ValueError(f"Failed to estimate fees: {e}") from e
|
||||
|
||||
|
||||
async def swap_to_primary_mint(
|
||||
token_obj: Token, token_wallet: Wallet
|
||||
) -> tuple[int, str, str]:
|
||||
logger.info(
|
||||
"swap_to_primary_mint",
|
||||
"swap_to_primary_mint: starting",
|
||||
extra={
|
||||
"mint": token_obj.mint,
|
||||
"amount": token_obj.amount,
|
||||
"foreign_mint": token_obj.mint,
|
||||
"token_amount": token_obj.amount,
|
||||
"unit": token_obj.unit,
|
||||
"primary_mint": settings.primary_mint,
|
||||
},
|
||||
)
|
||||
# Ensure amount is an integer
|
||||
@@ -81,24 +213,169 @@ async def swap_to_primary_mint(
|
||||
amount_msat = token_amount
|
||||
else:
|
||||
raise ValueError("Invalid unit")
|
||||
estimated_fee_sat = math.ceil(max(amount_msat // 1000 * 0.01, 2)) + 1
|
||||
amount_msat_after_fee = amount_msat - estimated_fee_sat * 1000
|
||||
primary_wallet = await get_wallet(settings.primary_mint, settings.primary_mint_unit)
|
||||
|
||||
if settings.primary_mint_unit == "sat":
|
||||
minted_amount = int(amount_msat_after_fee // 1000)
|
||||
else:
|
||||
minted_amount = int(amount_msat_after_fee)
|
||||
# If the token is already from the primary mint, we don't need to swap
|
||||
# and we definitely don't want to calculate or pay fees.
|
||||
if token_obj.mint == settings.primary_mint:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: token already on primary mint, skipping swap",
|
||||
extra={
|
||||
"mint": token_obj.mint,
|
||||
"amount": token_amount,
|
||||
"unit": token_obj.unit,
|
||||
},
|
||||
)
|
||||
await token_wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
return token_amount, token_obj.unit, token_obj.mint
|
||||
|
||||
minted_amount = await _calculate_swap_amount(
|
||||
amount_msat,
|
||||
token_obj.unit,
|
||||
token_obj.mint,
|
||||
token_wallet,
|
||||
primary_wallet,
|
||||
token_obj.proofs,
|
||||
)
|
||||
|
||||
mint_quote = await primary_wallet.request_mint(minted_amount)
|
||||
logger.info(
|
||||
"swap_to_primary_mint: mint quote received",
|
||||
extra={"mint_quote_id": mint_quote.quote},
|
||||
)
|
||||
|
||||
melt_quote = await token_wallet.melt_quote(mint_quote.request)
|
||||
_ = await token_wallet.melt(
|
||||
proofs=token_obj.proofs,
|
||||
invoice=mint_quote.request,
|
||||
fee_reserve_sat=melt_quote.fee_reserve,
|
||||
quote_id=melt_quote.quote,
|
||||
input_fees = token_wallet.get_fees_for_proofs(token_obj.proofs)
|
||||
total_needed = melt_quote.amount + melt_quote.fee_reserve + input_fees
|
||||
logger.info(
|
||||
"swap_to_primary_mint: melt quote received",
|
||||
extra={
|
||||
"melt_quote_id": melt_quote.quote,
|
||||
"melt_amount": melt_quote.amount,
|
||||
"melt_fee_reserve": melt_quote.fee_reserve,
|
||||
"input_fees": input_fees,
|
||||
"total_needed": total_needed,
|
||||
"token_amount": token_amount,
|
||||
},
|
||||
)
|
||||
|
||||
if total_needed > token_amount:
|
||||
logger.warning(
|
||||
"swap_to_primary_mint: insufficient token amount for melt fees",
|
||||
extra={
|
||||
"token_amount": token_amount,
|
||||
"melt_amount": melt_quote.amount,
|
||||
"melt_fee_reserve": melt_quote.fee_reserve,
|
||||
"input_fees": input_fees,
|
||||
"total_needed": total_needed,
|
||||
"shortfall": total_needed - token_amount,
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Token amount ({token_amount} {token_obj.unit}) is insufficient to cover "
|
||||
f"melt fees. Needed: {total_needed} {token_obj.unit} "
|
||||
f"(amount: {melt_quote.amount} + fee: {melt_quote.fee_reserve} + input_fees: {input_fees})"
|
||||
)
|
||||
|
||||
try:
|
||||
_ = await token_wallet.melt(
|
||||
proofs=token_obj.proofs,
|
||||
invoice=mint_quote.request,
|
||||
fee_reserve_sat=melt_quote.fee_reserve,
|
||||
quote_id=melt_quote.quote,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: melt failed",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"foreign_mint": token_obj.mint,
|
||||
"token_amount": token_amount,
|
||||
"melt_quote_id": melt_quote.quote,
|
||||
"total_needed": total_needed,
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Failed to melt token from foreign mint {token_obj.mint}: {e}"
|
||||
) from e
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: melt succeeded, minting on primary",
|
||||
extra={"minted_amount": minted_amount, "mint_quote_id": mint_quote.quote},
|
||||
)
|
||||
|
||||
await primary_wallet.load_proofs(reload=True)
|
||||
pre_mint_balance = primary_wallet.available_balance.amount
|
||||
try:
|
||||
_ = await primary_wallet.mint(minted_amount, quote_id=mint_quote.quote)
|
||||
except Exception as e:
|
||||
if "11003" in str(e) or "outputs already signed" in str(e).lower():
|
||||
# Previous mint call signed outputs at the mint but failed before
|
||||
# bump_secret_derivation ran locally. Recover orphaned proofs and
|
||||
# advance the counter so the next request derives fresh secrets.
|
||||
logger.warning(
|
||||
"swap_to_primary_mint: outputs already signed — recovering orphaned proofs",
|
||||
extra={"mint_quote_id": mint_quote.quote, "minted_amount": minted_amount},
|
||||
)
|
||||
try:
|
||||
for keyset_id in primary_wallet.keysets:
|
||||
await primary_wallet.restore_tokens_for_keyset(keyset_id, to=1, batch=25)
|
||||
await primary_wallet.load_proofs(reload=True)
|
||||
post_recovery_balance = primary_wallet.available_balance.amount
|
||||
balance_gained = post_recovery_balance - pre_mint_balance
|
||||
logger.info(
|
||||
"swap_to_primary_mint: recovery scan completed",
|
||||
extra={
|
||||
"pre_mint_balance": pre_mint_balance,
|
||||
"post_recovery_balance": post_recovery_balance,
|
||||
"balance_gained": balance_gained,
|
||||
"expected": minted_amount,
|
||||
},
|
||||
)
|
||||
if balance_gained < minted_amount:
|
||||
# Recovery scan ran but did NOT restore the orphaned proofs
|
||||
# (mint reports them as spent — they're stuck). Refuse to
|
||||
# credit the API key balance for proofs we don't actually hold.
|
||||
raise ValueError(
|
||||
f"Swap recovery failed: mint signed outputs but proofs are "
|
||||
f"unrecoverable (mint reports them spent). "
|
||||
f"Expected {minted_amount}, recovered {balance_gained}. "
|
||||
f"Local wallet DB ('.wallet/') state is corrupted — "
|
||||
f"the counter for keyset is stuck at a bad index range."
|
||||
)
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception as recovery_err:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: recovery failed",
|
||||
extra={"error": str(recovery_err)},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Mint on primary failed and recovery unsuccessful: {e}"
|
||||
) from e
|
||||
else:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: mint on primary failed after successful melt",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"minted_amount": minted_amount,
|
||||
"mint_quote_id": mint_quote.quote,
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: completed successfully",
|
||||
extra={
|
||||
"foreign_mint": token_obj.mint,
|
||||
"primary_mint": settings.primary_mint,
|
||||
"original_amount": token_amount,
|
||||
"minted_amount": minted_amount,
|
||||
"unit": settings.primary_mint_unit,
|
||||
},
|
||||
)
|
||||
_ = await primary_wallet.mint(minted_amount, quote_id=mint_quote.quote)
|
||||
|
||||
return int(minted_amount), settings.primary_mint_unit, settings.primary_mint
|
||||
|
||||
@@ -113,6 +390,8 @@ async def credit_balance(
|
||||
|
||||
try:
|
||||
amount, unit, mint_url = await recieve_token(cashu_token)
|
||||
original_amount = amount
|
||||
original_unit = unit
|
||||
logger.info(
|
||||
"credit_balance: Token redeemed successfully",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
@@ -124,6 +403,20 @@ async def credit_balance(
|
||||
"credit_balance: Converted to msat", extra={"amount_msat": amount}
|
||||
)
|
||||
|
||||
# Guard against zero/negative redemptions (empty or dust tokens, or
|
||||
# swap-to-primary-mint amounts that net to <= 0 after fees). Raising here
|
||||
# — before the UPDATE/commit below — leaves any freshly-created, still
|
||||
# uncommitted ApiKey row to be rolled back when the request session
|
||||
# closes, instead of persisting an orphan key with balance 0.
|
||||
if amount <= 0:
|
||||
logger.error(
|
||||
"credit_balance: Redeemed amount is zero or negative; refusing to credit",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Redeemed token amount must be positive, got {amount} msats"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"credit_balance: Updating balance",
|
||||
extra={"old_balance": key.balance, "credit_amount": amount},
|
||||
@@ -144,7 +437,20 @@ async def credit_balance(
|
||||
extra={"new_balance": key.balance},
|
||||
)
|
||||
|
||||
logger.info(
|
||||
try:
|
||||
await store_cashu_transaction(
|
||||
token=cashu_token,
|
||||
amount=original_amount,
|
||||
unit=original_unit,
|
||||
mint_url=mint_url,
|
||||
typ="in",
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.debug(
|
||||
"Cashu token successfully redeemed and stored",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
)
|
||||
@@ -246,7 +552,7 @@ async def fetch_all_balances(
|
||||
"unit": unit,
|
||||
"wallet_balance": proofs_balance,
|
||||
"user_balance": user_balance,
|
||||
"owner_balance": proofs_balance - user_balance,
|
||||
"owner_balance": proofs_balance - user_balance if proofs_balance != 0 else 0,
|
||||
}
|
||||
return result
|
||||
except Exception as e:
|
||||
@@ -294,7 +600,9 @@ async def fetch_all_balances(
|
||||
total_wallet_balance_sats += proofs_balance_sats
|
||||
total_user_balance_sats += user_balance_sats
|
||||
|
||||
owner_balance = total_wallet_balance_sats - total_user_balance_sats
|
||||
owner_balance = 0
|
||||
if total_wallet_balance_sats != 0:
|
||||
owner_balance = total_wallet_balance_sats - total_user_balance_sats
|
||||
|
||||
return (
|
||||
balance_details,
|
||||
@@ -305,11 +613,11 @@ async def fetch_all_balances(
|
||||
|
||||
|
||||
async def periodic_payout() -> None:
|
||||
if not settings.receive_ln_address:
|
||||
logger.error("RECEIVE_LN_ADDRESS is not set, skipping payout")
|
||||
return
|
||||
while True:
|
||||
await asyncio.sleep(60 * 15)
|
||||
await asyncio.sleep(settings.payout_interval_seconds)
|
||||
print(settings.payout_interval_seconds)
|
||||
if not settings.receive_ln_address:
|
||||
continue
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
for mint_url in settings.cashu_mints:
|
||||
@@ -327,7 +635,12 @@ async def periodic_payout() -> None:
|
||||
user_balance = user_balance // 1000
|
||||
proofs_balance = sum(proof.amount for proof in proofs)
|
||||
available_balance = proofs_balance - user_balance
|
||||
min_amount = 210 if unit == "sat" else 210000
|
||||
# Threshold is configured in sats; convert for msat wallets.
|
||||
min_amount = (
|
||||
settings.min_payout_sat
|
||||
if unit == "sat"
|
||||
else settings.min_payout_sat * 1000
|
||||
)
|
||||
if available_balance > min_amount:
|
||||
amount_received = await raw_send_to_lnurl(
|
||||
wallet,
|
||||
@@ -352,6 +665,101 @@ async def periodic_payout() -> None:
|
||||
)
|
||||
|
||||
|
||||
async def periodic_refund_sweep() -> None:
|
||||
while True:
|
||||
await asyncio.sleep(60 * 60) # every hour
|
||||
try:
|
||||
cutoff = int(time.time()) - settings.refund_sweep_ttl_seconds
|
||||
async with db.create_session() as session:
|
||||
stmt = select(db.CashuTransaction).where(
|
||||
db.CashuTransaction.type == "out",
|
||||
db.CashuTransaction.collected == False, # noqa: E712
|
||||
db.CashuTransaction.swept == False, # noqa: E712
|
||||
db.CashuTransaction.created_at < cutoff,
|
||||
)
|
||||
results = await session.exec(stmt)
|
||||
refunds = results.all()
|
||||
|
||||
for refund in refunds:
|
||||
try:
|
||||
await recieve_token(refund.token)
|
||||
refund.swept = True
|
||||
session.add(refund)
|
||||
logger.info(
|
||||
"Swept uncollected refund",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
"amount": refund.amount,
|
||||
"unit": refund.unit,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "already spent" in error_msg:
|
||||
refund.collected = True
|
||||
session.add(refund)
|
||||
logger.info(
|
||||
"Refund already spent (client collected), marking swept",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Failed to sweep refund",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in periodic refund sweep",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
|
||||
|
||||
async def periodic_routstr_fee_payout() -> None:
|
||||
from .auth import (
|
||||
ROUTSTR_FEE_DEFAULT_PAYOUT,
|
||||
ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS,
|
||||
ROUTSTR_LN_ADDRESS,
|
||||
)
|
||||
|
||||
if not ROUTSTR_LN_ADDRESS:
|
||||
logger.info("ROUTSTR_LN_ADDRESS not set, skipping fee payout")
|
||||
return
|
||||
while True:
|
||||
await asyncio.sleep(ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS)
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
fee = await db.get_routstr_fee(session)
|
||||
accumulated_sats = fee.accumulated_msats // 1000
|
||||
if accumulated_sats >= ROUTSTR_FEE_DEFAULT_PAYOUT:
|
||||
wallet = await get_wallet(settings.primary_mint, "sat")
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
wallet, settings.primary_mint, "sat", not_reserved=True
|
||||
)
|
||||
amount_received = await raw_send_to_lnurl(
|
||||
wallet, proofs, ROUTSTR_LN_ADDRESS, "sat", amount=accumulated_sats
|
||||
)
|
||||
paid_msats = accumulated_sats * 1000
|
||||
await db.reset_routstr_fee(session, paid_msats)
|
||||
logger.info(
|
||||
"Routstr fee payout sent",
|
||||
extra={
|
||||
"accumulated_sats": accumulated_sats,
|
||||
"amount_received": amount_received,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error in Routstr fee payout: {type(e).__name__}",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
|
||||
|
||||
async def send_to_lnurl(amount: int, unit: str, mint: str, address: str) -> int:
|
||||
wallet = await get_wallet(mint, unit)
|
||||
proofs = wallet._get_proofs_per_keyset(wallet.proofs)[wallet.keyset_id]
|
||||
|
||||
@@ -174,17 +174,20 @@ class TestmintWallet:
|
||||
"""Fallback method to create a basic test token"""
|
||||
import base64
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
import secrets
|
||||
|
||||
unique_id = int(time.time() * 1000000) + random.randint(1000, 9999)
|
||||
# Use a cryptographically random id/secret so every minted token is
|
||||
# guaranteed unique. Time/PRNG-based ids can collide on hosts with
|
||||
# coarse clock resolution, producing byte-identical tokens that hash to
|
||||
# the same api_key and silently dedupe (flaky concurrency tests).
|
||||
unique_id = secrets.token_hex(16)
|
||||
token_data = {
|
||||
"token": [
|
||||
{
|
||||
"mint": self.mint_url,
|
||||
"proofs": [
|
||||
{
|
||||
"id": f"009a1f293253e41e{unique_id % 100000000:08d}",
|
||||
"id": f"009a1f293253e41e{unique_id[:8]}",
|
||||
"amount": amount,
|
||||
"secret": f"test-secret-{amount}-{unique_id}",
|
||||
"C": "02194603ffa36356f4a56b7df9371fc3192472351453ec7398b8da8117e7c3e104",
|
||||
@@ -380,6 +383,13 @@ async def integration_session(
|
||||
yield session
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def patched_db_engine(integration_engine: Any) -> AsyncGenerator[None, None]:
|
||||
"""Patch the global db engine so create_session() uses the test engine."""
|
||||
with patch("routstr.core.db.engine", integration_engine):
|
||||
yield
|
||||
|
||||
|
||||
class DatabaseSnapshot:
|
||||
"""Utility to capture and compare database states"""
|
||||
|
||||
|
||||
80
tests/integration/test_admin_provider_balance.py
Normal file
80
tests/integration/test_admin_provider_balance.py
Normal file
@@ -0,0 +1,80 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import UpstreamProviderRow
|
||||
|
||||
|
||||
async def _create_routstr_provider() -> UpstreamProviderRow:
|
||||
return UpstreamProviderRow(
|
||||
provider_type="routstr",
|
||||
base_url="https://upstream.example",
|
||||
api_key="",
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
def _admin_headers() -> dict[str, str]:
|
||||
token = "test-admin-token"
|
||||
admin_sessions[token] = int(
|
||||
(datetime.now(timezone.utc) + timedelta(minutes=5)).timestamp()
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_provider_balance_timeout_returns_504(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
provider = await _create_routstr_provider()
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
request = httpx.Request("GET", f"{provider.base_url}/v1/balance/info")
|
||||
timeout_error = httpx.ConnectTimeout("Connect timeout", request=request)
|
||||
|
||||
with patch(
|
||||
"httpx.AsyncHTTPTransport.handle_async_request",
|
||||
new=AsyncMock(side_effect=timeout_error),
|
||||
):
|
||||
response = await integration_client.get(
|
||||
f"/admin/api/upstream-providers/{provider.id}/balance",
|
||||
headers=_admin_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 504
|
||||
assert response.json()["detail"] == "Timed out contacting upstream Routstr provider"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_provider_balance_request_error_returns_502(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
provider = await _create_routstr_provider()
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
request = httpx.Request("GET", f"{provider.base_url}/v1/balance/info")
|
||||
request_error = httpx.ConnectError("Connection failed", request=request)
|
||||
|
||||
with patch(
|
||||
"httpx.AsyncHTTPTransport.handle_async_request",
|
||||
new=AsyncMock(side_effect=request_error),
|
||||
):
|
||||
response = await integration_client.get(
|
||||
f"/admin/api/upstream-providers/{provider.id}/balance",
|
||||
headers=_admin_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 502
|
||||
assert response.json()["detail"] == "Failed to contact upstream Routstr provider"
|
||||
369
tests/integration/test_balance_negative_on_cost_overrun.py
Normal file
369
tests/integration/test_balance_negative_on_cost_overrun.py
Normal file
@@ -0,0 +1,369 @@
|
||||
"""
|
||||
Integration tests for the balance-goes-negative bug in adjust_payment_for_tokens.
|
||||
|
||||
Root cause: when actual token cost exceeds the discounted reservation
|
||||
(cost_difference > 0, caused by tolerance_percentage discounting the reservation),
|
||||
the finalization UPDATE had no WHERE guard on balance, allowing balance to go negative.
|
||||
|
||||
Fix: added `.where(col(ApiKey.balance) >= total_cost_msats)` so the UPDATE is a no-op
|
||||
when balance is insufficient, then falls back to charging only deducted_max_cost.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
|
||||
def _make_key(balance: int, reserved: int) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
total_requests=1,
|
||||
)
|
||||
|
||||
|
||||
async def _refresh(session: AsyncSession, key: ApiKey) -> ApiKey:
|
||||
await session.refresh(key)
|
||||
return key
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper: build a CostData where token cost > deducted_max_cost
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _cost_data(total_msats: int) -> CostData:
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=total_msats // 2,
|
||||
output_msats=total_msats - total_msats // 2,
|
||||
total_msats=total_msats,
|
||||
total_usd=0.0,
|
||||
input_tokens=100,
|
||||
output_tokens=100,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — exact reproduction of the bug
|
||||
#
|
||||
# Setup: balance == deducted_max_cost (user has just enough for the reservation,
|
||||
# nothing extra). Actual token cost is 1% higher (tolerance_percentage).
|
||||
#
|
||||
# Before fix: balance -= total_cost_msats → goes negative.
|
||||
# After fix: WHERE balance >= total_cost_msats fails → fallback charges
|
||||
# deducted_max_cost → balance reaches 0, never negative.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_never_negative_when_cost_exceeds_reservation(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Balance must not go negative when actual token cost > discounted reservation."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 990 # reserved (1% below true max of 1000)
|
||||
actual_token_cost = 1000 # actual cost at true max
|
||||
|
||||
# User has balance exactly equal to the reservation — tight budget
|
||||
key = _make_key(balance=deducted_max_cost, reserved=deducted_max_cost)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
|
||||
assert key.reserved_balance == 0, "Reservation must be fully released after finalization"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — balance is ZERO after the reservation is accounted for
|
||||
# (absolute floor case)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_floor_at_zero_on_overrun(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""When balance exactly covers deducted_max_cost and cost overruns, balance reaches 0 not negative."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 500
|
||||
actual_token_cost = 550 # 10% overrun
|
||||
|
||||
key = _make_key(balance=500, reserved=500)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance == 0, (
|
||||
f"Expected balance=0 (charged deducted_max_cost fallback), got {key.balance}"
|
||||
)
|
||||
assert key.reserved_balance == 0, f"Reserved balance should be 0, got {key.reserved_balance}"
|
||||
# Fallback charges deducted_max_cost
|
||||
assert key.total_spent == deducted_max_cost, (
|
||||
f"Expected total_spent={deducted_max_cost}, got {key.total_spent}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — balance has enough room: full token cost should be charged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_cost_charged_when_balance_sufficient_for_overrun(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""When balance covers total_cost_msats, the full amount is charged (not just deducted_max_cost)."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 990
|
||||
actual_token_cost = 1000
|
||||
|
||||
# User has extra balance beyond the reservation
|
||||
key = _make_key(balance=2000, reserved=990)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance == 0, f"Reservation not released: {key.reserved_balance}"
|
||||
assert key.total_spent == actual_token_cost, (
|
||||
f"Expected full charge of {actual_token_cost}, got {key.total_spent}"
|
||||
)
|
||||
assert key.balance == 2000 - actual_token_cost, (
|
||||
f"Expected balance={2000 - actual_token_cost}, got {key.balance}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — concurrent finalizations with cost overrun
|
||||
#
|
||||
# Multiple requests finish concurrently. Each has a small overrun.
|
||||
# None should drive balance negative.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_cost_overruns_never_negative(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Concurrent finalization with cost overruns must never produce negative balance."""
|
||||
import asyncio
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.core.db import create_session
|
||||
|
||||
deducted_max_cost = 990
|
||||
actual_token_cost = 1000
|
||||
n_requests = 5
|
||||
|
||||
# Fund the key with exactly enough for n_requests reservations + a tiny buffer
|
||||
starting_balance = deducted_max_cost * n_requests
|
||||
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=starting_balance,
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
# Reserve n_requests slots (sequentially, as pay_for_request is atomic)
|
||||
async with create_session() as session:
|
||||
key_to_reserve = await session.get(ApiKey, key_hash)
|
||||
assert key_to_reserve is not None
|
||||
for _ in range(n_requests):
|
||||
await pay_for_request(key_to_reserve, deducted_max_cost, session)
|
||||
await session.refresh(key_to_reserve)
|
||||
|
||||
# Now finalize all concurrently with cost overrun
|
||||
async def finalize() -> None:
|
||||
response_data = {
|
||||
"model": "test-model",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 100},
|
||||
}
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key, response_data, session, deducted_max_cost
|
||||
)
|
||||
|
||||
await asyncio.gather(*[finalize() for _ in range(n_requests)])
|
||||
|
||||
async with create_session() as session:
|
||||
final_key = await session.get(ApiKey, key_hash)
|
||||
assert final_key is not None
|
||||
|
||||
assert final_key.balance >= 0, (
|
||||
f"Balance went negative after concurrent overruns: {final_key.balance}"
|
||||
)
|
||||
assert final_key.reserved_balance == 0, (
|
||||
f"Reserved balance not fully released: {final_key.reserved_balance}"
|
||||
)
|
||||
assert final_key.total_spent <= starting_balance, (
|
||||
f"Total spent ({final_key.total_spent}) exceeds starting balance ({starting_balance})"
|
||||
)
|
||||
# Every request must have been charged at least deducted_max_cost — no free inference.
|
||||
assert final_key.total_spent == starting_balance, (
|
||||
f"Expected total_spent={starting_balance} (all {n_requests} reservations charged), "
|
||||
f"got {final_key.total_spent} — at least one request got free inference"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — overrun with no balance at all (reserved_balance == balance)
|
||||
# simulates a user who topped up to exactly the reservation floor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_free_balance_overrun_is_safe(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""User with zero free balance (all reserved) should never go negative on overrun."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 1000
|
||||
actual_token_cost = 1050
|
||||
|
||||
# balance == reserved_balance: zero free balance
|
||||
key = _make_key(balance=1000, reserved=1000)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — parallel requests: second finalization must not get free inference
|
||||
#
|
||||
# Root cause of the bug fixed in auth.py:
|
||||
# `.where(col(ApiKey.balance) >= total_cost_msats)` ignores other requests'
|
||||
# reservations, so after Request A charges total_cost_msats, balance can drop
|
||||
# below deducted_max_cost, causing Request B's fallback to release for free.
|
||||
#
|
||||
# Fix: use `balance - reserved_balance + deducted_max_cost >= total_cost_msats`
|
||||
# so the check accounts for concurrent reservations.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_requests_no_free_inference(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Second parallel finalization must be charged even when first depleted free balance."""
|
||||
import asyncio
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
from routstr.core.db import create_session
|
||||
|
||||
deducted_max_cost = 100
|
||||
actual_token_cost = 150 # overrun: 50 more than reserved
|
||||
|
||||
# Fund the key with exactly 2 * deducted_max_cost.
|
||||
# Both requests pre-reserved 100 each → balance=200, reserved=200, free=0.
|
||||
# Old check (balance >= total_cost_msats):
|
||||
# Request A: 200 >= 150 ✓ → charges 150 → balance=50, reserved=100
|
||||
# Request B: 50 >= 150 ✗ → fallback: 50 >= 100 ✗ → releases FREE
|
||||
# New check (balance - reserved + deducted >= total_cost_msats):
|
||||
# Both fall to fallback (0 free balance).
|
||||
# Both charge deducted_max_cost=100 → total_spent=200, balance=0.
|
||||
starting_balance = deducted_max_cost * 2
|
||||
key_hash = f"test_parallel_no_free_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=starting_balance,
|
||||
reserved_balance=deducted_max_cost * 2, # both slots pre-reserved
|
||||
total_spent=0,
|
||||
total_requests=2,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
async def finalize() -> None:
|
||||
response_data = {
|
||||
"model": "test-model",
|
||||
"usage": {"prompt_tokens": 50, "completion_tokens": 100},
|
||||
}
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key, response_data, session, deducted_max_cost
|
||||
)
|
||||
|
||||
await asyncio.gather(finalize(), finalize())
|
||||
|
||||
async with create_session() as session:
|
||||
final_key = await session.get(ApiKey, key_hash)
|
||||
assert final_key is not None
|
||||
|
||||
assert final_key.balance >= 0, f"Balance went negative: {final_key.balance}"
|
||||
assert final_key.reserved_balance == 0, (
|
||||
f"Reserved balance not released: {final_key.reserved_balance}"
|
||||
)
|
||||
# Both requests must have been charged — no free inference.
|
||||
assert final_key.total_spent == starting_balance, (
|
||||
f"Expected total_spent={starting_balance} (both reservations charged), "
|
||||
f"got {final_key.total_spent} — one request got free inference"
|
||||
)
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import secrets
|
||||
from typing import Any
|
||||
|
||||
@@ -7,7 +8,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.balance import ChildKeyRequest, create_child_key
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.db import ApiKey, create_session
|
||||
from routstr.core.settings import settings
|
||||
|
||||
|
||||
@@ -119,6 +120,54 @@ async def test_child_key_insufficient_balance(
|
||||
assert exc.value.status_code == 402
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_child_key_creation_is_atomic(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Two concurrent create_child_key() calls with balance for exactly one must
|
||||
result in exactly one success and one 402, with the parent balance deducted
|
||||
only once."""
|
||||
child_key_cost = 1000
|
||||
settings.child_key_cost = child_key_cost
|
||||
|
||||
parent_hash = f"parent_concurrent_{secrets.token_hex(8)}"
|
||||
async with create_session() as session:
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=child_key_cost)
|
||||
session.add(parent)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_parent = await session.get(ApiKey, parent_hash)
|
||||
assert fresh_parent is not None
|
||||
try:
|
||||
await create_child_key(ChildKeyRequest(count=1), fresh_parent, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt(), attempt())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, parent_hash)
|
||||
assert final is not None
|
||||
|
||||
assert final.balance == 0, (
|
||||
f"Balance should be fully deducted once: expected 0, got {final.balance}"
|
||||
)
|
||||
assert final.total_spent == child_key_cost, (
|
||||
f"total_spent should equal one deduction: expected {child_key_cost}, "
|
||||
f"got {final.total_spent}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_cannot_create_child(integration_session: AsyncSession) -> None:
|
||||
parent_key = ApiKey(
|
||||
|
||||
@@ -69,9 +69,14 @@ async def test_wallet_info_child_key_no_child_keys(
|
||||
info_response = await integration_client.get("/v1/wallet/info")
|
||||
assert info_response.status_code == 200
|
||||
info_data = info_response.json()
|
||||
parent_key = authenticated_client._test_api_key # type: ignore[attr-defined]
|
||||
parent_key_hash = parent_key.removeprefix("sk-")
|
||||
|
||||
assert info_data["is_child"] is True
|
||||
assert "child_keys" not in info_data
|
||||
assert "parent_key" not in info_data
|
||||
assert info_data["parent_key_preview"] == parent_key_hash[:8] + "..."
|
||||
assert info_data["parent_key_preview"] not in {parent_key, parent_key_hash}
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
|
||||
302
tests/integration/test_cli_tokens.py
Normal file
302
tests/integration/test_cli_tokens.py
Normal file
@@ -0,0 +1,302 @@
|
||||
"""Integration tests for CLI token management (/admin/api/cli-tokens).
|
||||
|
||||
Covers:
|
||||
- GET /admin/api/cli-tokens — list (preview only, no full token)
|
||||
- POST /admin/api/cli-tokens — create (returns full token once)
|
||||
- DELETE /admin/api/cli-tokens/{id} — revoke
|
||||
- Using a CLI token as Bearer auth against admin endpoints
|
||||
- Expiry enforcement (expired tokens are rejected by require_admin_api)
|
||||
- last_used_at bump on successful use
|
||||
- Auth failures: missing token, wrong token, revoked token
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
import time
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import AsyncSession, CliToken
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Fixtures
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def admin_session_token() -> AsyncGenerator[str, None]:
|
||||
"""Inject a short-lived admin session token into admin_sessions."""
|
||||
token = secrets.token_urlsafe(24)
|
||||
admin_sessions[token] = int(time.time()) + 3600
|
||||
yield token
|
||||
admin_sessions.pop(token, None)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def admin_client(
|
||||
integration_client: AsyncClient, admin_session_token: str
|
||||
) -> AsyncClient:
|
||||
"""An integration_client pre-authenticated with an admin session token."""
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_session_token}"
|
||||
return integration_client
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Creation
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_returns_full_token_once(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
"""POST /admin/api/cli-tokens returns the raw token only on creation."""
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "my-laptop"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
|
||||
assert body["name"] == "my-laptop"
|
||||
assert isinstance(body["id"], str) and body["id"]
|
||||
assert isinstance(body["token"], str) and len(body["token"]) >= 32
|
||||
assert body["expires_at"] is None
|
||||
assert isinstance(body["created_at"], int)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_with_expiry(admin_client: AsyncClient) -> None:
|
||||
"""expires_in_days sets expires_at ~= now + days * 86400."""
|
||||
before = int(time.time())
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "ci-runner", "expires_in_days": 7},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
|
||||
assert body["expires_at"] is not None
|
||||
delta = body["expires_at"] - before
|
||||
# Allow 10s jitter around 7 * 86400
|
||||
assert 7 * 86400 - 10 <= delta <= 7 * 86400 + 10
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_rejects_empty_name(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": " "}
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
"""No admin token / no bearer → 403."""
|
||||
resp = await integration_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "no-auth"}
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Listing
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_cli_tokens_returns_preview_not_full_token(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
"""Listing never leaks the raw token."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "secret-keeper"}
|
||||
)
|
||||
assert create.status_code == 200
|
||||
full_token = create.json()["token"]
|
||||
|
||||
resp = await admin_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 200
|
||||
items = resp.json()
|
||||
assert any(t["name"] == "secret-keeper" for t in items)
|
||||
|
||||
for t in items:
|
||||
# No 'token' field, only 'token_preview'
|
||||
assert "token" not in t
|
||||
assert "token_preview" in t
|
||||
assert full_token not in t["token_preview"]
|
||||
assert "..." in t["token_preview"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_cli_tokens_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Using a CLI token as admin auth
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_token_authorizes_admin_endpoints(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""A freshly-created CLI token can be used as Bearer on admin endpoints."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "cli-auth"}
|
||||
)
|
||||
assert create.status_code == 200
|
||||
cli_token = create.json()["token"]
|
||||
token_id = create.json()["id"]
|
||||
|
||||
# Use a NEW client to isolate the header from admin_session_token
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 200
|
||||
|
||||
# last_used_at should be populated after use
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is not None
|
||||
assert row.last_used_at is not None
|
||||
assert row.last_used_at >= row.created_at
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_cli_token_is_rejected(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""A CLI token with expires_at in the past → 403."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "will-expire", "expires_in_days": 1},
|
||||
)
|
||||
assert create.status_code == 200
|
||||
cli_token = create.json()["token"]
|
||||
token_id = create.json()["id"]
|
||||
|
||||
# Force-expire it in the DB
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is not None
|
||||
row.expires_at = int(time.time()) - 1
|
||||
integration_session.add(row)
|
||||
await integration_session.commit()
|
||||
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_bearer_token_is_rejected(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
integration_client.headers["Authorization"] = "Bearer not-a-real-token"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Revocation
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_cli_token_removes_auth(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""After DELETE, the token no longer authorizes."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "to-revoke"}
|
||||
)
|
||||
token_id = create.json()["id"]
|
||||
cli_token = create.json()["token"]
|
||||
|
||||
revoke = await admin_client.delete(f"/admin/api/cli-tokens/{token_id}")
|
||||
assert revoke.status_code == 200
|
||||
assert revoke.json() == {"ok": True, "deleted_id": token_id}
|
||||
|
||||
# Row is gone
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is None
|
||||
|
||||
# Can no longer be used for auth
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_unknown_cli_token_returns_404(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await admin_client.delete("/admin/api/cli-tokens/does-not-exist")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_cli_token_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await integration_client.delete("/admin/api/cli-tokens/anything")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Lifecycle / uniqueness
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_tokens_are_independent(
|
||||
admin_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Creating N tokens yields N unique tokens that all live in DB."""
|
||||
names = ["dev-a", "dev-b", "dev-c"]
|
||||
raw_tokens: list[str] = []
|
||||
ids: list[str] = []
|
||||
for name in names:
|
||||
r = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": name}
|
||||
)
|
||||
assert r.status_code == 200
|
||||
raw_tokens.append(r.json()["token"])
|
||||
ids.append(r.json()["id"])
|
||||
|
||||
# All unique
|
||||
assert len(set(raw_tokens)) == len(raw_tokens)
|
||||
assert len(set(ids)) == len(ids)
|
||||
|
||||
# All in DB
|
||||
result = await integration_session.exec(
|
||||
select(CliToken).where(CliToken.name.in_(names)) # type: ignore[attr-defined]
|
||||
)
|
||||
rows = result.all()
|
||||
assert {r.name for r in rows} == set(names)
|
||||
274
tests/integration/test_insufficient_balance.py
Normal file
274
tests/integration/test_insufficient_balance.py
Normal file
@@ -0,0 +1,274 @@
|
||||
"""
|
||||
Tests showing how a user hits "Insufficient balance: X mSats required for this model"
|
||||
when their balance is too low for the model's cost.
|
||||
|
||||
The log line that triggered this:
|
||||
WARNING Insufficient billing balance during validation
|
||||
ERROR Bearer token validation failed: HTTPException: 402:
|
||||
{'error': {'message': 'Insufficient balance: 622888 mSats required
|
||||
for this model. 20320 available.', ...}}
|
||||
|
||||
This happens in validate_bearer_key (auth.py) when:
|
||||
billing_key.total_balance < min_cost (model's max cost)
|
||||
|
||||
and also in pay_for_request when the atomic UPDATE finds no available balance.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
|
||||
def _key(balance: int, reserved: int = 0) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — simplest case: balance < model cost → pay_for_request raises 402
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_when_balance_too_low(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
User has 20_000 msats. Model costs 622_888 msats.
|
||||
pay_for_request must raise HTTP 402 with a clear message.
|
||||
"""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 622_888
|
||||
user_balance = 20_000
|
||||
|
||||
key = _key(balance=user_balance)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
detail = exc_info.value.detail
|
||||
assert isinstance(detail, dict)
|
||||
error = detail["error"]
|
||||
assert error["code"] == "insufficient_balance"
|
||||
assert str(model_cost) in error["message"]
|
||||
assert str(user_balance) in error["message"]
|
||||
|
||||
# Balance must be untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == user_balance
|
||||
assert key.reserved_balance == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — balance is zero
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_on_zero_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""User with zero balance cannot make any request."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
key = _key(balance=0)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, 1_000, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
detail = exc_info.value.detail
|
||||
assert isinstance(detail, dict)
|
||||
assert detail["error"]["code"] == "insufficient_balance"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — all balance is reserved (total_balance = balance - reserved = 0)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_when_all_balance_reserved(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
User has 50_000 msats balance but 50_000 is already reserved for in-flight
|
||||
requests. Free balance (total_balance) = 0. Should get 402.
|
||||
"""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
key = _key(balance=50_000, reserved=50_000)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, 1_000, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
# Balance and reserved must be untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == 50_000
|
||||
assert key.reserved_balance == 50_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_info_matches_chat_available_balance(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
Regression for /v1/balance/info showing gross funds while chat admission
|
||||
rejects with a negative available balance.
|
||||
"""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
key = _key(balance=4_404_339, reserved=4_410_636)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response = await integration_client.get(
|
||||
"/v1/balance/info",
|
||||
headers={"Authorization": f"Bearer sk-{key.hashed_key}"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["balance"] == -6_297
|
||||
assert body["reserved"] == 4_410_636
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, 1, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
detail = exc_info.value.detail
|
||||
assert isinstance(detail, dict)
|
||||
assert "-6297 available" in detail["error"]["message"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — balance just one msat below model cost
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_one_msat_short(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Off-by-one: balance is exactly model_cost - 1."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 10_000
|
||||
key = _key(balance=model_cost - 1)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == model_cost - 1 # untouched
|
||||
assert key.reserved_balance == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — balance exactly equal to model cost → succeeds
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_succeeds_when_balance_equals_cost(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Balance == model cost: the request should be reserved successfully."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 10_000
|
||||
key = _key(balance=model_cost)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Should not raise
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == model_cost
|
||||
assert key.balance == model_cost # balance unchanged, only reserved goes up
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — HTTP layer returns 402 JSON with the right shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_402_response_shape_on_insufficient_balance(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
End-to-end: POST /v1/chat/completions with a key whose balance is far below
|
||||
the mocked model cost returns HTTP 402 with the expected JSON error body.
|
||||
|
||||
Matches exactly the log snippet in the bug report:
|
||||
'Insufficient balance: X mSats required for this model. Y available.'
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
model_cost = 622_888
|
||||
user_balance = 20_320
|
||||
|
||||
key = _key(balance=user_balance)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Minimal model stub so proxy routing doesn't 400 before reaching balance check
|
||||
mock_model = MagicMock()
|
||||
mock_model.sats_pricing = None
|
||||
|
||||
# Upstream stub — never reached because balance check fires first
|
||||
mock_upstream = MagicMock()
|
||||
mock_upstream.prepare_headers = MagicMock(return_value={})
|
||||
|
||||
with (
|
||||
patch("routstr.proxy.get_model_instance", return_value=mock_model),
|
||||
patch("routstr.proxy.get_provider_for_model", return_value=[mock_upstream]),
|
||||
# Patch where it is used (proxy imports it at module level)
|
||||
patch(
|
||||
"routstr.proxy.get_max_cost_for_model",
|
||||
new=AsyncMock(return_value=model_cost),
|
||||
),
|
||||
):
|
||||
response = await integration_client.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer sk-{key.hashed_key}"},
|
||||
json={
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 402
|
||||
body = response.json()
|
||||
# FastAPI wraps HTTPException detail under "detail"
|
||||
error = body["detail"]["error"]
|
||||
assert error["code"] == "insufficient_balance"
|
||||
assert error["type"] == "insufficient_quota"
|
||||
assert str(model_cost) in error["message"]
|
||||
assert str(user_balance) in error["message"]
|
||||
|
||||
# Balance must be completely untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == user_balance
|
||||
assert key.reserved_balance == 0
|
||||
assert key.total_spent == 0
|
||||
@@ -1,12 +1,14 @@
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel import select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import pay_for_request
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.db import ApiKey, create_session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -120,6 +122,270 @@ async def test_periodic_key_reset_job(integration_session: AsyncSession) -> None
|
||||
assert key2.total_spent == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_limit_enforced_atomically_under_concurrency(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
parent_hash = "parent_limit_atomic"
|
||||
child_hash = "child_limit_atomic"
|
||||
cost = 300
|
||||
|
||||
async with create_session() as session:
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=10000)
|
||||
child = ApiKey(
|
||||
hashed_key=child_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
balance_limit=cost,
|
||||
total_spent=0,
|
||||
)
|
||||
session.add(parent)
|
||||
session.add(child)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_child = await session.get(ApiKey, child_hash)
|
||||
assert fresh_child is not None
|
||||
try:
|
||||
await pay_for_request(fresh_child, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt(), attempt())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final_child = await session.get(ApiKey, child_hash)
|
||||
assert final_child is not None
|
||||
|
||||
assert final_child.reserved_balance == cost, (
|
||||
f"Child reserved_balance should equal one reservation, "
|
||||
f"got {final_child.reserved_balance}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_payments_with_parent_and_child_key(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
parent_hash = "parent_parallel_mixed"
|
||||
child_hash = "child_parallel_mixed"
|
||||
cost = 300
|
||||
|
||||
async with create_session() as session:
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=10000)
|
||||
child = ApiKey(
|
||||
hashed_key=child_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
balance_limit=2 * cost,
|
||||
)
|
||||
session.add(parent)
|
||||
session.add(child)
|
||||
await session.commit()
|
||||
|
||||
async def attempt(key_hash: str) -> str:
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
try:
|
||||
await pay_for_request(fresh_key, cost, session)
|
||||
return "success"
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
return "blocked"
|
||||
|
||||
results = await asyncio.gather(attempt(parent_hash), attempt(child_hash))
|
||||
|
||||
assert results == ["success", "success"], (
|
||||
f"Both parent and child payments should succeed, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final_parent = await session.get(ApiKey, parent_hash)
|
||||
final_child = await session.get(ApiKey, child_hash)
|
||||
assert final_parent is not None
|
||||
assert final_child is not None
|
||||
|
||||
# Both requests bill the parent; only the child request reserves on the child.
|
||||
assert final_parent.reserved_balance == 2 * cost
|
||||
assert final_parent.total_requests == 2
|
||||
assert final_child.reserved_balance == cost
|
||||
assert final_child.total_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_limit_with_existing_total_spent_under_concurrency(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
parent_hash = "parent_total_spent"
|
||||
child_hash = "child_total_spent"
|
||||
cost = 300
|
||||
|
||||
async with create_session() as session:
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=10000)
|
||||
# 700 already spent against a 1000 limit: only one more 300 request fits.
|
||||
child = ApiKey(
|
||||
hashed_key=child_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
balance_limit=1000,
|
||||
total_spent=700,
|
||||
)
|
||||
session.add(parent)
|
||||
session.add(child)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_child = await session.get(ApiKey, child_hash)
|
||||
assert fresh_child is not None
|
||||
try:
|
||||
await pay_for_request(fresh_child, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt(), attempt())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final_child = await session.get(ApiKey, child_hash)
|
||||
assert final_child is not None
|
||||
|
||||
assert final_child.reserved_balance == cost
|
||||
assert final_child.total_spent == 700
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_limit_with_existing_reserved_balance(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
parent_hash = "parent_reserved_set"
|
||||
blocked_hash = "child_reserved_blocked"
|
||||
allowed_hash = "child_reserved_allowed"
|
||||
cost = 300
|
||||
|
||||
async with create_session() as session:
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=10000)
|
||||
# 800 already reserved against a 1000 limit: another 300 must be rejected.
|
||||
blocked_child = ApiKey(
|
||||
hashed_key=blocked_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
balance_limit=1000,
|
||||
reserved_balance=800,
|
||||
)
|
||||
# 500 reserved against a 1000 limit: another 300 still fits.
|
||||
allowed_child = ApiKey(
|
||||
hashed_key=allowed_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
balance_limit=1000,
|
||||
reserved_balance=500,
|
||||
)
|
||||
session.add(parent)
|
||||
session.add(blocked_child)
|
||||
session.add(allowed_child)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
fresh_blocked = await session.get(ApiKey, blocked_hash)
|
||||
assert fresh_blocked is not None
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(fresh_blocked, cost, session)
|
||||
assert exc_info.value.status_code == 402
|
||||
|
||||
async with create_session() as session:
|
||||
fresh_allowed = await session.get(ApiKey, allowed_hash)
|
||||
assert fresh_allowed is not None
|
||||
await pay_for_request(fresh_allowed, cost, session)
|
||||
|
||||
async with create_session() as session:
|
||||
final_blocked = await session.get(ApiKey, blocked_hash)
|
||||
final_allowed = await session.get(ApiKey, allowed_hash)
|
||||
final_parent = await session.get(ApiKey, parent_hash)
|
||||
assert final_blocked is not None
|
||||
assert final_allowed is not None
|
||||
assert final_parent is not None
|
||||
|
||||
assert final_blocked.reserved_balance == 800, "Rejected request must not reserve"
|
||||
assert final_blocked.total_requests == 0
|
||||
assert final_allowed.reserved_balance == 500 + cost
|
||||
assert final_allowed.total_requests == 1
|
||||
# Only the allowed request should have billed the parent.
|
||||
assert final_parent.reserved_balance == cost
|
||||
assert final_parent.total_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_reservation_discarded_when_parent_balance_depleted(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
parent_hash = "parent_depleted"
|
||||
child_hash = "child_depleted"
|
||||
cost = 300
|
||||
|
||||
async with create_session() as session:
|
||||
# Parent can only afford one request; child has no balance_limit.
|
||||
parent = ApiKey(hashed_key=parent_hash, balance=cost)
|
||||
child = ApiKey(
|
||||
hashed_key=child_hash,
|
||||
balance=0,
|
||||
parent_key_hash=parent_hash,
|
||||
)
|
||||
session.add(parent)
|
||||
session.add(child)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_child = await session.get(ApiKey, child_hash)
|
||||
assert fresh_child is not None
|
||||
try:
|
||||
await pay_for_request(fresh_child, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt(), attempt())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final_parent = await session.get(ApiKey, parent_hash)
|
||||
final_child = await session.get(ApiKey, child_hash)
|
||||
assert final_parent is not None
|
||||
assert final_child is not None
|
||||
|
||||
assert final_parent.reserved_balance == cost
|
||||
# The failed request must not leave a committed reservation on the child.
|
||||
assert final_child.reserved_balance == cost, (
|
||||
f"Child reserved_balance should reflect only the successful request, "
|
||||
f"got {final_child.reserved_balance}"
|
||||
)
|
||||
assert final_child.total_requests == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_does_not_delete_key(integration_session: AsyncSession) -> None:
|
||||
# This requires mocking the router call or testing the logic in balance.py
|
||||
|
||||
159
tests/integration/test_lightning_invoice_constraints.py
Normal file
159
tests/integration/test_lightning_invoice_constraints.py
Normal file
@@ -0,0 +1,159 @@
|
||||
"""Integration tests for Lightning invoice key constraint fields.
|
||||
|
||||
Covers two things:
|
||||
- The three constraint fields (balance_limit, balance_limit_reset, validity_date)
|
||||
are persisted on LightningInvoice and survive a DB round-trip.
|
||||
- create_api_key_from_invoice propagates those fields to the created ApiKey,
|
||||
so the constraints are actually enforced when the key is used.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey, LightningInvoice
|
||||
from routstr.lightning import create_api_key_from_invoice
|
||||
|
||||
|
||||
def _make_invoice(**kwargs: object) -> LightningInvoice:
|
||||
base = dict(
|
||||
id="inv_test_001",
|
||||
bolt11="lnbc1000n1test",
|
||||
amount_sats=1000,
|
||||
description="test invoice",
|
||||
payment_hash="deadbeef" * 8,
|
||||
status="paid",
|
||||
purpose="create",
|
||||
expires_at=int(time.time()) + 3600,
|
||||
paid_at=int(time.time()),
|
||||
)
|
||||
base.update(kwargs)
|
||||
return LightningInvoice(**base) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_wallet_mint() -> object:
|
||||
with patch("routstr.lightning.get_wallet") as mock_get_wallet:
|
||||
wallet = AsyncMock()
|
||||
wallet.mint = AsyncMock(return_value=[])
|
||||
mock_get_wallet.return_value = wallet
|
||||
yield mock_get_wallet
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_persists_balance_limit(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
invoice = _make_invoice(balance_limit=5000)
|
||||
integration_session.add(invoice)
|
||||
await integration_session.commit()
|
||||
|
||||
stored = await integration_session.get(LightningInvoice, invoice.id)
|
||||
assert stored is not None
|
||||
assert stored.balance_limit == 5000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_persists_balance_limit_reset(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
invoice = _make_invoice(balance_limit=5000, balance_limit_reset="daily")
|
||||
integration_session.add(invoice)
|
||||
await integration_session.commit()
|
||||
|
||||
stored = await integration_session.get(LightningInvoice, invoice.id)
|
||||
assert stored is not None
|
||||
assert stored.balance_limit_reset == "daily"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoice_persists_validity_date(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
expiry = int(time.time()) + 86400
|
||||
invoice = _make_invoice(validity_date=expiry)
|
||||
integration_session.add(invoice)
|
||||
await integration_session.commit()
|
||||
|
||||
stored = await integration_session.get(LightningInvoice, invoice.id)
|
||||
assert stored is not None
|
||||
assert stored.validity_date == expiry
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Propagation to ApiKey
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_created_key_receives_balance_limit(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
invoice = _make_invoice(balance_limit=8000)
|
||||
integration_session.add(invoice)
|
||||
await integration_session.flush()
|
||||
|
||||
api_key = await create_api_key_from_invoice(invoice, integration_session)
|
||||
await integration_session.commit()
|
||||
|
||||
stored_key = await integration_session.get(ApiKey, api_key.hashed_key)
|
||||
assert stored_key is not None
|
||||
assert stored_key.balance_limit == 8000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_created_key_receives_balance_limit_reset(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
invoice = _make_invoice(balance_limit=8000, balance_limit_reset="monthly")
|
||||
integration_session.add(invoice)
|
||||
await integration_session.flush()
|
||||
|
||||
api_key = await create_api_key_from_invoice(invoice, integration_session)
|
||||
await integration_session.commit()
|
||||
|
||||
stored_key = await integration_session.get(ApiKey, api_key.hashed_key)
|
||||
assert stored_key is not None
|
||||
assert stored_key.balance_limit_reset == "monthly"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_created_key_receives_validity_date(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
expiry = int(time.time()) + 86400
|
||||
invoice = _make_invoice(validity_date=expiry)
|
||||
integration_session.add(invoice)
|
||||
await integration_session.flush()
|
||||
|
||||
api_key = await create_api_key_from_invoice(invoice, integration_session)
|
||||
await integration_session.commit()
|
||||
|
||||
stored_key = await integration_session.get(ApiKey, api_key.hashed_key)
|
||||
assert stored_key is not None
|
||||
assert stored_key.validity_date == expiry
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_created_key_without_constraints_has_none_fields(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
invoice = _make_invoice()
|
||||
integration_session.add(invoice)
|
||||
await integration_session.flush()
|
||||
|
||||
api_key = await create_api_key_from_invoice(invoice, integration_session)
|
||||
await integration_session.commit()
|
||||
|
||||
stored_key = await integration_session.get(ApiKey, api_key.hashed_key)
|
||||
assert stored_key is not None
|
||||
assert stored_key.balance_limit is None
|
||||
assert stored_key.balance_limit_reset is None
|
||||
assert stored_key.validity_date is None
|
||||
198
tests/integration/test_lightning_invoice_rip08.py
Normal file
198
tests/integration/test_lightning_invoice_rip08.py
Normal file
@@ -0,0 +1,198 @@
|
||||
"""RIP-08 lightning invoice endpoint tests.
|
||||
|
||||
Verifies both the spec-compliant path (`POST /lightning/invoice` with
|
||||
`Authorization: Bearer sk-...`) and the legacy path
|
||||
(`POST /v1/balance/lightning/invoice` with `api_key` in body).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
RIP08_PATH = "/lightning/invoice"
|
||||
LEGACY_PATH = "/v1/balance/lightning/invoice"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def patch_invoice_generation() -> Any:
|
||||
"""Stub out `generate_lightning_invoice` so no mint round-trip is needed."""
|
||||
counter = {"n": 0}
|
||||
|
||||
async def fake_generate(amount_sats: int, description: str) -> tuple[str, str]:
|
||||
counter["n"] += 1
|
||||
return (
|
||||
f"lnbc{amount_sats}n1pfakeinvoice{counter['n']}",
|
||||
f"payment_hash_{counter['n']}",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"routstr.lightning.generate_lightning_invoice",
|
||||
side_effect=fake_generate,
|
||||
) as m:
|
||||
yield m
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def seeded_topup_key(integration_session: AsyncSession) -> str:
|
||||
"""Insert an ApiKey row and return the public `sk-...` form."""
|
||||
hashed = "0" * 64
|
||||
key = ApiKey(
|
||||
hashed_key=hashed,
|
||||
balance=0,
|
||||
refund_currency="sat",
|
||||
refund_mint_url="http://localhost:3338",
|
||||
)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
return f"sk-{hashed}"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_create_invoice_purpose_create(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
path: str,
|
||||
) -> None:
|
||||
"""`purpose=create` works on both paths and requires no auth."""
|
||||
resp = await integration_client.post(
|
||||
path,
|
||||
json={"amount_sats": 1000, "purpose": "create"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
assert body["amount_sats"] == 1000
|
||||
assert body["bolt11"].startswith("lnbc")
|
||||
assert body["invoice_id"]
|
||||
assert body["payment_hash"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_topup_with_authorization_header(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
seeded_topup_key: str,
|
||||
path: str,
|
||||
) -> None:
|
||||
"""RIP-08: topup using `Authorization: Bearer sk-...` header (no api_key in body)."""
|
||||
resp = await integration_client.post(
|
||||
path,
|
||||
json={"amount_sats": 500, "purpose": "topup"},
|
||||
headers={"Authorization": f"Bearer {seeded_topup_key}"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
body = resp.json()
|
||||
assert body["amount_sats"] == 500
|
||||
assert body["bolt11"].startswith("lnbc")
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_topup_with_legacy_api_key_in_body(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
seeded_topup_key: str,
|
||||
path: str,
|
||||
) -> None:
|
||||
"""Legacy: topup with `api_key` in body still accepted on both paths."""
|
||||
resp = await integration_client.post(
|
||||
path,
|
||||
json={
|
||||
"amount_sats": 250,
|
||||
"purpose": "topup",
|
||||
"api_key": seeded_topup_key,
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["amount_sats"] == 250
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_topup_missing_auth_returns_401(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
path: str,
|
||||
) -> None:
|
||||
"""Topup without any credential is rejected on both paths."""
|
||||
resp = await integration_client.post(
|
||||
path,
|
||||
json={"amount_sats": 100, "purpose": "topup"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_topup_unknown_api_key_returns_404(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
path: str,
|
||||
) -> None:
|
||||
resp = await integration_client.post(
|
||||
path,
|
||||
json={"amount_sats": 100, "purpose": "topup"},
|
||||
headers={"Authorization": "Bearer sk-deadbeef"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", [RIP08_PATH, LEGACY_PATH])
|
||||
async def test_invoice_status_404_for_unknown_id(
|
||||
integration_client: AsyncClient,
|
||||
path: str,
|
||||
) -> None:
|
||||
base = path.rsplit("/invoice", 1)[0] + "/invoice"
|
||||
resp = await integration_client.get(f"{base}/does-not-exist/status")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_purpose_defaults_to_create(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
) -> None:
|
||||
"""Per RIP-08, `purpose` may be omitted and defaults to `create`."""
|
||||
resp = await integration_client.post(
|
||||
RIP08_PATH,
|
||||
json={"amount_sats": 100},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["amount_sats"] == 100
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_header_overrides_body_api_key(
|
||||
integration_client: AsyncClient,
|
||||
patch_invoice_generation: Any,
|
||||
seeded_topup_key: str,
|
||||
) -> None:
|
||||
"""Header api_key wins over body api_key: bogus body must not cause 404."""
|
||||
resp = await integration_client.post(
|
||||
RIP08_PATH,
|
||||
json={
|
||||
"amount_sats": 100,
|
||||
"purpose": "topup",
|
||||
"api_key": "sk-" + "f" * 64, # bogus body key
|
||||
},
|
||||
headers={"Authorization": f"Bearer {seeded_topup_key}"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
224
tests/integration/test_model_test_endpoint_security.py
Normal file
224
tests/integration/test_model_test_endpoint_security.py
Normal file
@@ -0,0 +1,224 @@
|
||||
import time
|
||||
from types import TracebackType
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import ModelRow, UpstreamProviderRow
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_test_endpoint_requires_admin_auth(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
with patch("httpx.AsyncClient") as mock_async_client:
|
||||
response = await integration_client.post(
|
||||
"/api/models/test",
|
||||
json={
|
||||
"model_id": "model-a",
|
||||
"endpoint_type": "chat-completions",
|
||||
"request_data": {"messages": []},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
mock_async_client.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_test_endpoint_rejects_unsupported_endpoint_type(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
admin_token = "test-admin-token-model-test"
|
||||
admin_sessions[admin_token] = int(time.time()) + 3600
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
|
||||
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url="https://api.example.com/v1",
|
||||
api_key="sk-upstream-test",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
assert provider.id is not None
|
||||
|
||||
model = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=provider.id,
|
||||
enabled=True,
|
||||
)
|
||||
integration_session.add(model)
|
||||
await integration_session.commit()
|
||||
|
||||
try:
|
||||
with patch("httpx.AsyncClient") as mock_async_client:
|
||||
response = await integration_client.post(
|
||||
"/api/models/test",
|
||||
json={
|
||||
"model_id": "model-a",
|
||||
"endpoint_type": "../../abuse",
|
||||
"request_data": {"messages": []},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Unsupported endpoint_type"
|
||||
mock_async_client.assert_not_called()
|
||||
finally:
|
||||
admin_sessions.pop(admin_token, None)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_test_endpoint_rejects_oversized_request_data(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
admin_token = "test-admin-token-model-test-oversized"
|
||||
admin_sessions[admin_token] = int(time.time()) + 3600
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
|
||||
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url="https://api.example.com/v1",
|
||||
api_key="sk-upstream-test",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
assert provider.id is not None
|
||||
|
||||
model = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=provider.id,
|
||||
enabled=True,
|
||||
)
|
||||
integration_session.add(model)
|
||||
await integration_session.commit()
|
||||
|
||||
oversized = "x" * (64 * 1024 + 1)
|
||||
|
||||
try:
|
||||
with patch("httpx.AsyncClient") as mock_async_client:
|
||||
response = await integration_client.post(
|
||||
"/api/models/test",
|
||||
json={
|
||||
"model_id": "model-a",
|
||||
"endpoint_type": "chat-completions",
|
||||
"request_data": {"blob": oversized},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 413
|
||||
assert response.json()["detail"] == "request_data too large"
|
||||
mock_async_client.assert_not_called()
|
||||
finally:
|
||||
admin_sessions.pop(admin_token, None)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_test_endpoint_admin_uses_allowed_upstream_path(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
admin_token = "test-admin-token-model-test-success"
|
||||
admin_sessions[admin_token] = int(time.time()) + 3600
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
|
||||
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url="https://api.example.com/v1",
|
||||
api_key="sk-upstream-test",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
assert provider.id is not None
|
||||
|
||||
model = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=provider.id,
|
||||
enabled=True,
|
||||
forwarded_model_id="upstream-model-a",
|
||||
)
|
||||
integration_session.add(model)
|
||||
await integration_session.commit()
|
||||
|
||||
class MockResponse:
|
||||
status_code = 200
|
||||
text = '{"ok": true}'
|
||||
|
||||
def json(self) -> dict[str, bool]:
|
||||
return {"ok": True}
|
||||
|
||||
class MockAsyncClient:
|
||||
async def __aenter__(self) -> "MockAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
async def post(
|
||||
self, url: str, json: dict[str, Any], headers: dict[str, str]
|
||||
) -> MockResponse:
|
||||
assert url == "https://api.example.com/v1/chat/completions"
|
||||
assert json["model"] == "upstream-model-a"
|
||||
assert headers["Authorization"] == "Bearer sk-upstream-test"
|
||||
return MockResponse()
|
||||
|
||||
try:
|
||||
with patch("httpx.AsyncClient", return_value=MockAsyncClient()):
|
||||
response = await integration_client.post(
|
||||
"/api/models/test",
|
||||
json={
|
||||
"model_id": "model-a",
|
||||
"endpoint_type": "chat-completions",
|
||||
"request_data": {"messages": []},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"success": True,
|
||||
"data": {"ok": True},
|
||||
"status_code": 200,
|
||||
}
|
||||
finally:
|
||||
admin_sessions.pop(admin_token, None)
|
||||
143
tests/integration/test_provider_fee_enforcement.py
Normal file
143
tests/integration/test_provider_fee_enforcement.py
Normal file
@@ -0,0 +1,143 @@
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.core.db import AsyncSession, ModelRow, UpstreamProviderRow
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
from routstr.proxy import refresh_model_maps
|
||||
from routstr.upstream.base import BaseUpstreamProvider
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_lowest_provider_fee_for_same_url(
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
"""Test that the algorithm selects the provider with the lowest fee when URLs match."""
|
||||
|
||||
# 1. Create two providers with the same URL but different fees
|
||||
url = "https://api.example.com"
|
||||
p1 = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url=url,
|
||||
api_key="key1",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
p2 = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url=url,
|
||||
api_key="key2",
|
||||
enabled=True,
|
||||
provider_fee=1.05,
|
||||
)
|
||||
|
||||
integration_session.add(p1)
|
||||
integration_session.add(p2)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(p1)
|
||||
await integration_session.refresh(p2)
|
||||
|
||||
assert p1.id is not None
|
||||
assert p2.id is not None
|
||||
|
||||
# 2. Add a model for each provider
|
||||
m1 = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=p1.id,
|
||||
enabled=True,
|
||||
)
|
||||
m2 = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=p2.id,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
integration_session.add(m1)
|
||||
integration_session.add(m2)
|
||||
await integration_session.commit()
|
||||
|
||||
# 3. Create mock provider instances
|
||||
class MockProvider(BaseUpstreamProvider):
|
||||
db_id: int
|
||||
|
||||
def __init__(self, db_id: int, base_url: str, api_key: str, fee: float):
|
||||
super().__init__(base_url, api_key, fee)
|
||||
self.db_id = db_id
|
||||
self.provider_type = "custom"
|
||||
|
||||
def get_cached_models(self) -> list[Model]:
|
||||
return [
|
||||
Model(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture=Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="tiktoken",
|
||||
instruct_type="chat",
|
||||
),
|
||||
pricing=Pricing(prompt=1.0, completion=1.0),
|
||||
enabled=True,
|
||||
upstream_provider_id=self.db_id,
|
||||
)
|
||||
]
|
||||
|
||||
async def refresh_models_cache(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare_headers(self, request_headers: dict[str, str]) -> dict[str, str]:
|
||||
return request_headers
|
||||
|
||||
# 4. Inject mock providers into the proxy
|
||||
from routstr import proxy
|
||||
|
||||
assert p1.id is not None
|
||||
assert p2.id is not None
|
||||
|
||||
# Need to patch proxy._upstreams and proxy.create_session
|
||||
mp1: MockProvider = MockProvider(p1.id, url, "key1", 1.01)
|
||||
mp2: MockProvider = MockProvider(p2.id, url, "key2", 1.05)
|
||||
|
||||
with (
|
||||
patch("routstr.proxy._upstreams", [mp1, mp2]),
|
||||
patch("routstr.proxy.create_session") as mock_session_factory,
|
||||
):
|
||||
# Configure mock_session_factory to return a session that uses the test engine
|
||||
@asynccontextmanager
|
||||
async def mock_create_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
yield integration_session
|
||||
|
||||
mock_session_factory.return_value = mock_create_session()
|
||||
|
||||
await refresh_model_maps()
|
||||
|
||||
# 5. Check which provider is selected for 'model-a'
|
||||
provider_map = proxy.get_provider_for_model("model-a")
|
||||
|
||||
# Assertions
|
||||
assert provider_map is not None
|
||||
assert len(provider_map) >= 1
|
||||
|
||||
# Check the first one, cast to MockProvider to access db_id
|
||||
best_provider = cast(MockProvider, provider_map[0])
|
||||
assert best_provider.db_id == p1.id
|
||||
assert best_provider.provider_fee == 1.01
|
||||
@@ -3,12 +3,16 @@ Integration tests for provider management functionality.
|
||||
Tests GET /v1/providers/ endpoint for listing and managing providers.
|
||||
"""
|
||||
|
||||
import time
|
||||
from types import TracebackType
|
||||
from typing import Any, Generator
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import UpstreamProviderRow
|
||||
from routstr.nostr.discovery import _PROVIDERS_CACHE
|
||||
|
||||
from .utils import ResponseValidator
|
||||
@@ -678,3 +682,88 @@ async def test_no_database_changes_during_provider_operations(
|
||||
assert final_diff["api_keys"]["added"] == []
|
||||
assert final_diff["api_keys"]["modified"] == []
|
||||
assert final_diff["api_keys"]["removed"] == []
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_topup_retries_transient_upstream_failure(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
admin_token = "test-admin-token"
|
||||
admin_sessions[admin_token] = int(time.time()) + 3600
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
|
||||
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="routstr",
|
||||
base_url="https://node.example",
|
||||
api_key="sk-upstream-test",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, status_code: int, data: dict[str, Any] | None = None):
|
||||
self.status_code = status_code
|
||||
self._data = data or {}
|
||||
self.text = str(self._data)
|
||||
|
||||
def json(self) -> dict[str, Any]:
|
||||
return self._data
|
||||
|
||||
class MockAsyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def __aenter__(self) -> "MockAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
async def post(
|
||||
self, url: str, json: dict[str, Any], headers: dict[str, str]
|
||||
) -> MockResponse:
|
||||
self.calls += 1
|
||||
assert url == "https://node.example/v1/balance/lightning/invoice"
|
||||
assert json["amount_sats"] == 10
|
||||
assert json["purpose"] == "topup"
|
||||
assert json["api_key"] == "sk-upstream-test"
|
||||
assert headers["Authorization"] == "Bearer sk-upstream-test"
|
||||
|
||||
if self.calls == 1:
|
||||
return MockResponse(500, {"detail": "warmup failure"})
|
||||
|
||||
return MockResponse(
|
||||
200,
|
||||
{
|
||||
"bolt11": "lnbc1testinvoice",
|
||||
"invoice_id": "invoice-123",
|
||||
},
|
||||
)
|
||||
|
||||
mock_client = MockAsyncClient()
|
||||
|
||||
try:
|
||||
with patch("httpx.AsyncClient", return_value=mock_client):
|
||||
response = await integration_client.post(
|
||||
f"/admin/api/upstream-providers/{provider.id}/topup",
|
||||
json={"amount": 10},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["ok"] is True
|
||||
assert data["topup_data"]["payment_request"] == "lnbc1testinvoice"
|
||||
assert data["topup_data"]["invoice_id"] == "invoice-123"
|
||||
assert mock_client.calls == 2
|
||||
finally:
|
||||
admin_sessions.pop(admin_token, None)
|
||||
|
||||
@@ -171,10 +171,11 @@ async def test_proxy_get_unauthorized_access(integration_client: AsyncClient) ->
|
||||
assert response.status_code == 200 # GET requests are allowed
|
||||
|
||||
# Test 2: POST requests without auth should return 401
|
||||
# Note: Model validation happens before auth, so missing model returns 400
|
||||
response = await integration_client.post(
|
||||
"/v1/chat/completions", json={"test": "data"}
|
||||
)
|
||||
assert response.status_code == 401
|
||||
assert response.status_code in [400, 401] # Accept both for now
|
||||
|
||||
# Test 3: POST with invalid API key
|
||||
# Note: After refactor, model validation may happen before auth validation
|
||||
@@ -550,9 +551,6 @@ async def test_proxy_get_concurrent_requests(
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_get_response_format_preservation(
|
||||
|
||||
268
tests/integration/test_reservation_lifecycle.py
Normal file
268
tests/integration/test_reservation_lifecycle.py
Normal file
@@ -0,0 +1,268 @@
|
||||
"""
|
||||
Tests for the reservation lifecycle:
|
||||
|
||||
1. Reserve → reserved_balance increases, available (total_balance) decreases.
|
||||
2. Reserve → revert → reserved_balance restored, balance untouched.
|
||||
3. Reserve → finalise → reserved_balance released, balance charged.
|
||||
4. Two parallel reserves, only one fits → second blocked with 402.
|
||||
5. Three parallel reserves, two fit, third blocked with 402.
|
||||
6. Sequential reserves until balance exhausted → next request blocked.
|
||||
|
||||
Reservation invariant enforced by the atomic WHERE clause in pay_for_request:
|
||||
balance - reserved_balance >= cost_per_request
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import pay_for_request, revert_pay_for_request
|
||||
from routstr.core.db import ApiKey, create_session
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_key(balance: int, reserved: int = 0) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
|
||||
|
||||
async def _persist(session: AsyncSession, key: ApiKey) -> ApiKey:
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
return key
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — Reserve: reserved_balance increases, available balance decreases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_increases_reserved_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""pay_for_request must increment reserved_balance by cost_per_request."""
|
||||
cost = 100
|
||||
key = await _persist(integration_session, _make_key(balance=500))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == cost
|
||||
assert key.balance == 500 # balance column is NOT decremented on reserve
|
||||
assert key.total_balance == 500 - cost # available = balance - reserved
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — Revert: reserved_balance restored, balance untouched
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_releases_reservation(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""revert_pay_for_request must release the reservation without touching balance."""
|
||||
cost = 150
|
||||
key = await _persist(integration_session, _make_key(balance=300))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == cost
|
||||
|
||||
await revert_pay_for_request(key, integration_session, cost)
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == 0
|
||||
assert key.balance == 300 # balance unchanged after revert
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — Finalise: reservation released + balance charged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalise_releases_reservation_and_charges_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""adjust_payment_for_tokens must zero reserved_balance and deduct actual cost."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
cost = 100
|
||||
actual = 80 # actual < reserved → refund path
|
||||
key = await _persist(integration_session, _make_key(balance=500))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == cost
|
||||
|
||||
cost_data = CostData(
|
||||
base_msats=0,
|
||||
input_msats=40,
|
||||
output_msats=40,
|
||||
total_msats=actual,
|
||||
total_usd=0.0,
|
||||
input_tokens=50,
|
||||
output_tokens=50,
|
||||
)
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
|
||||
|
||||
with patch("routstr.auth.calculate_cost", return_value=cost_data):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, cost)
|
||||
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == 0
|
||||
assert key.balance == 500 - actual
|
||||
assert key.total_spent == actual
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — Concurrent: second parallel reserve blocked when balance exhausted
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_second_reserve_blocked_when_balance_exhausted(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""When two requests race for the same balance, only one succeeds; the other gets 402."""
|
||||
cost = 300
|
||||
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=300, # exactly enough for ONE reservation
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt_reserve() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
try:
|
||||
await pay_for_request(fresh_key, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt_reserve(), attempt_reserve())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, key_hash)
|
||||
assert final is not None
|
||||
|
||||
# reserved_balance must equal exactly one reservation (not two)
|
||||
assert final.reserved_balance == cost, (
|
||||
f"Expected reserved_balance={cost}, got {final.reserved_balance}"
|
||||
)
|
||||
assert final.balance == 300, "Balance column must not be modified by reservation"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — Concurrent: three requests, two fit, third blocked
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_three_parallel_reserves_third_blocked(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Balance covers two reservations exactly; the third concurrent request must be blocked."""
|
||||
cost = 100
|
||||
key_hash = f"test_three_parallel_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=200, # fits exactly 2 reservations of 100
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt_reserve() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
try:
|
||||
await pay_for_request(fresh_key, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(
|
||||
attempt_reserve(),
|
||||
attempt_reserve(),
|
||||
attempt_reserve(),
|
||||
)
|
||||
|
||||
successes = results.count("success")
|
||||
blocked = results.count("blocked")
|
||||
|
||||
assert successes == 2, f"Expected 2 successes, got {successes}: {results}"
|
||||
assert blocked == 1, f"Expected 1 blocked, got {blocked}: {results}"
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, key_hash)
|
||||
assert final is not None
|
||||
|
||||
assert final.reserved_balance == cost * 2, (
|
||||
f"Expected reserved_balance={cost * 2}, got {final.reserved_balance}"
|
||||
)
|
||||
assert final.balance == 200, "Balance column must not be modified by reservation"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — Sequential exhaustion: reserve until empty, next request blocked
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sequential_reserves_block_when_balance_exhausted(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Repeated reservations should block as soon as available balance drops below cost."""
|
||||
cost = 100
|
||||
key = await _persist(integration_session, _make_key(balance=250))
|
||||
|
||||
# First two succeed (100 + 100 = 200 ≤ 250)
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == 200
|
||||
assert key.total_balance == 50 # 250 - 200
|
||||
|
||||
# Third: only 50 available, need 100 → blocked
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == 200 # unchanged after failed reserve
|
||||
@@ -133,13 +133,16 @@ async def test_reserved_balance_with_successful_requests(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insufficient_reserved_balance_for_revert(
|
||||
async def test_revert_with_zero_reserved_balance_is_noop(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test revert_pay_for_request behavior with insufficient reserved balance."""
|
||||
"""Test that revert_pay_for_request is a no-op when reserved_balance is 0.
|
||||
|
||||
Previously this would drive reserved_balance negative. With the floor guard,
|
||||
it should return False and leave reserved_balance at 0.
|
||||
"""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
# Create key with zero reserved balance
|
||||
unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
@@ -149,17 +152,211 @@ async def test_insufficient_reserved_balance_for_revert(
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Try to revert more than available
|
||||
# Note: Current implementation allows reserved_balance to go negative
|
||||
await revert_pay_for_request(test_key, integration_session, 100)
|
||||
# Try to revert more than available — should be a no-op
|
||||
result = await revert_pay_for_request(test_key, integration_session, 100)
|
||||
|
||||
# Refresh to get updated values
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
# Current implementation allows negative reserved balance
|
||||
assert test_key.reserved_balance == -100, (
|
||||
f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}"
|
||||
assert result is False, "Revert should return False when reservation already released"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should remain 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == -1, (
|
||||
f"Expected total_requests to be -1, got: {test_key.total_requests}"
|
||||
assert test_key.total_requests == 0, (
|
||||
f"Total requests should remain 0, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_with_sufficient_reserved_balance_succeeds(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that revert_pay_for_request works correctly when there is enough reserved balance."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_revert_ok_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=5000,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
result = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result is True, "Revert should return True on success"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 2, (
|
||||
f"Total requests should be 2, got: {test_key.total_requests}"
|
||||
)
|
||||
assert test_key.balance == 5000, "Balance should not change on revert"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_partial_reserved_balance_is_noop(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that reverting more than the current reserved_balance is a no-op."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=5000,
|
||||
reserved_balance=50,
|
||||
total_requests=1,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Try to revert 500 when only 50 is reserved — should be no-op
|
||||
result = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result is False, "Revert should fail when cost > reserved_balance"
|
||||
assert test_key.reserved_balance == 50, (
|
||||
f"Reserved balance should stay at 50, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 1, (
|
||||
f"Total requests should stay at 1, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_double_revert_prevented(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that calling revert twice doesn't drive reserved_balance negative.
|
||||
|
||||
This simulates the double-revert scenario where both upstream/base.py
|
||||
and proxy.py attempt to revert the same reservation.
|
||||
"""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_double_revert_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=5,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# First revert — should succeed
|
||||
result1 = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result1 is True
|
||||
assert test_key.reserved_balance == 0
|
||||
assert test_key.total_requests == 4
|
||||
|
||||
# Second revert of the same amount — should be no-op
|
||||
result2 = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result2 is False, "Second revert should be a no-op"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should stay 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 4, (
|
||||
f"Total requests should stay 4, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sequential_reverts_never_go_negative(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that multiple reverts don't cause negative reserved_balance.
|
||||
|
||||
Simulates the double-revert scenario where multiple code paths
|
||||
attempt to revert the same reservation.
|
||||
"""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_multi_revert_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=5,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Run 5 sequential reverts for the same 500 reservation
|
||||
results = []
|
||||
for _ in range(5):
|
||||
r = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
results.append(r)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
# Exactly one should succeed, rest should be no-ops
|
||||
success_count = sum(1 for r in results if r is True)
|
||||
assert success_count == 1, (
|
||||
f"Exactly one revert should succeed, got {success_count} successes"
|
||||
)
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.reserved_balance >= 0, (
|
||||
f"Reserved balance went negative: {test_key.reserved_balance}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_revert_floor_guard(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that child key reserved_balance also has floor guard on revert."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
parent_key_hash = f"test_parent_{uuid.uuid4().hex[:8]}"
|
||||
child_key_hash = f"test_child_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
parent_key = ApiKey(
|
||||
hashed_key=parent_key_hash,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
)
|
||||
child_key = ApiKey(
|
||||
hashed_key=child_key_hash,
|
||||
balance=0,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
parent_key_hash=parent_key_hash,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
integration_session.add(child_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# First revert succeeds
|
||||
result1 = await revert_pay_for_request(child_key, integration_session, 500)
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
assert result1 is True
|
||||
assert parent_key.reserved_balance == 0
|
||||
assert child_key.reserved_balance == 0
|
||||
|
||||
# Second revert is a no-op for both parent and child
|
||||
result2 = await revert_pay_for_request(child_key, integration_session, 500)
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
assert result2 is False
|
||||
assert parent_key.reserved_balance == 0, (
|
||||
f"Parent reserved_balance should stay 0, got: {parent_key.reserved_balance}"
|
||||
)
|
||||
assert child_key.reserved_balance == 0, (
|
||||
f"Child reserved_balance should stay 0, got: {child_key.reserved_balance}"
|
||||
)
|
||||
|
||||
210
tests/integration/test_temporary_balances_api.py
Normal file
210
tests/integration/test_temporary_balances_api.py
Normal file
@@ -0,0 +1,210 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlmodel import col, update
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
|
||||
def _admin_headers() -> dict[str, str]:
|
||||
token = "test-admin-token"
|
||||
admin_sessions[token] = int(
|
||||
(datetime.now(timezone.utc) + timedelta(minutes=5)).timestamp()
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
|
||||
async def _add_key(
|
||||
session: AsyncSession,
|
||||
hashed_key: str,
|
||||
*,
|
||||
balance: int = 0,
|
||||
total_spent: int = 0,
|
||||
total_requests: int = 0,
|
||||
created_at: int | None = None,
|
||||
parent_key_hash: str | None = None,
|
||||
refund_address: str | None = None,
|
||||
) -> ApiKey:
|
||||
key = ApiKey(
|
||||
hashed_key=hashed_key,
|
||||
balance=balance,
|
||||
total_spent=total_spent,
|
||||
total_requests=total_requests,
|
||||
parent_key_hash=parent_key_hash,
|
||||
refund_address=refund_address,
|
||||
)
|
||||
key.created_at = created_at
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
# The model's default_factory is translated into a SQLAlchemy column
|
||||
# default that fires on INSERT whenever the value is None, so a true NULL
|
||||
# (a legacy row created before the column existed) can only be produced by
|
||||
# an explicit UPDATE after insert.
|
||||
if created_at is None:
|
||||
await session.exec(
|
||||
update(ApiKey) # type: ignore[call-overload]
|
||||
.where(col(ApiKey.hashed_key) == hashed_key)
|
||||
.values(created_at=None)
|
||||
)
|
||||
await session.commit()
|
||||
return key
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_envelope_and_created_at(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
await _add_key(integration_session, "key_a", balance=1000, created_at=1000)
|
||||
|
||||
response = await integration_client.get(
|
||||
"/admin/api/temporary-balances", headers=_admin_headers()
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert set(body.keys()) == {"balances", "total", "totals"}
|
||||
assert body["total"] == 1
|
||||
assert body["balances"][0]["hashed_key"] == "key_a"
|
||||
assert body["balances"][0]["created_at"] == 1000
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_sorted_latest_first_nulls_last(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
await _add_key(integration_session, "older", created_at=1000)
|
||||
await _add_key(integration_session, "newer", created_at=2000)
|
||||
await _add_key(integration_session, "legacy", created_at=None)
|
||||
|
||||
response = await integration_client.get(
|
||||
"/admin/api/temporary-balances", headers=_admin_headers()
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
order = [b["hashed_key"] for b in response.json()["balances"]]
|
||||
# Newest created first, NULL created_at (legacy) sorts last.
|
||||
assert order == ["newer", "older", "legacy"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_pagination(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
for i in range(5):
|
||||
await _add_key(integration_session, f"key_{i}", created_at=1000 + i)
|
||||
|
||||
headers = _admin_headers()
|
||||
page1 = (
|
||||
await integration_client.get(
|
||||
"/admin/api/temporary-balances?limit=2&offset=0", headers=headers
|
||||
)
|
||||
).json()
|
||||
page2 = (
|
||||
await integration_client.get(
|
||||
"/admin/api/temporary-balances?limit=2&offset=2", headers=headers
|
||||
)
|
||||
).json()
|
||||
|
||||
assert page1["total"] == 5
|
||||
assert page2["total"] == 5
|
||||
assert [b["hashed_key"] for b in page1["balances"]] == ["key_4", "key_3"]
|
||||
assert [b["hashed_key"] for b in page2["balances"]] == ["key_2", "key_1"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_totals_exclude_child_balance(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
await _add_key(
|
||||
integration_session,
|
||||
"parent",
|
||||
balance=5000,
|
||||
total_spent=100,
|
||||
total_requests=3,
|
||||
created_at=1000,
|
||||
)
|
||||
# Child draws from parent's balance, so its balance must NOT be summed,
|
||||
# but its spent/requests still count.
|
||||
await _add_key(
|
||||
integration_session,
|
||||
"child",
|
||||
balance=0,
|
||||
total_spent=200,
|
||||
total_requests=7,
|
||||
created_at=1001,
|
||||
parent_key_hash="parent",
|
||||
)
|
||||
|
||||
response = await integration_client.get(
|
||||
"/admin/api/temporary-balances", headers=_admin_headers()
|
||||
)
|
||||
|
||||
totals = response.json()["totals"]
|
||||
assert totals["total_balance"] == 5000
|
||||
assert totals["total_spent"] == 300
|
||||
assert totals["total_requests"] == 10
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_search_filters_total_and_totals(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
await _add_key(
|
||||
integration_session,
|
||||
"alpha",
|
||||
balance=1000,
|
||||
created_at=1000,
|
||||
refund_address="alice@ln.tld",
|
||||
)
|
||||
await _add_key(
|
||||
integration_session,
|
||||
"beta",
|
||||
balance=2000,
|
||||
created_at=1001,
|
||||
refund_address="bob@ln.tld",
|
||||
)
|
||||
|
||||
headers = _admin_headers()
|
||||
|
||||
# Match by hashed_key.
|
||||
by_key = (
|
||||
await integration_client.get(
|
||||
"/admin/api/temporary-balances?search=alpha", headers=headers
|
||||
)
|
||||
).json()
|
||||
assert by_key["total"] == 1
|
||||
assert by_key["balances"][0]["hashed_key"] == "alpha"
|
||||
# totals reflect only the filtered set.
|
||||
assert by_key["totals"]["total_balance"] == 1000
|
||||
|
||||
# Match by refund_address.
|
||||
by_addr = (
|
||||
await integration_client.get(
|
||||
"/admin/api/temporary-balances?search=bob@ln.tld", headers=headers
|
||||
)
|
||||
).json()
|
||||
assert by_addr["total"] == 1
|
||||
assert by_addr["balances"][0]["hashed_key"] == "beta"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporary_balances_requires_admin(
|
||||
integration_client: httpx.AsyncClient,
|
||||
) -> None:
|
||||
response = await integration_client.get("/admin/api/temporary-balances")
|
||||
assert response.status_code in (401, 403)
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@@ -356,6 +356,70 @@ async def test_concurrent_refund_requests(
|
||||
assert len(successful) + len(failed) == 5
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_rejects_concurrent_topup_on_same_key(
|
||||
authenticated_client: AsyncClient,
|
||||
testmint_wallet: Any,
|
||||
) -> None:
|
||||
"""Test refund returns 409 when a concurrent topup changes the balance first."""
|
||||
from routstr import balance as balance_module
|
||||
|
||||
wallet_response = await authenticated_client.get("/v1/wallet/")
|
||||
assert wallet_response.status_code == 200
|
||||
initial_balance = wallet_response.json()["balance"]
|
||||
|
||||
topup_amount_sat = 500
|
||||
topup_token = await testmint_wallet.mint_tokens(topup_amount_sat)
|
||||
|
||||
key_looked_up = asyncio.Event()
|
||||
allow_refund_to_continue = asyncio.Event()
|
||||
original_lookup = balance_module._lookup_key_no_create
|
||||
delayed_once = False
|
||||
|
||||
async def delayed_lookup_key_no_create(*args: Any, **kwargs: Any) -> ApiKey | None:
|
||||
nonlocal delayed_once
|
||||
key = await original_lookup(*args, **kwargs)
|
||||
if not delayed_once and key is not None:
|
||||
delayed_once = True
|
||||
key_looked_up.set()
|
||||
await allow_refund_to_continue.wait()
|
||||
return key
|
||||
|
||||
async def issue_refund() -> Any:
|
||||
return await authenticated_client.post("/v1/wallet/refund")
|
||||
|
||||
async def issue_topup() -> Any:
|
||||
await key_looked_up.wait()
|
||||
try:
|
||||
return await authenticated_client.post(
|
||||
"/v1/wallet/topup", params={"cashu_token": topup_token}
|
||||
)
|
||||
finally:
|
||||
allow_refund_to_continue.set()
|
||||
|
||||
with patch(
|
||||
"routstr.balance._lookup_key_no_create", new=delayed_lookup_key_no_create
|
||||
):
|
||||
refund_response, topup_response = await asyncio.gather(
|
||||
issue_refund(), issue_topup()
|
||||
)
|
||||
|
||||
assert topup_response.status_code == 200
|
||||
assert topup_response.json()["msats"] == topup_amount_sat * 1000
|
||||
assert refund_response.status_code == 409
|
||||
assert (
|
||||
refund_response.json()["detail"]
|
||||
== "Balance changed concurrently. Please retry the refund."
|
||||
)
|
||||
|
||||
final_balance_response = await authenticated_client.get("/v1/wallet/")
|
||||
assert final_balance_response.status_code == 200
|
||||
assert final_balance_response.json()["balance"] == (
|
||||
initial_balance + topup_amount_sat * 1000
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_during_active_usage(
|
||||
@@ -394,6 +458,42 @@ async def test_refund_during_active_usage(
|
||||
assert response.json()["balance"] == 0
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_wallet_history_returns_apikey_transactions(
|
||||
authenticated_client: AsyncClient,
|
||||
testmint_wallet: Any,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
wallet_response = await authenticated_client.get("/v1/wallet/")
|
||||
api_key = wallet_response.json()["api_key"]
|
||||
hashed_key = api_key[3:] if api_key.startswith("sk-") else api_key
|
||||
|
||||
topup_token = await testmint_wallet.mint_tokens(250)
|
||||
topup_response = await authenticated_client.post(
|
||||
"/v1/wallet/topup", params={"cashu_token": topup_token}
|
||||
)
|
||||
assert topup_response.status_code == 200
|
||||
|
||||
refund_response = await authenticated_client.post("/v1/wallet/refund")
|
||||
assert refund_response.status_code == 200
|
||||
|
||||
history_response = await authenticated_client.get("/v1/wallet/history")
|
||||
assert history_response.status_code == 200
|
||||
transactions = history_response.json()["transactions"]
|
||||
assert len(transactions) >= 2
|
||||
assert all("api_key_hashed_key" not in tx for tx in transactions)
|
||||
assert {tx["type"] for tx in transactions} >= {"in", "out"}
|
||||
|
||||
db_result = await integration_session.execute(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.api_key_hashed_key == hashed_key
|
||||
)
|
||||
)
|
||||
db_transactions = db_result.scalars().all()
|
||||
assert len(db_transactions) >= 2
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_mint_unavailability_handling(
|
||||
|
||||
@@ -11,7 +11,7 @@ import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
|
||||
from .utils import (
|
||||
CashuTokenGenerator,
|
||||
@@ -71,6 +71,16 @@ async def test_topup_with_valid_token( # type: ignore[no-untyped-def]
|
||||
assert db_key.balance == new_balance
|
||||
assert db_key.balance == initial_balance + (topup_amount * 1000)
|
||||
|
||||
tx_result = await integration_session.execute(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.token == token,
|
||||
CashuTransaction.type == "in",
|
||||
)
|
||||
)
|
||||
tx = tx_result.scalar_one()
|
||||
assert tx.api_key_hashed_key == hashed_key
|
||||
assert tx.source == "apikey"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
@@ -426,9 +436,9 @@ async def test_topup_with_zero_amount_token( # type: ignore[no-untyped-def]
|
||||
"/v1/wallet/topup", params={"cashu_token": token}
|
||||
)
|
||||
|
||||
# Should succeed but add 0 msats
|
||||
assert response.status_code == 200
|
||||
assert response.json()["msats"] == 0
|
||||
# Zero/negative redemptions are refused to avoid crediting empty
|
||||
# or dust tokens (and to prevent orphan zero-balance keys).
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
"""Tests for the model prioritization algorithm."""
|
||||
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
# Set required env vars before importing
|
||||
os.environ["UPSTREAM_BASE_URL"] = "http://test"
|
||||
os.environ["UPSTREAM_API_KEY"] = "test"
|
||||
|
||||
from routstr.algorithm import ( # noqa: E402
|
||||
calculate_model_cost_score,
|
||||
create_model_mappings,
|
||||
get_provider_penalty,
|
||||
)
|
||||
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
|
||||
@@ -45,11 +49,21 @@ def create_test_model(
|
||||
)
|
||||
|
||||
|
||||
def create_test_provider(name: str, base_url: str = "http://test.com") -> Mock:
|
||||
def create_test_provider(
|
||||
name: str,
|
||||
base_url: str = "http://test.com",
|
||||
*,
|
||||
db_id: int | None = None,
|
||||
models: list[Model] | None = None,
|
||||
upstream_name: str | None = None,
|
||||
) -> Mock:
|
||||
"""Helper to create a test provider mock."""
|
||||
provider = Mock()
|
||||
provider.provider_type = name
|
||||
provider.base_url = base_url
|
||||
provider.db_id = db_id
|
||||
provider.upstream_name = upstream_name or name
|
||||
provider.get_cached_models.return_value = models or []
|
||||
return provider
|
||||
|
||||
|
||||
@@ -99,3 +113,80 @@ def test_get_provider_penalty_openrouter() -> None:
|
||||
provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1")
|
||||
penalty = get_provider_penalty(provider)
|
||||
assert penalty == 1.001
|
||||
|
||||
|
||||
def test_create_model_mappings_includes_db_override_for_missing_cached_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Model overrides should still map when provider discovery misses the model."""
|
||||
provider = create_test_provider(
|
||||
"azure",
|
||||
"https://example.openai.azure.com/openai/v1",
|
||||
db_id=7,
|
||||
models=[],
|
||||
)
|
||||
override_model = create_test_model("azure/gpt-4o")
|
||||
override_model.canonical_slug = "azure-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=7, enabled=True)
|
||||
|
||||
model_instances, provider_map, unique_models = create_model_mappings(
|
||||
upstreams=[provider],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
assert "azure/gpt-4o" in model_instances
|
||||
assert provider_map["azure/gpt-4o"] == [provider]
|
||||
assert "gpt-4o" in unique_models
|
||||
|
||||
|
||||
def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Different provider instances of same type should both survive dedupe."""
|
||||
provider_a_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.01, completion_price=0.01
|
||||
)
|
||||
provider_a = create_test_provider(
|
||||
"azure",
|
||||
"https://a.openai.azure.com/openai/v1",
|
||||
db_id=1,
|
||||
models=[provider_a_model],
|
||||
upstream_name="azure-a",
|
||||
)
|
||||
provider_b = create_test_provider(
|
||||
"azure",
|
||||
"https://b.openai.azure.com/openai/v1",
|
||||
db_id=2,
|
||||
models=[],
|
||||
upstream_name="azure-b",
|
||||
)
|
||||
|
||||
override_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.001, completion_price=0.001
|
||||
)
|
||||
override_model.canonical_slug = "azure-b-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=2, enabled=True)
|
||||
|
||||
_, provider_map, _ = create_model_mappings(
|
||||
upstreams=[provider_a, provider_b],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
providers_for_alias = provider_map["azure/gpt-4o"]
|
||||
assert provider_a in providers_for_alias
|
||||
assert provider_b in providers_for_alias
|
||||
assert len(providers_for_alias) == 2
|
||||
|
||||
403
tests/unit/test_balance.py
Normal file
403
tests/unit/test_balance.py
Normal file
@@ -0,0 +1,403 @@
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from routstr.balance import refund_wallet_endpoint
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
from routstr.wallet import credit_balance
|
||||
|
||||
|
||||
def _make_cashu_tx(
|
||||
token: str,
|
||||
amount: int,
|
||||
unit: str,
|
||||
type: str = "out",
|
||||
request_id: str | None = "req-abc",
|
||||
swept: bool = False,
|
||||
collected: bool = False,
|
||||
) -> CashuTransaction:
|
||||
tx = CashuTransaction(token=token, amount=amount, unit=unit, type=type, request_id=request_id)
|
||||
tx.swept = swept
|
||||
tx.collected = collected
|
||||
return tx
|
||||
|
||||
|
||||
def _exec_result(tx: CashuTransaction | None) -> MagicMock:
|
||||
result = MagicMock()
|
||||
result.first.return_value = tx
|
||||
return result
|
||||
|
||||
|
||||
def _update_result(rowcount: int) -> MagicMock:
|
||||
result = MagicMock()
|
||||
result.rowcount = rowcount
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_returns_token() -> None:
|
||||
x_cashu_token = "cashuAtest_token_value"
|
||||
in_tx = _make_cashu_tx(token=x_cashu_token, amount=0, unit="msat", type="in", request_id="req-abc")
|
||||
out_tx = _make_cashu_tx(token="cashuArefund_token", amount=1000, unit="msat", type="out", request_id="req-abc")
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu=x_cashu_token,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, JSONResponse)
|
||||
body = json.loads(result.body)
|
||||
assert body["token"] == "cashuArefund_token"
|
||||
assert body["msats"] == "1000"
|
||||
assert result.headers["X-Cashu"] == "cashuArefund_token"
|
||||
assert out_tx.collected is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_sat_unit() -> None:
|
||||
x_cashu_token = "cashuAsat_token"
|
||||
in_tx = _make_cashu_tx(token=x_cashu_token, amount=0, unit="sat", type="in", request_id="req-sat")
|
||||
out_tx = _make_cashu_tx(token="cashuArefund_sat", amount=500, unit="sat", type="out", request_id="req-sat")
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu=x_cashu_token,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, JSONResponse)
|
||||
body = json.loads(result.body)
|
||||
assert body["token"] == "cashuArefund_sat"
|
||||
assert body["sats"] == "500"
|
||||
assert "msats" not in body
|
||||
assert result.headers["X-Cashu"] == "cashuArefund_sat"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_not_found_raises_404() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_exec_result(None))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu="cashuAmissing_token",
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_swept_raises_410() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
in_tx = _make_cashu_tx(token="cashuAswept_token", amount=0, unit="msat", type="in", request_id="req-swept")
|
||||
out_tx = _make_cashu_tx(token="cashuAswept", amount=100, unit="msat", type="out", request_id="req-swept", swept=True)
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu="cashuAswept_token",
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 410
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# source field defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_cashu_transaction_source_defaults_to_x_cashu() -> None:
|
||||
tx = CashuTransaction(token="cashuAtest", amount=100, unit="msat")
|
||||
assert tx.source == "x-cashu"
|
||||
|
||||
|
||||
def test_cashu_transaction_source_can_be_apikey() -> None:
|
||||
tx = CashuTransaction(token="cashuAtest", amount=100, unit="msat", source="apikey")
|
||||
assert tx.source == "apikey"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apikey-based refund: token logging and CashuTransaction storage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_api_key(
|
||||
balance: int = 5000,
|
||||
refund_currency: str | None = "sat",
|
||||
refund_mint_url: str | None = "https://mint.example.com",
|
||||
refund_address: str | None = None,
|
||||
parent_key_hash: str | None = None,
|
||||
) -> ApiKey:
|
||||
key = ApiKey(hashed_key="testhash")
|
||||
key.balance = balance
|
||||
key.reserved_balance = 0
|
||||
key.refund_currency = refund_currency
|
||||
key.refund_mint_url = refund_mint_url
|
||||
key.refund_address = refund_address
|
||||
key.parent_key_hash = parent_key_hash
|
||||
key.total_spent = 0
|
||||
key.total_requests = 0
|
||||
return key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_stores_cashu_transaction_with_apikey_source() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuArefund_apikey_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=key)
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()) as mock_store,
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
):
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert result["token"] == refund_token
|
||||
|
||||
mock_store.assert_awaited_once()
|
||||
call_kwargs = mock_store.call_args.kwargs
|
||||
assert call_kwargs["source"] == "apikey"
|
||||
assert call_kwargs["token"] == refund_token
|
||||
assert call_kwargs["typ"] == "out"
|
||||
assert call_kwargs["api_key_hashed_key"] == key.hashed_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_logs_token() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuAlogged_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=key)
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger") as mock_logger,
|
||||
):
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
calls = [str(c) for c in mock_logger.info.call_args_list]
|
||||
assert any("cashu token issued" in c for c in calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_log_includes_path() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuApath_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=key)
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger") as mock_logger,
|
||||
):
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
# Find the "cashu token issued" call and verify extra contains the path
|
||||
token_issued_calls = [
|
||||
c for c in mock_logger.info.call_args_list
|
||||
if c.args and "cashu token issued" in c.args[0]
|
||||
]
|
||||
assert len(token_issued_calls) == 1
|
||||
extra = token_issued_calls[0].kwargs.get("extra", {})
|
||||
assert extra.get("path") == "/v1/wallet/refund"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_rejects_on_concurrent_balance_change() -> None:
|
||||
"""When the debit CAS fails (rowcount=0), no token is minted and 409 is returned."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=key)
|
||||
# Debit returns rowcount=0 → balance changed concurrently
|
||||
session.exec = AsyncMock(return_value=_update_result(0))
|
||||
session.commit = AsyncMock()
|
||||
|
||||
mock_send_token = AsyncMock(return_value="cashuAshould_not_be_minted")
|
||||
|
||||
with (
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", mock_send_token),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
# Crucially: send_token must NOT have been called
|
||||
mock_send_token.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credit_balance_stores_apikey_transaction_history() -> None:
|
||||
key = _make_api_key(balance=1000)
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.commit = AsyncMock()
|
||||
session.refresh = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"routstr.wallet.recieve_token",
|
||||
AsyncMock(return_value=(100, "sat", "https://mint.example")),
|
||||
),
|
||||
patch("routstr.wallet.store_cashu_transaction", AsyncMock()) as mock_store,
|
||||
):
|
||||
amount = await credit_balance("cashuAtopup_token", key, session)
|
||||
|
||||
assert amount == 100_000
|
||||
mock_store.assert_awaited_once()
|
||||
call_kwargs = mock_store.call_args.kwargs
|
||||
assert call_kwargs["typ"] == "in"
|
||||
assert call_kwargs["source"] == "apikey"
|
||||
assert call_kwargs["api_key_hashed_key"] == key.hashed_key
|
||||
assert call_kwargs["amount"] == 100
|
||||
assert call_kwargs["unit"] == "sat"
|
||||
assert call_kwargs["token"] == "cashuAtopup_token"
|
||||
assert call_kwargs["mint_url"] == "https://mint.example"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_restores_balance_on_mint_failure() -> None:
|
||||
"""When debit succeeds but minting fails, balance must be restored."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
|
||||
# First exec call = debit (succeeds), second = restore
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=key)
|
||||
session.exec = AsyncMock(side_effect=[_update_result(1), _update_result(1)])
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(side_effect=Exception("mint down"))),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger"),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 503
|
||||
# Verify two exec calls: debit + restore
|
||||
assert session.exec.await_count == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# no-create guarantee: fresh Cashu/unknown sk- tokens must not create API keys
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_fresh_cashu_bearer_returns_401() -> None:
|
||||
"""Fresh Cashu token not in DB must get 401, never create a new ApiKey."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer cashuAfresh_never_deposited_token",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
session.get.assert_awaited_once()
|
||||
# No add/commit → no key was persisted
|
||||
session.add.assert_not_called()
|
||||
session.commit.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_unknown_sk_bearer_returns_401() -> None:
|
||||
"""Unknown sk- key not in DB must get 401."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
session = MagicMock()
|
||||
session.get = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-unknownhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
session.get.assert_awaited_once()
|
||||
189
tests/unit/test_cache_breakpoints.py
Normal file
189
tests/unit/test_cache_breakpoints.py
Normal file
@@ -0,0 +1,189 @@
|
||||
"""Tests for Anthropic cache-breakpoint injection on forwarded requests.
|
||||
|
||||
Anthropic prompt caching is explicit; a client that doesn't recognise a routstr
|
||||
URL as Anthropic-backed never sends ``cache_control`` markers, so caching never
|
||||
engages over routstr. ``prepare_request_body`` must stamp the standard
|
||||
breakpoints for Anthropic-family models while always deferring to client-set
|
||||
markers and never touching automatic-cache providers.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
from routstr.upstream.cache_breakpoints import (
|
||||
body_has_cache_control,
|
||||
inject_anthropic_cache_breakpoints,
|
||||
is_explicit_cache_model,
|
||||
)
|
||||
|
||||
|
||||
def _chat_body() -> dict:
|
||||
return {
|
||||
"model": "anthropic/claude-sonnet-4.5",
|
||||
"stream": True,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are concise."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
],
|
||||
"tools": [
|
||||
{"type": "function", "function": {"name": "a"}},
|
||||
{"type": "function", "function": {"name": "b"}},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id,expected",
|
||||
[
|
||||
("anthropic/claude-sonnet-4.5", True),
|
||||
("claude-haiku-4-5-20251001", True),
|
||||
# Alibaba explicit-cache models share Anthropic's wire format
|
||||
("qwen/qwen3-max", True),
|
||||
("qwen/qwen3-coder-plus", True),
|
||||
("deepseek/deepseek-v3.2", True),
|
||||
# Automatic-cache providers need no markers
|
||||
("openai/gpt-4o", False),
|
||||
("google/gemini-2.5-flash", False),
|
||||
("deepseek/deepseek-chat", False),
|
||||
("qwen/qwen3.5-plus-02-15", False), # snapshot, no explicit caching
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def test_is_explicit_cache_model(model_id: str | None, expected: bool) -> None:
|
||||
assert is_explicit_cache_model(model_id) is expected
|
||||
|
||||
|
||||
def test_is_explicit_cache_model_uses_fallbacks() -> None:
|
||||
# routstr id is opaque but a forwarded/canonical alias reveals the family.
|
||||
assert is_explicit_cache_model("model-xyz", None, "anthropic/claude-opus-4.1")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Breakpoint placement
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_injects_three_breakpoints() -> None:
|
||||
data = _chat_body()
|
||||
assert inject_anthropic_cache_breakpoints(data) is True
|
||||
|
||||
# system prompt promoted to array form with a marker
|
||||
system = data["messages"][0]["content"]
|
||||
assert system == [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are concise.",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
# last tool marked
|
||||
assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in data["tools"][0]
|
||||
# last user message marked
|
||||
user = data["messages"][1]["content"]
|
||||
assert user[-1]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_defers_to_client_supplied_cache_control() -> None:
|
||||
data = _chat_body()
|
||||
data["messages"][1]["content"] = [
|
||||
{"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}}
|
||||
]
|
||||
assert body_has_cache_control(data) is True
|
||||
# No additional stamping when the client already controls caching.
|
||||
assert inject_anthropic_cache_breakpoints(data) is False
|
||||
assert "cache_control" not in data["tools"][-1]
|
||||
|
||||
|
||||
def test_marks_last_text_part_of_array_content() -> None:
|
||||
data = _chat_body()
|
||||
data["messages"][1]["content"] = [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "image_url", "image_url": {"url": "x"}},
|
||||
{"type": "text", "text": "last"},
|
||||
]
|
||||
inject_anthropic_cache_breakpoints(data)
|
||||
parts = data["messages"][1]["content"]
|
||||
assert parts[2]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in parts[0]
|
||||
|
||||
|
||||
def test_noop_without_messages() -> None:
|
||||
assert inject_anthropic_cache_breakpoints({"prompt": "x"}) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# prepare_request_body integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _model(model_id: str): # type: ignore[no-untyped-def]
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="",
|
||||
context_length=200000,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="Claude",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=Pricing(prompt=0.0, completion=0.0),
|
||||
)
|
||||
|
||||
|
||||
def _openrouter_provider() -> "GenericUpstreamProvider":
|
||||
# OpenRouter endpoint via the generic provider — recognised by base URL.
|
||||
return GenericUpstreamProvider(base_url="https://openrouter.ai/api/v1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id", ["anthropic/claude-sonnet-4.5", "qwen/qwen3-max", "deepseek/deepseek-v3.2"]
|
||||
)
|
||||
def test_prepare_request_body_injects_for_explicit_models(model_id: str) -> None:
|
||||
provider = _openrouter_provider()
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model(model_id))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is True
|
||||
assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
|
||||
def test_prepare_request_body_skips_for_automatic_provider_model() -> None:
|
||||
provider = _openrouter_provider()
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model("openai/gpt-4o"))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is False
|
||||
|
||||
|
||||
def test_prepare_request_body_skips_when_upstream_rejects_markers() -> None:
|
||||
# Claude id but a non-OpenRouter/Anthropic upstream → must NOT inject,
|
||||
# since the markers could be rejected by an upstream that doesn't accept them.
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
|
||||
provider = GenericUpstreamProvider(base_url="https://some-gateway.example/v1")
|
||||
body = json.dumps(_chat_body()).encode()
|
||||
out = provider.prepare_request_body(body, _model("anthropic/claude-sonnet-4.5"))
|
||||
assert out is not None
|
||||
data = json.loads(out)
|
||||
assert body_has_cache_control(data) is False
|
||||
243
tests/unit/test_cache_pricing.py
Normal file
243
tests/unit/test_cache_pricing.py
Normal file
@@ -0,0 +1,243 @@
|
||||
"""Tests for cache-aware pricing of cached input tokens.
|
||||
|
||||
Specifies two things:
|
||||
|
||||
1. ``backfill_cache_pricing`` — when the OpenRouter model feed omits cache
|
||||
rates (it does for most DeepSeek models and e.g. openai/gpt-4o), they are
|
||||
filled from litellm's bundled cost map instead of silently billing cache
|
||||
reads at the full input rate. Existing OpenRouter values are never
|
||||
overwritten, and provider fees apply to backfilled rates like any other.
|
||||
2. ``calculate_cost`` — cached tokens are billed at the cache rates from the
|
||||
model's sats_pricing; the full input rate remains only as the documented
|
||||
last resort when no cache rate could be resolved anywhere.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.core.settings import settings
|
||||
from routstr.payment.cost_calculation import CostData, calculate_cost
|
||||
from routstr.payment.models import (
|
||||
Architecture,
|
||||
Model,
|
||||
Pricing,
|
||||
backfill_cache_pricing,
|
||||
)
|
||||
from routstr.upstream import GenericUpstreamProvider
|
||||
|
||||
|
||||
def _make_model(model_id: str, pricing: Pricing) -> Model:
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="",
|
||||
context_length=64000,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="Other",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=pricing,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# backfill_cache_pricing — litellm as fallback source for missing cache rates
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def test_backfill_deepseek_cache_read_from_litellm() -> None:
|
||||
"""deepseek/deepseek-chat has no input_cache_read on OpenRouter; litellm
|
||||
knows the real rate (10x cheaper than input)."""
|
||||
pricing = Pricing(prompt=2.8e-07, completion=4.2e-07)
|
||||
|
||||
result = backfill_cache_pricing("deepseek/deepseek-chat", pricing)
|
||||
|
||||
expected = litellm.model_cost["deepseek/deepseek-chat"][
|
||||
"cache_read_input_token_cost"
|
||||
]
|
||||
assert result.input_cache_read == expected
|
||||
assert result.input_cache_read < pricing.prompt # sanity: it's a discount
|
||||
|
||||
|
||||
def test_backfill_strips_vendor_prefix_for_litellm_lookup() -> None:
|
||||
"""OpenRouter ids are vendor-prefixed (openai/gpt-4o); litellm keys most
|
||||
non-DeepSeek models without the prefix (gpt-4o)."""
|
||||
pricing = Pricing(prompt=2.5e-06, completion=1e-05)
|
||||
|
||||
result = backfill_cache_pricing("openai/gpt-4o", pricing)
|
||||
|
||||
expected = litellm.model_cost["gpt-4o"]["cache_read_input_token_cost"]
|
||||
assert result.input_cache_read == expected
|
||||
|
||||
|
||||
def test_backfill_fills_cache_write_rate() -> None:
|
||||
"""Anthropic cache writes cost more than input (1.25x); billing them at
|
||||
the input rate undercharges. litellm carries the write rate."""
|
||||
pricing = Pricing(prompt=3e-06, completion=1.5e-05)
|
||||
|
||||
result = backfill_cache_pricing("anthropic/claude-sonnet-4-5", pricing)
|
||||
|
||||
expected = litellm.model_cost["claude-sonnet-4-5"][
|
||||
"cache_creation_input_token_cost"
|
||||
]
|
||||
assert result.input_cache_write == expected
|
||||
assert result.input_cache_write > pricing.prompt # sanity: write premium
|
||||
|
||||
|
||||
def test_backfill_never_overwrites_openrouter_rates() -> None:
|
||||
"""When OpenRouter provides a cache rate, it is authoritative."""
|
||||
pricing = Pricing(
|
||||
prompt=2.1e-07, completion=7.9e-07, input_cache_read=1.3e-07
|
||||
)
|
||||
|
||||
result = backfill_cache_pricing("deepseek/deepseek-chat", pricing)
|
||||
|
||||
assert result.input_cache_read == 1.3e-07
|
||||
|
||||
|
||||
def test_backfill_unknown_model_unchanged() -> None:
|
||||
"""Models litellm doesn't know stay untouched (last-resort fallback to
|
||||
the input rate happens later, at billing time)."""
|
||||
pricing = Pricing(prompt=1e-06, completion=2e-06)
|
||||
|
||||
result = backfill_cache_pricing("artificial-dumbness/dumb-1", pricing)
|
||||
|
||||
assert result.input_cache_read == 0.0
|
||||
assert result.input_cache_write == 0.0
|
||||
|
||||
|
||||
def test_provider_fee_applies_to_backfilled_cache_rates() -> None:
|
||||
"""Backfill happens before the provider fee, so cache rates carry the
|
||||
same markup as every other price component."""
|
||||
provider = GenericUpstreamProvider(
|
||||
base_url="http://upstream.example", provider_fee=2.0
|
||||
)
|
||||
model = _make_model(
|
||||
"deepseek/deepseek-chat", Pricing(prompt=2.8e-07, completion=4.2e-07)
|
||||
)
|
||||
|
||||
adjusted = provider._apply_provider_fee_to_model(model)
|
||||
|
||||
litellm_read = litellm.model_cost["deepseek/deepseek-chat"][
|
||||
"cache_read_input_token_cost"
|
||||
]
|
||||
assert adjusted.pricing.input_cache_read == pytest.approx(litellm_read * 2.0)
|
||||
assert adjusted.pricing.prompt == pytest.approx(2.8e-07 * 2.0)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# calculate_cost — cached tokens billed at cache rates
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_sats_usd_price() -> None: # type: ignore[misc]
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_pricing(monkeypatch: pytest.MonkeyPatch) -> Mock:
|
||||
"""Model-based pricing: 1 msat per input token, 2 per output token,
|
||||
0.1 per cache-read token, 1.25 per cache-write token."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", False)
|
||||
model = Mock()
|
||||
model.sats_pricing = Pricing(
|
||||
prompt=0.001,
|
||||
completion=0.002,
|
||||
input_cache_read=0.0001,
|
||||
input_cache_write=0.00125,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hits_billed_at_cache_rate(model_pricing: Mock) -> None:
|
||||
"""The reported overcharge scenario: a 10k-token prompt with 90% cache
|
||||
hits costs 2900 msats at honest rates, not the 11000 msats that billing
|
||||
every prompt token at the full input rate would charge."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model_pricing):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 1000 input @ 1 msat + 9000 cache reads @ 0.1 msat + 500 output @ 2 msat.
|
||||
# Client-visible input_msats includes cache-read input cost for display,
|
||||
# while cache_read_msats keeps the detailed breakdown.
|
||||
assert result.input_msats == 1900
|
||||
assert result.cache_read_msats == 900
|
||||
assert result.output_msats == 1000
|
||||
assert result.total_msats == 2900
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_write_billed_at_write_rate(model_pricing: Mock) -> None:
|
||||
"""Cache writes carry their premium rate (1.25x input here), instead of
|
||||
being silently billed at the plain input rate."""
|
||||
response = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"usage": {
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"cache_read_input_tokens": 500,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model_pricing):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 300 @ 1 + 500 @ 0.1 + 2000 @ 1.25 + 100 @ 2
|
||||
assert result.cache_read_msats == 50
|
||||
assert result.cache_creation_msats == 2500
|
||||
assert result.total_msats == 3050
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_cache_rate_falls_back_to_input_rate(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Documented last resort: when no cache rate could be resolved anywhere
|
||||
(OpenRouter and litellm both silent), cache reads bill at the input rate —
|
||||
never cheaper, never free."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", False)
|
||||
model = Mock()
|
||||
model.sats_pricing = Pricing(prompt=0.001, completion=0.002)
|
||||
|
||||
response = {
|
||||
"model": "dumb-1",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
|
||||
with patch("routstr.proxy.get_model_instance", return_value=model):
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# 1000 @ 1 + 9000 @ 1 (fallback) + 500 @ 2
|
||||
assert result.total_msats == 11000
|
||||
433
tests/unit/test_cost_calculation_caching.py
Normal file
433
tests/unit/test_cost_calculation_caching.py
Normal file
@@ -0,0 +1,433 @@
|
||||
"""Tests for cache token handling in cost calculation.
|
||||
|
||||
Covers OpenAI, Anthropic and DeepSeek caching formats, dialect precedence,
|
||||
edge cases, and billing accuracy.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.core.settings import settings
|
||||
from routstr.payment.cost_calculation import CostData, MaxCostData, calculate_cost
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Mock settings and price to use fixed pricing."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", True)
|
||||
monkeypatch.setattr(settings, "fixed_per_1k_input_tokens", 0.001)
|
||||
monkeypatch.setattr(settings, "fixed_per_1k_output_tokens", 0.001)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_sats_usd_price() -> None: # type: ignore[misc]
|
||||
"""Patch sats_usd_price to avoid initialization issues."""
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5):
|
||||
yield
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 1: OpenAI Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_cache_subtraction() -> None:
|
||||
"""OpenAI includes cached_tokens in prompt_tokens, subtract them."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 2000, # ← Includes 1000 cached
|
||||
"completion_tokens": 100,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 1000 # ← Extracted separately
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 2000 - 1000
|
||||
assert result.cache_read_input_tokens == 1000
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 2: Anthropic Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_additive(mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache tokens are separate (additive) from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500, # ← Regular input only
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 1500, # ← Additive, not included above
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
assert result.cache_creation_input_tokens == 1500
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 3: Invalid Cache (Edge Case)
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_read_exceeds_prompt_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle buggy upstream reporting cached > prompt_tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 150 # ← Invalid! Greater than prompt
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Should not go negative
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0 # max(0, 100 - 150)
|
||||
assert result.cache_read_input_tokens == 150
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 4: Malformed Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_malformed_cache_tokens_coerce_to_zero(mock_fixed_pricing: None) -> None:
|
||||
"""Handle non-numeric cache token values."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"cache_read_input_tokens": "-50", # ← String, negative
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": "invalid" # ← Non-numeric string
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Both should coerce to 0
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.input_tokens == 100 # No subtraction if cache_read = 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 5: Anthropic Cache Not Subtracted
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_not_subtracted(mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache fields should NOT be subtracted from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"completion_tokens": 100,
|
||||
"cache_read_input_tokens": 200, # ← Additive, don't subtract
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
# Anthropic: input_tokens stays as-is
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500 # NOT 300
|
||||
assert result.cache_read_input_tokens == 200
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 6: Only Cache Read, No Regular Input
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_read_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache read tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 1000
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0 # max(0, 0 - 1000)
|
||||
assert result.cache_read_input_tokens == 1000
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 7: Only Cache Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_creation_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache creation tokens (Anthropic)."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
assert result.cache_creation_input_tokens == 2000
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 8: Both Cache Read and Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_both_cache_read_and_creation(mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with both cache read and creation."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
"cache_read_input_tokens": 500,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 300
|
||||
assert result.cache_creation_input_tokens == 2000
|
||||
assert result.cache_read_input_tokens == 500
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 9: Token Field Fallback
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_field_fallback_order(mock_fixed_pricing: None) -> None:
|
||||
"""Verify fallback order for token extraction."""
|
||||
# When prompt_tokens is not present, fall back to input_tokens
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"input_tokens": 250,
|
||||
"completion_tokens": 50,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 250
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 10: Float Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_float_token_values_coerced_to_int(mock_fixed_pricing: None) -> None:
|
||||
"""Handle float token values by converting to int."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100.7, # Float
|
||||
"completion_tokens": 50.3, # Float
|
||||
"prompt_tokens_details": {"cached_tokens": 25.9}, # Float
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
# cached_tokens are part of prompt_tokens (OpenAI dialect) → subtracted: 100 - 25
|
||||
assert result.input_tokens == 75 # Floored
|
||||
assert result.output_tokens == 50 # Floored
|
||||
assert result.cache_read_input_tokens == 25 # Floored
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 11: Boolean Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_boolean_cache_tokens_coerced_to_zero(mock_fixed_pricing: None) -> None:
|
||||
"""Handle boolean cache token values by coercing to zero."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"cache_read_input_tokens": True, # Boolean
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0 # Boolean coerced to 0
|
||||
assert result.input_tokens == 100 # No subtraction
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 12: Zero Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cache_tokens(mock_fixed_pricing: None) -> None:
|
||||
"""Handle explicit zero cache tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.input_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# DeepSeek Cache Format
|
||||
# DeepSeek emits neither OpenAI's prompt_tokens_details nor Anthropic's
|
||||
# cache_read_input_tokens — only prompt_cache_hit_tokens and
|
||||
# prompt_cache_miss_tokens, with the documented guarantee
|
||||
# prompt_tokens = hit + miss. Hits are ~10x cheaper upstream, so billing
|
||||
# them as regular input is a large overcharge.
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_cache_hit_tokens_extracted() -> None:
|
||||
"""DeepSeek cache hits are extracted and removed from regular input.
|
||||
|
||||
Payload shape verbatim from the DeepSeek API reference (usage object).
|
||||
"""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000, # = hit + miss
|
||||
"completion_tokens": 500,
|
||||
"total_tokens": 10500,
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # only the cache misses
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
assert result.output_tokens == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_all_tokens_cached() -> None:
|
||||
"""A fully cached DeepSeek prompt bills zero regular input tokens."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 5000,
|
||||
"completion_tokens": 100,
|
||||
"prompt_cache_hit_tokens": 5000,
|
||||
"prompt_cache_miss_tokens": 0,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 5000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dialect_precedence_never_double_subtracts() -> None:
|
||||
"""If a vendor emits both OpenAI-style and DeepSeek-style cache fields for
|
||||
the same cached tokens, they are counted once, not subtracted twice."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 10000,
|
||||
"completion_tokens": 500,
|
||||
"prompt_tokens_details": {"cached_tokens": 9000},
|
||||
"prompt_cache_hit_tokens": 9000,
|
||||
"prompt_cache_miss_tokens": 1000,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 10000 - 9000, applied exactly once
|
||||
assert result.cache_read_input_tokens == 9000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_malformed_hit_tokens_coerce_to_zero() -> None:
|
||||
"""Malformed DeepSeek cache fields degrade to billing all input at full
|
||||
rate instead of crashing or going negative."""
|
||||
response = {
|
||||
"model": "deepseek-chat",
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 50,
|
||||
"prompt_cache_hit_tokens": "garbage",
|
||||
"prompt_cache_miss_tokens": -5,
|
||||
},
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000
|
||||
assert result.cache_read_input_tokens == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 13: Missing Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_usage_block(mock_fixed_pricing: None) -> None:
|
||||
"""When usage is missing, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "choices": [{"message": {"content": "test"}}]}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 14: Null Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_usage_block(mock_fixed_pricing: None) -> None:
|
||||
"""When usage is null, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "usage": None}
|
||||
result = await calculate_cost(response, max_cost=100000)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 0
|
||||
177
tests/unit/test_count_tokens_local.py
Normal file
177
tests/unit/test_count_tokens_local.py
Normal file
@@ -0,0 +1,177 @@
|
||||
"""Unit tests for the local count_tokens shim.
|
||||
|
||||
The shim runs whenever an upstream that does not support Anthropic's
|
||||
``/v1/messages`` endpoint is asked for a token count. It must always
|
||||
return a 200 JSON ``{"input_tokens": N}`` response and must never raise.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
from routstr.upstream import count_tokens as count_tokens_module
|
||||
from routstr.upstream.count_tokens import count_tokens_locally
|
||||
|
||||
|
||||
def _make_model(model_id: str = "anthropic/claude-3-5-sonnet") -> Model:
|
||||
pricing = Pricing(prompt=0.000003, completion=0.000015)
|
||||
architecture = Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="cl100k_base",
|
||||
instruct_type=None,
|
||||
)
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="",
|
||||
context_length=200_000,
|
||||
architecture=architecture,
|
||||
pricing=pricing,
|
||||
)
|
||||
|
||||
|
||||
def _body(payload: dict[str, Any]) -> bytes:
|
||||
return json.dumps(payload).encode()
|
||||
|
||||
|
||||
def _read_payload(response: Any) -> dict[str, Any]:
|
||||
body = response.body if isinstance(response.body, bytes) else bytes(response.body)
|
||||
return json.loads(body.decode())
|
||||
|
||||
|
||||
def test_returns_input_tokens_for_simple_messages() -> None:
|
||||
model = _make_model()
|
||||
request_body = _body(
|
||||
{
|
||||
"model": model.id,
|
||||
"messages": [{"role": "user", "content": "hello world"}],
|
||||
}
|
||||
)
|
||||
|
||||
response = count_tokens_locally(request_body, model)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.media_type == "application/json"
|
||||
payload = _read_payload(response)
|
||||
assert "input_tokens" in payload
|
||||
assert isinstance(payload["input_tokens"], int)
|
||||
assert payload["input_tokens"] >= 0
|
||||
|
||||
|
||||
def test_falls_back_to_estimator_when_litellm_raises() -> None:
|
||||
model = _make_model()
|
||||
request_body = _body(
|
||||
{
|
||||
"model": model.id,
|
||||
"messages": [{"role": "user", "content": "this is a longer message"}],
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
count_tokens_module,
|
||||
"_count_with_litellm",
|
||||
side_effect=RuntimeError("boom"),
|
||||
):
|
||||
response = count_tokens_locally(request_body, model)
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = _read_payload(response)
|
||||
assert payload["input_tokens"] >= 1
|
||||
|
||||
|
||||
def test_handles_missing_model_object() -> None:
|
||||
request_body = _body(
|
||||
{
|
||||
"model": "anthropic/claude-3-5-sonnet",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
)
|
||||
|
||||
response = count_tokens_locally(request_body, None)
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = _read_payload(response)
|
||||
assert payload["input_tokens"] >= 0
|
||||
|
||||
|
||||
def test_handles_empty_request_body() -> None:
|
||||
response = count_tokens_locally(b"", _make_model())
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = _read_payload(response)
|
||||
assert payload["input_tokens"] >= 0
|
||||
|
||||
|
||||
def test_handles_malformed_json() -> None:
|
||||
response = count_tokens_locally(b"not-json", _make_model())
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = _read_payload(response)
|
||||
assert payload["input_tokens"] >= 0
|
||||
|
||||
|
||||
def test_includes_system_prompt_in_count() -> None:
|
||||
model = _make_model()
|
||||
short = _body(
|
||||
{
|
||||
"model": model.id,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
)
|
||||
with_system = _body(
|
||||
{
|
||||
"model": model.id,
|
||||
"system": "You are a helpful assistant with a long preamble " * 10,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
)
|
||||
|
||||
short_count = _read_payload(count_tokens_locally(short, model))["input_tokens"]
|
||||
long_count = _read_payload(count_tokens_locally(with_system, model))["input_tokens"]
|
||||
|
||||
assert long_count > short_count
|
||||
|
||||
|
||||
def test_supports_anthropic_system_block_list() -> None:
|
||||
model = _make_model()
|
||||
request_body = _body(
|
||||
{
|
||||
"model": model.id,
|
||||
"system": [{"type": "text", "text": "be terse" * 50}],
|
||||
"messages": [{"role": "user", "content": "ok"}],
|
||||
}
|
||||
)
|
||||
|
||||
response = count_tokens_locally(request_body, model)
|
||||
|
||||
payload = _read_payload(response)
|
||||
assert payload["input_tokens"] > 0
|
||||
|
||||
|
||||
def test_uses_forwarded_model_id_when_present() -> None:
|
||||
model = _make_model("anthropic/claude-3-5-sonnet")
|
||||
model.forwarded_model_id = "claude-3-5-sonnet-20241022"
|
||||
request_body = _body(
|
||||
{
|
||||
"model": "ignored",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
)
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def _capture(model_name: str, body: dict[str, Any]) -> int:
|
||||
captured["model"] = model_name
|
||||
return 7
|
||||
|
||||
with patch.object(count_tokens_module, "_count_with_litellm", side_effect=_capture):
|
||||
response = count_tokens_locally(request_body, model)
|
||||
|
||||
assert captured["model"] == "claude-3-5-sonnet-20241022"
|
||||
assert _read_payload(response)["input_tokens"] == 7
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user