mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-22 12:22:20 +00:00
Compare commits
389 Commits
v0.4.0-ui-
...
remove-req
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ebb2267844 | ||
|
|
c48693973b | ||
|
|
407f4745a0 | ||
|
|
806771df13 | ||
|
|
3c1b6d5f17 | ||
|
|
a9d5a52b32 | ||
|
|
b36bb56a30 | ||
|
|
97172641f8 | ||
|
|
46bbba7027 | ||
|
|
65ccf83dbd | ||
|
|
2baea4149e | ||
|
|
64d9460711 | ||
|
|
f55c3e9c8c | ||
|
|
e3ac06342f | ||
|
|
e2aa02307b | ||
|
|
66bc1260a4 | ||
|
|
9dc4c979b8 | ||
|
|
ec0fd1143b | ||
|
|
ae9748db02 | ||
|
|
ccee76b31e | ||
|
|
9f04525c82 | ||
|
|
c81de0de0a | ||
|
|
12c4c1f030 | ||
|
|
81658d0ba6 | ||
|
|
ee5ced10c7 | ||
|
|
2a27fb239e | ||
|
|
ffde661d93 | ||
|
|
7f49ba1771 | ||
|
|
949dc433f1 | ||
|
|
b7fcf000af | ||
|
|
a1af1383e1 | ||
|
|
5bedbd129f | ||
|
|
f588147b41 | ||
|
|
17bc949597 | ||
|
|
7705656016 | ||
|
|
434283c58d | ||
|
|
9fcf870e3f | ||
|
|
782b233d40 | ||
|
|
0c7373675f | ||
|
|
a33ea5da07 | ||
|
|
b509581950 | ||
|
|
38d7075652 | ||
|
|
ed27ad0dea | ||
|
|
6b1b4285fa | ||
|
|
7b56dcb902 | ||
|
|
2131cfc7f9 | ||
|
|
d6d9af91e4 | ||
|
|
cbd97e8e3a | ||
|
|
e4b26dd27b | ||
|
|
1e50afbd82 | ||
|
|
4d6e2cdf59 | ||
|
|
d079f45ba3 | ||
|
|
d08f07d6fb | ||
|
|
af0796689f | ||
|
|
f408785409 | ||
|
|
90f1ba89e5 | ||
|
|
e29c9241ce | ||
|
|
5a3774a414 | ||
|
|
e972b62758 | ||
|
|
b57f2d408e | ||
|
|
70d9e1a1ce | ||
|
|
c0a34a391f | ||
|
|
a97ea2995a | ||
|
|
5b50a78d95 | ||
|
|
ccab5e4216 | ||
|
|
ff9c645bb1 | ||
|
|
d4318dcc07 | ||
|
|
75ed865a5f | ||
|
|
7969a8da55 | ||
|
|
830c299130 | ||
|
|
e9dced8148 | ||
|
|
ecd46975b4 | ||
|
|
8f5f3d9738 | ||
|
|
355e3f19ef | ||
|
|
439ac48216 | ||
|
|
bf05c96dfc | ||
|
|
ed6e0c0189 | ||
|
|
412bfe479c | ||
|
|
861fda21b7 | ||
|
|
919dcf5535 | ||
|
|
cbc424e8e7 | ||
|
|
068fb3572f | ||
|
|
eaf74edbba | ||
|
|
eb6dc5189c | ||
|
|
cbdbe4b899 | ||
|
|
7870db9af1 | ||
|
|
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 | ||
|
|
8a373276ce |
@@ -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
|
||||
|
||||
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.
|
||||
|
||||
@@ -113,42 +113,109 @@ All errors follow a consistent JSON structure:
|
||||
**Status:** 402
|
||||
**Resolution:** Top up API key balance
|
||||
|
||||
#### Invalid Token
|
||||
### Cashu Token Redemption Errors
|
||||
|
||||
These errors are returned when a Cashu token you pay with cannot be redeemed.
|
||||
They apply to every endpoint that accepts a token:
|
||||
|
||||
- **Per-request payment** via the `X-Cashu` header (chat completions + Responses API).
|
||||
- **API key top-up** via `POST /v1/wallet/topup`.
|
||||
- **Minting an API key** from a token sent in `Authorization: Bearer <cashu-token>`.
|
||||
|
||||
All three share one classifier, so the same failure yields the same HTTP status
|
||||
and sanitized message everywhere. Structured error envelopes (`X-Cashu` and
|
||||
`Authorization: Bearer <cashu-token>`) also expose the same `type` and `code` —
|
||||
branch on `type` (or `code` for finer granularity). `POST /v1/wallet/topup`
|
||||
keeps its existing plain-string `detail` envelope, so branch on status there.
|
||||
|
||||
| `type` | Status | `code` | Retryable | Meaning |
|
||||
|--------|--------|--------|-----------|---------|
|
||||
| `token_already_spent` | 400 | `cashu_token_already_spent` | No | The token was already redeemed. |
|
||||
| `invalid_token` | 400 | `invalid_cashu_token` | No | The token is malformed or cannot be decoded. |
|
||||
| `mint_error` | 422 | `cashu_token_swap_fees_exceed_amount` | No | Token value is too small to cover the mint's swap/melt fees. |
|
||||
| `mint_error` | 422 | `cashu_foreign_mint_swap_failed` | No | Swapping the token from a foreign mint to the primary mint failed. |
|
||||
| `mint_unreachable` | 503 | `cashu_mint_unreachable` | **Yes** | The mint could not be reached (DNS failure, refused/reset connection, timeout). The token is fine — retry once the mint recovers. |
|
||||
| `cashu_error` | 400 | `cashu_token_redemption_failed` | No | The token could not be redeemed for another expected reason. |
|
||||
| `cashu_error` | 400 | `cashu_token_zero_value` | No | The token redeemed to zero (empty/dust token, or value fully consumed by fees). |
|
||||
| `token_consumed` | 500 | `cashu_token_consumed` | No | The token was **spent** (melted/redeemed) but crediting it then failed. Do not retry — the token is gone; contact support to reconcile. |
|
||||
| `api_error` | 500 | `internal_error` | Maybe | Unexpected server-side fault during redemption. |
|
||||
|
||||
!!! important "Retry only `mint_unreachable`"
|
||||
Only `mint_unreachable` (503) means the same token will work again later —
|
||||
everything else is a permanent property of the token and must not be
|
||||
blindly retried. Use exponential backoff for the 503. In particular, a
|
||||
`token_consumed` 500 means the mint already spent the token, so a retry
|
||||
would fail as `token_already_spent`.
|
||||
|
||||
#### Mint Unreachable (retryable)
|
||||
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"type": "payment_error",
|
||||
"message": "Invalid Cashu token",
|
||||
"code": "invalid_token",
|
||||
"details": {
|
||||
"reason": "Token already spent"
|
||||
}
|
||||
"type": "mint_unreachable",
|
||||
"message": "Cashu mint is unreachable",
|
||||
"code": "cashu_mint_unreachable"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Status:** 400
|
||||
**Resolution:** Use a valid, unspent token
|
||||
**Status:** 503
|
||||
|
||||
#### Mint Unavailable
|
||||
**Resolution:** The token is valid — the mint is temporarily down. Retry with
|
||||
backoff, or pay with a token from a different mint.
|
||||
|
||||
#### Token Already Spent
|
||||
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"type": "payment_error",
|
||||
"message": "Cannot connect to Cashu mint",
|
||||
"code": "mint_unavailable",
|
||||
"details": {
|
||||
"mint_url": "https://mint.example.com",
|
||||
"retry_after": 60
|
||||
}
|
||||
"type": "token_already_spent",
|
||||
"message": "Cashu token already spent",
|
||||
"code": "cashu_token_already_spent"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Status:** 503
|
||||
**Resolution:** Try again later or use different mint
|
||||
**Status:** 400
|
||||
|
||||
**Resolution:** Use a fresh, unspent token. Do not retry with the same token.
|
||||
|
||||
#### Response envelope differs by endpoint
|
||||
|
||||
The `error` object above is identical everywhere, but the surrounding envelope
|
||||
depends on how you paid:
|
||||
|
||||
- **`X-Cashu` header payments** (chat + Responses API) return the object at the
|
||||
top level, alongside a `request_id`:
|
||||
|
||||
```json
|
||||
{
|
||||
"error": { "type": "mint_unreachable", "message": "Cashu mint is unreachable", "code": "cashu_mint_unreachable" },
|
||||
"request_id": "req-abc123"
|
||||
}
|
||||
```
|
||||
|
||||
The original token is echoed back in the `X-Cashu` **response header only when
|
||||
it is still spendable** (e.g. `mint_unreachable`, `invalid_cashu_token`, fee
|
||||
errors) so you can recover/retry it. It is **not** echoed for spent/consumed
|
||||
tokens (`cashu_token_already_spent`, `cashu_token_consumed`,
|
||||
`cashu_token_zero_value`, `internal_error`) — retrying those can never succeed.
|
||||
|
||||
- **`Authorization: Bearer <cashu-token>`** (API key minting) wraps it in
|
||||
FastAPI's `detail` field:
|
||||
|
||||
```json
|
||||
{ "detail": { "error": { "type": "mint_unreachable", "message": "Cashu mint is unreachable", "code": "cashu_mint_unreachable" } } }
|
||||
```
|
||||
|
||||
- **`POST /v1/wallet/topup`** returns a plain string message under `detail` —
|
||||
it carries the shared HTTP **status** and **message** (e.g. `503` for an
|
||||
unreachable mint) but not the structured `type`/`code`, so branch on the
|
||||
status code here:
|
||||
|
||||
```json
|
||||
{ "detail": "Cashu mint is unreachable" }
|
||||
```
|
||||
|
||||
### Validation Errors
|
||||
|
||||
@@ -347,7 +414,7 @@ class ErrorHandler:
|
||||
'rate_limit',
|
||||
'upstream_timeout',
|
||||
'model_overloaded',
|
||||
'mint_unavailable'
|
||||
'cashu_mint_unreachable'
|
||||
}
|
||||
|
||||
# Errors requiring user action
|
||||
|
||||
@@ -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,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")
|
||||
@@ -0,0 +1,88 @@
|
||||
"""add slug to upstream_providers
|
||||
|
||||
Revision ID: c6d7e8f9a0b1
|
||||
Revises: b5e7c9d1f3a2
|
||||
Create Date: 2026-06-29 00:00:00.000000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from routstr.core.provider_slugs import provider_slug_base, provider_slug_candidate
|
||||
|
||||
revision = "c6d7e8f9a0b1"
|
||||
down_revision = "b5e7c9d1f3a2"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _allocate_backfill_slug(provider_type: str, reserved_slugs: set[str]) -> str:
|
||||
base = provider_slug_base(provider_type)
|
||||
suffix_number = 1
|
||||
while True:
|
||||
candidate = provider_slug_candidate(base, suffix_number)
|
||||
if candidate not in reserved_slugs:
|
||||
reserved_slugs.add(candidate)
|
||||
return candidate
|
||||
suffix_number += 1
|
||||
|
||||
|
||||
def _backfill_provider_slugs(conn: sa.Connection) -> None:
|
||||
existing_rows = conn.execute(
|
||||
sa.text(
|
||||
"SELECT slug FROM upstream_providers "
|
||||
"WHERE slug IS NOT NULL AND slug != ''"
|
||||
)
|
||||
)
|
||||
reserved_slugs = {str(row.slug).lower() for row in existing_rows}
|
||||
|
||||
rows_to_backfill = conn.execute(
|
||||
sa.text(
|
||||
"SELECT id, provider_type FROM upstream_providers "
|
||||
"WHERE slug IS NULL OR slug = '' "
|
||||
"ORDER BY id"
|
||||
)
|
||||
)
|
||||
for row in rows_to_backfill:
|
||||
slug = _allocate_backfill_slug(str(row.provider_type), reserved_slugs)
|
||||
conn.execute(
|
||||
sa.text("UPDATE upstream_providers SET slug = :slug WHERE id = :id"),
|
||||
{"slug": slug, "id": row.id},
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = {c["name"] for c in inspector.get_columns("upstream_providers")}
|
||||
|
||||
if "slug" not in columns:
|
||||
op.add_column(
|
||||
"upstream_providers",
|
||||
sa.Column("slug", sa.String(), nullable=True),
|
||||
)
|
||||
|
||||
_backfill_provider_slugs(conn)
|
||||
|
||||
existing_indexes = {idx["name"] for idx in inspector.get_indexes("upstream_providers")}
|
||||
if "ix_upstream_providers_slug" not in existing_indexes:
|
||||
op.create_index(
|
||||
"ix_upstream_providers_slug",
|
||||
"upstream_providers",
|
||||
["slug"],
|
||||
unique=True,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
existing_indexes = {idx["name"] for idx in inspector.get_indexes("upstream_providers")}
|
||||
if "ix_upstream_providers_slug" in existing_indexes:
|
||||
op.drop_index("ix_upstream_providers_slug", table_name="upstream_providers")
|
||||
|
||||
columns = {c["name"] for c in inspector.get_columns("upstream_providers")}
|
||||
if "slug" in columns:
|
||||
op.drop_column("upstream_providers", "slug")
|
||||
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.4"
|
||||
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" }
|
||||
|
||||
@@ -126,11 +126,24 @@ def create_model_mappings(
|
||||
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
|
||||
@@ -182,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(
|
||||
@@ -204,6 +218,10 @@ 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)
|
||||
@@ -255,18 +273,19 @@ def create_model_mappings(
|
||||
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 base_id not in unique_models:
|
||||
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[base_id] = unique_model
|
||||
unique_models[unique_key] = unique_model
|
||||
|
||||
try:
|
||||
aliases = resolve_model_alias(
|
||||
@@ -292,6 +311,10 @@ 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)
|
||||
|
||||
for alias in aliases:
|
||||
_add_candidate(alias, model_to_use, upstream_for_override)
|
||||
seen_model_provider.add(dedupe_key)
|
||||
@@ -301,7 +324,26 @@ def create_model_mappings(
|
||||
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
|
||||
|
||||
506
routstr/auth.py
506
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,
|
||||
@@ -19,9 +20,20 @@ from .payment.cost_calculation import (
|
||||
MaxCostData,
|
||||
calculate_cost,
|
||||
)
|
||||
from .wallet import credit_balance, deserialize_token_from_string
|
||||
from .wallet import (
|
||||
classify_redemption_error,
|
||||
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)
|
||||
@@ -70,6 +82,37 @@ async def check_and_reset_limit(key: ApiKey, session: AsyncSession) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def redemption_error_to_http_exception(error: Exception) -> HTTPException:
|
||||
"""Map a Cashu token redemption failure to a sanitized client-facing error.
|
||||
|
||||
Thin wrapper over the shared :func:`classify_redemption_error` so the bearer
|
||||
path stays identical to the X-Cashu and top-up paths.
|
||||
"""
|
||||
classified = classify_redemption_error(error)
|
||||
if classified is None:
|
||||
return HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Internal error during token redemption",
|
||||
"type": "api_error",
|
||||
"code": "internal_error",
|
||||
}
|
||||
},
|
||||
)
|
||||
error_type, status_code, message, error_code = classified
|
||||
return HTTPException(
|
||||
status_code=status_code,
|
||||
detail={
|
||||
"error": {
|
||||
"message": message,
|
||||
"type": error_type,
|
||||
"code": error_code,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def validate_bearer_key(
|
||||
bearer_key: str,
|
||||
session: AsyncSession,
|
||||
@@ -208,7 +251,17 @@ async def validate_bearer_key(
|
||||
|
||||
try:
|
||||
hashed_key = hashlib.sha256(bearer_key.encode()).hexdigest()
|
||||
token_obj = deserialize_token_from_string(bearer_key)
|
||||
try:
|
||||
token_obj = deserialize_token_from_string(bearer_key)
|
||||
except Exception as decode_error:
|
||||
# A malformed token is a bad token (400 invalid_cashu_token via
|
||||
# the shared taxonomy), not an auth failure (401) — otherwise it
|
||||
# would fall through to the generic "Invalid API key" handler.
|
||||
raise redemption_error_to_http_exception(
|
||||
ValueError(
|
||||
f"Invalid Cashu token: could not decode token ({decode_error})"
|
||||
)
|
||||
) from decode_error
|
||||
logger.debug(
|
||||
"Generated token hash", extra={"hash_preview": hashed_key[:16] + "..."}
|
||||
)
|
||||
@@ -309,13 +362,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:
|
||||
@@ -326,14 +379,32 @@ async def validate_bearer_key(
|
||||
"error_type": type(credit_error).__name__,
|
||||
},
|
||||
)
|
||||
raise credit_error
|
||||
await session.rollback()
|
||||
raise redemption_error_to_http_exception(credit_error) from credit_error
|
||||
|
||||
if msats <= 0:
|
||||
logger.error(
|
||||
"Token redemption returned zero or negative amount",
|
||||
extra={"msats": msats, "key_hash": hashed_key[:8] + "..."},
|
||||
)
|
||||
raise Exception("Token redemption failed")
|
||||
# Defense-in-depth: credit_balance already raises
|
||||
# ValueError("Redeemed token amount must be positive…") before
|
||||
# returning (wallet.py), so this branch is only reachable if a
|
||||
# zero/negative row was somehow persisted; drop it so we never
|
||||
# leave an orphan zero-balance key. Reuse the shared taxonomy
|
||||
# (cashu_error) so the envelope matches the mapper above.
|
||||
await session.delete(new_key)
|
||||
await session.commit()
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Failed to redeem Cashu token: token yielded no value",
|
||||
"type": "cashu_error",
|
||||
"code": "cashu_token_zero_value",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
await session.refresh(new_key)
|
||||
await session.commit()
|
||||
@@ -351,6 +422,7 @@ async def validate_bearer_key(
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
await session.rollback()
|
||||
logger.error(
|
||||
"Cashu token redemption failed",
|
||||
extra={
|
||||
@@ -365,19 +437,19 @@ async def validate_bearer_key(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Invalid or expired Cashu key: {str(e)}",
|
||||
"message": "Invalid or expired Cashu key",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_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),
|
||||
},
|
||||
)
|
||||
@@ -386,7 +458,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",
|
||||
}
|
||||
@@ -521,31 +593,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",
|
||||
@@ -557,7 +617,6 @@ async def pay_for_request(
|
||||
},
|
||||
)
|
||||
|
||||
# Another concurrent request spent the balance first
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
@@ -569,6 +628,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)
|
||||
@@ -584,6 +681,18 @@ 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
|
||||
|
||||
@@ -595,12 +704,19 @@ async def revert_pay_for_request(
|
||||
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,
|
||||
)
|
||||
)
|
||||
@@ -616,6 +732,7 @@ async def revert_pay_for_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]
|
||||
@@ -635,16 +752,33 @@ async def revert_pay_for_request(
|
||||
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")
|
||||
@@ -716,7 +850,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)",
|
||||
@@ -728,11 +873,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,
|
||||
)
|
||||
@@ -741,13 +907,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]
|
||||
@@ -776,10 +946,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:
|
||||
@@ -799,6 +989,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,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -811,12 +1003,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,
|
||||
)
|
||||
@@ -825,13 +1037,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]
|
||||
@@ -840,77 +1056,99 @@ 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).
|
||||
# CASE expressions keep this atomic and safe even when the
|
||||
# stale-reservation sweeper has already released the reservation.
|
||||
chargeable = case(
|
||||
(col(ApiKey.balance) >= total_cost_msats, total_cost_msats),
|
||||
else_=col(ApiKey.balance),
|
||||
)
|
||||
overrun_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,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=overrun_safe_reserved,
|
||||
balance=col(ApiKey.balance) - chargeable,
|
||||
total_spent=col(ApiKey.total_spent) + chargeable,
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
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)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
reserved_balance=overrun_safe_reserved,
|
||||
total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats),
|
||||
)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
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)
|
||||
|
||||
logger.info(
|
||||
"Finalized payment with additional charge",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": total_cost_msats,
|
||||
"new_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Failed to finalize additional charge - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"attempted_charge": total_cost_msats,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await release_reservation_only()
|
||||
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={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": total_cost_msats,
|
||||
"new_balance": billing_key.balance,
|
||||
"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:
|
||||
# Refund some of the base cost
|
||||
refund = abs(cost_difference)
|
||||
@@ -925,12 +1163,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,
|
||||
)
|
||||
@@ -939,13 +1198,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]
|
||||
@@ -982,6 +1245,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()
|
||||
|
||||
@@ -1071,3 +1353,53 @@ async def periodic_key_reset() -> None:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in periodic_key_reset: {e}")
|
||||
|
||||
|
||||
async def periodic_dead_key_prune() -> None:
|
||||
"""Periodically prune dead API keys. Interval <= 0 disables it.
|
||||
|
||||
See ``prune_dead_api_keys`` for eligibility.
|
||||
"""
|
||||
from .core.db import create_session, prune_dead_api_keys
|
||||
|
||||
interval = settings.dead_key_prune_interval_seconds
|
||||
if interval <= 0:
|
||||
logger.info("Dead-key pruning disabled (interval <= 0)")
|
||||
return
|
||||
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(interval)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
|
||||
try:
|
||||
async with create_session() as session:
|
||||
await prune_dead_api_keys(
|
||||
session, settings.dead_key_min_age_seconds
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in periodic_dead_key_prune: {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,15 +5,29 @@ 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
|
||||
from .wallet import credit_balance, recieve_token, send_to_lnurl, send_token
|
||||
from .wallet import (
|
||||
classify_redemption_error,
|
||||
credit_balance,
|
||||
is_mint_connection_error,
|
||||
recieve_token,
|
||||
send_to_lnurl,
|
||||
send_token,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
balance_router = APIRouter(prefix="/v1/balance")
|
||||
@@ -38,10 +52,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 +62,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)
|
||||
@@ -148,16 +163,18 @@ async def topup_wallet_endpoint(
|
||||
raise HTTPException(status_code=400, detail="Invalid token format")
|
||||
try:
|
||||
amount_msats = await credit_balance(cashu_token, billing_key, session)
|
||||
except ValueError as e:
|
||||
error_msg = str(e)
|
||||
if "already spent" in error_msg.lower():
|
||||
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")
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Failed to redeem token")
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
except Exception as e:
|
||||
# Shared taxonomy so top-up matches the bearer/X-Cashu paths (503 for an
|
||||
# unreachable mint, 422 for fee/swap failures, 400 for token faults).
|
||||
classified = classify_redemption_error(e)
|
||||
if classified is None:
|
||||
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")
|
||||
_type, status_code, message, _code = classified
|
||||
raise HTTPException(status_code=status_code, detail=message)
|
||||
return {"msats": amount_msats}
|
||||
|
||||
|
||||
@@ -190,20 +207,108 @@ 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:
|
||||
# The "in" row exists with a request_id, but the "out" (refund)
|
||||
# row hasn't been written yet — the upstream request is still in
|
||||
# flight and the refund will be minted once it completes. Tell the
|
||||
# client to retry instead of 404ing permanently (race condition
|
||||
# where /v1/wallet/refund is polled before the refund exists).
|
||||
logger.debug(
|
||||
"refund_wallet_endpoint: refund pending (in row exists, out row not yet created)",
|
||||
extra={"request_id": in_tx.request_id},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=425,
|
||||
detail="Refund is pending; retry shortly.",
|
||||
headers={"Retry-After": "2"},
|
||||
)
|
||||
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 = await validate_bearer_key(bearer_value, session)
|
||||
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.",
|
||||
)
|
||||
|
||||
if key.total_balance <= 0:
|
||||
if cached := await _refund_cache_get(bearer_value):
|
||||
@@ -215,6 +320,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":
|
||||
@@ -227,22 +365,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}
|
||||
|
||||
@@ -251,32 +418,105 @@ 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)
|
||||
if (
|
||||
"mint" in error_msg.lower()
|
||||
or "connection" in error_msg.lower()
|
||||
or isinstance(e, Exception)
|
||||
and "ConnectError" in str(type(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 is_mint_connection_error(e):
|
||||
raise HTTPException(status_code=503, detail="Mint service unavailable")
|
||||
else:
|
||||
raise HTTPException(status_code=500, detail="Refund failed")
|
||||
|
||||
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:
|
||||
@@ -327,10 +567,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
|
||||
@@ -355,6 +609,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,
|
||||
@@ -401,6 +656,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"],
|
||||
@@ -413,7 +669,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)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,14 +2,18 @@ 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 sqlalchemy import UniqueConstraint
|
||||
from alembic.util.exc import CommandError
|
||||
from sqlalchemy import UniqueConstraint, delete
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio.engine import create_async_engine
|
||||
from sqlmodel import Field, Relationship, SQLModel, func, select, update
|
||||
from sqlalchemy.orm import aliased
|
||||
from sqlmodel import Field, Relationship, SQLModel, col, func, select, update
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from .logging import get_logger
|
||||
@@ -30,6 +34,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 +54,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 +96,91 @@ 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
|
||||
|
||||
|
||||
async def prune_dead_api_keys(session: AsyncSession, min_age_seconds: int) -> int:
|
||||
"""Delete dead parentless API keys; return the count removed.
|
||||
|
||||
Dead = 0 balance/reservation/spend/requests, older than the grace period,
|
||||
no parent, no children, no pending invoice. Cashu rows are unlinked (not
|
||||
deleted) first to keep the audit trail.
|
||||
"""
|
||||
cutoff = int(time.time()) - min_age_seconds
|
||||
|
||||
child = aliased(ApiKey)
|
||||
has_children = (
|
||||
select(child.hashed_key).where(
|
||||
col(child.parent_key_hash) == col(ApiKey.hashed_key)
|
||||
)
|
||||
).exists()
|
||||
pending_invoice = (
|
||||
select(LightningInvoice.id)
|
||||
.where(col(LightningInvoice.api_key_hash) == col(ApiKey.hashed_key))
|
||||
.where(col(LightningInvoice.status) == "pending")
|
||||
).exists()
|
||||
|
||||
eligible_hashes = (
|
||||
select(ApiKey.hashed_key)
|
||||
.where(col(ApiKey.balance) == 0)
|
||||
.where(col(ApiKey.reserved_balance) == 0)
|
||||
.where(col(ApiKey.total_spent) == 0)
|
||||
.where(col(ApiKey.total_requests) == 0)
|
||||
.where(col(ApiKey.parent_key_hash).is_(None))
|
||||
.where(
|
||||
(col(ApiKey.created_at).is_(None)) | (col(ApiKey.created_at) < cutoff)
|
||||
)
|
||||
.where(~pending_invoice)
|
||||
.where(~has_children)
|
||||
)
|
||||
|
||||
# Unlink transactions rather than cascade-deleting them, so the financial
|
||||
# audit trail survives. The eligibility predicate is re-evaluated inside both
|
||||
# statements so a key that gained balance mid-run is left untouched.
|
||||
await session.exec( # type: ignore[call-overload]
|
||||
update(CashuTransaction)
|
||||
.where(col(CashuTransaction.api_key_hashed_key).in_(eligible_hashes))
|
||||
.values(api_key_hashed_key=None)
|
||||
)
|
||||
|
||||
result = await session.exec( # type: ignore[call-overload]
|
||||
delete(ApiKey).where(col(ApiKey.hashed_key).in_(eligible_hashes))
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
pruned = int(result.rowcount or 0)
|
||||
logger.info(
|
||||
"Pruned dead API keys",
|
||||
extra={"pruned_keys": pruned, "min_age_seconds": min_age_seconds},
|
||||
)
|
||||
return pruned
|
||||
|
||||
|
||||
class ModelRow(SQLModel, table=True): # type: ignore
|
||||
@@ -103,6 +203,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 +230,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
|
||||
@@ -136,6 +319,12 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
),
|
||||
)
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
slug: str | None = Field(
|
||||
default=None,
|
||||
unique=True,
|
||||
index=True,
|
||||
description="Stable external slug used for updates via API key.",
|
||||
)
|
||||
provider_type: str = Field(
|
||||
description="Provider type: custom, openai, anthropic, azure, openrouter, etc."
|
||||
)
|
||||
@@ -157,6 +346,66 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
@@ -227,6 +476,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:
|
||||
@@ -248,8 +508,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")
|
||||
|
||||
|
||||
@@ -7,11 +7,26 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class UpstreamError(Exception):
|
||||
"""Exception raised when an upstream provider fails."""
|
||||
"""Exception raised when an upstream provider fails.
|
||||
|
||||
def __init__(self, message: str, status_code: int = 502):
|
||||
``code`` carries a stable, machine-readable classification (e.g.
|
||||
``UPSTREAM_RATE_LIMIT``) so callers can distinguish failure kinds without
|
||||
string-matching the message. ``details`` holds optional structured,
|
||||
redaction-safe context. Both default to ``None`` for backwards
|
||||
compatibility.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
status_code: int = 502,
|
||||
code: str | None = None,
|
||||
details: dict[str, object] | None = None,
|
||||
):
|
||||
self.message = message
|
||||
self.status_code = status_code
|
||||
self.code = code
|
||||
self.details = details
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
@@ -22,16 +37,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,
|
||||
|
||||
@@ -1,17 +1,73 @@
|
||||
import json
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from heapq import heappush, heapreplace
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator
|
||||
from threading import Lock
|
||||
from typing import Any, Callable, Iterator, TypeVar
|
||||
|
||||
from .logging import get_logger
|
||||
from .usage_analytics_store import UsageAnalyticsStore
|
||||
|
||||
logger = get_logger(__name__)
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class LogManager:
|
||||
def __init__(self, logs_dir: Path = Path("logs")):
|
||||
self.logs_dir = logs_dir
|
||||
self._usage_store = UsageAnalyticsStore(logs_dir=logs_dir)
|
||||
self._analytics_cache_ttl_seconds = 30.0
|
||||
self._analytics_cache: dict[tuple[Any, ...], tuple[float, Any]] = {}
|
||||
self._analytics_cache_lock = Lock()
|
||||
self._cache_miss = object()
|
||||
|
||||
def _get_cached(self, key: tuple[Any, ...]) -> Any:
|
||||
now = time.time()
|
||||
with self._analytics_cache_lock:
|
||||
cached = self._analytics_cache.get(key)
|
||||
if cached is None:
|
||||
return self._cache_miss
|
||||
|
||||
expires_at, value = cached
|
||||
if expires_at <= now:
|
||||
self._analytics_cache.pop(key, None)
|
||||
return self._cache_miss
|
||||
|
||||
return value
|
||||
|
||||
def _set_cached(
|
||||
self, key: tuple[Any, ...], value: Any, ttl_seconds: float | None = None
|
||||
) -> None:
|
||||
ttl = (
|
||||
self._analytics_cache_ttl_seconds
|
||||
if ttl_seconds is None
|
||||
else max(1.0, ttl_seconds)
|
||||
)
|
||||
expires_at = time.time() + ttl
|
||||
with self._analytics_cache_lock:
|
||||
self._analytics_cache[key] = (expires_at, value)
|
||||
|
||||
def _cache_call(
|
||||
self,
|
||||
key: tuple[Any, ...],
|
||||
compute: Callable[[], T],
|
||||
ttl_seconds: float | None = None,
|
||||
) -> T:
|
||||
cached = self._get_cached(key)
|
||||
if cached is not self._cache_miss:
|
||||
return cached
|
||||
|
||||
value = compute()
|
||||
self._set_cached(key, value, ttl_seconds=ttl_seconds)
|
||||
return value
|
||||
|
||||
def _get_cached_entries(self, hours: int) -> list[dict[str, Any]]:
|
||||
return self._cache_call(
|
||||
("usage_entries", hours),
|
||||
lambda: list(self._yield_log_entries(hours_back=hours)),
|
||||
)
|
||||
|
||||
def _yield_log_entries(
|
||||
self,
|
||||
@@ -19,7 +75,6 @@ class LogManager:
|
||||
specific_date: str | None = None,
|
||||
reverse_files: bool = False,
|
||||
max_files: int | None = None,
|
||||
window_center: datetime | None = None,
|
||||
) -> Iterator[dict[str, Any]]:
|
||||
"""
|
||||
Yields log entries from files.
|
||||
@@ -29,7 +84,6 @@ class LogManager:
|
||||
specific_date: specific date string (YYYY-MM-DD) to look at.
|
||||
reverse_files: if True, process files in reverse order (newest first).
|
||||
max_files: maximum number of log files to process (most recent if reverse_files is True).
|
||||
window_center: datetime object to center a 5-month window around.
|
||||
"""
|
||||
if not self.logs_dir.exists():
|
||||
return
|
||||
@@ -44,36 +98,6 @@ class LogManager:
|
||||
log_files.append(log_file)
|
||||
else:
|
||||
log_files = sorted(self.logs_dir.glob("app_*.log"))
|
||||
|
||||
if window_center:
|
||||
# Calculate the 5 months: [center-2, center-1, center, center+1, center+2]
|
||||
allowed_month_years = []
|
||||
cur_m = window_center.month
|
||||
cur_y = window_center.year
|
||||
|
||||
for offset in range(-2, 3):
|
||||
m = cur_m + offset
|
||||
y = cur_y
|
||||
while m <= 0:
|
||||
m += 12
|
||||
y -= 1
|
||||
while m > 12:
|
||||
m -= 12
|
||||
y += 1
|
||||
allowed_month_years.append(f"{y}-{m:02d}")
|
||||
|
||||
filtered_files = []
|
||||
for log_path in log_files:
|
||||
try:
|
||||
# Stem is "app_YYYY-MM-DD"
|
||||
file_date_str = log_path.stem.split("_")[1]
|
||||
file_month_year = file_date_str[:7] # YYYY-MM
|
||||
if file_month_year in allowed_month_years:
|
||||
filtered_files.append(log_path)
|
||||
except Exception:
|
||||
continue
|
||||
log_files = filtered_files
|
||||
|
||||
if reverse_files:
|
||||
log_files.reverse()
|
||||
|
||||
@@ -303,128 +327,208 @@ class LogManager:
|
||||
return 0
|
||||
|
||||
def get_usage_summary(self, hours: int = 24) -> dict:
|
||||
entries = list(
|
||||
self._yield_log_entries(
|
||||
hours_back=hours, window_center=datetime.now(timezone.utc)
|
||||
)
|
||||
def compute() -> dict:
|
||||
try:
|
||||
return self._usage_store.get_summary(hours_back=hours)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Usage analytics index failed, falling back to log scan: {e}"
|
||||
)
|
||||
return self._calculate_summary_stats(self._get_cached_entries(hours))
|
||||
|
||||
return self._cache_call(
|
||||
("usage_summary", hours),
|
||||
compute,
|
||||
)
|
||||
return self._calculate_summary_stats(entries)
|
||||
|
||||
def get_usage_metrics(self, interval: int = 15, hours: int = 24) -> dict:
|
||||
entries = list(
|
||||
self._yield_log_entries(
|
||||
hours_back=hours, window_center=datetime.now(timezone.utc)
|
||||
)
|
||||
def compute() -> dict:
|
||||
try:
|
||||
return self._usage_store.get_metrics(
|
||||
interval_minutes=interval,
|
||||
hours_back=hours,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Usage analytics index failed, falling back to log scan: {e}"
|
||||
)
|
||||
return self._aggregate_metrics_by_time(
|
||||
self._get_cached_entries(hours), interval, hours
|
||||
)
|
||||
|
||||
return self._cache_call(
|
||||
("usage_metrics", interval, hours),
|
||||
compute,
|
||||
)
|
||||
|
||||
def get_usage_dashboard(
|
||||
self,
|
||||
interval: int = 15,
|
||||
hours: int = 24,
|
||||
error_limit: int = 100,
|
||||
model_limit: int = 20,
|
||||
) -> dict:
|
||||
# Large ranges are expensive to scan; keep cached longer.
|
||||
if hours <= 24:
|
||||
cache_ttl = 60.0
|
||||
elif hours <= 7 * 24:
|
||||
cache_ttl = 300.0
|
||||
elif hours <= 30 * 24:
|
||||
cache_ttl = 1800.0
|
||||
elif hours <= 90 * 24:
|
||||
cache_ttl = 7200.0
|
||||
else:
|
||||
cache_ttl = 21600.0
|
||||
|
||||
def compute() -> dict:
|
||||
try:
|
||||
return self._usage_store.get_dashboard(
|
||||
interval_minutes=interval,
|
||||
hours_back=hours,
|
||||
error_limit=error_limit,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Usage analytics index failed, falling back to log scan: {e}"
|
||||
)
|
||||
return self._aggregate_dashboard(
|
||||
interval_minutes=interval,
|
||||
hours_back=hours,
|
||||
error_limit=error_limit,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
return self._cache_call(
|
||||
("usage_dashboard", interval, hours, error_limit, model_limit),
|
||||
compute,
|
||||
ttl_seconds=cache_ttl,
|
||||
)
|
||||
return self._aggregate_metrics_by_time(entries, interval, hours)
|
||||
|
||||
def get_error_details(self, hours: int = 24, limit: int = 100) -> dict:
|
||||
errors: list[dict[str, Any]] = []
|
||||
def compute() -> dict:
|
||||
try:
|
||||
return self._usage_store.get_error_details(hours_back=hours, limit=limit)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Usage analytics index failed, falling back to log scan: {e}"
|
||||
)
|
||||
|
||||
for entry in self._yield_log_entries(hours_back=hours):
|
||||
if str(entry.get("levelname", "")).upper() != "ERROR":
|
||||
continue
|
||||
errors: list[dict] = []
|
||||
for entry in self._get_cached_entries(hours):
|
||||
if str(entry.get("levelname", "")).upper() == "ERROR":
|
||||
timestamp_str = entry.get("asctime", "")
|
||||
errors.append(
|
||||
{
|
||||
"timestamp": timestamp_str,
|
||||
"message": entry.get("message", ""),
|
||||
"error_type": entry.get("error_type", "unknown"),
|
||||
"pathname": entry.get("pathname", ""),
|
||||
"lineno": entry.get("lineno", 0),
|
||||
"request_id": entry.get("request_id", ""),
|
||||
}
|
||||
)
|
||||
|
||||
errors.append(
|
||||
{
|
||||
"timestamp": entry.get("asctime", ""),
|
||||
"message": entry.get("message", ""),
|
||||
"error_type": entry.get("error_type", "unknown"),
|
||||
"pathname": entry.get("pathname", ""),
|
||||
"lineno": entry.get("lineno", 0),
|
||||
"request_id": entry.get("request_id", ""),
|
||||
}
|
||||
)
|
||||
errors.sort(key=lambda x: x["timestamp"], reverse=True)
|
||||
return {"errors": errors[:limit], "total_count": len(errors)}
|
||||
|
||||
errors.sort(key=lambda x: str(x["timestamp"]), reverse=True)
|
||||
return {"errors": errors[:limit], "total_count": len(errors)}
|
||||
return self._cache_call(("error_details", hours, limit), compute)
|
||||
|
||||
def get_revenue_by_model(self, hours: int = 24, limit: int = 20) -> dict:
|
||||
entries = list(
|
||||
self._yield_log_entries(
|
||||
hours_back=hours, window_center=datetime.now(timezone.utc)
|
||||
)
|
||||
)
|
||||
|
||||
model_stats: dict[str, dict[str, int | float]] = defaultdict(
|
||||
lambda: {
|
||||
"revenue_msats": 0,
|
||||
"refunds_msats": 0,
|
||||
"requests": 0,
|
||||
"successful": 0,
|
||||
"failed": 0,
|
||||
}
|
||||
)
|
||||
|
||||
for entry in entries:
|
||||
def compute() -> dict:
|
||||
try:
|
||||
model = entry.get("model", "unknown")
|
||||
if not isinstance(model, str):
|
||||
model = "unknown"
|
||||
|
||||
message = str(entry.get("message", "")).lower()
|
||||
|
||||
completed, revenue_msats, _, _ = self._extract_success_metrics(
|
||||
entry, message
|
||||
return self._usage_store.get_revenue_by_model(
|
||||
hours_back=hours, limit=limit
|
||||
)
|
||||
if completed:
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["successful"] += 1
|
||||
if revenue_msats > 0:
|
||||
model_stats[model]["revenue_msats"] += revenue_msats
|
||||
|
||||
failed = (
|
||||
"revert payment" in message or "upstream request failed" in message
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Usage analytics index failed, falling back to log scan: {e}"
|
||||
)
|
||||
if failed:
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["failed"] += 1
|
||||
if "revert payment" in message:
|
||||
max_cost = entry.get("max_cost_for_model", 0)
|
||||
if isinstance(max_cost, (int, float)) and max_cost > 0:
|
||||
model_stats[model]["refunds_msats"] += max_cost
|
||||
|
||||
except Exception:
|
||||
continue
|
||||
entries = self._get_cached_entries(hours)
|
||||
|
||||
models: list[dict[str, Any]] = []
|
||||
total_revenue = 0.0
|
||||
|
||||
for model, stats in model_stats.items():
|
||||
revenue_msats = float(stats["revenue_msats"])
|
||||
refunds_msats = float(stats["refunds_msats"])
|
||||
|
||||
revenue_sats = revenue_msats / 1000
|
||||
refunds_sats = refunds_msats / 1000
|
||||
net_revenue_sats = revenue_sats - refunds_sats
|
||||
|
||||
total_revenue += net_revenue_sats
|
||||
|
||||
requests = int(stats["requests"])
|
||||
successful = int(stats["successful"])
|
||||
|
||||
models.append(
|
||||
{
|
||||
"model": model,
|
||||
"revenue_sats": revenue_sats,
|
||||
"refunds_sats": refunds_sats,
|
||||
"net_revenue_sats": net_revenue_sats,
|
||||
"requests": requests,
|
||||
"successful": successful,
|
||||
"failed": int(stats["failed"]),
|
||||
"avg_revenue_per_request": (
|
||||
revenue_sats / successful if successful > 0 else 0
|
||||
),
|
||||
model_stats: dict[str, dict[str, int | float]] = defaultdict(
|
||||
lambda: {
|
||||
"revenue_msats": 0,
|
||||
"refunds_msats": 0,
|
||||
"requests": 0,
|
||||
"successful": 0,
|
||||
"failed": 0,
|
||||
}
|
||||
)
|
||||
|
||||
models.sort(key=lambda x: float(x["net_revenue_sats"]), reverse=True)
|
||||
for entry in entries:
|
||||
try:
|
||||
model = entry.get("model", "unknown")
|
||||
if not isinstance(model, str):
|
||||
model = "unknown"
|
||||
|
||||
return {
|
||||
"models": models[:limit],
|
||||
"total_revenue_sats": total_revenue,
|
||||
"total_models": len(models),
|
||||
}
|
||||
message = str(entry.get("message", "")).lower()
|
||||
|
||||
completed, revenue_msats, _, _ = self._extract_success_metrics(
|
||||
entry, message
|
||||
)
|
||||
if completed:
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["successful"] += 1
|
||||
if revenue_msats > 0:
|
||||
model_stats[model]["revenue_msats"] += revenue_msats
|
||||
|
||||
failed = (
|
||||
"revert payment" in message
|
||||
or "upstream request failed" in message
|
||||
)
|
||||
if failed:
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["failed"] += 1
|
||||
if "revert payment" in message:
|
||||
max_cost = entry.get("max_cost_for_model", 0)
|
||||
if isinstance(max_cost, (int, float)) and max_cost > 0:
|
||||
model_stats[model]["refunds_msats"] += max_cost
|
||||
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
models: list[dict[str, Any]] = []
|
||||
total_revenue = 0.0
|
||||
|
||||
for model, stats in model_stats.items():
|
||||
revenue_msats = float(stats["revenue_msats"])
|
||||
refunds_msats = float(stats["refunds_msats"])
|
||||
|
||||
revenue_sats = revenue_msats / 1000
|
||||
refunds_sats = refunds_msats / 1000
|
||||
net_revenue_sats = revenue_sats - refunds_sats
|
||||
|
||||
total_revenue += net_revenue_sats
|
||||
|
||||
requests = int(stats["requests"])
|
||||
successful = int(stats["successful"])
|
||||
|
||||
models.append(
|
||||
{
|
||||
"model": model,
|
||||
"revenue_sats": revenue_sats,
|
||||
"refunds_sats": refunds_sats,
|
||||
"net_revenue_sats": net_revenue_sats,
|
||||
"requests": requests,
|
||||
"successful": successful,
|
||||
"failed": int(stats["failed"]),
|
||||
"avg_revenue_per_request": (
|
||||
revenue_sats / successful if successful > 0 else 0
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
models.sort(key=lambda x: float(x["net_revenue_sats"]), reverse=True)
|
||||
|
||||
return {
|
||||
"models": models[:limit],
|
||||
"total_revenue_sats": total_revenue,
|
||||
"total_models": len(models),
|
||||
}
|
||||
|
||||
return self._cache_call(("revenue_by_model", hours, limit), compute)
|
||||
|
||||
def _build_summary_response(self, stats: dict[str, Any]) -> dict[str, Any]:
|
||||
revenue_sats = stats["revenue_msats"] / 1000
|
||||
@@ -555,6 +659,344 @@ class LogManager:
|
||||
|
||||
return self._build_summary_response(stats)
|
||||
|
||||
def _aggregate_dashboard(
|
||||
self,
|
||||
interval_minutes: int,
|
||||
hours_back: int,
|
||||
error_limit: int,
|
||||
model_limit: int,
|
||||
) -> dict[str, Any]:
|
||||
time_buckets: dict[str, dict[str, Any]] = defaultdict(
|
||||
lambda: {
|
||||
"total_requests": 0,
|
||||
"successful_chat_completions": 0,
|
||||
"failed_requests": 0,
|
||||
"errors": 0,
|
||||
"warnings": 0,
|
||||
"payment_processed": 0,
|
||||
"upstream_errors": 0,
|
||||
"revenue_msats": 0.0,
|
||||
"refunds_msats": 0.0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
)
|
||||
summary_stats: dict[str, Any] = {
|
||||
"total_entries": 0,
|
||||
"total_requests": 0,
|
||||
"successful_chat_completions": 0,
|
||||
"failed_requests": 0,
|
||||
"total_errors": 0,
|
||||
"total_warnings": 0,
|
||||
"payment_processed": 0,
|
||||
"upstream_errors": 0,
|
||||
"unique_models": set(),
|
||||
"error_types": defaultdict(int),
|
||||
"revenue_msats": 0.0,
|
||||
"refunds_msats": 0.0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
model_stats: dict[str, dict[str, int | float]] = defaultdict(
|
||||
lambda: {
|
||||
"revenue_msats": 0,
|
||||
"refunds_msats": 0,
|
||||
"requests": 0,
|
||||
"successful": 0,
|
||||
"failed": 0,
|
||||
}
|
||||
)
|
||||
model_mix_buckets: dict[str, dict[str, int]] = defaultdict(
|
||||
lambda: defaultdict(int)
|
||||
)
|
||||
model_mix_revenue_buckets: dict[str, dict[str, float]] = defaultdict(
|
||||
lambda: defaultdict(float)
|
||||
)
|
||||
model_mix_token_buckets: dict[str, dict[str, int]] = defaultdict(
|
||||
lambda: defaultdict(int)
|
||||
)
|
||||
model_mix_totals: dict[str, int] = defaultdict(int)
|
||||
model_mix_revenue_totals: dict[str, float] = defaultdict(float)
|
||||
model_mix_token_totals: dict[str, int] = defaultdict(int)
|
||||
latest_errors_heap: list[tuple[str, dict[str, Any]]] = []
|
||||
total_error_count = 0
|
||||
|
||||
for entry in self._yield_log_entries(hours_back=hours_back):
|
||||
try:
|
||||
summary_stats["total_entries"] += 1
|
||||
|
||||
timestamp_str = entry.get("asctime", "")
|
||||
message = str(entry.get("message", "")).lower()
|
||||
level = str(entry.get("levelname", "")).upper()
|
||||
model = entry.get("model", "unknown")
|
||||
if not isinstance(model, str):
|
||||
model = "unknown"
|
||||
|
||||
bucket_key = (
|
||||
self._bucket_key_for_timestamp(timestamp_str, interval_minutes)
|
||||
if isinstance(timestamp_str, str)
|
||||
else None
|
||||
)
|
||||
bucket = time_buckets[bucket_key] if bucket_key else None
|
||||
|
||||
if level == "ERROR":
|
||||
summary_stats["total_errors"] += 1
|
||||
if bucket:
|
||||
bucket["errors"] += 1
|
||||
if "error_type" in entry:
|
||||
summary_stats["error_types"][str(entry["error_type"])] += 1
|
||||
|
||||
total_error_count += 1
|
||||
error_item = {
|
||||
"timestamp": timestamp_str,
|
||||
"message": entry.get("message", ""),
|
||||
"error_type": entry.get("error_type", "unknown"),
|
||||
"pathname": entry.get("pathname", ""),
|
||||
"lineno": entry.get("lineno", 0),
|
||||
"request_id": entry.get("request_id", ""),
|
||||
}
|
||||
if len(latest_errors_heap) < error_limit:
|
||||
heappush(latest_errors_heap, (timestamp_str, error_item))
|
||||
elif timestamp_str > latest_errors_heap[0][0]:
|
||||
heapreplace(latest_errors_heap, (timestamp_str, error_item))
|
||||
elif level == "WARNING":
|
||||
summary_stats["total_warnings"] += 1
|
||||
if bucket:
|
||||
bucket["warnings"] += 1
|
||||
|
||||
completed, revenue_msats, input_tokens, output_tokens = (
|
||||
self._extract_success_metrics(entry, message)
|
||||
)
|
||||
if completed:
|
||||
summary_stats["total_requests"] += 1
|
||||
summary_stats["successful_chat_completions"] += 1
|
||||
summary_stats["input_tokens"] += input_tokens
|
||||
summary_stats["output_tokens"] += output_tokens
|
||||
summary_stats["total_tokens"] += input_tokens + output_tokens
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["successful"] += 1
|
||||
model_mix_totals[model] += 1
|
||||
if bucket:
|
||||
bucket["total_requests"] += 1
|
||||
bucket["successful_chat_completions"] += 1
|
||||
bucket["input_tokens"] += input_tokens
|
||||
bucket["output_tokens"] += output_tokens
|
||||
bucket["total_tokens"] += input_tokens + output_tokens
|
||||
if bucket_key:
|
||||
model_mix_buckets[bucket_key][model] += 1
|
||||
if revenue_msats > 0:
|
||||
model_mix_revenue_buckets[bucket_key][model] += revenue_msats
|
||||
model_mix_revenue_totals[model] += revenue_msats
|
||||
if input_tokens > 0 or output_tokens > 0:
|
||||
token_total = input_tokens + output_tokens
|
||||
model_mix_token_buckets[bucket_key][model] += token_total
|
||||
model_mix_token_totals[model] += token_total
|
||||
|
||||
if revenue_msats > 0:
|
||||
summary_stats["revenue_msats"] += revenue_msats
|
||||
model_stats[model]["revenue_msats"] += revenue_msats
|
||||
if bucket:
|
||||
bucket["revenue_msats"] += revenue_msats
|
||||
|
||||
failed = (
|
||||
"upstream request failed" in message
|
||||
or "revert payment" in message
|
||||
)
|
||||
if failed:
|
||||
summary_stats["total_requests"] += 1
|
||||
summary_stats["failed_requests"] += 1
|
||||
model_stats[model]["requests"] += 1
|
||||
model_stats[model]["failed"] += 1
|
||||
if bucket:
|
||||
bucket["total_requests"] += 1
|
||||
bucket["failed_requests"] += 1
|
||||
|
||||
if "payment processed successfully" in message:
|
||||
summary_stats["payment_processed"] += 1
|
||||
if bucket:
|
||||
bucket["payment_processed"] += 1
|
||||
|
||||
if "upstream" in message and level == "ERROR":
|
||||
summary_stats["upstream_errors"] += 1
|
||||
if bucket:
|
||||
bucket["upstream_errors"] += 1
|
||||
|
||||
if model != "unknown":
|
||||
summary_stats["unique_models"].add(model)
|
||||
|
||||
if "revert payment" in message:
|
||||
max_cost = entry.get("max_cost_for_model", 0)
|
||||
if isinstance(max_cost, (int, float)) and max_cost > 0:
|
||||
max_cost_float = float(max_cost)
|
||||
summary_stats["refunds_msats"] += max_cost_float
|
||||
model_stats[model]["refunds_msats"] += max_cost_float
|
||||
if bucket:
|
||||
bucket["refunds_msats"] += max_cost_float
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
metrics_result = []
|
||||
for bucket_key in sorted(time_buckets.keys()):
|
||||
bucket = dict(time_buckets[bucket_key])
|
||||
bucket["requests"] = bucket["total_requests"]
|
||||
metrics_result.append({"timestamp": bucket_key, **bucket})
|
||||
|
||||
models: list[dict[str, Any]] = []
|
||||
total_revenue = 0.0
|
||||
for model_name, stats in model_stats.items():
|
||||
revenue_msats = float(stats["revenue_msats"])
|
||||
refunds_msats = float(stats["refunds_msats"])
|
||||
revenue_sats = revenue_msats / 1000
|
||||
refunds_sats = refunds_msats / 1000
|
||||
net_revenue_sats = revenue_sats - refunds_sats
|
||||
total_revenue += net_revenue_sats
|
||||
|
||||
successful = int(stats["successful"])
|
||||
models.append(
|
||||
{
|
||||
"model": model_name,
|
||||
"revenue_sats": revenue_sats,
|
||||
"refunds_sats": refunds_sats,
|
||||
"net_revenue_sats": net_revenue_sats,
|
||||
"requests": int(stats["requests"]),
|
||||
"successful": successful,
|
||||
"failed": int(stats["failed"]),
|
||||
"avg_revenue_per_request": (
|
||||
revenue_sats / successful if successful > 0 else 0
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
models.sort(key=lambda x: float(x["net_revenue_sats"]), reverse=True)
|
||||
latest_errors = [
|
||||
item
|
||||
for _, item in sorted(
|
||||
latest_errors_heap, key=lambda x: x[0], reverse=True
|
||||
)
|
||||
]
|
||||
top_model_limit = max(1, min(model_limit, 20))
|
||||
top_models_requests = [
|
||||
model_name
|
||||
for model_name, _ in sorted(
|
||||
(
|
||||
(name, count)
|
||||
for name, count in model_mix_totals.items()
|
||||
if name != "unknown"
|
||||
),
|
||||
key=lambda item: item[1],
|
||||
reverse=True,
|
||||
)[:top_model_limit]
|
||||
]
|
||||
top_models_revenue = [
|
||||
model_name
|
||||
for model_name, _ in sorted(
|
||||
(
|
||||
(name, amount)
|
||||
for name, amount in model_mix_revenue_totals.items()
|
||||
if name != "unknown"
|
||||
),
|
||||
key=lambda item: item[1],
|
||||
reverse=True,
|
||||
)[:top_model_limit]
|
||||
]
|
||||
top_models_tokens = [
|
||||
model_name
|
||||
for model_name, _ in sorted(
|
||||
(
|
||||
(name, token_count)
|
||||
for name, token_count in model_mix_token_totals.items()
|
||||
if name != "unknown"
|
||||
),
|
||||
key=lambda item: item[1],
|
||||
reverse=True,
|
||||
)[:top_model_limit]
|
||||
]
|
||||
selected_models: list[str] = []
|
||||
for model in top_models_requests + top_models_revenue + top_models_tokens:
|
||||
if model not in selected_models:
|
||||
selected_models.append(model)
|
||||
top_model_set = set(selected_models)
|
||||
|
||||
model_usage_mix_metrics: list[dict[str, Any]] = []
|
||||
mix_bucket_keys = sorted(
|
||||
set(model_mix_buckets.keys())
|
||||
| set(model_mix_revenue_buckets.keys())
|
||||
| set(model_mix_token_buckets.keys())
|
||||
)
|
||||
for bucket_key in mix_bucket_keys:
|
||||
counts = model_mix_buckets.get(bucket_key, {})
|
||||
revenue_counts = model_mix_revenue_buckets.get(bucket_key, {})
|
||||
token_counts = model_mix_token_buckets.get(bucket_key, {})
|
||||
others = 0
|
||||
others_revenue_msats = 0.0
|
||||
others_tokens = 0
|
||||
model_counts: dict[str, int] = {}
|
||||
model_revenue_msats: dict[str, float] = {}
|
||||
model_tokens: dict[str, int] = {}
|
||||
for model_name, successful_count in counts.items():
|
||||
if model_name in top_model_set:
|
||||
model_counts[model_name] = int(successful_count)
|
||||
else:
|
||||
others += int(successful_count)
|
||||
for model_name, revenue_value in revenue_counts.items():
|
||||
if model_name in top_model_set:
|
||||
model_revenue_msats[model_name] = float(revenue_value)
|
||||
else:
|
||||
others_revenue_msats += float(revenue_value)
|
||||
for model_name, token_value in token_counts.items():
|
||||
if model_name in top_model_set:
|
||||
model_tokens[model_name] = int(token_value)
|
||||
else:
|
||||
others_tokens += int(token_value)
|
||||
|
||||
model_usage_mix_metrics.append(
|
||||
{
|
||||
"timestamp": bucket_key,
|
||||
"total_successful": int(sum(counts.values())),
|
||||
"total_revenue_msats": float(sum(revenue_counts.values())),
|
||||
"total_tokens": int(sum(token_counts.values())),
|
||||
"others": others,
|
||||
"others_revenue_msats": others_revenue_msats,
|
||||
"others_tokens": others_tokens,
|
||||
"model_counts": model_counts,
|
||||
"model_revenue_msats": model_revenue_msats,
|
||||
"model_tokens": model_tokens,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"metrics": {
|
||||
"metrics": metrics_result,
|
||||
"interval_minutes": interval_minutes,
|
||||
"hours_back": hours_back,
|
||||
"total_buckets": len(metrics_result),
|
||||
},
|
||||
"summary": self._build_summary_response(summary_stats),
|
||||
"error_details": {
|
||||
"errors": latest_errors,
|
||||
"total_count": total_error_count,
|
||||
},
|
||||
"revenue_by_model": {
|
||||
"models": models[:model_limit],
|
||||
"total_revenue_sats": total_revenue,
|
||||
"total_models": len(models),
|
||||
},
|
||||
"model_usage_mix": {
|
||||
"top_models": top_models_requests,
|
||||
"top_models_by_metric": {
|
||||
"requests": top_models_requests,
|
||||
"revenue": top_models_revenue,
|
||||
"tokens": top_models_tokens,
|
||||
},
|
||||
"metrics": model_usage_mix_metrics,
|
||||
"interval_minutes": interval_minutes,
|
||||
"hours_back": hours_back,
|
||||
"total_buckets": len(model_usage_mix_metrics),
|
||||
},
|
||||
}
|
||||
|
||||
def _aggregate_metrics_by_time(
|
||||
self, entries: list[dict], interval_minutes: int, hours_back: int
|
||||
) -> dict:
|
||||
|
||||
@@ -3,50 +3,63 @@ 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
|
||||
|
||||
from .redaction import redact_obj, redact_org_ids
|
||||
|
||||
# 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")
|
||||
@@ -169,6 +182,37 @@ class RequestIdFilter(logging.Filter):
|
||||
return True
|
||||
|
||||
|
||||
# Standard ``LogRecord`` attributes that are never user-supplied ``extra``
|
||||
# fields; skipped when redacting structured extras (``msg``/``message`` are
|
||||
# handled separately above).
|
||||
_NON_EXTRA_RECORD_ATTRS = frozenset(
|
||||
{
|
||||
"name",
|
||||
"msg",
|
||||
"args",
|
||||
"levelname",
|
||||
"levelno",
|
||||
"pathname",
|
||||
"filename",
|
||||
"module",
|
||||
"exc_info",
|
||||
"exc_text",
|
||||
"stack_info",
|
||||
"lineno",
|
||||
"funcName",
|
||||
"created",
|
||||
"msecs",
|
||||
"relativeCreated",
|
||||
"thread",
|
||||
"threadName",
|
||||
"processName",
|
||||
"process",
|
||||
"taskName",
|
||||
"message",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class SecurityFilter(logging.Filter):
|
||||
"""Filter to remove sensitive information from logs."""
|
||||
|
||||
@@ -192,6 +236,7 @@ class SecurityFilter(logging.Filter):
|
||||
"""Filter out sensitive information from log records."""
|
||||
try:
|
||||
message = record.getMessage()
|
||||
message = redact_org_ids(message)
|
||||
standalone_patterns = [
|
||||
r"Bearer\s+([a-zA-Z0-9_\-\.]{10,})", # Bearer token (must be 10 characters or more to reduce false-positives)
|
||||
r"cashu[A-Z]+([a-zA-Z0-9_\-\.=/+]+)", # Cashu tokens
|
||||
@@ -213,6 +258,16 @@ class SecurityFilter(logging.Filter):
|
||||
record.msg = message
|
||||
record.args = ()
|
||||
|
||||
# Structured `extra={...}` fields are emitted by the JSON formatter
|
||||
# straight from the record dict and never pass through the message
|
||||
# formatting above. Redact organization IDs from any string-valued
|
||||
# extra so they cannot leak via structured logs.
|
||||
for attr, value in list(record.__dict__.items()):
|
||||
if attr in _NON_EXTRA_RECORD_ATTRS:
|
||||
continue
|
||||
if isinstance(value, (str, dict, list, tuple)):
|
||||
record.__dict__[attr] = redact_obj(value)
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -259,6 +314,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 +343,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 +354,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,33 +8,43 @@ 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_dead_key_prune,
|
||||
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 ..upstream.auto_topup import periodic_auto_topup
|
||||
from ..wallet import periodic_payout
|
||||
from ..upstream.deepseek_v4_pricing_shim import register_deepseek_v4_pricing
|
||||
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]:
|
||||
@@ -45,13 +54,28 @@ 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
|
||||
dead_key_prune_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()
|
||||
|
||||
# TEMPORARY: backfill DeepSeek V4 pricing missing from litellm's cost
|
||||
# map (BerriAI/litellm#30430). Remove this call and
|
||||
# deepseek_v4_pricing_shim.py once litellm ships these models.
|
||||
register_deepseek_v4_pricing()
|
||||
|
||||
# Run database migrations on startup
|
||||
run_migrations()
|
||||
|
||||
@@ -96,17 +120,28 @@ 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()
|
||||
)
|
||||
dead_key_prune_task = asyncio.create_task(periodic_dead_key_prune())
|
||||
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
|
||||
|
||||
@@ -130,6 +165,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:
|
||||
@@ -138,8 +175,18 @@ 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 dead_key_prune_task is not None:
|
||||
dead_key_prune_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 = []
|
||||
@@ -151,6 +198,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:
|
||||
@@ -159,8 +208,18 @@ 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 dead_key_prune_task is not None:
|
||||
tasks_to_wait.append(dead_key_prune_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)
|
||||
@@ -172,6 +231,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)
|
||||
|
||||
|
||||
@@ -181,7 +257,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
|
||||
@@ -218,7 +294,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",
|
||||
)
|
||||
|
||||
@@ -226,100 +302,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"
|
||||
@@ -331,12 +377,12 @@ 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"
|
||||
"UI dist directory not found at %s; serving API only. Run `make ui-build` "
|
||||
"to build the static UI served from here, or `make ui-dev` for the Next.js "
|
||||
"dev server with hot reload on :3000 (it targets this backend on :8000).",
|
||||
UI_DIST_PATH,
|
||||
)
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
@@ -353,6 +399,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)
|
||||
65
routstr/core/provider_slugs.py
Normal file
65
routstr/core/provider_slugs.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from itertools import count
|
||||
from typing import Collection
|
||||
|
||||
from sqlmodel import select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from .db import UpstreamProviderRow
|
||||
|
||||
_SLUG_BASE_PATTERN = re.compile(r"[^a-z0-9]+")
|
||||
_MAX_SLUG_LENGTH = 64
|
||||
|
||||
|
||||
def provider_slug_base(provider_type: str) -> str:
|
||||
"""Return a deterministic slug base for a provider type."""
|
||||
base = _SLUG_BASE_PATTERN.sub("-", provider_type.lower()).strip("-")
|
||||
if not base:
|
||||
base = "provider"
|
||||
elif base.isdigit():
|
||||
base = f"provider-{base}"
|
||||
elif len(base) < 3:
|
||||
base = f"{base}-provider"
|
||||
|
||||
if len(base) > _MAX_SLUG_LENGTH:
|
||||
base = base[:_MAX_SLUG_LENGTH].rstrip("-") or "provider"
|
||||
return base
|
||||
|
||||
|
||||
def provider_slug_candidate(base: str, suffix_number: int) -> str:
|
||||
if suffix_number == 1:
|
||||
return base
|
||||
|
||||
suffix = f"-{suffix_number}"
|
||||
max_base_length = _MAX_SLUG_LENGTH - len(suffix)
|
||||
return f"{base[:max_base_length].rstrip('-')}{suffix}"
|
||||
|
||||
|
||||
async def allocate_unique_provider_slug(
|
||||
session: AsyncSession,
|
||||
provider_type: str,
|
||||
reserved_slugs: Collection[str] = (),
|
||||
) -> str:
|
||||
"""Allocate a stable, deterministic provider slug.
|
||||
|
||||
The first provider of a type gets ``openai``; later collisions get
|
||||
``openai-2``, ``openai-3``, etc. ``reserved_slugs`` covers rows staged in
|
||||
memory but not flushed yet, such as settings/env seeding.
|
||||
"""
|
||||
base = provider_slug_base(provider_type)
|
||||
reserved = {slug.lower() for slug in reserved_slugs}
|
||||
|
||||
for suffix_number in count(1):
|
||||
candidate = provider_slug_candidate(base, suffix_number)
|
||||
if candidate in reserved:
|
||||
continue
|
||||
|
||||
result = await session.exec(
|
||||
select(UpstreamProviderRow).where(UpstreamProviderRow.slug == candidate)
|
||||
)
|
||||
if result.first() is None:
|
||||
return candidate
|
||||
|
||||
raise RuntimeError("unreachable")
|
||||
51
routstr/core/redaction.py
Normal file
51
routstr/core/redaction.py
Normal file
@@ -0,0 +1,51 @@
|
||||
"""Redaction helpers for sensitive provider identifiers.
|
||||
|
||||
Single source of truth for stripping account-scoped identifiers (e.g. OpenAI
|
||||
organization IDs) from any text before it is logged, returned to a caller, or
|
||||
written to an audit entry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
# OpenAI-style organization identifiers look like ``org-<base62>``. Require at
|
||||
# least 6 trailing chars so the already-redacted literal ``org-[REDACTED]`` is
|
||||
# never re-matched (``[`` is not in the character class).
|
||||
_ORG_ID_PATTERN = re.compile(r"\borg-[A-Za-z0-9]{6,}\b")
|
||||
|
||||
ORG_ID_PLACEHOLDER = "org-[REDACTED]"
|
||||
|
||||
|
||||
def redact_org_ids(text: str) -> str:
|
||||
"""Replace OpenAI-style organization IDs with ``org-[REDACTED]``.
|
||||
|
||||
Args:
|
||||
text: Arbitrary text that may embed an ``org-*`` identifier.
|
||||
|
||||
Returns:
|
||||
The text with every organization ID replaced. Non-string input is
|
||||
returned unchanged after coercion to ``str``.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
return _ORG_ID_PATTERN.sub(ORG_ID_PLACEHOLDER, text)
|
||||
|
||||
|
||||
def redact_obj(obj: Any) -> Any:
|
||||
"""Recursively redact organization IDs in arbitrary nested structures.
|
||||
|
||||
Strings are redacted in place; dicts and lists/tuples are walked so that
|
||||
identifiers nested inside structured payloads (e.g. log ``extra`` fields or
|
||||
error ``details``) are also stripped. Other types are returned unchanged.
|
||||
"""
|
||||
if isinstance(obj, str):
|
||||
return redact_org_ids(obj)
|
||||
if isinstance(obj, dict):
|
||||
return {key: redact_obj(value) for key, value in obj.items()}
|
||||
if isinstance(obj, list):
|
||||
return [redact_obj(value) for value in obj]
|
||||
if isinstance(obj, tuple):
|
||||
return tuple(redact_obj(value) for value in obj)
|
||||
return obj
|
||||
@@ -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
|
||||
@@ -58,6 +67,20 @@ class Settings(BaseSettings):
|
||||
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"
|
||||
)
|
||||
# Background prune of dead (zero balance, never used) API keys.
|
||||
# Interval 0 disables it; min-age is a grace period (default 1 week).
|
||||
dead_key_prune_interval_seconds: int = Field(
|
||||
default=3600, env="DEAD_KEY_PRUNE_INTERVAL_SECONDS"
|
||||
)
|
||||
dead_key_min_age_seconds: int = Field(
|
||||
default=604_800, env="DEAD_KEY_MIN_AGE_SECONDS"
|
||||
)
|
||||
|
||||
# Network
|
||||
cors_origins: list[str] = Field(default_factory=lambda: ["*"], env="CORS_ORIGINS")
|
||||
@@ -74,6 +97,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 +116,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 +269,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 +291,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 +314,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 +348,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.4"
|
||||
|
||||
_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,33 @@ class CostDataError(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
async def calculate_cost( # todo: can be sync
|
||||
response_data: dict, max_cost: int, session: AsyncSession
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
def _empty_cost(cls: type[CostData] = CostData) -> CostData:
|
||||
"""Build an all-zero cost object — a full refund for an empty response.
|
||||
|
||||
Shared by the two paths that must not bill: an upstream response with no
|
||||
usage data at all, and one that reports a USD cost but carries zero tokens
|
||||
in every bucket.
|
||||
"""
|
||||
Calculate the cost of an API request based on token usage.
|
||||
return cls(
|
||||
base_msats=0,
|
||||
input_msats=0,
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
async def calculate_cost(
|
||||
response_data: dict,
|
||||
max_cost: int,
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
"""Calculate the cost of an API request based on token usage.
|
||||
|
||||
Args:
|
||||
response_data: Response data containing usage information
|
||||
@@ -39,6 +75,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,61 +88,90 @@ 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(
|
||||
base_msats=0,
|
||||
input_msats=0,
|
||||
output_msats=0,
|
||||
total_msats=0,
|
||||
total_usd=0.0,
|
||||
)
|
||||
return _empty_cost(MaxCostData)
|
||||
|
||||
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",
|
||||
truly_empty = (
|
||||
input_tokens == 0
|
||||
and output_tokens == 0
|
||||
and cache_read_tokens == 0
|
||||
and cache_creation_tokens == 0
|
||||
)
|
||||
if truly_empty:
|
||||
logger.warning(
|
||||
"Upstream reported a USD cost but the response carries no "
|
||||
"tokens at all (input, output, cache-read and cache-creation "
|
||||
"are all zero) — refunding in full rather than billing the "
|
||||
"USD-derived cost for an empty response.",
|
||||
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,
|
||||
return _empty_cost()
|
||||
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={
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"usd_cost": usd_cost,
|
||||
"usage_keys": sorted(usage_data.keys())
|
||||
if isinstance(usage_data, dict)
|
||||
else None,
|
||||
},
|
||||
)
|
||||
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 +182,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 +218,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 +437,39 @@ 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"),
|
||||
},
|
||||
)
|
||||
|
||||
# Fold the cache-read/write cost into the visible ``input_msats`` so a
|
||||
# dashboard that renders I / O / T sees ``input + output == total``
|
||||
# exactly. This mirrors ``_fold_cache_into_input_tokens`` (which rolls the
|
||||
# cache token counts into the visible prompt total). The standalone
|
||||
# ``cache_read_msats`` / ``cache_creation_msats`` fields stay populated for
|
||||
# clients that want the breakdown; nothing sums the components to derive
|
||||
# ``total_msats`` (it is computed independently above), so this is
|
||||
# display-only and does not change what is billed.
|
||||
visible_output_msats = int(calc_output_msats)
|
||||
visible_input_msats = token_based_cost - visible_output_msats
|
||||
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=int(input_msats),
|
||||
output_msats=int(output_msats),
|
||||
input_msats=visible_input_msats,
|
||||
output_msats=visible_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),
|
||||
)
|
||||
|
||||
@@ -11,6 +11,8 @@ from PIL import Image
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.exceptions import UpstreamError
|
||||
from ..core.redaction import redact_org_ids
|
||||
from ..core.settings import settings
|
||||
from ..wallet import deserialize_token_from_string
|
||||
|
||||
@@ -418,16 +420,27 @@ def create_error_response(
|
||||
status_code: int,
|
||||
request: Request,
|
||||
token: str | None = None,
|
||||
code: str | int | None = None,
|
||||
details: dict[str, object] | None = None,
|
||||
) -> Response:
|
||||
"""Create a standardized error response."""
|
||||
"""Create a standardized error response.
|
||||
|
||||
``code`` is a stable, machine-readable classification (e.g.
|
||||
``UPSTREAM_RATE_LIMIT``); when omitted it defaults to the HTTP status code
|
||||
for backwards compatibility. ``details`` carries optional structured,
|
||||
redaction-safe context.
|
||||
"""
|
||||
error_obj: dict[str, object] = {
|
||||
"message": redact_org_ids(message),
|
||||
"type": error_type,
|
||||
"code": code if code is not None else status_code,
|
||||
}
|
||||
if details is not None:
|
||||
error_obj["details"] = details
|
||||
return Response(
|
||||
content=json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"message": message,
|
||||
"type": error_type,
|
||||
"code": status_code,
|
||||
},
|
||||
"error": error_obj,
|
||||
"request_id": getattr(request.state, "request_id", "unknown"),
|
||||
}
|
||||
),
|
||||
@@ -435,3 +448,20 @@ def create_error_response(
|
||||
media_type="application/json",
|
||||
headers={"X-Cashu": token} if token else {},
|
||||
)
|
||||
|
||||
|
||||
def create_upstream_error_response(
|
||||
error: UpstreamError,
|
||||
request: Request,
|
||||
fallback_status: int = 502,
|
||||
) -> Response:
|
||||
"""Build an error response from an :class:`UpstreamError`, preserving its
|
||||
structured ``code``, ``details``, and original ``status_code``."""
|
||||
return create_error_response(
|
||||
"upstream_error",
|
||||
str(error),
|
||||
error.status_code or fallback_status,
|
||||
request=request,
|
||||
code=getattr(error, "code", None),
|
||||
details=getattr(error, "details", None),
|
||||
)
|
||||
|
||||
@@ -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,70 @@ 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 litellm_cost_entry(model_id: str) -> dict | None:
|
||||
"""Look up ``model_id`` in litellm's bundled cost map.
|
||||
|
||||
litellm ships per-model USD rates keyed by the exact OpenRouter id
|
||||
(``deepseek/deepseek-chat``) or the bare model name (``gpt-4o``,
|
||||
``claude-sonnet-4-5``), so both spellings are tried. Keys are lowercase, so
|
||||
a mixed-case upstream id (``deepseek-ai/DeepSeek-V4-Flash``) is retried via
|
||||
a case-insensitive scan. Returns the matched cost dict, or ``None``.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
candidates = (model_id, model_id.split("/", 1)[-1])
|
||||
for key in candidates:
|
||||
info = litellm.model_cost.get(key)
|
||||
if isinstance(info, dict):
|
||||
return info
|
||||
|
||||
lowered = {c.lower() for c in candidates}
|
||||
for key, info in litellm.model_cost.items():
|
||||
if isinstance(key, str) and key.lower() in lowered and isinstance(info, dict):
|
||||
return info
|
||||
return None
|
||||
|
||||
|
||||
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). The lookup (see ``litellm_cost_entry``) tries both
|
||||
id spellings and a case-insensitive fallback.
|
||||
|
||||
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
|
||||
|
||||
info = litellm_cost_entry(model_id)
|
||||
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", {})
|
||||
@@ -153,13 +231,29 @@ def _row_to_model(
|
||||
)
|
||||
top_provider_dict = json.loads(row.top_provider) if row.top_provider else None
|
||||
|
||||
if apply_provider_fee and isinstance(pricing, dict):
|
||||
pricing = {k: float(v) * provider_fee for k, v in pricing.items()}
|
||||
|
||||
if isinstance(pricing, dict) and float(pricing.get("request", 0.0)) <= 0.0:
|
||||
pricing["request"] = max(pricing.get("request", 0.0), 0.0)
|
||||
|
||||
parsed_pricing = Pricing.parse_obj(pricing)
|
||||
|
||||
# Fill missing cache-read/write rates from litellm's cost map BEFORE applying
|
||||
# the provider fee, so they carry the same markup as every other component.
|
||||
# DB-stored override pricing (e.g. generic providers) omits cache rates;
|
||||
# without this, ``_row_to_model`` bills cache reads at the full input rate —
|
||||
# the ``_apply_provider_fee_to_model`` path backfills, but the override path
|
||||
# used for admin-configured providers did not.
|
||||
#
|
||||
# Key on ``forwarded_model_id`` (the actual upstream model name litellm
|
||||
# prices) when set: an alias row (id="local-alias",
|
||||
# forwarded_model_id="deepseek-v4-flash") would otherwise look up the alias
|
||||
# and miss the cache rate.
|
||||
pricing_model_id = getattr(row, "forwarded_model_id", None) or row.id
|
||||
parsed_pricing = backfill_cache_pricing(pricing_model_id, parsed_pricing)
|
||||
|
||||
if apply_provider_fee:
|
||||
parsed_pricing = Pricing.parse_obj(
|
||||
{k: float(v) * provider_fee for k, v in parsed_pricing.dict().items()}
|
||||
)
|
||||
model = Model(
|
||||
id=row.id,
|
||||
name=row.name,
|
||||
@@ -177,6 +271,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 +424,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 +443,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 +455,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 +504,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,
|
||||
)
|
||||
278
routstr/proxy.py
278
routstr/proxy.py
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
@@ -17,16 +18,19 @@ 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,
|
||||
check_token_balance,
|
||||
create_error_response,
|
||||
create_upstream_error_response,
|
||||
get_max_cost_for_model,
|
||||
)
|
||||
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 +73,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,16 +154,70 @@ 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:
|
||||
# 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)
|
||||
|
||||
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_upstream_error_response(e, 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:
|
||||
@@ -192,22 +268,29 @@ 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
|
||||
continue
|
||||
|
||||
if last_error is not None:
|
||||
return create_upstream_error_response(last_error, request)
|
||||
return create_error_response(
|
||||
"upstream_error",
|
||||
str(last_error) if last_error else "All upstreams failed",
|
||||
502,
|
||||
request=request,
|
||||
"upstream_error", "All upstreams failed", 502, request=request
|
||||
)
|
||||
|
||||
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:
|
||||
@@ -258,9 +341,7 @@ async def proxy(
|
||||
except UpstreamError as e:
|
||||
logger.warning(f"Upstream {upstream.provider_type} failed (GET): {e}")
|
||||
if i == len(upstreams) - 1:
|
||||
last_error_response = create_error_response(
|
||||
"upstream_error", str(e), 502, request=request
|
||||
)
|
||||
last_error_response = create_upstream_error_response(e, request)
|
||||
continue
|
||||
return last_error_response or create_error_response(
|
||||
"upstream_error", "All upstreams failed", 502, request=request
|
||||
@@ -269,50 +350,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 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,
|
||||
},
|
||||
)
|
||||
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 +454,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,34 +469,59 @@ 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
|
||||
if i == len(upstreams) - 1:
|
||||
await revert_pay_for_request(key, session, max_cost_for_model)
|
||||
return create_error_response(
|
||||
"upstream_error", str(e), 502, request=request
|
||||
)
|
||||
return create_upstream_error_response(e, request)
|
||||
|
||||
# Otherwise loop continues to next provider
|
||||
continue
|
||||
@@ -387,7 +533,12 @@ 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."""
|
||||
parts = auth.split()
|
||||
@@ -455,15 +606,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
|
||||
|
||||
@@ -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__(
|
||||
@@ -22,7 +24,7 @@ class AnthropicUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "AnthropicUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -97,6 +97,8 @@ async def _check_and_topup(row: UpstreamProviderRow) -> None:
|
||||
|
||||
# Instantiate provider and check balance
|
||||
provider = RoutstrUpstreamProvider.from_db_row(row)
|
||||
if provider is None:
|
||||
return
|
||||
balance = await provider.get_balance()
|
||||
|
||||
if balance is None:
|
||||
|
||||
@@ -13,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,
|
||||
@@ -37,7 +38,7 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
self.api_version = api_version
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "AzureUpstreamProvider | None":
|
||||
if not provider_row.api_version:
|
||||
|
||||
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",
|
||||
)
|
||||
73
routstr/upstream/deepseek_v4_pricing_shim.py
Normal file
73
routstr/upstream/deepseek_v4_pricing_shim.py
Normal file
@@ -0,0 +1,73 @@
|
||||
"""TEMPORARY: local DeepSeek V4 pricing shim.
|
||||
|
||||
litellm's bundled cost map does not yet ship ``deepseek-v4-flash`` /
|
||||
``deepseek-v4-pro``. Without an entry, ``backfill_cache_pricing`` cannot find a
|
||||
``cache_read_input_token_cost`` and cache reads fall back to the full input
|
||||
rate — a large overcharge on cache hits (DeepSeek V4 hits are ~0.008-0.02x
|
||||
input, i.e. cached tokens cost 50-120x less than regular input).
|
||||
|
||||
This module injects the missing entries into ``litellm.model_cost`` at startup
|
||||
so the existing backfill path resolves them. Rates mirror the canonical
|
||||
``deepseek`` provider entries now in litellm's ``model_prices`` map
|
||||
(``input_cost_per_token`` is the cache-*miss* rate;
|
||||
``cache_read_input_token_cost`` is the cache-*hit* rate), sourced from
|
||||
https://api-docs.deepseek.com/quick_start/pricing via
|
||||
https://github.com/BerriAI/litellm/pull/26380 (issue
|
||||
https://github.com/BerriAI/litellm/issues/30430).
|
||||
|
||||
=== REMOVAL (once litellm ships these models) ===
|
||||
Delete this file and the single ``register_deepseek_v4_pricing()`` call in
|
||||
``routstr/core/main.py``. Nothing else depends on it. Entries are only added
|
||||
when absent, so a stale shim is harmless after upstream lands — but remove it.
|
||||
"""
|
||||
|
||||
import litellm
|
||||
|
||||
from ..core import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# USD per token. Mirrors the canonical ``deepseek`` provider entries in
|
||||
# litellm's model_prices map (source: DeepSeek API pricing docs). Keep these in
|
||||
# sync with ``litellm.model_cost["deepseek/deepseek-v4-*"]``.
|
||||
_DEEPSEEK_V4_RATES: dict[str, dict[str, float]] = {
|
||||
"deepseek-v4-flash": {
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"input_cost_per_token_cache_hit": 2.8e-09,
|
||||
},
|
||||
"deepseek-v4-pro": {
|
||||
"input_cost_per_token": 4.35e-07,
|
||||
"output_cost_per_token": 8.7e-07,
|
||||
"cache_read_input_token_cost": 3.625e-09,
|
||||
"cache_creation_input_token_cost": 0.0,
|
||||
"input_cost_per_token_cache_hit": 3.625e-09,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def register_deepseek_v4_pricing() -> None:
|
||||
"""Inject DeepSeek V4 pricing into ``litellm.model_cost`` if absent.
|
||||
|
||||
Idempotent and non-destructive: a key already present in the cost map
|
||||
(e.g. once litellm ships it) is left untouched. Registers both the bare
|
||||
(``deepseek-v4-flash``) and prefixed (``deepseek/deepseek-v4-flash``)
|
||||
spellings since ``backfill_cache_pricing`` tries both.
|
||||
"""
|
||||
added = []
|
||||
for bare, rates in _DEEPSEEK_V4_RATES.items():
|
||||
for key in (bare, f"deepseek/{bare}"):
|
||||
if key in litellm.model_cost:
|
||||
continue
|
||||
entry: dict[str, object] = dict(rates)
|
||||
entry["litellm_provider"] = "deepseek"
|
||||
entry["mode"] = "chat"
|
||||
litellm.model_cost[key] = entry
|
||||
added.append(key)
|
||||
if added:
|
||||
logger.info(
|
||||
"Registered temporary DeepSeek V4 pricing shim",
|
||||
extra={"models": added},
|
||||
)
|
||||
@@ -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__(
|
||||
@@ -19,7 +20,7 @@ class FireworksUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "FireworksUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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,13 +45,13 @@ 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
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "GeminiUpstreamProvider":
|
||||
return cls(
|
||||
@@ -65,248 +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
|
||||
)
|
||||
await session.refresh(key)
|
||||
remaining_balance_msats = key.balance
|
||||
openai_format_response["cost"] = cost_data
|
||||
openai_format_response["cost"]["sats_cost"] = (
|
||||
cost_data.get("total_msats", 0) // 1000
|
||||
)
|
||||
openai_format_response["cost"]["remaining_balance_msats"] = (
|
||||
remaining_balance_msats
|
||||
)
|
||||
|
||||
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
|
||||
@@ -5,6 +5,12 @@ from typing import TYPE_CHECKING
|
||||
import httpx
|
||||
|
||||
from .base import BaseUpstreamProvider
|
||||
from .pricing_resolver import (
|
||||
FallbackPricingResolver,
|
||||
ResolvedPricing,
|
||||
_as_float,
|
||||
estimate_context_length,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
@@ -45,7 +51,7 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "GenericUpstreamProvider":
|
||||
return cls(
|
||||
@@ -64,6 +70,40 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
"platform_url": cls.platform_url,
|
||||
}
|
||||
|
||||
def _native_pricing(
|
||||
self, model_id: str, model_spec: dict
|
||||
) -> ResolvedPricing | None:
|
||||
"""Read pricing/metadata from Venice's bespoke ``model_spec`` schema.
|
||||
|
||||
Returns ``None`` when the upstream reported no *usable* native price —
|
||||
absent, non-numeric, negative, or both-zero — so the caller falls
|
||||
through to the shared resolution chain instead of fabricating a number
|
||||
or trusting a bogus one. This mirrors the money-safety guards the
|
||||
litellm and OpenRouter rungs already apply: a both-zero price would
|
||||
serve the model free, a negative one would credit the caller, and a
|
||||
non-numeric string would otherwise throw and drop the whole catalog.
|
||||
"""
|
||||
pricing_info = model_spec.get("pricing", {})
|
||||
input_usd = _as_float(pricing_info.get("input", {}).get("usd"))
|
||||
output_usd = _as_float(pricing_info.get("output", {}).get("usd"))
|
||||
if input_usd is None or output_usd is None:
|
||||
return None
|
||||
if input_usd < 0 or output_usd < 0 or (input_usd == 0 and output_usd == 0):
|
||||
return None
|
||||
|
||||
capabilities = model_spec.get("capabilities", {})
|
||||
input_modalities = ["text"]
|
||||
if capabilities.get("supportsVision", False):
|
||||
input_modalities.append("image")
|
||||
|
||||
return ResolvedPricing(
|
||||
prompt=input_usd / 1_000_000,
|
||||
completion=output_usd / 1_000_000,
|
||||
context_length=model_spec.get("availableContextTokens"),
|
||||
source="native",
|
||||
input_modalities=input_modalities,
|
||||
)
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
"""Fetch models from upstream API using /models endpoint."""
|
||||
from ..payment.models import Architecture, Model, Pricing, TopProvider
|
||||
@@ -78,6 +118,7 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
resolver = FallbackPricingResolver()
|
||||
models_list = []
|
||||
for model_data in data.get("data", []):
|
||||
model_id = model_data.get("id", "")
|
||||
@@ -89,41 +130,44 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
owned_by = model_data.get("owned_by", "unknown")
|
||||
model_spec = model_data.get("model_spec", {})
|
||||
|
||||
context_length = 4096
|
||||
if model_spec.get("availableContextTokens"):
|
||||
context_length = model_spec["availableContextTokens"]
|
||||
elif any(
|
||||
pattern in model_id.lower() for pattern in ["32k", "32000"]
|
||||
):
|
||||
context_length = 32768
|
||||
elif any(
|
||||
pattern in model_id.lower() for pattern in ["16k", "16000"]
|
||||
):
|
||||
context_length = 16384
|
||||
elif any(pattern in model_id.lower() for pattern in ["8k", "8000"]):
|
||||
context_length = 8192
|
||||
elif "gpt-4" in model_id.lower():
|
||||
context_length = 8192
|
||||
elif "claude" in model_id.lower():
|
||||
context_length = 200000
|
||||
resolved = self._native_pricing(model_id, model_spec)
|
||||
if resolved is None:
|
||||
resolved = await resolver.resolve(model_id)
|
||||
|
||||
pricing_info = model_spec.get("pricing", {})
|
||||
input_pricing = pricing_info.get("input", {})
|
||||
output_pricing = pricing_info.get("output", {})
|
||||
if resolved is None:
|
||||
# Fail closed: never invent a price. Import the model
|
||||
# disabled with a warning so the operator can price it
|
||||
# (the admin UI surfaces disabled remote models).
|
||||
logger.warning(
|
||||
f"No pricing source resolved for '{model_id}' from "
|
||||
f"{self.upstream_name}; importing it disabled",
|
||||
extra={"model_id": model_id, "base_url": self.base_url},
|
||||
)
|
||||
resolved = ResolvedPricing(
|
||||
prompt=0.0,
|
||||
completion=0.0,
|
||||
context_length=None,
|
||||
source="unresolved",
|
||||
)
|
||||
enabled = False
|
||||
else:
|
||||
enabled = True
|
||||
|
||||
prompt_price = input_pricing.get("usd", 0.001) / 1000000
|
||||
completion_price = output_pricing.get("usd", 0.001) / 1000000
|
||||
# Prefer the source's own modality string (OpenRouter ships
|
||||
# one, e.g. "text+image->text"); otherwise derive it from the
|
||||
# captured input/output modalities in the same "in->out" shape
|
||||
# rather than flattening vision models to "text->text".
|
||||
modality = resolved.modality or (
|
||||
f"{'+'.join(resolved.input_modalities)}"
|
||||
f"->{'+'.join(resolved.output_modalities)}"
|
||||
)
|
||||
|
||||
capabilities = model_spec.get("capabilities", {})
|
||||
input_modalities = ["text"]
|
||||
output_modalities = ["text"]
|
||||
|
||||
if capabilities.get("supportsVision", False):
|
||||
input_modalities.append("image")
|
||||
|
||||
modality = "text"
|
||||
if capabilities.get("supportsVision", False):
|
||||
modality = "text->text"
|
||||
# A source can carry a price but no context (e.g. a litellm
|
||||
# entry missing max_input_tokens); fall back to an id-based
|
||||
# estimate so we never persist a zero-length window.
|
||||
context_length = resolved.context_length or estimate_context_length(
|
||||
model_id
|
||||
)
|
||||
|
||||
spec_name = model_spec.get("name", model_name)
|
||||
description = f"{spec_name}"
|
||||
@@ -139,30 +183,33 @@ class GenericUpstreamProvider(BaseUpstreamProvider):
|
||||
context_length=context_length,
|
||||
architecture=Architecture(
|
||||
modality=modality,
|
||||
input_modalities=input_modalities,
|
||||
output_modalities=output_modalities,
|
||||
tokenizer="unknown",
|
||||
instruct_type=None,
|
||||
input_modalities=resolved.input_modalities,
|
||||
output_modalities=resolved.output_modalities,
|
||||
tokenizer=resolved.tokenizer,
|
||||
instruct_type=resolved.instruct_type,
|
||||
),
|
||||
pricing=Pricing(
|
||||
prompt=prompt_price,
|
||||
completion=completion_price,
|
||||
prompt=resolved.prompt,
|
||||
completion=resolved.completion,
|
||||
request=0.0,
|
||||
image=0.0,
|
||||
web_search=0.0,
|
||||
internal_reasoning=0.0,
|
||||
max_prompt_cost=0.001,
|
||||
max_completion_cost=0.001,
|
||||
max_cost=0.001,
|
||||
input_cache_read=resolved.input_cache_read,
|
||||
input_cache_write=resolved.input_cache_write,
|
||||
),
|
||||
sats_pricing=None,
|
||||
per_request_limits=None,
|
||||
top_provider=TopProvider(
|
||||
context_length=context_length,
|
||||
max_completion_tokens=context_length // 2,
|
||||
is_moderated=False,
|
||||
max_completion_tokens=(
|
||||
resolved.max_completion_tokens
|
||||
if resolved.max_completion_tokens is not None
|
||||
else context_length // 2
|
||||
),
|
||||
is_moderated=bool(resolved.is_moderated),
|
||||
),
|
||||
enabled=True,
|
||||
enabled=enabled,
|
||||
upstream_provider_id=None,
|
||||
canonical_slug=None,
|
||||
)
|
||||
|
||||
@@ -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__(
|
||||
@@ -19,7 +20,7 @@ class GroqUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(cls, provider_row: "UpstreamProviderRow") -> "GroqUpstreamProvider":
|
||||
def _build_from_row(cls, provider_row: "UpstreamProviderRow") -> "GroqUpstreamProvider":
|
||||
return cls(
|
||||
api_key=provider_row.api_key,
|
||||
provider_fee=provider_row.provider_fee,
|
||||
|
||||
@@ -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
|
||||
@@ -12,6 +12,7 @@ from sqlmodel import select
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.db import AsyncSession, ModelRow, UpstreamProviderRow, create_session
|
||||
from ..core.provider_slugs import allocate_unique_provider_slug
|
||||
from ..payment.models import Model
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
@@ -122,12 +123,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 +144,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 +198,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,9 +216,6 @@ 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",
|
||||
@@ -240,6 +248,7 @@ async def _seed_providers_from_settings(
|
||||
|
||||
providers_to_add: list[UpstreamProviderRow] = []
|
||||
seeded_provider_keys: set[tuple[str, str]] = set()
|
||||
reserved_slugs: set[str] = set()
|
||||
|
||||
provider_classes_by_type = {
|
||||
cls.provider_type: cls
|
||||
@@ -269,8 +278,13 @@ async def _seed_providers_from_settings(
|
||||
)
|
||||
)
|
||||
if not result.first():
|
||||
slug = await allocate_unique_provider_slug(
|
||||
session, provider_type, reserved_slugs
|
||||
)
|
||||
reserved_slugs.add(slug)
|
||||
providers_to_add.append(
|
||||
UpstreamProviderRow(
|
||||
slug=slug,
|
||||
provider_type=provider_type,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
@@ -289,8 +303,13 @@ async def _seed_providers_from_settings(
|
||||
)
|
||||
)
|
||||
if not result.first():
|
||||
slug = await allocate_unique_provider_slug(
|
||||
session, "ollama", reserved_slugs
|
||||
)
|
||||
reserved_slugs.add(slug)
|
||||
providers_to_add.append(
|
||||
UpstreamProviderRow(
|
||||
slug=slug,
|
||||
provider_type="ollama",
|
||||
base_url=ollama_base_url,
|
||||
api_key=ollama_api_key,
|
||||
@@ -310,8 +329,13 @@ async def _seed_providers_from_settings(
|
||||
)
|
||||
)
|
||||
if not result.first():
|
||||
slug = await allocate_unique_provider_slug(
|
||||
session, "azure", reserved_slugs
|
||||
)
|
||||
reserved_slugs.add(slug)
|
||||
providers_to_add.append(
|
||||
UpstreamProviderRow(
|
||||
slug=slug,
|
||||
provider_type="azure",
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
@@ -332,8 +356,13 @@ async def _seed_providers_from_settings(
|
||||
)
|
||||
)
|
||||
if not result.first():
|
||||
slug = await allocate_unique_provider_slug(
|
||||
session, "custom", reserved_slugs
|
||||
)
|
||||
reserved_slugs.add(slug)
|
||||
providers_to_add.append(
|
||||
UpstreamProviderRow(
|
||||
slug=slug,
|
||||
provider_type="custom",
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
@@ -346,7 +375,7 @@ async def _seed_providers_from_settings(
|
||||
session.add(provider)
|
||||
logger.info(
|
||||
f"Seeding {provider.provider_type} provider", # type: ignore[str-format]
|
||||
extra={"base_url": provider.base_url},
|
||||
extra={"base_url": provider.base_url, "slug": provider.slug},
|
||||
)
|
||||
|
||||
|
||||
@@ -381,9 +410,7 @@ def _instantiate_provider(
|
||||
return provider
|
||||
|
||||
if provider_row.provider_type == "custom":
|
||||
return BaseUpstreamProvider(
|
||||
provider_row.base_url, provider_row.api_key, provider_row.provider_fee
|
||||
)
|
||||
return BaseUpstreamProvider.from_db_row(provider_row)
|
||||
|
||||
logger.error(
|
||||
f"Unknown provider type: {provider_row.provider_type}",
|
||||
|
||||
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
|
||||
575
routstr/upstream/messages_dispatch.py
Normal file
575
routstr/upstream/messages_dispatch.py
Normal file
@@ -0,0 +1,575 @@
|
||||
"""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 ..core.redaction import redact_org_ids
|
||||
from ..payment.models import Model
|
||||
from .rate_limit import classify_rate_limit
|
||||
|
||||
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:
|
||||
raw_message = getattr(exc, "message", None) or str(exc) or repr(exc)
|
||||
# Redact provider account identifiers before the message reaches logs
|
||||
# or the surfaced error.
|
||||
exc_message = redact_org_ids(raw_message)
|
||||
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 = redact_org_ids(
|
||||
getattr(exc_response, "text", str(exc_response))
|
||||
)
|
||||
except Exception:
|
||||
response_text = "<unreadable>"
|
||||
status_for_classify = exc_status if isinstance(exc_status, int) else 502
|
||||
rate_limit = classify_rate_limit(
|
||||
status_for_classify, exc_message, getattr(exc, "headers", None)
|
||||
)
|
||||
logger.error(
|
||||
"litellm dispatch failed",
|
||||
extra={
|
||||
"error": exc_message,
|
||||
"error_type": type(exc).__name__,
|
||||
"status_code": exc_status,
|
||||
"error_code": rate_limit.code if rate_limit else None,
|
||||
"llm_provider": getattr(exc, "llm_provider", None),
|
||||
"body": redact_org_ids(str(getattr(exc, "body", "") or "")) or None,
|
||||
"response_text": response_text,
|
||||
"model": litellm_model,
|
||||
"api_base": base_url,
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Upstream error via litellm: {exc_message}",
|
||||
status_code=status_for_classify,
|
||||
code=rate_limit.code if rate_limit else None,
|
||||
details=rate_limit.as_details() if rate_limit else None,
|
||||
) 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
|
||||
@@ -21,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,
|
||||
@@ -42,7 +43,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OllamaUpstreamProvider":
|
||||
return cls(
|
||||
@@ -184,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)},
|
||||
|
||||
@@ -20,7 +20,7 @@ class OpenAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OpenAIUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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.
|
||||
@@ -28,7 +57,7 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "OpenRouterUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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__(
|
||||
@@ -22,7 +23,7 @@ class PerplexityUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "PerplexityUpstreamProvider":
|
||||
return cls(
|
||||
|
||||
@@ -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):
|
||||
@@ -44,7 +46,7 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "PPQAIUpstreamProvider":
|
||||
return cls(
|
||||
@@ -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",
|
||||
@@ -204,17 +229,14 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
f"Disabling PPQ.AI provider ({self.base_url}) due to insufficient balance",
|
||||
extra={"error": error_message},
|
||||
)
|
||||
from sqlmodel import select
|
||||
|
||||
from ..core.db import UpstreamProviderRow, create_session
|
||||
|
||||
async with create_session() as session:
|
||||
statement = select(UpstreamProviderRow).where(
|
||||
UpstreamProviderRow.base_url == self.base_url,
|
||||
UpstreamProviderRow.api_key == self.api_key,
|
||||
provider = (
|
||||
await session.get(UpstreamProviderRow, self.db_id)
|
||||
if self.db_id is not None
|
||||
else None
|
||||
)
|
||||
result = await session.exec(statement)
|
||||
provider = result.first()
|
||||
|
||||
if provider:
|
||||
provider.enabled = False
|
||||
|
||||
203
routstr/upstream/pricing_resolver.py
Normal file
203
routstr/upstream/pricing_resolver.py
Normal file
@@ -0,0 +1,203 @@
|
||||
"""Shared price/metadata resolution chain for upstream model discovery.
|
||||
|
||||
Most OpenAI-compatible ``/models`` responses carry no pricing. Rather than let
|
||||
a provider fabricate one, this module resolves a model through decreasingly
|
||||
trustworthy sources — litellm's bundled cost map (curated list prices, mirrors
|
||||
provider docs), then the OpenRouter feed (resale prices, broader coverage) —
|
||||
and returns ``None`` when none of them know the model, so the caller can fail
|
||||
closed instead of inventing a number.
|
||||
|
||||
Provider-native pricing (a gateway's own ``/models`` schema, e.g. Venice's
|
||||
``model_spec``) is authoritative and handled by the provider before this chain
|
||||
is consulted; only the shared fallback lives here so a later refactor can hoist
|
||||
it into the base provider unchanged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResolvedPricing:
|
||||
"""Per-token pricing plus whatever metadata the answering source carried.
|
||||
|
||||
Prices are USD per token. ``source`` records provenance
|
||||
(``native``/``litellm``/``openrouter``/``unresolved``) so later work can
|
||||
surface where each price came from.
|
||||
"""
|
||||
|
||||
prompt: float
|
||||
completion: float
|
||||
context_length: int | None
|
||||
source: str
|
||||
modality: str | None = None
|
||||
max_completion_tokens: int | None = None
|
||||
input_cache_read: float = 0.0
|
||||
input_cache_write: float = 0.0
|
||||
input_modalities: list[str] = field(default_factory=lambda: ["text"])
|
||||
output_modalities: list[str] = field(default_factory=lambda: ["text"])
|
||||
tokenizer: str = "unknown"
|
||||
instruct_type: str | None = None
|
||||
is_moderated: bool | None = None
|
||||
|
||||
|
||||
def estimate_context_length(model_id: str) -> int:
|
||||
"""Best-effort context window from a model id when no source reports one.
|
||||
|
||||
The last rung of the fallback chain, reached only for a model whose price
|
||||
resolved but whose context did not (or that imported disabled). Context is
|
||||
not a billing input, so a rough id-based guess is acceptable here where a
|
||||
guessed *price* never would be.
|
||||
"""
|
||||
lowered = model_id.lower()
|
||||
if any(pattern in lowered for pattern in ["32k", "32000"]):
|
||||
return 32768
|
||||
if any(pattern in lowered for pattern in ["16k", "16000"]):
|
||||
return 16384
|
||||
if any(pattern in lowered for pattern in ["8k", "8000"]):
|
||||
return 8192
|
||||
if "gpt-4" in lowered:
|
||||
return 8192
|
||||
if "claude" in lowered:
|
||||
return 200000
|
||||
return 4096
|
||||
|
||||
|
||||
def _as_float(value: object) -> float | None:
|
||||
"""OpenRouter reports prices as strings; coerce, ``None`` if unparseable."""
|
||||
try:
|
||||
return float(value) # type: ignore[arg-type]
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _as_int(value: object) -> int | None:
|
||||
"""Coerce an already-numeric token count to ``int``, else ``None``."""
|
||||
return int(value) if isinstance(value, (int, float)) else None
|
||||
|
||||
|
||||
def _from_litellm(model_id: str) -> ResolvedPricing | None:
|
||||
# Lazy import so the resolver stays import-light and shares the exact
|
||||
# lookup semantics used by cache-rate backfill.
|
||||
from ..payment.models import litellm_cost_entry
|
||||
|
||||
info = litellm_cost_entry(model_id)
|
||||
if info is None:
|
||||
return None
|
||||
|
||||
prompt = info.get("input_cost_per_token")
|
||||
completion = info.get("output_cost_per_token")
|
||||
if not isinstance(prompt, (int, float)) or not isinstance(completion, (int, float)):
|
||||
return None
|
||||
# A both-zero entry is litellm listing a model without a real price (free
|
||||
# moderation/rerank tiers do this) — treating 0/0 as resolved would serve
|
||||
# the model for free. Reject it (and any negative) so the caller falls
|
||||
# through, mirroring async_fetch_openrouter_models' _has_valid_pricing.
|
||||
if prompt < 0 or completion < 0 or (prompt == 0 and completion == 0):
|
||||
return None
|
||||
|
||||
input_modalities = ["text"]
|
||||
if info.get("supports_vision"):
|
||||
input_modalities.append("image")
|
||||
|
||||
return ResolvedPricing(
|
||||
prompt=float(prompt),
|
||||
completion=float(completion),
|
||||
# max_input_tokens is the context window; max_tokens is litellm's
|
||||
# completion cap (it tracks max_output_tokens for ~94% of models), so
|
||||
# it is never a context source. A missing window falls to the id-based
|
||||
# estimate downstream rather than borrowing the output cap.
|
||||
context_length=_as_int(info.get("max_input_tokens")),
|
||||
source="litellm",
|
||||
max_completion_tokens=_as_int(info.get("max_output_tokens")),
|
||||
input_cache_read=float(info.get("cache_read_input_token_cost") or 0.0),
|
||||
input_cache_write=float(info.get("cache_creation_input_token_cost") or 0.0),
|
||||
input_modalities=input_modalities,
|
||||
)
|
||||
|
||||
|
||||
def _match_openrouter(model_id: str, feed: list[dict]) -> dict | None:
|
||||
"""Find ``model_id`` in the OpenRouter feed, exact id before bare tail.
|
||||
|
||||
Bare-tail matching (``deepseek-chat`` ↔ ``deepseek/deepseek-chat``) is a
|
||||
looser, lower-trust match — OpenRouter fans a model out across resellers —
|
||||
so an exact id match always wins first. When several entries share the bare
|
||||
tail, the one with the highest *combined* (prompt + completion) per-token
|
||||
cost wins: the choice must be deterministic (not feed-order-dependent) and
|
||||
money-safe whichever way traffic leans, since undercharging is the hazard.
|
||||
Ranking on prompt alone could pick an entry that is cheap on input but dear
|
||||
on output. The live feed has no such collisions today; this only governs
|
||||
the latent case.
|
||||
"""
|
||||
bare = model_id.split("/", 1)[-1]
|
||||
exact = next((m for m in feed if m.get("id") == model_id), None)
|
||||
if exact is not None:
|
||||
return exact
|
||||
matches = [m for m in feed if m.get("id", "").split("/", 1)[-1] == bare]
|
||||
if not matches:
|
||||
return None
|
||||
|
||||
def _combined_cost(m: dict) -> float:
|
||||
pricing = m.get("pricing", {})
|
||||
return (_as_float(pricing.get("prompt")) or 0.0) + (
|
||||
_as_float(pricing.get("completion")) or 0.0
|
||||
)
|
||||
|
||||
return max(matches, key=_combined_cost)
|
||||
|
||||
|
||||
def _from_openrouter(model_id: str, feed: list[dict]) -> ResolvedPricing | None:
|
||||
entry = _match_openrouter(model_id, feed)
|
||||
if entry is None:
|
||||
return None
|
||||
|
||||
pricing = entry.get("pricing", {})
|
||||
prompt = _as_float(pricing.get("prompt"))
|
||||
completion = _as_float(pricing.get("completion"))
|
||||
if prompt is None or completion is None:
|
||||
return None
|
||||
|
||||
architecture = entry.get("architecture", {})
|
||||
top_provider = entry.get("top_provider", {})
|
||||
|
||||
return ResolvedPricing(
|
||||
prompt=prompt,
|
||||
completion=completion,
|
||||
context_length=_as_int(entry.get("context_length")),
|
||||
source="openrouter",
|
||||
modality=architecture.get("modality"),
|
||||
max_completion_tokens=_as_int(top_provider.get("max_completion_tokens")),
|
||||
input_cache_read=_as_float(pricing.get("input_cache_read")) or 0.0,
|
||||
input_cache_write=_as_float(pricing.get("input_cache_write")) or 0.0,
|
||||
input_modalities=architecture.get("input_modalities") or ["text"],
|
||||
output_modalities=architecture.get("output_modalities") or ["text"],
|
||||
tokenizer=architecture.get("tokenizer") or "unknown",
|
||||
instruct_type=architecture.get("instruct_type"),
|
||||
is_moderated=top_provider.get("is_moderated"),
|
||||
)
|
||||
|
||||
|
||||
class FallbackPricingResolver:
|
||||
"""Resolves models via litellm → OpenRouter for one discovery pass.
|
||||
|
||||
The OpenRouter catalog is fetched at most once and only when a model
|
||||
actually misses litellm, so a provider full of litellm-known models never
|
||||
touches the network. Instantiate one per ``fetch_models`` call.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._openrouter_feed: list[dict] | None = None
|
||||
|
||||
async def resolve(self, model_id: str) -> ResolvedPricing | None:
|
||||
"""Resolve ``model_id``; ``None`` if no source knows it."""
|
||||
resolved = _from_litellm(model_id)
|
||||
if resolved is not None:
|
||||
return resolved
|
||||
|
||||
if self._openrouter_feed is None:
|
||||
# Lazy import so tests can patch the feed at its source.
|
||||
from ..payment.models import async_fetch_openrouter_models
|
||||
|
||||
self._openrouter_feed = await async_fetch_openrouter_models()
|
||||
return _from_openrouter(model_id, self._openrouter_feed)
|
||||
136
routstr/upstream/rate_limit.py
Normal file
136
routstr/upstream/rate_limit.py
Normal file
@@ -0,0 +1,136 @@
|
||||
"""Detection and parsing of upstream provider rate-limit errors.
|
||||
|
||||
Upstream OpenAI-compatible providers signal rate limits via HTTP 429 and/or a
|
||||
human-readable message such as::
|
||||
|
||||
Rate limit reached for gpt-5.5-2026-04-23 (for limit gpt-5.5) in organization
|
||||
org-XXXX on tokens per min (TPM): Limit 180000000, Used 180000000,
|
||||
Requested 8929. Please try again in 2ms.
|
||||
|
||||
This module classifies those failures into a stable :data:`UPSTREAM_RATE_LIMIT`
|
||||
code and extracts useful debugging fields. All retained text is redacted of
|
||||
organization IDs first.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
from ..core.redaction import redact_org_ids
|
||||
|
||||
# Stable error code callers can switch on to distinguish upstream rate limits
|
||||
# from generic request failures. The literal value matches the identifier named
|
||||
# in issue #555 ("UPSTREAM_RATE_LIMIT") so the public API contract is exact.
|
||||
UPSTREAM_RATE_LIMIT = "UPSTREAM_RATE_LIMIT"
|
||||
|
||||
# Message fragments that indicate a rate-limit even when the status code is not
|
||||
# 429 (some providers wrap it in a 400/500 envelope).
|
||||
_RATE_LIMIT_MARKERS = (
|
||||
"rate limit reached",
|
||||
"rate_limit_exceeded",
|
||||
"rate limit exceeded",
|
||||
"too many requests",
|
||||
)
|
||||
|
||||
_MODEL_RE = re.compile(r"Rate limit reached for ([^\s(]+)", re.IGNORECASE)
|
||||
_LIMIT_NAME_RE = re.compile(r"\(for limit ([^)]+)\)", re.IGNORECASE)
|
||||
_METRIC_RE = re.compile(r"on ([a-z ]+\((?:TPM|RPM|TPD|RPD|IPM)\))", re.IGNORECASE)
|
||||
_LIMIT_RE = re.compile(r"Limit (\d+)", re.IGNORECASE)
|
||||
_USED_RE = re.compile(r"Used (\d+)", re.IGNORECASE)
|
||||
_REQUESTED_RE = re.compile(r"Requested (\d+)", re.IGNORECASE)
|
||||
_RETRY_RE = re.compile(r"try again in ([\d.]+)\s*(ms|s)", re.IGNORECASE)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitInfo:
|
||||
"""Structured, redaction-safe view of an upstream rate-limit error."""
|
||||
|
||||
code: str
|
||||
message: str
|
||||
model: str | None = None
|
||||
limit_name: str | None = None
|
||||
metric: str | None = None
|
||||
limit: int | None = None
|
||||
used: int | None = None
|
||||
requested: int | None = None
|
||||
retry_after_seconds: float | None = None
|
||||
|
||||
def as_details(self) -> dict[str, object]:
|
||||
"""Return a JSON-serialisable dict for embedding in an error envelope."""
|
||||
return {k: v for k, v in asdict(self).items() if v is not None}
|
||||
|
||||
|
||||
def _looks_like_rate_limit(status_code: int, message: str) -> bool:
|
||||
if status_code == 429:
|
||||
return True
|
||||
lowered = message.lower()
|
||||
return any(marker in lowered for marker in _RATE_LIMIT_MARKERS)
|
||||
|
||||
|
||||
def _parse_retry_after_header(headers: dict[str, str] | None) -> float | None:
|
||||
"""Parse a ``Retry-After`` header (delta-seconds form) into seconds."""
|
||||
if not headers:
|
||||
return None
|
||||
raw = headers.get("retry-after") or headers.get("Retry-After")
|
||||
if raw is None:
|
||||
return None
|
||||
try:
|
||||
return float(str(raw).strip())
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _int_or_none(match: re.Match[str] | None) -> int | None:
|
||||
if match is None:
|
||||
return None
|
||||
try:
|
||||
return int(match.group(1))
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def classify_rate_limit(
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> RateLimitInfo | None:
|
||||
"""Classify an upstream error as a rate-limit and extract its fields.
|
||||
|
||||
Args:
|
||||
status_code: HTTP status code from the upstream response.
|
||||
message: Upstream error message (may contain sensitive identifiers).
|
||||
headers: Optional upstream response headers, used for ``Retry-After``.
|
||||
|
||||
Returns:
|
||||
A :class:`RateLimitInfo` when the error is a rate-limit, else ``None``.
|
||||
"""
|
||||
message = message or ""
|
||||
if not _looks_like_rate_limit(status_code, message):
|
||||
return None
|
||||
|
||||
redacted = redact_org_ids(message)
|
||||
|
||||
model_match = _MODEL_RE.search(redacted)
|
||||
metric_match = _METRIC_RE.search(redacted)
|
||||
|
||||
retry_after = _parse_retry_after_header(headers)
|
||||
if retry_after is None:
|
||||
retry_match = _RETRY_RE.search(redacted)
|
||||
if retry_match is not None:
|
||||
value = float(retry_match.group(1))
|
||||
retry_after = value / 1000.0 if retry_match.group(2).lower() == "ms" else value
|
||||
|
||||
limit_name_match = _LIMIT_NAME_RE.search(redacted)
|
||||
|
||||
return RateLimitInfo(
|
||||
code=UPSTREAM_RATE_LIMIT,
|
||||
message=redacted,
|
||||
model=model_match.group(1) if model_match else None,
|
||||
limit_name=limit_name_match.group(1).strip() if limit_name_match else None,
|
||||
metric=metric_match.group(1).strip() if metric_match else None,
|
||||
limit=_int_or_none(_LIMIT_RE.search(redacted)),
|
||||
used=_int_or_none(_USED_RE.search(redacted)),
|
||||
requested=_int_or_none(_REQUESTED_RE.search(redacted)),
|
||||
retry_after_seconds=retry_after,
|
||||
)
|
||||
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
|
||||
@@ -18,6 +18,10 @@ class RoutstrUpstreamProvider(BaseUpstreamProvider):
|
||||
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,
|
||||
@@ -43,8 +47,15 @@ class RoutstrUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
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(
|
||||
def _build_from_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "RoutstrUpstreamProvider":
|
||||
import json
|
||||
@@ -83,7 +94,7 @@ class RoutstrUpstreamProvider(BaseUpstreamProvider):
|
||||
Balance in satoshis, or None if failed
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/info"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
|
||||
@@ -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__(
|
||||
@@ -20,7 +21,7 @@ class XAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_db_row(cls, provider_row: "UpstreamProviderRow") -> "XAIUpstreamProvider":
|
||||
def _build_from_row(cls, provider_row: "UpstreamProviderRow") -> "XAIUpstreamProvider":
|
||||
return cls(
|
||||
api_key=provider_row.api_key,
|
||||
provider_fee=provider_row.provider_fee,
|
||||
|
||||
@@ -1,24 +1,204 @@
|
||||
import asyncio
|
||||
import math
|
||||
import re
|
||||
import socket
|
||||
import time
|
||||
import typing
|
||||
from typing import TypedDict
|
||||
|
||||
import httpx
|
||||
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__)
|
||||
|
||||
|
||||
class MintConnectionError(Exception):
|
||||
"""The mint could not be reached (network transport failure).
|
||||
|
||||
Maps to a 503, not a 4xx: the token is fine, the mint is just unavailable.
|
||||
"""
|
||||
|
||||
|
||||
class TokenConsumedError(Exception):
|
||||
"""A failure that happened AFTER the token's proofs were spent (melt
|
||||
succeeded, or redemption already returned) — e.g. minting on the primary
|
||||
mint or the DB credit then failed.
|
||||
|
||||
Non-retryable: the same token will not work again. Seals the cause chain so
|
||||
a transport error underneath is never re-surfaced as a retryable
|
||||
mint_unreachable.
|
||||
"""
|
||||
|
||||
|
||||
# httpx base classes cover their subclasses. HTTPStatusError is excluded on
|
||||
# purpose — that means the mint answered, just with an error status.
|
||||
_TRANSPORT_EXC_TYPES: tuple[type[BaseException], ...] = (
|
||||
httpx.NetworkError,
|
||||
httpx.TimeoutException,
|
||||
ConnectionError, # refused/reset/aborted
|
||||
socket.gaierror, # DNS failure
|
||||
asyncio.TimeoutError,
|
||||
)
|
||||
|
||||
|
||||
def is_mint_connection_error(error: BaseException) -> bool:
|
||||
"""True if ``error`` (or anything in its cause/context chain) is a mint
|
||||
transport failure. Walks the chain because some sites re-raise transport
|
||||
errors wrapped in ValueError/MintConnectionError; matches on TYPE, not text.
|
||||
"""
|
||||
seen: set[int] = set()
|
||||
current: BaseException | None = error
|
||||
while current is not None and id(current) not in seen:
|
||||
seen.add(id(current))
|
||||
if isinstance(current, TokenConsumedError):
|
||||
# Sealed: the token was already spent, so whatever transport error
|
||||
# sits underneath must not make this look retryable.
|
||||
return False
|
||||
if isinstance(current, MintConnectionError):
|
||||
return True
|
||||
if isinstance(current, _TRANSPORT_EXC_TYPES):
|
||||
return True
|
||||
current = current.__cause__ or current.__context__
|
||||
return False
|
||||
|
||||
|
||||
# Redemption ``code`` values whose token is spent/consumed/unusable — the
|
||||
# X-Cashu path must NOT echo the original token for these (echoing invites a
|
||||
# retry with a token that can never succeed again).
|
||||
SPENT_TOKEN_CODES: frozenset[str] = frozenset(
|
||||
{
|
||||
"cashu_token_already_spent",
|
||||
"cashu_token_consumed",
|
||||
"cashu_token_zero_value",
|
||||
"internal_error",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def classify_redemption_error(
|
||||
error: Exception,
|
||||
) -> tuple[str, int, str, str] | None:
|
||||
"""Map a token-redemption failure to ``(type, status, message, code)``.
|
||||
|
||||
Single source of truth for every endpoint that redeems a token (bearer,
|
||||
X-Cashu, top-up) so the same failure yields the same taxonomy everywhere.
|
||||
``type`` and ``code`` are stable client contract; ``message`` is sanitized
|
||||
(raw error text stays in logs). Returns None for an unclassified internal
|
||||
fault — the caller emits a generic 500.
|
||||
"""
|
||||
if isinstance(error, TokenConsumedError):
|
||||
return (
|
||||
"token_consumed",
|
||||
500,
|
||||
"Token was redeemed but could not be credited; do not retry",
|
||||
"cashu_token_consumed",
|
||||
)
|
||||
if is_mint_connection_error(error):
|
||||
return (
|
||||
"mint_unreachable",
|
||||
503,
|
||||
"Cashu mint is unreachable",
|
||||
"cashu_mint_unreachable",
|
||||
)
|
||||
lowered = str(error).lower()
|
||||
if "already spent" in lowered:
|
||||
return (
|
||||
"token_already_spent",
|
||||
400,
|
||||
"Cashu token already spent",
|
||||
"cashu_token_already_spent",
|
||||
)
|
||||
if (
|
||||
"insufficient" in lowered
|
||||
or "melt fee" in lowered
|
||||
or "exceed token amount" in lowered
|
||||
or "estimate fees" in lowered
|
||||
):
|
||||
return (
|
||||
"mint_error",
|
||||
422,
|
||||
"Token value is too small to cover swap fees",
|
||||
"cashu_token_swap_fees_exceed_amount",
|
||||
)
|
||||
if "failed to melt" in lowered:
|
||||
return (
|
||||
"mint_error",
|
||||
422,
|
||||
"Failed to swap token from foreign mint",
|
||||
"cashu_foreign_mint_swap_failed",
|
||||
)
|
||||
if ("invalid" in lowered or "decode" in lowered) and "token" in lowered:
|
||||
# Anchored to "token" so internal faults whose text merely contains
|
||||
# "invalid"/"decode" fall through to the 500 branch, not a token error.
|
||||
return (
|
||||
"invalid_token",
|
||||
400,
|
||||
"Invalid Cashu token",
|
||||
"invalid_cashu_token",
|
||||
)
|
||||
if "must be positive" in lowered or "yielded no value" in lowered:
|
||||
# Redeemed to <= 0 (empty/dust token, or value fully consumed by fees).
|
||||
# Consumed, so non-retryable, but its own code — not the generic bucket.
|
||||
return (
|
||||
"cashu_error",
|
||||
400,
|
||||
"Failed to redeem Cashu token: token yielded no value",
|
||||
"cashu_token_zero_value",
|
||||
)
|
||||
if isinstance(error, ValueError):
|
||||
return (
|
||||
"cashu_error",
|
||||
400,
|
||||
"Failed to redeem Cashu token",
|
||||
"cashu_token_redemption_failed",
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def get_balance(unit: str) -> int:
|
||||
wallet = await get_wallet(settings.primary_mint, unit)
|
||||
return wallet.available_balance.amount
|
||||
|
||||
|
||||
async def _redeem_same_mint(
|
||||
wallet: Wallet, token_obj: Token
|
||||
) -> tuple[int, str, str]: # amount, unit, mint_url
|
||||
"""Redeem proofs at their own issuing mint (no cross-mint swap).
|
||||
|
||||
split() re-mints the incoming proofs into fresh ones we own so the sender
|
||||
can't double-spend them. With include_fees=True the mint deducts its NUT-02
|
||||
per-proof input fee, so we end up holding only `amount - input_fees`. Credit
|
||||
that, not the face value, or routstr over-credits the user and its wallet
|
||||
drifts insolvent.
|
||||
"""
|
||||
await wallet.load_mint(keyset_id=token_obj.keysets[0])
|
||||
wallet.verify_proofs_dleq(token_obj.proofs)
|
||||
input_fees = wallet.get_fees_for_proofs(token_obj.proofs)
|
||||
await wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
return int(token_obj.amount) - input_fees, token_obj.unit, token_obj.mint
|
||||
|
||||
|
||||
async def recieve_token(
|
||||
token: str,
|
||||
) -> tuple[int, str, str]: # amount, unit, mint_url
|
||||
@@ -32,24 +212,61 @@ async def recieve_token(
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
|
||||
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
|
||||
return await _redeem_same_mint(wallet, token_obj)
|
||||
|
||||
|
||||
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 +275,149 @@ async def send_token(amount: int, unit: str, mint_url: str | None = None) -> str
|
||||
return token
|
||||
|
||||
|
||||
# A foreign mint's fee_reserve is a non-binding estimate (NUT-05): the mint may
|
||||
# demand more when re-quoting or at melt execution. Instead of padding the
|
||||
# estimate with a safety buffer (which strands the margin at the foreign mint
|
||||
# on every swap), the swap retries with the amount recomputed from the fees the
|
||||
# mint actually demands, up to this many attempts.
|
||||
_MAX_SWAP_ATTEMPTS = 3
|
||||
|
||||
_MINT_ERROR_CODE_RE = re.compile(r"\(Code: (\d+)\)")
|
||||
_MELT_SHORTFALL_RE = re.compile(r"Provided: (\d+), needed: (\d+)")
|
||||
|
||||
# Insufficient-melt-inputs failures differ across mint implementations. 11005 is
|
||||
# the registered "Transaction is not balanced" code (cdk), specific enough to
|
||||
# trust on the code alone. 11000 is nutshell's generic, unregistered
|
||||
# TransactionError covering many unrelated failures, so it only counts as a fee
|
||||
# shortfall alongside the "not enough inputs" detail text. With no code suffix at
|
||||
# all, that same text is the only signal.
|
||||
|
||||
|
||||
def _net_minted_amount(amount_msat: int, token_unit: str, fees: int) -> int:
|
||||
"""
|
||||
Convert the token value minus fees (given in the token unit) into an
|
||||
amount in the primary mint's unit.
|
||||
"""
|
||||
fee_msat = fees * 1000 if token_unit == "sat" else fees
|
||||
remaining_msat = amount_msat - fee_msat
|
||||
if settings.primary_mint_unit == "sat":
|
||||
return int(remaining_msat // 1000)
|
||||
return int(remaining_msat)
|
||||
|
||||
|
||||
def _melt_insufficient_shortfall(error: Exception) -> int | None:
|
||||
"""
|
||||
Classify a melt failure: return the observed shortfall (in the token unit)
|
||||
when the mint rejected the inputs as insufficient, or None when the failure
|
||||
is unrelated to fees and must not be retried (e.g. a Lightning payment
|
||||
failure, where a smaller invoice would not help).
|
||||
|
||||
Cashu errors carry no structured amounts (NUT-00 defines only detail/code,
|
||||
flattened to "Mint Error: <detail> (Code: <code>)" by cashu-py), so the
|
||||
classification uses the code and the shortfall must be inferred: the
|
||||
"Provided: X, needed: Y" amounts are nutshell-specific free text and only
|
||||
refine the shortfall when present; otherwise shrink one unit at a time.
|
||||
"""
|
||||
message = str(error)
|
||||
code_match = _MINT_ERROR_CODE_RE.search(message)
|
||||
code = code_match.group(1) if code_match is not None else None
|
||||
has_shortfall_text = "not enough inputs" in message.lower()
|
||||
|
||||
match code:
|
||||
case "11005": # registered TransactionUnbalanced: trust the code
|
||||
pass
|
||||
case "11000" if has_shortfall_text: # generic nutshell error: needs the text
|
||||
pass
|
||||
case None if has_shortfall_text: # no code suffix: text is the only signal
|
||||
pass
|
||||
case _: # other codes, a bare 11000, or no signal: must not retry
|
||||
return None
|
||||
|
||||
amounts = _MELT_SHORTFALL_RE.search(message)
|
||||
if amounts is not None:
|
||||
provided, needed = int(amounts.group(1)), int(amounts.group(2))
|
||||
if needed > provided:
|
||||
return needed - provided
|
||||
return 1
|
||||
|
||||
|
||||
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)
|
||||
total_fees = fee_reserve + input_fees
|
||||
minted_amount = _net_minted_amount(amount_msat, token_unit, total_fees)
|
||||
|
||||
if minted_amount <= 0:
|
||||
raise ValueError(f"Fees ({total_fees} {token_unit}) exceed token amount")
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: fee estimation result",
|
||||
extra={
|
||||
"token_amount_sat": amount_msat // 1000,
|
||||
"estimated_fee": total_fees,
|
||||
"estimated_fee_unit": token_unit,
|
||||
"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)},
|
||||
)
|
||||
if is_mint_connection_error(e):
|
||||
raise MintConnectionError("Cashu mint is unreachable") from 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 +432,228 @@ 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
|
||||
# If the token is already from the primary mint, we don't need a cross-mint
|
||||
# swap — redeem it same-mint. There's no melt/Lightning fee, but the mint's
|
||||
# NUT-02 input fee still applies; _redeem_same_mint accounts for it.
|
||||
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,
|
||||
},
|
||||
)
|
||||
return await _redeem_same_mint(token_wallet, token_obj)
|
||||
|
||||
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)
|
||||
mint_quote = await primary_wallet.request_mint(minted_amount)
|
||||
|
||||
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,
|
||||
minted_amount = await _calculate_swap_amount(
|
||||
amount_msat,
|
||||
token_obj.unit,
|
||||
token_obj.mint,
|
||||
token_wallet,
|
||||
primary_wallet,
|
||||
token_obj.proofs,
|
||||
)
|
||||
|
||||
# The estimate above is non-binding: the mint may demand a higher fee on the
|
||||
# real quote or reject the melt outright. Retry the quote/melt cycle with the
|
||||
# amount recomputed from the fees the mint actually demands.
|
||||
observed_extra_fee = 0
|
||||
attempt = 0
|
||||
while True:
|
||||
attempt += 1
|
||||
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, "attempt": attempt},
|
||||
)
|
||||
|
||||
melt_quote = await token_wallet.melt_quote(mint_quote.request)
|
||||
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,
|
||||
"attempt": attempt,
|
||||
},
|
||||
)
|
||||
|
||||
if total_needed > token_amount:
|
||||
recomputed = _net_minted_amount(
|
||||
amount_msat,
|
||||
token_obj.unit,
|
||||
melt_quote.fee_reserve + input_fees + observed_extra_fee,
|
||||
)
|
||||
if attempt >= _MAX_SWAP_ATTEMPTS or recomputed <= 0:
|
||||
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,
|
||||
"attempts": attempt,
|
||||
},
|
||||
)
|
||||
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})"
|
||||
)
|
||||
logger.warning(
|
||||
"swap_to_primary_mint: melt quote exceeds token amount, retrying",
|
||||
extra={
|
||||
"total_needed": total_needed,
|
||||
"token_amount": token_amount,
|
||||
"retry_minted_amount": recomputed,
|
||||
"attempt": attempt,
|
||||
},
|
||||
)
|
||||
minted_amount = recomputed
|
||||
continue
|
||||
|
||||
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:
|
||||
# A down mint won't fix itself by retrying with a smaller amount.
|
||||
if is_mint_connection_error(e):
|
||||
logger.error(
|
||||
"swap_to_primary_mint: melt failed — mint unreachable",
|
||||
extra={"error": str(e), "foreign_mint": token_obj.mint},
|
||||
)
|
||||
raise MintConnectionError("Cashu mint is unreachable") from e
|
||||
shortfall = _melt_insufficient_shortfall(e)
|
||||
recomputed = 0
|
||||
if shortfall is not None:
|
||||
observed_extra_fee += shortfall
|
||||
recomputed = _net_minted_amount(
|
||||
amount_msat,
|
||||
token_obj.unit,
|
||||
melt_quote.fee_reserve + input_fees + observed_extra_fee,
|
||||
)
|
||||
if shortfall is None or attempt >= _MAX_SWAP_ATTEMPTS or recomputed <= 0:
|
||||
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,
|
||||
"attempts": attempt,
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Failed to melt token from foreign mint {token_obj.mint}: {e}"
|
||||
) from e
|
||||
logger.warning(
|
||||
"swap_to_primary_mint: mint demanded more than quoted at melt, retrying",
|
||||
extra={
|
||||
"shortfall": shortfall,
|
||||
"retry_minted_amount": recomputed,
|
||||
"attempt": attempt,
|
||||
},
|
||||
)
|
||||
minted_amount = recomputed
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
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 TokenConsumedError(
|
||||
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 TokenConsumedError:
|
||||
raise
|
||||
except Exception as recovery_err:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: recovery failed",
|
||||
extra={"error": str(recovery_err)},
|
||||
)
|
||||
raise TokenConsumedError(
|
||||
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,
|
||||
},
|
||||
)
|
||||
# Foreign proofs already melted (spent) — non-retryable.
|
||||
raise TokenConsumedError(
|
||||
"Mint on primary failed after successful melt"
|
||||
) from e
|
||||
|
||||
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 +668,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,27 +681,72 @@ 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},
|
||||
)
|
||||
|
||||
# Use atomic SQL UPDATE to prevent race conditions during concurrent topups
|
||||
stmt = (
|
||||
update(db.ApiKey)
|
||||
.where(col(db.ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(balance=(db.ApiKey.balance) + amount)
|
||||
)
|
||||
await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
# The token is already redeemed (spent) here, so any crediting failure
|
||||
# is post-redemption and non-retryable — surface it as TokenConsumedError
|
||||
# (a key that vanished mid-flight, or an unexpected DB fault), never a
|
||||
# retryable/token-error taxonomy.
|
||||
try:
|
||||
# Atomic UPDATE to prevent race conditions during concurrent topups.
|
||||
stmt = (
|
||||
update(db.ApiKey)
|
||||
.where(col(db.ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(balance=(db.ApiKey.balance) + amount)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
# If pruning removed this key after redemption, do not commit a no-op
|
||||
# balance update and pretend the top-up succeeded.
|
||||
if (getattr(result, "rowcount", 0) or 0) == 0:
|
||||
raise TokenConsumedError(
|
||||
"Token redeemed but the API key disappeared before the "
|
||||
"credit could be recorded"
|
||||
)
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
except TokenConsumedError:
|
||||
raise
|
||||
except Exception as db_error:
|
||||
raise TokenConsumedError(
|
||||
"Token redeemed but crediting the balance failed"
|
||||
) from db_error
|
||||
|
||||
logger.info(
|
||||
"credit_balance: Balance updated successfully",
|
||||
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 +848,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:
|
||||
@@ -261,11 +863,19 @@ async def fetch_all_balances(
|
||||
}
|
||||
return error_result
|
||||
|
||||
# Build the set of mints to inspect. Received tokens are stored against
|
||||
# ``primary_mint`` (which defaults to a real mint even when ``cashu_mints``
|
||||
# is empty), so include it as a fallback — otherwise a node that accepts
|
||||
# payments would still report empty balances when ``cashu_mints`` is unset.
|
||||
mint_urls: list[str] = list(settings.cashu_mints)
|
||||
if settings.primary_mint and settings.primary_mint not in mint_urls:
|
||||
mint_urls.append(settings.primary_mint)
|
||||
|
||||
# Create tasks for all mint/unit combinations
|
||||
async with db.create_session() as session:
|
||||
tasks = [
|
||||
fetch_balance(session, mint_url, unit)
|
||||
for mint_url in settings.cashu_mints
|
||||
for mint_url in mint_urls
|
||||
for unit in units
|
||||
]
|
||||
|
||||
@@ -294,7 +904,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 +917,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 +939,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 +969,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"""
|
||||
|
||||
|
||||
229
tests/integration/test_admin_model_cache_pricing.py
Normal file
229
tests/integration/test_admin_model_cache_pricing.py
Normal file
@@ -0,0 +1,229 @@
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import ModelRow, UpstreamProviderRow
|
||||
from routstr.payment.cost_calculation import CostData, calculate_cost
|
||||
from routstr.proxy import get_model_instance, reinitialize_upstreams
|
||||
|
||||
|
||||
def _admin_headers() -> dict[str, str]:
|
||||
token = "test-admin-cache-pricing-token"
|
||||
admin_sessions[token] = int(
|
||||
(datetime.now(timezone.utc) + timedelta(minutes=5)).timestamp()
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
|
||||
def _model_payload(
|
||||
provider_id: int,
|
||||
*,
|
||||
cache_read: float,
|
||||
cache_write: float,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"id": "custom-cache-model",
|
||||
"name": "Custom Cache Model",
|
||||
"description": "custom model with explicit cache pricing",
|
||||
"created": 0,
|
||||
"context_length": 128000,
|
||||
"architecture": {
|
||||
"modality": "text",
|
||||
"input_modalities": ["text"],
|
||||
"output_modalities": ["text"],
|
||||
"tokenizer": "unknown",
|
||||
"instruct_type": None,
|
||||
},
|
||||
"pricing": {
|
||||
"prompt": 1.4e-7,
|
||||
"completion": 2.8e-7,
|
||||
"input_cache_read": cache_read,
|
||||
"input_cache_write": cache_write,
|
||||
"request": 0.0,
|
||||
"image": 0.0,
|
||||
"web_search": 0.0,
|
||||
"internal_reasoning": 0.0,
|
||||
},
|
||||
"per_request_limits": None,
|
||||
"top_provider": None,
|
||||
"upstream_provider_id": provider_id,
|
||||
"canonical_slug": None,
|
||||
"alias_ids": [],
|
||||
"enabled": True,
|
||||
"forwarded_model_id": "custom-cache-model",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_provider_model_api_persists_cache_pricing_on_create_and_update(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="generic",
|
||||
base_url="https://custom-upstream.example/v1",
|
||||
api_key="test-key",
|
||||
provider_fee=1.0,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
assert provider.id is not None
|
||||
await reinitialize_upstreams()
|
||||
|
||||
headers = _admin_headers()
|
||||
create_payload = _model_payload(
|
||||
provider.id,
|
||||
cache_read=2.8e-9,
|
||||
cache_write=3.5e-9,
|
||||
)
|
||||
|
||||
with patch("routstr.payment.models.sats_usd_price", return_value=1e-6):
|
||||
create_response = await integration_client.post(
|
||||
f"/admin/api/upstream-providers/{provider.id}/models",
|
||||
headers=headers,
|
||||
json=create_payload,
|
||||
)
|
||||
|
||||
assert create_response.status_code == 200
|
||||
create_body = create_response.json()
|
||||
assert create_body["pricing"]["input_cache_read"] == pytest.approx(2.8e-9)
|
||||
assert create_body["pricing"]["input_cache_write"] == pytest.approx(3.5e-9)
|
||||
|
||||
row = await integration_session.get(ModelRow, ("custom-cache-model", provider.id))
|
||||
assert row is not None
|
||||
stored_pricing = json.loads(row.pricing)
|
||||
assert stored_pricing["input_cache_read"] == pytest.approx(2.8e-9)
|
||||
assert stored_pricing["input_cache_write"] == pytest.approx(3.5e-9)
|
||||
|
||||
update_payload = _model_payload(
|
||||
provider.id,
|
||||
cache_read=1.25e-9,
|
||||
cache_write=4.5e-9,
|
||||
)
|
||||
|
||||
with patch("routstr.payment.models.sats_usd_price", return_value=1e-6):
|
||||
update_response = await integration_client.post(
|
||||
f"/admin/api/upstream-providers/{provider.id}/models",
|
||||
headers=headers,
|
||||
json=update_payload,
|
||||
)
|
||||
|
||||
assert update_response.status_code == 200
|
||||
update_body = update_response.json()
|
||||
assert update_body["pricing"]["input_cache_read"] == pytest.approx(1.25e-9)
|
||||
assert update_body["pricing"]["input_cache_write"] == pytest.approx(4.5e-9)
|
||||
|
||||
await integration_session.refresh(row)
|
||||
updated_pricing = json.loads(row.pricing)
|
||||
assert updated_pricing["input_cache_read"] == pytest.approx(1.25e-9)
|
||||
assert updated_pricing["input_cache_write"] == pytest.approx(4.5e-9)
|
||||
|
||||
model = get_model_instance("custom-cache-model")
|
||||
assert model is not None
|
||||
assert model.sats_pricing is not None
|
||||
assert model.sats_pricing.input_cache_read == pytest.approx(0.00125)
|
||||
assert model.sats_pricing.input_cache_write == pytest.approx(0.0045)
|
||||
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=1e-6):
|
||||
cost = await calculate_cost(
|
||||
{
|
||||
"model": "custom-cache-model",
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 100,
|
||||
"prompt_tokens_details": {"cached_tokens": 800},
|
||||
},
|
||||
},
|
||||
max_cost=1_000_000,
|
||||
)
|
||||
|
||||
assert isinstance(cost, CostData)
|
||||
assert cost.input_tokens == 200
|
||||
assert cost.cache_read_input_tokens == 800
|
||||
assert cost.cache_read_msats == 1000
|
||||
assert cost.output_msats == 28000
|
||||
assert cost.input_msats == 29000
|
||||
assert cost.total_msats == 57000
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_response_cost_uses_model_cache_pricing(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""A model's configured cache price must discount upstream cached-token usage."""
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="generic",
|
||||
base_url="https://cache-priced-upstream.example/v1",
|
||||
api_key="test-key",
|
||||
provider_fee=1.0,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
assert provider.id is not None
|
||||
|
||||
row = ModelRow(
|
||||
id="cache-priced-model",
|
||||
name="Cache Priced Model",
|
||||
description="model seeded with explicit cache pricing",
|
||||
created=0,
|
||||
context_length=128000,
|
||||
architecture=json.dumps(
|
||||
{
|
||||
"modality": "text",
|
||||
"input_modalities": ["text"],
|
||||
"output_modalities": ["text"],
|
||||
"tokenizer": "unknown",
|
||||
"instruct_type": None,
|
||||
}
|
||||
),
|
||||
pricing=json.dumps(
|
||||
{
|
||||
"prompt": 1.4e-7,
|
||||
"completion": 2.8e-7,
|
||||
"input_cache_read": 1.25e-9,
|
||||
"input_cache_write": 4.5e-9,
|
||||
}
|
||||
),
|
||||
upstream_provider_id=provider.id,
|
||||
enabled=True,
|
||||
forwarded_model_id="cache-priced-model",
|
||||
)
|
||||
integration_session.add(row)
|
||||
await integration_session.commit()
|
||||
|
||||
with patch("routstr.payment.models.sats_usd_price", return_value=1e-6):
|
||||
await reinitialize_upstreams()
|
||||
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=1e-6):
|
||||
cost = await calculate_cost(
|
||||
{
|
||||
"model": "cache-priced-model",
|
||||
"usage": {
|
||||
"prompt_tokens": 1000,
|
||||
"completion_tokens": 100,
|
||||
"prompt_tokens_details": {"cached_tokens": 800},
|
||||
},
|
||||
},
|
||||
max_cost=1_000_000,
|
||||
)
|
||||
|
||||
assert isinstance(cost, CostData)
|
||||
# prompt 1.4e-7 USD/token -> 0.14 sats/token -> 140 msats/token
|
||||
# cache 1.25e-9 USD/token -> 0.00125 sats/token -> 1.25 msats/token
|
||||
# completion 2.8e-7 USD/token -> 0.28 sats/token -> 280 msats/token
|
||||
assert cost.input_tokens == 200
|
||||
assert cost.cache_read_input_tokens == 800
|
||||
assert cost.cache_read_msats == 1000
|
||||
assert cost.input_msats == 29000
|
||||
assert cost.output_msats == 28000
|
||||
assert cost.total_msats == 57000
|
||||
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)
|
||||
146
tests/integration/test_free_response_stale_reservation.py
Normal file
146
tests/integration/test_free_response_stale_reservation.py
Normal file
@@ -0,0 +1,146 @@
|
||||
"""Regression tests for charging after stale reservation cleanup."""
|
||||
|
||||
import time
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overrun_charges_after_reservation_swept(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Overrun finalize must charge even when the reservation was already released."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 990 # discounted reservation
|
||||
actual_token_cost = 1000 # actual cost overruns the reservation
|
||||
|
||||
# Sweeper has zeroed reserved_balance but left balance untouched.
|
||||
key = _make_key(balance=1000, reserved=0)
|
||||
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 integration_session.refresh(key)
|
||||
|
||||
assert key.total_spent == actual_token_cost, (
|
||||
f"Request was not billed (total_spent={key.total_spent}) — free response bug"
|
||||
)
|
||||
assert key.balance == 1000 - actual_token_cost, (
|
||||
f"Balance not charged: {key.balance}"
|
||||
)
|
||||
assert key.balance >= 0
|
||||
assert key.reserved_balance == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_free_response_path_closed_end_to_end(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""A reservation released by the real sweeper must not yield a free response."""
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.core.db import create_session, release_stale_reservations
|
||||
|
||||
deducted_max_cost = 990
|
||||
actual_token_cost = 1000
|
||||
key_hash = f"test_sweep_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
session.add(
|
||||
ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=1000,
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
# Reserve the request, then backdate reserved_at so the sweeper treats it as
|
||||
# stale (simulates a stream that outlived stale_reservation_timeout_seconds).
|
||||
async with create_session() as session:
|
||||
key = await session.get(ApiKey, key_hash)
|
||||
assert key is not None
|
||||
await pay_for_request(key, deducted_max_cost, session)
|
||||
await session.refresh(key)
|
||||
assert key.reserved_balance == deducted_max_cost
|
||||
key.reserved_at = int(time.time()) - 10_000
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
# Sweeper releases the stale reservation without charging.
|
||||
async with create_session() as session:
|
||||
released = await release_stale_reservations(session, max_age_seconds=300)
|
||||
assert released == 1
|
||||
|
||||
async with create_session() as session:
|
||||
key = await session.get(ApiKey, key_hash)
|
||||
assert key is not None
|
||||
assert key.reserved_balance == 0, "Precondition: sweeper zeroed the reservation"
|
||||
|
||||
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, session, deducted_max_cost
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, key_hash)
|
||||
assert final is not None
|
||||
|
||||
assert final.total_spent == actual_token_cost, (
|
||||
f"Free response: total_spent={final.total_spent}, expected {actual_token_cost}"
|
||||
)
|
||||
assert final.balance == 1000 - actual_token_cost, (
|
||||
f"Balance not charged after sweep: {final.balance}"
|
||||
)
|
||||
assert final.balance >= 0
|
||||
assert final.reserved_balance == 0
|
||||
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)
|
||||
|
||||
128
tests/integration/test_provider_self_lookup.py
Normal file
128
tests/integration/test_provider_self_lookup.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""A live upstream provider resolves its OWN database row by stable identity
|
||||
(its primary key), not by its mutable/secret ``api_key``.
|
||||
|
||||
Today ``from_db_row`` drops ``provider_row.id`` and the two self-referential
|
||||
paths — PPQ.AI's insufficient-balance self-disable and the base
|
||||
``refresh_models_cache`` — re-find their own row with
|
||||
``WHERE base_url == self.base_url AND api_key == self.api_key``. That uses a
|
||||
rotatable secret as a self-handle: the moment the row's key changes underneath a
|
||||
live object (a rotation racing an in-flight request), the object can no longer
|
||||
find itself. These tests pin the invariant that a provider looks itself up by
|
||||
identity, so the lookup survives a key change (and, later, key encryption).
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.core.db import UpstreamProviderRow
|
||||
from routstr.upstream.ppqai import PPQAIUpstreamProvider
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_object_carries_its_persistent_identity(
|
||||
integration_session: object,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""``from_db_row`` gives the in-memory object its row's identity (``db_id``)."""
|
||||
row = UpstreamProviderRow(
|
||||
provider_type="ppqai",
|
||||
base_url="https://api.ppq.ai",
|
||||
api_key="sk-original",
|
||||
enabled=True,
|
||||
provider_fee=1.0,
|
||||
)
|
||||
integration_session.add(row) # type: ignore[attr-defined]
|
||||
await integration_session.commit() # type: ignore[attr-defined]
|
||||
await integration_session.refresh(row) # type: ignore[attr-defined]
|
||||
|
||||
provider = PPQAIUpstreamProvider.from_db_row(row)
|
||||
assert provider is not None
|
||||
|
||||
assert provider.db_id == row.id
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_self_disable_targets_own_row_after_key_rotation(
|
||||
integration_session: object,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""PPQ.AI self-disable must disable *its* row even after the key rotated.
|
||||
|
||||
RED (current): the object holds the pre-rotation key, so the
|
||||
``(base_url, api_key)`` lookup misses the row → the provider is never
|
||||
disabled. GREEN: lookup by ``id`` finds it and disables it.
|
||||
"""
|
||||
row = UpstreamProviderRow(
|
||||
provider_type="ppqai",
|
||||
base_url="https://api.ppq.ai",
|
||||
api_key="sk-original",
|
||||
enabled=True,
|
||||
provider_fee=1.0,
|
||||
)
|
||||
integration_session.add(row) # type: ignore[attr-defined]
|
||||
await integration_session.commit() # type: ignore[attr-defined]
|
||||
await integration_session.refresh(row) # type: ignore[attr-defined]
|
||||
|
||||
provider = PPQAIUpstreamProvider.from_db_row(row) # captures sk-original
|
||||
assert provider is not None
|
||||
|
||||
# Key is rotated in the DB while `provider` is still live.
|
||||
row.api_key = "sk-rotated"
|
||||
integration_session.add(row) # type: ignore[attr-defined]
|
||||
await integration_session.commit() # type: ignore[attr-defined]
|
||||
|
||||
with patch("routstr.proxy.reinitialize_upstreams", new=AsyncMock()):
|
||||
await provider.on_upstream_error_redirect(402, "Insufficient balance")
|
||||
|
||||
await integration_session.refresh(row) # type: ignore[attr-defined]
|
||||
assert row.enabled is False
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_models_cache_finds_own_row_after_key_rotation(
|
||||
integration_session: object,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""``refresh_models_cache`` must resolve its own row after a key rotation.
|
||||
|
||||
``refresh_models_cache`` swallows every exception (it only logs), so the
|
||||
observable proof it found its row is that it reaches ``list_models`` — which
|
||||
is called with the row's ``id`` only *after* the row is resolved. RED
|
||||
(current): the stale-key ``(base_url, api_key)`` lookup returns nothing, the
|
||||
method raises ``404`` internally and returns before ``list_models`` is ever
|
||||
called. GREEN: lookup by ``id`` finds the row and ``list_models`` runs for
|
||||
that ``id``.
|
||||
"""
|
||||
row = UpstreamProviderRow(
|
||||
provider_type="ppqai",
|
||||
base_url="https://api.ppq.ai",
|
||||
api_key="sk-original",
|
||||
enabled=True,
|
||||
provider_fee=1.0,
|
||||
)
|
||||
integration_session.add(row) # type: ignore[attr-defined]
|
||||
await integration_session.commit() # type: ignore[attr-defined]
|
||||
await integration_session.refresh(row) # type: ignore[attr-defined]
|
||||
row_id = row.id
|
||||
|
||||
provider = PPQAIUpstreamProvider.from_db_row(row) # captures sk-original
|
||||
assert provider is not None
|
||||
|
||||
row.api_key = "sk-rotated"
|
||||
integration_session.add(row) # type: ignore[attr-defined]
|
||||
await integration_session.commit() # type: ignore[attr-defined]
|
||||
|
||||
list_models_mock = AsyncMock(return_value=[])
|
||||
with (
|
||||
patch.object(provider, "fetch_models", new=AsyncMock(return_value=[])),
|
||||
patch("routstr.upstream.base.list_models", new=list_models_mock),
|
||||
):
|
||||
await provider.refresh_models_cache()
|
||||
|
||||
list_models_mock.assert_awaited_once()
|
||||
assert list_models_mock.await_args is not None
|
||||
assert list_models_mock.await_args.kwargs["upstream_id"] == row_id
|
||||
283
tests/integration/test_prune_dead_api_keys.py
Normal file
283
tests/integration/test_prune_dead_api_keys.py
Normal file
@@ -0,0 +1,283 @@
|
||||
"""
|
||||
Tests for prune_dead_api_keys — the janitor that removes provably-dead 0/0/0
|
||||
API keys (funded keys fully refunded/expired without ever being used, plus bare
|
||||
orphans), while protecting keys that are still meaningful.
|
||||
"""
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.sql.dml import Update
|
||||
from sqlmodel import col, update
|
||||
|
||||
from routstr.core.db import (
|
||||
ApiKey,
|
||||
CashuTransaction,
|
||||
LightningInvoice,
|
||||
create_session,
|
||||
prune_dead_api_keys,
|
||||
)
|
||||
|
||||
OLD = 100 # min_age_seconds used by the tests
|
||||
NOW = int(time.time())
|
||||
LONG_AGO = NOW - 10_000 # well past the grace period
|
||||
|
||||
|
||||
async def _exists(key_hash: str) -> bool:
|
||||
async with create_session() as session:
|
||||
return (await session.get(ApiKey, key_hash)) is not None
|
||||
|
||||
|
||||
def _dead_key(created_at: int | None) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"dead_{uuid.uuid4().hex}",
|
||||
balance=0,
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
created_at=created_at,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prunes_old_refunded_zero_key(patched_db_engine: None) -> None:
|
||||
"""A funded-then-refunded key (0/0/0, NULL parent, old) is pruned."""
|
||||
key = _dead_key(LONG_AGO)
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 1
|
||||
assert not await _exists(key.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grace_period_protects_fresh_key(patched_db_engine: None) -> None:
|
||||
"""A dead-looking but recently created key is protected by the grace period."""
|
||||
key = _dead_key(int(time.time()))
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 0
|
||||
assert await _exists(key.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_used_key_never_pruned(patched_db_engine: None) -> None:
|
||||
"""Keys with any spend/requests or live balance are never pruned."""
|
||||
spent = _dead_key(LONG_AGO)
|
||||
spent.total_spent = 1
|
||||
requested = _dead_key(LONG_AGO)
|
||||
requested.total_requests = 1
|
||||
funded = _dead_key(LONG_AGO)
|
||||
funded.balance = 1000
|
||||
reserved = _dead_key(LONG_AGO)
|
||||
reserved.reserved_balance = 500
|
||||
|
||||
async with create_session() as session:
|
||||
for k in (spent, requested, funded, reserved):
|
||||
session.add(k)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 0
|
||||
for k in (spent, requested, funded, reserved):
|
||||
assert await _exists(k.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_and_child_keys_are_not_pruned(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Pruning must not orphan child keys or delete valid children."""
|
||||
parent = _dead_key(LONG_AGO)
|
||||
child = ApiKey(
|
||||
hashed_key=f"child_{uuid.uuid4().hex}",
|
||||
balance=0,
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
created_at=LONG_AGO,
|
||||
parent_key_hash=parent.hashed_key,
|
||||
)
|
||||
async with create_session() as session:
|
||||
session.add(parent)
|
||||
session.add(child)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 0
|
||||
assert await _exists(parent.hashed_key)
|
||||
assert await _exists(child.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_invoice_protects_key(patched_db_engine: None) -> None:
|
||||
"""A key referenced by a pending topup invoice is never pruned mid-topup."""
|
||||
key = _dead_key(LONG_AGO)
|
||||
invoice = LightningInvoice(
|
||||
id=f"inv_{uuid.uuid4().hex}",
|
||||
bolt11=f"lnbc_{uuid.uuid4().hex}",
|
||||
amount_sats=10,
|
||||
description="topup",
|
||||
payment_hash=uuid.uuid4().hex,
|
||||
status="pending",
|
||||
api_key_hash=key.hashed_key,
|
||||
purpose="topup",
|
||||
expires_at=NOW + 10_000,
|
||||
)
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
session.add(invoice)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 0
|
||||
assert await _exists(key.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_paid_invoice_does_not_protect_key(patched_db_engine: None) -> None:
|
||||
"""A settled (non-pending) invoice does not keep a dead key alive."""
|
||||
key = _dead_key(LONG_AGO)
|
||||
invoice = LightningInvoice(
|
||||
id=f"inv_{uuid.uuid4().hex}",
|
||||
bolt11=f"lnbc_{uuid.uuid4().hex}",
|
||||
amount_sats=10,
|
||||
description="topup",
|
||||
payment_hash=uuid.uuid4().hex,
|
||||
status="paid",
|
||||
api_key_hash=key.hashed_key,
|
||||
purpose="topup",
|
||||
expires_at=NOW - 1,
|
||||
)
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
session.add(invoice)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 1
|
||||
assert not await _exists(key.hashed_key)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_that_becomes_meaningful_during_prune_survives(
|
||||
patched_db_engine: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Revalidate before unlink/delete so a late top-up cannot be pruned."""
|
||||
key = _dead_key(LONG_AGO)
|
||||
txn = CashuTransaction(
|
||||
id=uuid.uuid4().hex,
|
||||
token="cashuABC",
|
||||
amount=21,
|
||||
unit="sat",
|
||||
type="in",
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
session.add(txn)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
original_exec = cast(Callable[..., Awaitable[Any]], session.exec)
|
||||
topped_up = False
|
||||
|
||||
async def exec_with_late_topup(
|
||||
statement: Any, *args: Any, **kwargs: Any
|
||||
) -> Any:
|
||||
nonlocal topped_up
|
||||
if (
|
||||
not topped_up
|
||||
and isinstance(statement, Update)
|
||||
and getattr(statement.table, "name", None) == "cashu_transactions"
|
||||
):
|
||||
topped_up = True
|
||||
async with create_session() as topup_session:
|
||||
await topup_session.exec( # type: ignore[call-overload]
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(balance=42)
|
||||
)
|
||||
await topup_session.commit()
|
||||
return await original_exec(statement, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(session, "exec", exec_with_late_topup)
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 0
|
||||
assert await _exists(key.hashed_key)
|
||||
async with create_session() as session:
|
||||
surviving = await session.get(CashuTransaction, txn.id)
|
||||
assert surviving is not None
|
||||
assert surviving.api_key_hashed_key == key.hashed_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transaction_audit_trail_preserved(patched_db_engine: None) -> None:
|
||||
"""Pruning a refunded key keeps its cashu_transactions, unlinked from the key."""
|
||||
key = _dead_key(LONG_AGO)
|
||||
txn = CashuTransaction(
|
||||
id=uuid.uuid4().hex,
|
||||
token="cashuABC",
|
||||
amount=21,
|
||||
unit="sat",
|
||||
type="in",
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
async with create_session() as session:
|
||||
session.add(key)
|
||||
session.add(txn)
|
||||
await session.commit()
|
||||
|
||||
async with create_session() as session:
|
||||
pruned = await prune_dead_api_keys(session, OLD)
|
||||
|
||||
assert pruned == 1
|
||||
assert not await _exists(key.hashed_key)
|
||||
|
||||
async with create_session() as session:
|
||||
surviving = await session.get(CashuTransaction, txn.id)
|
||||
assert surviving is not None, "Financial audit row must survive key deletion"
|
||||
assert surviving.api_key_hashed_key is None, "Link must be nulled, not dangling"
|
||||
assert surviving.amount == 21
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_periodic_prune_disabled_returns_immediately(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Non-positive intervals disable the janitor."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from routstr import auth
|
||||
from routstr.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "dead_key_prune_interval_seconds", 0)
|
||||
sleep_mock = AsyncMock()
|
||||
monkeypatch.setattr(auth.asyncio, "sleep", sleep_mock)
|
||||
|
||||
await auth.periodic_dead_key_prune()
|
||||
|
||||
sleep_mock.assert_not_called()
|
||||
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
|
||||
194
tests/integration/test_swap_fee_retry.py
Normal file
194
tests/integration/test_swap_fee_retry.py
Normal file
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Integration tests for reactive swap fee retries via the wallet topup endpoint.
|
||||
|
||||
Foreign-mint tokens are swapped to the primary mint using the foreign mint's
|
||||
melt quote, whose fee_reserve is a non-binding estimate (NUT-05): the mint may
|
||||
demand more when re-quoting or at melt execution. These tests cover the
|
||||
endpoint behaviour in those cases:
|
||||
|
||||
1. The mint demands one sat more at melt time than every quote reported
|
||||
(the mint.cubabitcoin.org incident): the swap retries with a smaller
|
||||
invoice and the topup succeeds, crediting the recomputed amount.
|
||||
2. The real melt quote reports a higher fee_reserve than the estimate: the
|
||||
swap re-quotes from the observed fee and the topup succeeds.
|
||||
3. The mint escalates its fee demands on every attempt: the retry budget is
|
||||
exhausted and the endpoint returns 400 with a clear error (never 500),
|
||||
without ever executing a melt.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient, Response
|
||||
|
||||
from routstr.core.settings import settings
|
||||
|
||||
# Captured at collection time, before the integration_app fixture replaces it
|
||||
# with the testmint stub that bypasses swapping (see conftest.py).
|
||||
from routstr.wallet import recieve_token as _real_recieve_token
|
||||
|
||||
PRIMARY_MINT = "http://primary:3338"
|
||||
|
||||
|
||||
def _make_swap_mocks(
|
||||
token_amount: int,
|
||||
fee_reserves: list[int],
|
||||
input_fees: int = 0,
|
||||
mint_url: str = "http://foreign-mint:3338",
|
||||
) -> tuple[Mock, Mock, Mock]:
|
||||
"""Return (token, token_wallet, primary_wallet) mocks that act like a mint.
|
||||
|
||||
Mint quotes pass the requested amount through their ``request`` field and
|
||||
melt quotes echo that amount back, so the mocks stay consistent for
|
||||
whatever amounts the implementation requests. ``fee_reserves`` supplies the
|
||||
fee_reserve of each successive melt quote (the first serves the estimation
|
||||
pass); requesting more quotes than provided fails the test.
|
||||
"""
|
||||
mock_token = Mock()
|
||||
mock_token.mint = mint_url
|
||||
mock_token.unit = "sat"
|
||||
mock_token.amount = token_amount
|
||||
mock_token.keysets = ["keyset1"]
|
||||
mock_token.proofs = [Mock(amount=token_amount)]
|
||||
|
||||
mock_token_wallet = Mock()
|
||||
mock_token_wallet.load_mint = AsyncMock()
|
||||
mock_token_wallet.load_proofs = AsyncMock()
|
||||
mock_token_wallet.get_fees_for_proofs = Mock(return_value=input_fees)
|
||||
|
||||
mock_primary_wallet = Mock()
|
||||
mock_primary_wallet.load_mint = AsyncMock()
|
||||
mock_primary_wallet.load_proofs = AsyncMock()
|
||||
mock_primary_wallet.available_balance = Mock(amount=0)
|
||||
mock_primary_wallet.mint = AsyncMock(return_value=Mock())
|
||||
|
||||
fees = iter(fee_reserves)
|
||||
|
||||
def _next_fee() -> int:
|
||||
try:
|
||||
return next(fees)
|
||||
except StopIteration:
|
||||
raise AssertionError(
|
||||
"more melt quotes requested than fee_reserves provided"
|
||||
) from None
|
||||
|
||||
mock_primary_wallet.request_mint = AsyncMock(
|
||||
side_effect=lambda amount: Mock(quote=f"mint_quote_{amount}", request=amount)
|
||||
)
|
||||
mock_token_wallet.melt_quote = AsyncMock(
|
||||
side_effect=lambda invoice: Mock(
|
||||
quote=f"melt_quote_{invoice}", amount=invoice, fee_reserve=_next_fee()
|
||||
)
|
||||
)
|
||||
mock_token_wallet.melt = AsyncMock(return_value=Mock())
|
||||
|
||||
return mock_token, mock_token_wallet, mock_primary_wallet
|
||||
|
||||
|
||||
def _wallet_router(primary_wallet: Mock, token_wallet: Mock) -> Callable[..., Mock]:
|
||||
"""Route get_wallet calls to the primary or foreign wallet mock by URL."""
|
||||
|
||||
def fake_get_wallet(mint_url: str, unit: str = "sat", load: bool = True) -> Mock:
|
||||
return primary_wallet if mint_url == PRIMARY_MINT else token_wallet
|
||||
|
||||
return fake_get_wallet
|
||||
|
||||
|
||||
async def _post_topup(
|
||||
client: AsyncClient,
|
||||
mock_token: Mock,
|
||||
token_wallet: Mock,
|
||||
primary_wallet: Mock,
|
||||
) -> Response:
|
||||
"""POST /v1/wallet/topup with the swap layer mocked at the mint boundary.
|
||||
|
||||
The conftest's testmint stub for recieve_token is swapped back for the
|
||||
real implementation so the request exercises the actual swap path.
|
||||
"""
|
||||
with patch("routstr.wallet.recieve_token", _real_recieve_token):
|
||||
with patch(
|
||||
"routstr.wallet.deserialize_token_from_string", return_value=mock_token
|
||||
):
|
||||
with patch(
|
||||
"routstr.wallet.get_wallet",
|
||||
side_effect=_wallet_router(primary_wallet, token_wallet),
|
||||
):
|
||||
with patch.object(settings, "primary_mint", PRIMARY_MINT):
|
||||
with patch.object(settings, "primary_mint_unit", "sat"):
|
||||
with patch.object(settings, "cashu_mints", [PRIMARY_MINT]):
|
||||
return await client.post(
|
||||
"/v1/wallet/topup",
|
||||
params={"cashu_token": "cashuAtest_foreign_token"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_topup_retries_when_melt_demands_more_than_quoted(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""A 179-sat token where every quote reports fee_reserve=1 but the mint
|
||||
rejects the first melt demanding 180. The retry shrinks the invoice to 177
|
||||
and the topup credits 177 sats (177_000 msats)."""
|
||||
mock_token, token_wallet, primary_wallet = _make_swap_mocks(
|
||||
179, fee_reserves=[1, 1, 1], mint_url="http://mint.cubabitcoin.org"
|
||||
)
|
||||
token_wallet.melt.side_effect = [
|
||||
Exception(
|
||||
"Mint Error: not enough inputs provided for melt. "
|
||||
"Provided: 179, needed: 180 (Code: 11000)"
|
||||
),
|
||||
Mock(),
|
||||
]
|
||||
|
||||
response = await _post_topup(
|
||||
authenticated_client, mock_token, token_wallet, primary_wallet
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["msats"] == 177_000
|
||||
assert token_wallet.melt.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_topup_retries_when_quote_fee_exceeds_estimate(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""A 1000-sat token estimated at fee 20, but the real quote demands 23.
|
||||
The retry recomputes 1000 - 23 = 977, which fits, and the topup credits
|
||||
977 sats (977_000 msats) with a single melt."""
|
||||
mock_token, token_wallet, primary_wallet = _make_swap_mocks(
|
||||
1000, fee_reserves=[20, 23, 23]
|
||||
)
|
||||
|
||||
response = await _post_topup(
|
||||
authenticated_client, mock_token, token_wallet, primary_wallet
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["msats"] == 977_000
|
||||
assert token_wallet.melt.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_topup_returns_422_when_retries_exhausted(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""A mint that escalates fee_reserve on every re-quote (1 → 10 → 25 → 50)
|
||||
exhausts the retry budget: clean 422 mint_error/too-small taxonomy, melt
|
||||
never executed."""
|
||||
mock_token, token_wallet, primary_wallet = _make_swap_mocks(
|
||||
1000, fee_reserves=[1, 10, 25, 50]
|
||||
)
|
||||
|
||||
response = await _post_topup(
|
||||
authenticated_client, mock_token, token_wallet, primary_wallet
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
assert "too small to cover swap fees" in response.json()["detail"]
|
||||
assert token_wallet.melt_quote.call_count == 4 # estimation + 3 attempts
|
||||
token_wallet.melt.assert_not_called()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user