chore: initial commit
This commit is contained in:
@@ -0,0 +1,10 @@
|
|||||||
|
.git
|
||||||
|
.mimocode
|
||||||
|
config.yaml
|
||||||
|
.env
|
||||||
|
data
|
||||||
|
ai-router
|
||||||
|
web/node_modules
|
||||||
|
web/dist
|
||||||
|
output
|
||||||
|
.playwright-cli
|
||||||
+12
@@ -0,0 +1,12 @@
|
|||||||
|
/config.yaml
|
||||||
|
/.env
|
||||||
|
/data/
|
||||||
|
/ai-router
|
||||||
|
/web/node_modules/
|
||||||
|
/web/dist/
|
||||||
|
*.db
|
||||||
|
*.db-shm
|
||||||
|
*.db-wal
|
||||||
|
/output/
|
||||||
|
/.playwright-cli/
|
||||||
|
/.playwright/
|
||||||
+23
@@ -0,0 +1,23 @@
|
|||||||
|
FROM node:24-alpine AS frontend
|
||||||
|
WORKDIR /src/web
|
||||||
|
COPY web/package*.json ./
|
||||||
|
RUN npm ci
|
||||||
|
COPY web/ ./
|
||||||
|
RUN npm run build
|
||||||
|
|
||||||
|
FROM golang:1.26-alpine AS backend
|
||||||
|
WORKDIR /src
|
||||||
|
COPY go.mod go.sum ./
|
||||||
|
RUN go mod download
|
||||||
|
COPY . .
|
||||||
|
COPY --from=frontend /src/web/dist ./web/dist
|
||||||
|
RUN CGO_ENABLED=0 go build -trimpath -ldflags='-s -w' -o /ai-router .
|
||||||
|
|
||||||
|
FROM alpine:3.23
|
||||||
|
RUN apk add --no-cache ca-certificates && addgroup -g 10001 router && adduser -D -u 10001 -G router router
|
||||||
|
WORKDIR /app
|
||||||
|
RUN mkdir /app/data && chown router:router /app/data
|
||||||
|
COPY --from=backend /ai-router /usr/local/bin/ai-router
|
||||||
|
USER router
|
||||||
|
EXPOSE 8080
|
||||||
|
ENTRYPOINT ["ai-router", "-config", "/app/config.yaml"]
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
.PHONY: build frontend test dev
|
||||||
|
frontend:
|
||||||
|
cd web && npm ci && npm run build
|
||||||
|
build: frontend
|
||||||
|
go build -trimpath -o ai-router .
|
||||||
|
test: frontend
|
||||||
|
go test -race ./...
|
||||||
|
dev:
|
||||||
|
go run . -config config.yaml
|
||||||
@@ -0,0 +1,149 @@
|
|||||||
|
# AI Router
|
||||||
|
|
||||||
|
个人使用的 Go AI API 网关，React 管理台内嵌在单个二进制中，数据保存在 SQLite。
|
||||||
|
|
||||||
|
## 构建与启动
|
||||||
|
|
||||||
|
需要 Go 1.26+、Node.js 22+ 和 npm（Node.js 仅用于构建前端）。
|
||||||
|
|
||||||
|
```sh
|
||||||
|
make build
|
||||||
|
cp config.example.yaml config.yaml
|
||||||
|
chmod 600 config.yaml
|
||||||
|
# 编辑 config.yaml，替换 admin.password（至少 12 个字符）
|
||||||
|
./ai-router -config config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
打开 http://127.0.0.1:8080，使用配置中的密码登录，无需用户名。示例密码不能用于启动。默认监听本机；远程使用应配置 HTTPS 反向代理，禁止代理缓冲 SSE。
|
||||||
|
|
||||||
|
HTTPS 反向代理部署时，在配置中设置 `server.public_origin: "https://router.example.com"`（仅 origin，不含路径或末尾斜线）。这用于管理请求的同源检查并启用 Secure Cookie，比较不区分大小写；不会盲目信任客户端传入的转发头。
|
||||||
|
|
||||||
|
数据库路径相对于进程工作目录。备份时建议停止服务后复制整个 data 目录（包含可能存在的 WAL 文件）。上游凭据保存在数据库中，务必保护数据库及备份；不要提交 config.yaml 到 Git。下游 API Key 仅保存哈希，创建时完整显示一次。
|
||||||
|
|
||||||
|
## Docker Compose 部署
|
||||||
|
|
||||||
|
需要 Docker Engine 和 Docker Compose v2。Compose 使用现有多阶段 Dockerfile
|
||||||
|
构建前后端，容器以 UID/GID 10001 运行，SQLite 保存在命名卷 router-data 中。
|
||||||
|
|
||||||
|
首次部署时创建配置（如果已有 config.yaml，请备份并编辑，不要覆盖）：
|
||||||
|
|
||||||
|
```sh
|
||||||
|
cp config.docker.example.yaml config.yaml
|
||||||
|
openssl rand -hex 32
|
||||||
|
# 将生成的随机值填入 config.yaml 的 admin.password
|
||||||
|
sudo chown 10001:10001 config.yaml
|
||||||
|
sudo chmod 600 config.yaml
|
||||||
|
docker compose config --quiet
|
||||||
|
docker compose up -d --build
|
||||||
|
docker compose ps
|
||||||
|
```
|
||||||
|
|
||||||
|
配置只读挂载到容器，不会打入镜像。容器进程必须有权限读取该文件；
|
||||||
|
上面的权限命令适用于普通 Linux Docker Engine，rootless/user namespace 环境
|
||||||
|
需按实际 UID 映射设置。不要将真实密码提交到版本库。
|
||||||
|
|
||||||
|
默认访问 http://127.0.0.1:8080。容器内监听必须是 0.0.0.0:8080，
|
||||||
|
不能沿用本地运行示例的 127.0.0.1。宿主机端口已被占用时，例如已有本地服务：
|
||||||
|
|
||||||
|
```sh
|
||||||
|
AI_ROUTER_PORT=8081 docker compose up -d --build
|
||||||
|
```
|
||||||
|
|
||||||
|
后续 Compose 命令也应使用相同的 AI_ROUTER_PORT 值，或将它写入本地 .env。
|
||||||
|
默认仅发布到宿主机回环地址，不直接暴露公网。远程访问建议在宿主机配置
|
||||||
|
HTTPS 反向代理，同时设置 server.public_origin 并禁用 SSE 代理缓冲。
|
||||||
|
如果反向代理也在容器中，应加入同一 Docker 网络，连接 ai-router:8080。
|
||||||
|
|
||||||
|
常用命令：
|
||||||
|
|
||||||
|
```sh
|
||||||
|
docker compose logs -f --tail=100
|
||||||
|
docker compose up -d --build # 更新源码后重新构建和部署
|
||||||
|
docker compose restart # 修改配置文件后重启
|
||||||
|
docker compose down # 停止并删除容器，保留数据库卷
|
||||||
|
```
|
||||||
|
|
||||||
|
不要执行 docker compose down -v，除非明确要删除数据库卷。
|
||||||
|
首次启动使用独立的新数据库，不会自动导入本地运行的数据。
|
||||||
|
迁移时先停止旧服务，备份原数据库目录（含可能存在的 WAL/SHM 文件），
|
||||||
|
再将数据复制到命名卷并确保 UID/GID 10001 可读写；不要让两套服务同时写同一数据库。
|
||||||
|
|
||||||
|
## 渠道与模型映射
|
||||||
|
|
||||||
|
在管理台创建渠道，选择协议并填写 Base URL。OpenAI Chat/Responses 填 API 前缀，例如 https://api.openai.com/v1，程序分别追加 /chat/completions、/responses；Anthropic 填服务根地址，例如 https://api.anthropic.com（不带 /v1），程序追加 /v1/messages，获取模型追加 /v1/models。自定义路径前缀会保留。不要填写完整接口端点。不兼容 Anthropic 旧的带 /v1 配置，也不会自动去重或迁移；已有渠道需手动修改 Base URL。下游接口路径不变。
|
||||||
|
|
||||||
|
点击“获取模型列表”读取上游模型（新建渠道无需先保存，已有渠道密钥留空时使用原密钥），随后弹窗选择要加入列表的模型。支持搜索、多选、全选搜索结果和清空选择；已有模型标记为已加入，不会重复添加或覆盖展示名称。确认后只将所选模型加入编辑草稿，取消不改变原列表。可修改展示名称；上游不支持模型列表时，可手动添加。点击“保存渠道”后，勾选展示的模型才会保存并汇总到“模型”页面。展示名称即下面的模型别名，修改后下游调用名也随之改变。
|
||||||
|
|
||||||
|
管理台模型页汇总所有渠道已保存的展示模型，并标注渠道是否停用；下游 `/v1/models` 进一步按渠道启用状态及当前 Key 权限过滤。只允许调用已配置的模型，未映射模型不能调用；历史渠道的透传设置也不再生效。
|
||||||
|
|
||||||
|
下游模型格式为 `渠道名/模型别名`，按第一个 `/` 拆分。别名在渠道内查询一次，得到的真实模型名原样发给上游：
|
||||||
|
|
||||||
|
| 渠道 | 模型别名 | 上游真实模型 | 下游调用 |
|
||||||
|
| --- | --- | --- | --- |
|
||||||
|
| abc | gpt-4 | abc/gpt-4 | abc/gpt-4 |
|
||||||
|
| abc | gpt-5 | gpt-4 | abc/gpt-5 |
|
||||||
|
|
||||||
|
不会再次拼接渠道前缀，也不会递归路由。未配置的别名一律拒绝，不支持未映射模型透传。创建下游 API Key 时可授权全部渠道或选择渠道白名单。空白名单不允许任何渠道。
|
||||||
|
|
||||||
|
```sh
|
||||||
|
curl http://127.0.0.1:8080/v1/chat/completions \
|
||||||
|
-H "Authorization: Bearer YOUR_ROUTER_KEY" \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d '{"model":"abc/gpt-5","messages":[{"role":"user","content":"hello"}],"stream":true}'
|
||||||
|
```
|
||||||
|
|
||||||
|
## API 与兼容范围
|
||||||
|
|
||||||
|
模型页面支持“手动测试”和“一键测试全部”：手动测试可填写最多 2000 字符的提示词，一键测试按顺序检查当前所有启用渠道的已展示模型，跳过停用渠道。每次发送一条限制输出 Token 的非流式请求，最长等待 30 秒，展示成功/失败、耗时及上游 HTTP 状态。支持停止测试，离开模型页也会停止后续请求；已发送的请求仍可能计费。
|
||||||
|
|
||||||
|
测试属于管理员诊断，使用已保存的渠道密钥和模型映射，不创建下游 Key，不计入正式调用日志及统计，也不验证某个下游 Key 的授权。结果仅保留在当前页面，不返回上游响应正文；成功表示接口返回有效协议响应，不代表回答质量、所有协议转换或流式能力均通过验证。测试可能产生上游 Token 费用，一键测试开始前会要求确认。
|
||||||
|
|
||||||
|
- POST /v1/chat/completions — OpenAI Chat Completions。
|
||||||
|
- POST /v1/responses — OpenAI Responses。
|
||||||
|
- POST /v1/messages — Anthropic Messages，支持 x-api-key。同协议透传时，客户端提供的 `anthropic-version`/`anthropic-beta` 只有通过校验（短且仅含字母数字与 `-`、`_`、`.`、`,`、空格的 token 列表）才会转发上游，其余一律丢弃并由服务端默认版本兜底。
|
||||||
|
- GET /v1/models — 当前 Key 有权限访问的已配置模型别名。
|
||||||
|
|
||||||
|
同协议保留原生字段；跨协议转换通用文本、图片、工具调用、工具结果及 SSE。Responses 输出中的 phase 元数据不会阻断转换；转换为 Chat/Claude 时省略 reasoning 项及相关推理事件，不将其混入回答正文，不伪造 Claude thinking 签名，正文、工具调用和用量继续正常转换。需要保留原生推理内容时应使用 Responses 接口。其它厂商独有的服务端工具、思考块、会话引用等不保证可跨协议转换，不支持的内容应明确报错。强制工具选择是否可用取决于上游支持，路由器不会静默降级为自动选择。编程代理建议优先选择其原生协议渠道。此项目的本地自动化测试不等于所有 Claude Code/Codex 版本的真实上游兼容认证；使用自己的渠道后仍需联调。
|
||||||
|
|
||||||
|
跨协议转为 Anthropic 时必须提供 token 上限（Chat 的 `max_tokens` / `max_completion_tokens`，或 Responses 的 `max_output_tokens`）；不会猜测额度。详细限制见 [协议说明](internal/protocol/README.md)。不提供 Responses 的服务端会话存储、WebSocket 或 Anthropic token-counting 扩展端点。
|
||||||
|
|
||||||
|
## 按 Key 统计
|
||||||
|
|
||||||
|
统计包含请求数、成功/失败数、输入、输出、缓存读取、缓存写入 Token，支持按 Key 和时间范围筛选，以及按天、渠道/模型汇总。调用记录保存请求模型与真实模型，不保存提示词、响应正文或工具参数。
|
||||||
|
|
||||||
|
输入 Token 统一为包含缓存的总输入，缓存读取/写入为细分项，不能再重复相加。OpenAI 的缓存通常包含在输入内；Anthropic 普通输入与缓存字段相加得到总输入。缺失 usage 标记未知；流中断时已收到的用量可能不完整。撤销 Key 不删除历史记录。
|
||||||
|
|
||||||
|
后端统一保存 UTC 时间。页面记录、Key 创建/到期时间按浏览器本地时区展示；时间筛选、今天/近 7/30 天和有效期输入按浏览器本地时间解释，并转换为 UTC 提交。无需时区配置。每日用量也按浏览器时区的自然日汇总：前端自动传递 IANA 时区名称，后端仅在查询时调整分组边界，支持夏令时且不改变 UTC 存储。统计 API 的 `timezone` 参数省略时默认 UTC，非法时区返回 400。
|
||||||
|
|
||||||
|
汇总仅累计完整已知用量，未知/中断请求另行计数。统计由数据库聚合，不再将全部原始日志载入应用内存；查询有超时及结果规模限制，大量历史数据建议按时间范围筛选。可通过 `storage.log_retention_days` 显式开启日志清理，默认关闭。Key 创建后支持停用与撤销，修改授权需重新创建 Key。
|
||||||
|
|
||||||
|
### 异步日志持久化
|
||||||
|
|
||||||
|
模型请求完成后将元数据交给后台写入队列，不再等待该条日志的数据库提交。队列最多 1024 条，后台每批最多 100 条，或每 100ms 提交一次事务。统计页面采用最终一致性，刚完成的请求可能稍后才出现；数据库繁忙时延迟可能超过 100ms。队列满时请求处理会等待可用位置，不静默丢弃日志。
|
||||||
|
|
||||||
|
正常退出会先等待在途请求结束，再排空日志队列后关闭数据库。写入失败进行有限重试，并输出不含请求内容的错误提示；持续数据库故障仍可能导致日志丢失，不能保证恰好一次持久化。进程崩溃、SIGKILL 或断电时，队列及尚未提交批次可能丢失。此队列是内存缓冲，不是持久化消息队列。
|
||||||
|
|
||||||
|
## 开发与验证
|
||||||
|
|
||||||
|
```sh
|
||||||
|
make test
|
||||||
|
# 前端开发服务器（后端须先启动）
|
||||||
|
cd web && npm run dev
|
||||||
|
```
|
||||||
|
|
||||||
|
管理写接口要求同源请求和 `X-Requested-With: ai-router` 请求头；不要对管理接口开放宽泛 CORS。生产构建使用内嵌前端，不需要单独启动 Vite。
|
||||||
|
|
||||||
|
## 生产安全配置
|
||||||
|
|
||||||
|
- 管理台建议仅通过 VPN 或 IP 白名单访问。公网入口必须使用 HTTPS，设置正确的 `server.public_origin`，并保持后端端口不直接对公网开放。
|
||||||
|
- `server.trusted_proxies` 默认是空列表，不信任任何客户端转发头。使用反向代理时，填写**应用实际看到的代理 IP/CIDR**，例如宿主机本地代理可用 `["127.0.0.1/32", "::1/128"]`；Docker 场景须按实际网络填写，不能照抄。代理必须覆盖或正确追加 `X-Forwarded-For`。应用从右向左剥离可信代理，按首个非可信地址执行登录限流；未配置时仍按直接连接地址限流，多个用户可能共享额度。不要配置 `0.0.0.0/0` 或 `::/0`。
|
||||||
|
- `limits.requests_per_minute`、`limits.per_key_concurrency`、`limits.global_concurrency` 分别控制每个 Key 的请求速率、并发数和全部 Key 的总并发，默认是 120、8、64，必须为正整数。超过限制返回 `429` 和 `Retry-After`。这些限制保存在单进程内存中，重启会重置，多副本需要网关侧统一限流。
|
||||||
|
- 登录限速在**校验密码之前**执行：同一 IP 每分钟 10 次失败、或全部 IP 合计每分钟 300 次失败后，锁定期间**一切登录尝试（包括正确密码）统一返回 429** 且不做密码比对——响应不泄露候选密码的正确性，爆破速率被硬性封顶（每 IP 10 次/分、全局 300 次/分）。代价是攻击进行中管理员也可能收到 429：已有会话（24 小时有效）与数据面不受影响，稍候重试或换网络即可，这也是建议管理面仅经 VPN 或 IP 白名单访问的原因之一。成功登录不消耗失败配额，并会清空本 IP 的失败计数，自己的手误不会累积锁定。
|
||||||
|
- `/v1` 数据面在鉴权**之前**按 IP 对鉴权失败限速：同一 IP 每分钟 60 次「无效 Key」之后，该 IP 的后续请求在查询数据库前即返回 429，用于保护单一 SQLite 连接不被无效 Key 洪泛挤占。注意该门限发生在 Key 校验之前，因此与管理员不同出口共享 IP 的合法 Key 也可能被连坐至多一分钟。
|
||||||
|
- 请求体读取限时 30 秒；登录请求体上限 16 KiB，模型请求体上限 16 MiB（超限返回 413，无法读取返回 400）。上游模型请求最长 10 分钟。反向代理也应设置请求体读取超时、连接数和请求大小限制，同时关闭 SSE 响应缓冲。
|
||||||
|
- 请求与响应的 JSON 解析有节点数上限（262144 个 token）：由海量微型节点构成的敌意 JSON 会在解码前被拒绝，防止内存放大耗尽进程；正常的对话请求与响应远低于该上限。
|
||||||
|
- SSE 单事件最多 16 MiB，跨协议累计保留内容最多 16 MiB、块/索引最多 4096 个，累计 usage 最多 16 KiB；同协议透传模式另有整条流 64 MiB 总量上限，超限都会中断流。所有下游响应的单次写入/刷新限时 30 秒，防止停止读取的客户端长期占用并发槽。上游请求不发送 `Accept-Encoding: gzip`（禁用透明解压），避免被控上游以压缩流消耗代理 CPU。
|
||||||
|
- 所有响应默认携带 `Content-Security-Policy`（`default-src 'none'` + 按需放开 connect/script/style，管理台脚本全部同源、无内联 script）、`X-Frame-Options: DENY`、`Referrer-Policy: no-referrer` 与 `X-Content-Type-Options: nosniff`；反代不应移除这些头。
|
||||||
|
- `storage.log_retention_days` 默认 0，即不自动删除历史记录。需要自动清理时先备份，再设置正数（例如 90）；**启用后会分批删除早于保留期限的日志，历史统计也会相应减少**。删除并不保证 SQLite 文件立即缩小，仍需监控数据卷空间。
|
||||||
|
- 请求限流不是费用硬额度。仍需在上游供应商设置预算、费用告警和密钥权限，并在密钥泄露时及时撤销。上游渠道只配置可信目标，公网目标使用 HTTPS；必要时通过网络出口规则禁止访问内网和云元数据地址。
|
||||||
|
- 上游密钥保存在数据库中。配置文件、数据卷和备份应限制访问权限；备份应加密，不要将开发输出或真实配置上传到公开仓库。
|
||||||
@@ -0,0 +1,35 @@
|
|||||||
|
services:
|
||||||
|
ai-router:
|
||||||
|
build:
|
||||||
|
context: .
|
||||||
|
image: ai-router:local
|
||||||
|
restart: unless-stopped
|
||||||
|
ports:
|
||||||
|
- "127.0.0.1:${AI_ROUTER_PORT:-8080}:8080"
|
||||||
|
volumes:
|
||||||
|
- type: bind
|
||||||
|
source: ./config.yaml
|
||||||
|
target: /app/config.yaml
|
||||||
|
read_only: true
|
||||||
|
bind:
|
||||||
|
create_host_path: false
|
||||||
|
- router-data:/app/data
|
||||||
|
security_opt:
|
||||||
|
- no-new-privileges:true
|
||||||
|
cap_drop:
|
||||||
|
- ALL
|
||||||
|
stop_grace_period: 30s
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD", "wget", "-q", "-O", "/dev/null", "http://127.0.0.1:8080/admin/api/session"]
|
||||||
|
interval: 30s
|
||||||
|
timeout: 5s
|
||||||
|
start_period: 10s
|
||||||
|
retries: 3
|
||||||
|
logging:
|
||||||
|
driver: json-file
|
||||||
|
options:
|
||||||
|
max-size: "10m"
|
||||||
|
max-file: "3"
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
router-data:
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
server:
|
||||||
|
# 容器内必须监听所有接口；Compose 默认仅发布到宿主机回环地址。
|
||||||
|
listen: "0.0.0.0:8080"
|
||||||
|
# 使用 HTTPS 反向代理时填写实际访问 origin，不含路径或末尾斜线。
|
||||||
|
# public_origin: "https://router.example.com"
|
||||||
|
# 填实际代理的地址/CIDR，不要信任整个公网或不受控的 Docker 网络。
|
||||||
|
trusted_proxies: []
|
||||||
|
admin:
|
||||||
|
# 启动前必须替换，至少 12 字符；建议使用 32 字节随机值。
|
||||||
|
password: "CHANGE-ME-use-a-long-random-password"
|
||||||
|
storage:
|
||||||
|
path: "/app/data/router.db"
|
||||||
|
# 0 不自动删除；启用正数前备份，过期日志会被分批删除。
|
||||||
|
log_retention_days: 0
|
||||||
|
limits:
|
||||||
|
requests_per_minute: 120
|
||||||
|
per_key_concurrency: 8
|
||||||
|
global_concurrency: 64
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
server:
|
||||||
|
listen: "127.0.0.1:8080"
|
||||||
|
# HTTPS 反向代理部署时填写浏览器访问的 origin，不含路径或末尾斜线。
|
||||||
|
# public_origin: "https://router.example.com"
|
||||||
|
# 仅填写实际反向代理的地址/CIDR；默认不信任任何转发头。
|
||||||
|
trusted_proxies: []
|
||||||
|
admin:
|
||||||
|
password: "CHANGE-ME-use-a-long-random-password"
|
||||||
|
storage:
|
||||||
|
path: "./data/router.db"
|
||||||
|
# 0 禁用自动删除；设为 90 将分批删除超过 90 天的日志，请先备份。
|
||||||
|
log_retention_days: 0
|
||||||
|
limits:
|
||||||
|
requests_per_minute: 120
|
||||||
|
per_key_concurrency: 8
|
||||||
|
global_concurrency: 64
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
module ai-router
|
||||||
|
|
||||||
|
go 1.26.0
|
||||||
|
|
||||||
|
require (
|
||||||
|
gopkg.in/yaml.v3 v3.0.1
|
||||||
|
modernc.org/sqlite v1.57.0
|
||||||
|
)
|
||||||
|
|
||||||
|
require (
|
||||||
|
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||||
|
github.com/google/uuid v1.6.0 // indirect
|
||||||
|
github.com/mattn/go-isatty v0.0.24 // indirect
|
||||||
|
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||||
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||||
|
golang.org/x/sys v0.47.0 // indirect
|
||||||
|
modernc.org/libc v1.74.4 // indirect
|
||||||
|
modernc.org/mathutil v1.7.1 // indirect
|
||||||
|
modernc.org/memory v1.11.0 // indirect
|
||||||
|
)
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||||
|
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||||
|
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo=
|
||||||
|
github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk=
|
||||||
|
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||||
|
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||||
|
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||||
|
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||||
|
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
|
||||||
|
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
|
||||||
|
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||||
|
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||||
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||||
|
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||||
|
golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ=
|
||||||
|
golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0=
|
||||||
|
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
|
||||||
|
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||||
|
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||||
|
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||||
|
golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q=
|
||||||
|
golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA=
|
||||||
|
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||||
|
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||||
|
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||||
|
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||||
|
modernc.org/cc/v4 v4.29.1 h1:MKgdCV3WykTSPqpVrnxdEDS0HEd2FHpKZDzxzU5LyeI=
|
||||||
|
modernc.org/cc/v4 v4.29.1/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI=
|
||||||
|
modernc.org/ccgo/v4 v4.34.6 h1:sBgfIwyN0TQ9C5hwIeuqyeAKyMWnbvj2fvpF4L11uzU=
|
||||||
|
modernc.org/ccgo/v4 v4.34.6/go.mod h1:SZ8YcN9NG7XVsQYdm6jYBvi8PQP1qi+kqB6OhjqI3Fk=
|
||||||
|
modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM=
|
||||||
|
modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU=
|
||||||
|
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||||
|
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||||
|
modernc.org/gc/v3 v3.1.4 h1:2g65LGVSmFQrXeITAw97x7hCRvZFcyE1uDP+7Vng7JI=
|
||||||
|
modernc.org/gc/v3 v3.1.4/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||||
|
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||||
|
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||||
|
modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k=
|
||||||
|
modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co=
|
||||||
|
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||||
|
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||||
|
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||||
|
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||||
|
modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg=
|
||||||
|
modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||||
|
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||||
|
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||||
|
modernc.org/sqlite v1.57.0 h1:qNQP6xnx5M0ISNtlnxoOX0+cD5bJ0/gr9aMmndFczzg=
|
||||||
|
modernc.org/sqlite v1.57.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ=
|
||||||
|
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||||
|
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||||
|
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||||
|
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/netip"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"gopkg.in/yaml.v3"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Config struct {
|
||||||
|
Server struct {
|
||||||
|
TrustedProxies []string `yaml:"trusted_proxies"`
|
||||||
|
Listen string `yaml:"listen"`
|
||||||
|
PublicOrigin string `yaml:"public_origin"`
|
||||||
|
} `yaml:"server"`
|
||||||
|
Admin struct {
|
||||||
|
Password string `yaml:"password"`
|
||||||
|
} `yaml:"admin"`
|
||||||
|
Storage struct {
|
||||||
|
LogRetentionDays int `yaml:"log_retention_days"`
|
||||||
|
Path string `yaml:"path"`
|
||||||
|
} `yaml:"storage"`
|
||||||
|
Limits struct {
|
||||||
|
RequestsPerMinute int `yaml:"requests_per_minute"`
|
||||||
|
PerKeyConcurrency int `yaml:"per_key_concurrency"`
|
||||||
|
GlobalConcurrency int `yaml:"global_concurrency"`
|
||||||
|
} `yaml:"limits"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func Load(path string) (Config, error) {
|
||||||
|
var c Config
|
||||||
|
c.Limits.RequestsPerMinute = 120
|
||||||
|
c.Limits.PerKeyConcurrency = 8
|
||||||
|
c.Limits.GlobalConcurrency = 64
|
||||||
|
c.Server.Listen = "127.0.0.1:8080"
|
||||||
|
c.Storage.Path = "./data/router.db"
|
||||||
|
f, err := os.Open(path)
|
||||||
|
if err != nil {
|
||||||
|
return c, err
|
||||||
|
}
|
||||||
|
defer f.Close()
|
||||||
|
dec := yaml.NewDecoder(f)
|
||||||
|
dec.KnownFields(true)
|
||||||
|
if err = dec.Decode(&c); err != nil {
|
||||||
|
return c, fmt.Errorf("invalid config: %w", err)
|
||||||
|
}
|
||||||
|
var extra any
|
||||||
|
if err = dec.Decode(&extra); err != io.EOF {
|
||||||
|
return c, fmt.Errorf("config must contain exactly one YAML document")
|
||||||
|
}
|
||||||
|
if len(c.Admin.Password) < 12 || strings.HasPrefix(c.Admin.Password, "CHANGE-ME") {
|
||||||
|
return c, fmt.Errorf("admin.password must be a non-placeholder password of at least 12 characters")
|
||||||
|
}
|
||||||
|
if c.Server.Listen == "" || c.Storage.Path == "" {
|
||||||
|
return c, fmt.Errorf("server.listen and storage.path must not be empty")
|
||||||
|
}
|
||||||
|
if c.Limits.RequestsPerMinute <= 0 || c.Limits.PerKeyConcurrency <= 0 || c.Limits.GlobalConcurrency <= 0 {
|
||||||
|
return c, fmt.Errorf("limits must be positive integers")
|
||||||
|
}
|
||||||
|
if c.Storage.LogRetentionDays < 0 {
|
||||||
|
return c, fmt.Errorf("storage.log_retention_days must be nonnegative")
|
||||||
|
}
|
||||||
|
for _, cidr := range c.Server.TrustedProxies {
|
||||||
|
p, err := netip.ParsePrefix(cidr)
|
||||||
|
if err != nil {
|
||||||
|
return c, fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
|
||||||
|
}
|
||||||
|
if p.Bits() == 0 {
|
||||||
|
return c, fmt.Errorf("trusted proxy CIDR must not trust all addresses")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return c, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
package config
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestLoad(t *testing.T) {
|
||||||
|
for _, tt := range []struct {
|
||||||
|
name, body string
|
||||||
|
wantErr bool
|
||||||
|
}{
|
||||||
|
{"valid", "admin:\n password: a-long-test-password\n", false},
|
||||||
|
{"short", "admin:\n password: short\n", true},
|
||||||
|
{"multiple documents", "admin:\n password: a-long-test-password\n---\nadmin:\n password: ignored-password\n", true},
|
||||||
|
{"placeholder", "admin:\n password: CHANGE-ME-use-a-long-random-password\n", true},
|
||||||
|
{"unknown", "admin:\n password: a-long-test-password\nunknown: true\n", true},
|
||||||
|
} {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
path := filepath.Join(t.TempDir(), "config.yaml")
|
||||||
|
if err := os.WriteFile(path, []byte(tt.body), 0600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
c, err := Load(path)
|
||||||
|
if (err != nil) != tt.wantErr {
|
||||||
|
t.Fatalf("error=%v, wantErr=%v", err, tt.wantErr)
|
||||||
|
}
|
||||||
|
if err == nil && (c.Server.Listen != "127.0.0.1:8080" || c.Storage.Path != "./data/router.db") {
|
||||||
|
t.Fatal("defaults missing")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSecurityConfig(t *testing.T) {
|
||||||
|
for _, tt := range []struct {
|
||||||
|
body string
|
||||||
|
invalid bool
|
||||||
|
}{
|
||||||
|
{"", false},
|
||||||
|
{"limits:\n requests_per_minute: 0\n", true},
|
||||||
|
{"limits:\n per_key_concurrency: -1\n", true},
|
||||||
|
{"limits:\n global_concurrency: 0\n", true},
|
||||||
|
{"limits:\n requests_per_minute: nope\n", true},
|
||||||
|
{"server:\n trusted_proxies: [not-a-cidr]\n", true},
|
||||||
|
{"server:\n trusted_proxies: [0.0.0.0/0]\n", true},
|
||||||
|
{"server:\n trusted_proxies: ['::/0']\n", true},
|
||||||
|
{"server:\n trusted_proxies: [10.0.0.0/8, '::1/128']\n", false},
|
||||||
|
{"storage:\n log_retention_days: -1\n", true},
|
||||||
|
} {
|
||||||
|
t.Run(tt.body, func(t *testing.T) {
|
||||||
|
p := filepath.Join(t.TempDir(), "config.yaml")
|
||||||
|
if err := os.WriteFile(p, []byte("admin:\n password: a-long-test-password\n"+tt.body), 0600); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
c, err := Load(p)
|
||||||
|
if (err != nil) != tt.invalid {
|
||||||
|
t.Fatalf("error=%v", err)
|
||||||
|
}
|
||||||
|
if err == nil && (c.Limits.RequestsPerMinute != 120 || c.Limits.PerKeyConcurrency != 8 || c.Limits.GlobalConcurrency != 64 || c.Storage.LogRetentionDays != 0) {
|
||||||
|
t.Fatal("security defaults missing")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,71 @@
|
|||||||
|
# Protocol conversion
|
||||||
|
|
||||||
|
Public entry points are `ConvertRequest`, `ConvertResponse`, `Stream`, and
|
||||||
|
`ExtractUsage`. No external dependencies are needed. The service should bypass
|
||||||
|
conversion for same-protocol non-streaming traffic. `Stream` preserves original
|
||||||
|
SSE bytes for same-protocol traffic, including comments and CRLF, and returns
|
||||||
|
upstream errors instead of silently swallowing them.
|
||||||
|
|
||||||
|
Cross-protocol support includes text, user URL/base64 images, system messages,
|
||||||
|
function schemas and selection, multiple function calls, tool results, multi-turn
|
||||||
|
history, and incremental tool argument JSON. Responses output item/content part/
|
||||||
|
argument start, delta, done and terminal events are emitted. Anthropic content
|
||||||
|
block and message lifecycles are emitted. Each write flushes when supported.
|
||||||
|
|
||||||
|
## Deliberate limits
|
||||||
|
|
||||||
|
- Anthropic output thinking/redacted_thinking blocks and their thinking/signature
|
||||||
|
stream deltas are omitted when converting to Chat/Responses, never appended to
|
||||||
|
visible text. Block indices and lifecycle are checked. Native Anthropic output
|
||||||
|
bytes remain unchanged; request-side thinking is not translated.
|
||||||
|
|
||||||
|
- Chat output message/delta reasoning_content (string or null) is omitted only
|
||||||
|
during cross-protocol conversion; it is never appended to visible text or
|
||||||
|
fabricated as Anthropic thinking. Native Chat responses/SSE remain unchanged.
|
||||||
|
Optional indices on complete Chat tool calls must match their array positions.
|
||||||
|
|
||||||
|
- Responses output message `phase` is accepted as metadata. Cross-protocol
|
||||||
|
conversion deliberately omits Responses output reasoning, summaries and
|
||||||
|
encrypted content, including their recognized stream events. They are never
|
||||||
|
appended to visible text or converted to Anthropic thinking/signatures.
|
||||||
|
Same-protocol responses and SSE retain these bytes unchanged. Reasoning
|
||||||
|
item/part indices, identities and lifecycles are checked; unrelated unknown
|
||||||
|
content and event types remain errors. Empty output-text logprobs arrays
|
||||||
|
are accepted, but actual logprobs remain unsupported.
|
||||||
|
- Only one choice and function tools are supported. Other reasoning/thinking blocks,
|
||||||
|
audio/video/documents, hosted tools, citations, refusal content, logprobs,
|
||||||
|
structured-output modes, prompt-cache controls and Responses server-side
|
||||||
|
conversation state are rejected rather than approximated. Unknown request
|
||||||
|
fields and unknown stream event types are rejected.
|
||||||
|
- Anthropic targets require an explicit token limit. Developer roles,
|
||||||
|
mid-conversation system messages, strict tool schemas and explicit image
|
||||||
|
detail levels cannot be translated to Anthropic. The implementation currently
|
||||||
|
conservatively rejects explicit image detail for all cross-protocol pairs.
|
||||||
|
- Tool results must be text and cannot carry Anthropic's `is_error=true`.
|
||||||
|
Stop sequences cannot target Responses; stop-sequence termination cannot
|
||||||
|
be represented losslessly across protocols.
|
||||||
|
- Responses streaming supports one output-text content part per message.
|
||||||
|
Final content must agree with deltas; nonincremental/final-only variants
|
||||||
|
are rejected. Tool identity must be complete before argument streaming;
|
||||||
|
later ID/name fragments are rejected. Parallel tool argument fragments are
|
||||||
|
tracked independently.
|
||||||
|
- `Usage.Input` includes cache-read and cache-write tokens. Anthropic input
|
||||||
|
counters exclude these, so they are added on extraction and subtracted on
|
||||||
|
encoding. `Usage.Raw` holds upstream usage JSON; Anthropic stream snapshots
|
||||||
|
are merged so final output-only usage does not erase initial input counts.
|
||||||
|
Unknown usage remains `Known=false`; Anthropic wire messages still carry
|
||||||
|
required zero counters until actual usage is available. OpenAI wire schemas
|
||||||
|
have no separate Anthropic cache-write counter; it remains in returned Usage.
|
||||||
|
- SSE events are limited to 16 MiB; generated content is retained for final
|
||||||
|
Responses snapshots and argument validation. LF and CRLF framing is supported
|
||||||
|
(not standalone CR). Same-protocol streams allow EOF without a terminal marker;
|
||||||
|
cross-protocol streams report truncation. Terminal markers finish immediately.
|
||||||
|
- Context is checked between reads and writes. A generic `io.Reader` cannot be
|
||||||
|
interrupted safely: the service must bind upstream HTTP requests to the same
|
||||||
|
context or close their body on cancellation. Writers/readers remain caller-owned.
|
||||||
|
- Response transport metadata and echoed Responses request settings are not
|
||||||
|
round-tripped. Generated IDs use cryptographic randomness if upstream omits
|
||||||
|
an ID. No external provider or SDK integration tests have been performed.
|
||||||
|
|
||||||
|
Run `go test ./internal/protocol`, `go test -race -cover ./internal/protocol`,
|
||||||
|
and `go vet ./internal/protocol`.
|
||||||
@@ -0,0 +1,101 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAnthropicThinkingOutput(t *testing.T) {
|
||||||
|
for _, tool := range []bool{false, true} {
|
||||||
|
visible := `{"type":"text","text":"OK"}`
|
||||||
|
initial := `{"type":"text","text":""}`
|
||||||
|
delta := `{"type":"text_delta","text":"OK"}`
|
||||||
|
finish := "end_turn"
|
||||||
|
if tool {
|
||||||
|
visible = `{"type":"tool_use","id":"call_1","name":"weather","input":{}}`
|
||||||
|
initial = visible
|
||||||
|
delta = `{"type":"input_json_delta","partial_json":"{}"}`
|
||||||
|
finish = "tool_use"
|
||||||
|
}
|
||||||
|
hidden := `{"type":"thinking","thinking":"PRIVATE","signature":"SIGNATURE"}`
|
||||||
|
redacted := `{"type":"redacted_thinking","data":"CIPHER"}`
|
||||||
|
body := `{"id":"m","type":"message","role":"assistant","model":"m","content":[` + hidden + "," + redacted + "," + visible + `],"stop_reason":"` + finish + `","usage":{"input_tokens":9,"output_tokens":7}}`
|
||||||
|
events := []string{
|
||||||
|
`{"type":"message_start","message":{"id":"m","content":[],"usage":{"input_tokens":9,"output_tokens":0}}}`,
|
||||||
|
`{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}`,
|
||||||
|
`{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"PRIVATE"}}`,
|
||||||
|
`{"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"SIGNATURE"}}`,
|
||||||
|
`{"type":"content_block_stop","index":0}`,
|
||||||
|
`{"type":"content_block_start","index":1,"content_block":` + redacted + "}",
|
||||||
|
`{"type":"content_block_stop","index":1}`,
|
||||||
|
`{"type":"content_block_start","index":2,"content_block":` + initial + "}",
|
||||||
|
`{"type":"content_block_delta","index":2,"delta":` + delta + "}",
|
||||||
|
`{"type":"content_block_stop","index":2}`,
|
||||||
|
`{"type":"message_delta","delta":{"stop_reason":"` + finish + `"},"usage":{"output_tokens":7}}`,
|
||||||
|
`{"type":"message_stop"}`,
|
||||||
|
}
|
||||||
|
src := reasoningSSE(events...)
|
||||||
|
for _, target := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
out, err := ConvertResponse(Anthropic, target, []byte(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var dst bytes.Buffer
|
||||||
|
u, err := Stream(context.Background(), &dst, strings.NewReader(src), Anthropic, target, "m")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !u.Known || u.Input != 9 || u.Output != 7 {
|
||||||
|
t.Fatalf("usage %+v", u)
|
||||||
|
}
|
||||||
|
if target == Anthropic {
|
||||||
|
if string(out) != body || dst.String() != src {
|
||||||
|
t.Fatal("native bytes changed")
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, result := range []string{string(out), dst.String()} {
|
||||||
|
for _, secret := range []string{"PRIVATE", "SIGNATURE", "CIPHER", "thinking"} {
|
||||||
|
if strings.Contains(result, secret) {
|
||||||
|
t.Fatal("thinking leaked")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
want := "OK"
|
||||||
|
if tool {
|
||||||
|
want = "weather"
|
||||||
|
}
|
||||||
|
if !strings.Contains(result, want) {
|
||||||
|
t.Fatal("lost visible output")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
terminal := "[DONE]"
|
||||||
|
if target == Responses {
|
||||||
|
terminal = "response.completed"
|
||||||
|
}
|
||||||
|
if !strings.Contains(dst.String(), terminal) {
|
||||||
|
t.Fatal("no terminal")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, bad := range []string{
|
||||||
|
strings.Replace(src, `{"type":"content_block_stop","index":0}`, `{"type":"ping"}`, 1),
|
||||||
|
strings.Replace(src, `"type":"thinking_delta"`, `"type":"text_delta"`, 1),
|
||||||
|
strings.Replace(src, `"index":2,"content_block"`, `"index":0,"content_block"`, 1),
|
||||||
|
} {
|
||||||
|
var dst bytes.Buffer
|
||||||
|
if _, err := Stream(context.Background(), &dst, strings.NewReader(bad), Anthropic, Chat, "m"); err == nil {
|
||||||
|
t.Fatal("accepted invalid thinking lifecycle")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, block := range []object{
|
||||||
|
{"type": "thinking", "thinking": true, "signature": ""},
|
||||||
|
{"type": "thinking", "thinking": "", "signature": "", "unknown": true},
|
||||||
|
{"type": "redacted_thinking", "data": nil},
|
||||||
|
} {
|
||||||
|
if _, err := anthropicThinking(block); err == nil {
|
||||||
|
t.Fatal("accepted malformed thinking")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestChatOutputReasoning(t *testing.T) {
|
||||||
|
for _, tool := range []bool{false, true} {
|
||||||
|
message := object{"role": "assistant", "content": "OK", "reasoning_content": "PRIVATE"}
|
||||||
|
delta := object{"content": "OK", "reasoning_content": "PRIVATE"}
|
||||||
|
finish := "stop"
|
||||||
|
if tool {
|
||||||
|
call := object{"index": 0, "id": "call_1", "type": "function", "function": object{"name": "weather", "arguments": "{}"}}
|
||||||
|
message["content"] = nil
|
||||||
|
message["tool_calls"] = []any{call}
|
||||||
|
delta = object{"tool_calls": []any{call}, "reasoning_content": nil}
|
||||||
|
finish = "tool_calls"
|
||||||
|
}
|
||||||
|
usage := object{"prompt_tokens": 9, "completion_tokens": 7}
|
||||||
|
response := object{"id": "c", "object": "chat.completion", "model": "m", "choices": []any{object{"index": 0, "message": message, "finish_reason": finish}}, "usage": usage}
|
||||||
|
body, _ := json.Marshal(response)
|
||||||
|
chunk, _ := json.Marshal(object{"id": "c", "choices": []any{object{"index": 0, "delta": delta, "finish_reason": finish}}, "usage": usage})
|
||||||
|
reasoning, _ := json.Marshal(object{"id": "c", "choices": []any{object{"index": 0, "delta": object{"role": "assistant", "reasoning_content": "PRIVATE"}, "finish_reason": nil}}})
|
||||||
|
sep := string([]byte{10, 10})
|
||||||
|
src := "data: " + string(reasoning) + sep + "data: " + string(chunk) + sep + "data: [DONE]" + sep
|
||||||
|
for _, target := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
out, err := ConvertResponse(Chat, target, body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var dst bytes.Buffer
|
||||||
|
u, err := Stream(context.Background(), &dst, strings.NewReader(src), Chat, target, "m")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !u.Known || u.Input != 9 || u.Output != 7 {
|
||||||
|
t.Fatalf("usage: %+v", u)
|
||||||
|
}
|
||||||
|
if target == Chat {
|
||||||
|
if !bytes.Equal(out, body) || dst.String() != src {
|
||||||
|
t.Fatal("native bytes changed")
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, text := range []string{string(out), dst.String()} {
|
||||||
|
if strings.Contains(text, "PRIVATE") || strings.Contains(text, "reasoning_content") || strings.Contains(text, "thinking") {
|
||||||
|
t.Fatal("reasoning leaked")
|
||||||
|
}
|
||||||
|
want := "OK"
|
||||||
|
if tool {
|
||||||
|
want = "weather"
|
||||||
|
}
|
||||||
|
if !strings.Contains(text, want) {
|
||||||
|
t.Fatal("lost output")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
terminal := "response.completed"
|
||||||
|
if target == Anthropic {
|
||||||
|
terminal = "message_stop"
|
||||||
|
}
|
||||||
|
if !strings.Contains(dst.String(), terminal) {
|
||||||
|
t.Fatal("no terminal")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestChatOutputMetadataValidation(t *testing.T) {
|
||||||
|
for _, val := range []any{nil, "", "private"} {
|
||||||
|
if _, err := chatOutput(object{"reasoning_content": val}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, val := range []any{true, object{}, []any{}} {
|
||||||
|
if _, err := chatOutput(object{"reasoning_content": val}); err == nil {
|
||||||
|
t.Fatal("accepted malformed reasoning")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, index := range []any{json.Number("1"), json.Number("-1"), json.Number("0.5"), "0", nil} {
|
||||||
|
if _, err := chatResponseMessage(object{"tool_calls": []any{object{"index": index}}}); err == nil {
|
||||||
|
t.Fatal("accepted malformed index")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
body := []byte(`{"id":"c","choices":[{"message":{"role":"assistant","content":"OK","reasoning_content":"hidden","unknown_field":true},"finish_reason":"stop"}]}`)
|
||||||
|
if _, err := ConvertResponse(Chat, Responses, body); err == nil {
|
||||||
|
t.Fatal("accepted unknown field")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,516 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newID() string { return fmt.Sprintf("resp_%x", randomID()) }
|
||||||
|
func randomID() []byte {
|
||||||
|
b := make([]byte, 16)
|
||||||
|
if _, e := rand.Read(b); e != nil {
|
||||||
|
panic(e)
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func contentBlock(k Kind, p part, role string) (object, error) {
|
||||||
|
switch p.typ {
|
||||||
|
case "text":
|
||||||
|
typ := "text"
|
||||||
|
if k == Responses {
|
||||||
|
typ = "input_text"
|
||||||
|
if role == "assistant" {
|
||||||
|
typ = "output_text"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return object{"type": typ, "text": p.text}, nil
|
||||||
|
case "image":
|
||||||
|
if role != "user" {
|
||||||
|
return nil, fmt.Errorf("images only supported in user messages")
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
return object{"type": "image_url", "image_url": object{"url": p.url}}, nil
|
||||||
|
case Responses:
|
||||||
|
return object{"type": "input_image", "image_url": p.url}, nil
|
||||||
|
default:
|
||||||
|
src := object{"type": "url", "url": p.url}
|
||||||
|
if strings.HasPrefix(p.url, "data:") {
|
||||||
|
a, b, ok := strings.Cut(strings.TrimPrefix(p.url, "data:"), ";base64,")
|
||||||
|
if !ok || a == "" || b == "" {
|
||||||
|
return nil, fmt.Errorf("unsupported image data URL")
|
||||||
|
}
|
||||||
|
src = object{"type": "base64", "media_type": a, "data": b}
|
||||||
|
}
|
||||||
|
return object{"type": "image", "source": src}, nil
|
||||||
|
}
|
||||||
|
case "call":
|
||||||
|
if role != "assistant" {
|
||||||
|
return nil, fmt.Errorf("tool call outside assistant message")
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
return object{"id": p.id, "type": "function", "function": object{"name": p.name, "arguments": p.args}}, nil
|
||||||
|
case Responses:
|
||||||
|
return object{"type": "function_call", "call_id": p.id, "name": p.name, "arguments": p.args}, nil
|
||||||
|
default:
|
||||||
|
args, e := decode([]byte(p.args))
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
return object{"type": "tool_use", "id": p.id, "name": p.name, "input": args}, nil
|
||||||
|
}
|
||||||
|
case "result":
|
||||||
|
if role != "tool" && role != "user" {
|
||||||
|
return nil, fmt.Errorf("tool result outside user/tool message")
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
return object{"role": "tool", "tool_call_id": p.id, "content": p.text}, nil
|
||||||
|
case Responses:
|
||||||
|
return object{"type": "function_call_output", "call_id": p.id, "output": p.text}, nil
|
||||||
|
default:
|
||||||
|
return object{"type": "tool_result", "tool_use_id": p.id, "content": p.text}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("unsupported part")
|
||||||
|
}
|
||||||
|
func writeRequest(k Kind, r request) (object, error) {
|
||||||
|
out := object{}
|
||||||
|
for key, val := range r.options {
|
||||||
|
switch key {
|
||||||
|
case "max":
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
key = "max_completion_tokens"
|
||||||
|
case Responses:
|
||||||
|
key = "max_output_tokens"
|
||||||
|
case Anthropic:
|
||||||
|
key = "max_tokens"
|
||||||
|
}
|
||||||
|
case "stop":
|
||||||
|
if k == Responses {
|
||||||
|
return nil, fmt.Errorf("Responses does not support stop sequences")
|
||||||
|
}
|
||||||
|
if k == Anthropic {
|
||||||
|
key = "stop_sequences"
|
||||||
|
}
|
||||||
|
case "parallel_tool_calls":
|
||||||
|
if k == Anthropic {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out[key] = val
|
||||||
|
}
|
||||||
|
if k == Anthropic {
|
||||||
|
if _, ok := out["max_tokens"]; !ok {
|
||||||
|
return nil, fmt.Errorf("Anthropic requires explicit max_tokens")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ms := []any{}
|
||||||
|
sys := []any{}
|
||||||
|
seenNonSystem := false
|
||||||
|
for _, m := range r.messages {
|
||||||
|
if k == Anthropic && (m.role == "system" || m.role == "developer") {
|
||||||
|
if m.role == "developer" {
|
||||||
|
return nil, fmt.Errorf("developer role cannot be represented by Anthropic")
|
||||||
|
}
|
||||||
|
if seenNonSystem {
|
||||||
|
return nil, fmt.Errorf("mid-conversation system message cannot be represented by Anthropic")
|
||||||
|
}
|
||||||
|
for _, p := range m.parts {
|
||||||
|
if p.typ != "text" {
|
||||||
|
return nil, fmt.Errorf("non-text system content")
|
||||||
|
}
|
||||||
|
sys = append(sys, object{"type": "text", "text": p.text})
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seenNonSystem = true
|
||||||
|
role := m.role
|
||||||
|
if k == Anthropic && role == "tool" {
|
||||||
|
role = "user"
|
||||||
|
}
|
||||||
|
cs := []any{}
|
||||||
|
calls := []any{}
|
||||||
|
flush := func() {
|
||||||
|
if len(cs) > 0 {
|
||||||
|
ms = append(ms, object{"role": role, "content": cs})
|
||||||
|
cs = []any{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, p := range m.parts {
|
||||||
|
b, e := contentBlock(k, p, m.role)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case k == Chat && p.typ == "call":
|
||||||
|
calls = append(calls, b)
|
||||||
|
case (k == Chat && p.typ == "result") || (k == Responses && (p.typ == "call" || p.typ == "result")):
|
||||||
|
flush()
|
||||||
|
ms = append(ms, b)
|
||||||
|
default:
|
||||||
|
cs = append(cs, b)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Chat && len(calls) > 0 {
|
||||||
|
x := object{"role": role, "tool_calls": calls, "content": nil}
|
||||||
|
if len(cs) > 0 {
|
||||||
|
x["content"] = cs
|
||||||
|
}
|
||||||
|
ms = append(ms, x)
|
||||||
|
} else {
|
||||||
|
flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Anthropic {
|
||||||
|
merged := []any{}
|
||||||
|
for _, v := range ms {
|
||||||
|
m := obj(v)
|
||||||
|
if len(merged) > 0 && obj(merged[len(merged)-1])["role"] == m["role"] {
|
||||||
|
last := obj(merged[len(merged)-1])
|
||||||
|
last["content"] = append(arr(last["content"]), arr(m["content"])...)
|
||||||
|
} else {
|
||||||
|
merged = append(merged, m)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ms = merged
|
||||||
|
if len(sys) > 0 {
|
||||||
|
out["system"] = sys
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Responses {
|
||||||
|
out["input"] = ms
|
||||||
|
out["store"] = false
|
||||||
|
} else {
|
||||||
|
out["messages"] = ms
|
||||||
|
}
|
||||||
|
if len(r.tools) > 0 {
|
||||||
|
ts := []any{}
|
||||||
|
for _, f := range r.tools {
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
ts = append(ts, object{"type": "function", "function": f})
|
||||||
|
case Responses:
|
||||||
|
t := object{"type": "function"}
|
||||||
|
for key, v := range f {
|
||||||
|
t[key] = v
|
||||||
|
}
|
||||||
|
ts = append(ts, t)
|
||||||
|
case Anthropic:
|
||||||
|
t := object{}
|
||||||
|
for key, v := range f {
|
||||||
|
if key == "strict" {
|
||||||
|
if v != false {
|
||||||
|
return nil, fmt.Errorf("strict tool schema cannot be enforced by Anthropic")
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if key == "parameters" {
|
||||||
|
key = "input_schema"
|
||||||
|
}
|
||||||
|
t[key] = v
|
||||||
|
}
|
||||||
|
ts = append(ts, t)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out["tools"] = ts
|
||||||
|
}
|
||||||
|
if r.choice != nil {
|
||||||
|
if s, ok := r.choice.(string); ok {
|
||||||
|
if k == Anthropic {
|
||||||
|
if s == "required" {
|
||||||
|
s = "any"
|
||||||
|
}
|
||||||
|
out["tool_choice"] = object{"type": s}
|
||||||
|
} else {
|
||||||
|
out["tool_choice"] = s
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
name := obj(r.choice)["name"]
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
out["tool_choice"] = object{"type": "function", "function": object{"name": name}}
|
||||||
|
case Responses:
|
||||||
|
out["tool_choice"] = object{"type": "function", "name": name}
|
||||||
|
case Anthropic:
|
||||||
|
out["tool_choice"] = object{"type": "tool", "name": name}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Anthropic {
|
||||||
|
if v, ok := r.options["parallel_tool_calls"]; ok {
|
||||||
|
b, ok := v.(bool)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid parallel_tool_calls")
|
||||||
|
}
|
||||||
|
tc := obj(out["tool_choice"])
|
||||||
|
if tc == nil {
|
||||||
|
tc = object{"type": "auto"}
|
||||||
|
out["tool_choice"] = tc
|
||||||
|
}
|
||||||
|
tc["disable_parallel_tool_use"] = !b
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Chat && out["stream"] == true {
|
||||||
|
out["stream_options"] = object{"include_usage": true}
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type response struct {
|
||||||
|
id, model, finish string
|
||||||
|
parts []part
|
||||||
|
usage Usage
|
||||||
|
}
|
||||||
|
|
||||||
|
func finishReason(k Kind, s string) (string, error) {
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
switch s {
|
||||||
|
case "stop":
|
||||||
|
return "stop", nil
|
||||||
|
case "tool_calls":
|
||||||
|
return "tool", nil
|
||||||
|
case "length":
|
||||||
|
return "length", nil
|
||||||
|
}
|
||||||
|
case Anthropic:
|
||||||
|
switch s {
|
||||||
|
case "end_turn":
|
||||||
|
return "stop", nil
|
||||||
|
case "tool_use":
|
||||||
|
return "tool", nil
|
||||||
|
case "max_tokens":
|
||||||
|
return "length", nil
|
||||||
|
}
|
||||||
|
case Responses:
|
||||||
|
if s == "completed" {
|
||||||
|
return "stop", nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("unsupported finish reason %q", s)
|
||||||
|
}
|
||||||
|
func wireFinish(k Kind, s string) string {
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
switch s {
|
||||||
|
case "tool":
|
||||||
|
return "tool_calls"
|
||||||
|
case "length":
|
||||||
|
return "length"
|
||||||
|
}
|
||||||
|
return "stop"
|
||||||
|
case Anthropic:
|
||||||
|
switch s {
|
||||||
|
case "tool":
|
||||||
|
return "tool_use"
|
||||||
|
case "length":
|
||||||
|
return "max_tokens"
|
||||||
|
}
|
||||||
|
return "end_turn"
|
||||||
|
default:
|
||||||
|
if s == "length" {
|
||||||
|
return "incomplete"
|
||||||
|
}
|
||||||
|
return "completed"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func ConvertResponse(from, to Kind, body []byte) ([]byte, error) {
|
||||||
|
if !valid(from) || !valid(to) {
|
||||||
|
return nil, fmt.Errorf("unknown protocol")
|
||||||
|
}
|
||||||
|
m, e := decode(body)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if from == to {
|
||||||
|
return body, nil
|
||||||
|
}
|
||||||
|
r, e := readResponse(from, m)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
out, e := writeResponse(to, r)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
return json.Marshal(out)
|
||||||
|
}
|
||||||
|
func readResponse(k Kind, m object) (response, error) {
|
||||||
|
r := response{id: str(m["id"]), model: str(m["model"]), usage: usage(k, obj(m["usage"]))}
|
||||||
|
if m["error"] != nil {
|
||||||
|
return r, fmt.Errorf("upstream error: %v", m["error"])
|
||||||
|
}
|
||||||
|
var e error
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
if e = fields(m, "id", "object", "created", "model", "choices", "usage", "system_fingerprint", "service_tier"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
cs := arr(m["choices"])
|
||||||
|
if len(cs) != 1 {
|
||||||
|
return r, fmt.Errorf("expected one choice")
|
||||||
|
}
|
||||||
|
c := obj(cs[0])
|
||||||
|
if e = fields(c, "index", "message", "finish_reason", "logprobs"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if c["logprobs"] != nil {
|
||||||
|
return r, fmt.Errorf("logprobs unsupported")
|
||||||
|
}
|
||||||
|
message, e := chatResponseMessage(obj(c["message"]))
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
mm, e := readMessage(k, message)
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.parts = mm.parts
|
||||||
|
r.finish, e = finishReason(k, str(c["finish_reason"]))
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
case Anthropic:
|
||||||
|
if e = fields(m, "id", "type", "role", "model", "content", "stop_reason", "stop_sequence", "usage"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if m["stop_sequence"] != nil {
|
||||||
|
return r, fmt.Errorf("stop sequence cannot be preserved")
|
||||||
|
}
|
||||||
|
content, e := anthropicOutput(m["content"])
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.parts, e = readContent(k, content)
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.finish, e = finishReason(k, str(m["stop_reason"]))
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
case Responses:
|
||||||
|
// Responses echoes request configuration; only output items carry generated content.
|
||||||
|
for _, v := range arr(m["output"]) {
|
||||||
|
x := obj(v)
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "reasoning":
|
||||||
|
if e = validateReasoning(x); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
case "message":
|
||||||
|
x = responseMessage(x)
|
||||||
|
mm, e := readMessage(k, x)
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.parts = append(r.parts, mm.parts...)
|
||||||
|
case "function_call":
|
||||||
|
if e = fields(x, "type", "id", "call_id", "name", "arguments", "status"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
p, e := callPart(x, "call_id", "arguments")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.parts = append(r.parts, p)
|
||||||
|
default:
|
||||||
|
return r, fmt.Errorf("unsupported output item %q", str(x["type"]))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if str(m["status"]) == "incomplete" {
|
||||||
|
if str(obj(m["incomplete_details"])["reason"]) != "max_output_tokens" {
|
||||||
|
return r, fmt.Errorf("unsupported incomplete reason")
|
||||||
|
}
|
||||||
|
r.finish = "length"
|
||||||
|
} else {
|
||||||
|
r.finish, e = finishReason(k, str(m["status"]))
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
for _, p := range r.parts {
|
||||||
|
if p.typ == "call" {
|
||||||
|
r.finish = "tool"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return r, nil
|
||||||
|
}
|
||||||
|
func writeResponse(k Kind, r response) (object, error) {
|
||||||
|
if r.id == "" {
|
||||||
|
r.id = newID()
|
||||||
|
}
|
||||||
|
out := object{"id": r.id, "model": r.model}
|
||||||
|
cs := []any{}
|
||||||
|
calls := []any{}
|
||||||
|
text := ""
|
||||||
|
for i, p := range r.parts {
|
||||||
|
if p.typ != "text" && p.typ != "call" {
|
||||||
|
return nil, fmt.Errorf("unsupported response content")
|
||||||
|
}
|
||||||
|
b, e := contentBlock(k, p, "assistant")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
if p.typ == "call" {
|
||||||
|
calls = append(calls, b)
|
||||||
|
} else {
|
||||||
|
text += p.text
|
||||||
|
}
|
||||||
|
case Anthropic:
|
||||||
|
cs = append(cs, b)
|
||||||
|
case Responses:
|
||||||
|
if p.typ == "call" {
|
||||||
|
b["id"] = fmt.Sprintf("fc_%s_%d", r.id, i)
|
||||||
|
b["status"] = "completed"
|
||||||
|
cs = append(cs, b)
|
||||||
|
} else {
|
||||||
|
b["annotations"] = []any{}
|
||||||
|
cs = append(cs, object{"type": "message", "id": fmt.Sprintf("msg_%s_%d", r.id, i), "role": "assistant", "status": "completed", "content": []any{b}})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
m := object{"role": "assistant", "content": text}
|
||||||
|
if len(calls) > 0 {
|
||||||
|
m["tool_calls"] = calls
|
||||||
|
if text == "" {
|
||||||
|
m["content"] = nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out["object"] = "chat.completion"
|
||||||
|
out["created"] = 0
|
||||||
|
out["choices"] = []any{object{"index": 0, "message": m, "finish_reason": wireFinish(k, r.finish)}}
|
||||||
|
case Anthropic:
|
||||||
|
out["type"] = "message"
|
||||||
|
out["role"] = "assistant"
|
||||||
|
out["content"] = cs
|
||||||
|
out["stop_reason"] = wireFinish(k, r.finish)
|
||||||
|
out["stop_sequence"] = nil
|
||||||
|
case Responses:
|
||||||
|
out["object"] = "response"
|
||||||
|
out["created_at"] = time.Now().Unix()
|
||||||
|
out["status"] = wireFinish(k, r.finish)
|
||||||
|
out["output"] = cs
|
||||||
|
out["error"] = nil
|
||||||
|
out["incomplete_details"] = nil
|
||||||
|
if r.finish == "length" {
|
||||||
|
out["incomplete_details"] = object{"reason": "max_output_tokens"}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if r.usage.Known || k == Anthropic {
|
||||||
|
out["usage"] = usageJSON(k, r.usage)
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,669 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Kind string
|
||||||
|
|
||||||
|
const (
|
||||||
|
Chat Kind = "openai-chat"
|
||||||
|
Responses Kind = "openai-responses"
|
||||||
|
Anthropic Kind = "anthropic"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Usage struct {
|
||||||
|
Input int64 `json:"input_tokens"`
|
||||||
|
Output int64 `json:"output_tokens"`
|
||||||
|
CacheRead int64 `json:"cache_read_tokens"`
|
||||||
|
CacheWrite int64 `json:"cache_write_tokens"`
|
||||||
|
Known bool `json:"known"`
|
||||||
|
Raw json.RawMessage `json:"raw,omitempty"`
|
||||||
|
}
|
||||||
|
type object = map[string]any
|
||||||
|
|
||||||
|
// maxJSONNodes bounds how many JSON tokens one body may contain before it is
|
||||||
|
// materialized into maps. Without it a 16 MiB body of micro-values expands
|
||||||
|
// into hundreds of MB of Go objects per request.
|
||||||
|
const maxJSONNodes = 1 << 18
|
||||||
|
|
||||||
|
func valid(k Kind) bool { return k == Chat || k == Responses || k == Anthropic }
|
||||||
|
|
||||||
|
// scanNodeCount walks b with the streaming decoder and fails once more than
|
||||||
|
// limit tokens have been seen. Malformed JSON is not an error here; the
|
||||||
|
// materializing pass reports its usual message.
|
||||||
|
func scanNodeCount(b []byte, limit int) error {
|
||||||
|
d := json.NewDecoder(bytes.NewReader(b))
|
||||||
|
n := 0
|
||||||
|
for {
|
||||||
|
_, err := d.Token()
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if n++; n > limit {
|
||||||
|
return fmt.Errorf("JSON exceeds %d nodes", limit)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func decode(b []byte) (object, error) {
|
||||||
|
if err := scanNodeCount(b, maxJSONNodes); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
d := json.NewDecoder(bytes.NewReader(b))
|
||||||
|
d.UseNumber()
|
||||||
|
var v object
|
||||||
|
if err := d.Decode(&v); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if v == nil {
|
||||||
|
return nil, fmt.Errorf("expected JSON object")
|
||||||
|
}
|
||||||
|
var extra any
|
||||||
|
if err := d.Decode(&extra); err != io.EOF {
|
||||||
|
return nil, fmt.Errorf("trailing JSON data")
|
||||||
|
}
|
||||||
|
return v, nil
|
||||||
|
}
|
||||||
|
func obj(v any) object { m, _ := v.(map[string]any); return m }
|
||||||
|
func str(v any) string { s, _ := v.(string); return s }
|
||||||
|
func arr(v any) []any { a, _ := v.([]any); return a }
|
||||||
|
func num(v any) int64 {
|
||||||
|
switch n := v.(type) {
|
||||||
|
case json.Number:
|
||||||
|
i, _ := n.Int64()
|
||||||
|
return i
|
||||||
|
case int64:
|
||||||
|
return n
|
||||||
|
case int:
|
||||||
|
return int64(n)
|
||||||
|
case float64:
|
||||||
|
return int64(n)
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
func fields(m object, names ...string) error {
|
||||||
|
allowed := map[string]bool{}
|
||||||
|
for _, n := range names {
|
||||||
|
allowed[n] = true
|
||||||
|
}
|
||||||
|
for k := range m {
|
||||||
|
if !allowed[k] {
|
||||||
|
return fmt.Errorf("unsupported field %q", k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func required(m object, k string) (string, error) {
|
||||||
|
s, ok := m[k].(string)
|
||||||
|
if !ok || s == "" {
|
||||||
|
return "", fmt.Errorf("missing or invalid %s", k)
|
||||||
|
}
|
||||||
|
return s, nil
|
||||||
|
}
|
||||||
|
func ExtractUsage(kind Kind, body []byte) Usage {
|
||||||
|
m, e := decode(body)
|
||||||
|
if e != nil {
|
||||||
|
return Usage{}
|
||||||
|
}
|
||||||
|
if r := obj(m["response"]); r != nil {
|
||||||
|
m = r
|
||||||
|
}
|
||||||
|
if r := obj(m["message"]); r != nil {
|
||||||
|
m = r
|
||||||
|
}
|
||||||
|
return usage(kind, obj(m["usage"]))
|
||||||
|
}
|
||||||
|
func usage(k Kind, m object) Usage {
|
||||||
|
if m == nil {
|
||||||
|
return Usage{}
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(m)
|
||||||
|
u := Usage{Raw: raw}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
_, a := m["prompt_tokens"]
|
||||||
|
_, b := m["completion_tokens"]
|
||||||
|
u.Known = a || b
|
||||||
|
u.Input = num(m["prompt_tokens"])
|
||||||
|
u.Output = num(m["completion_tokens"])
|
||||||
|
u.CacheRead = num(obj(m["prompt_tokens_details"])["cached_tokens"])
|
||||||
|
case Responses:
|
||||||
|
_, a := m["input_tokens"]
|
||||||
|
_, b := m["output_tokens"]
|
||||||
|
u.Known = a || b
|
||||||
|
u.Input = num(m["input_tokens"])
|
||||||
|
u.Output = num(m["output_tokens"])
|
||||||
|
u.CacheRead = num(obj(m["input_tokens_details"])["cached_tokens"])
|
||||||
|
case Anthropic:
|
||||||
|
_, a := m["input_tokens"]
|
||||||
|
_, b := m["output_tokens"]
|
||||||
|
u.Known = a || b
|
||||||
|
u.CacheRead = num(m["cache_read_input_tokens"])
|
||||||
|
u.CacheWrite = num(m["cache_creation_input_tokens"])
|
||||||
|
u.Input = num(m["input_tokens"]) + u.CacheRead + u.CacheWrite
|
||||||
|
u.Output = num(m["output_tokens"])
|
||||||
|
}
|
||||||
|
return u
|
||||||
|
}
|
||||||
|
func usageJSON(k Kind, u Usage) object {
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
return object{"prompt_tokens": u.Input, "completion_tokens": u.Output, "total_tokens": u.Input + u.Output, "prompt_tokens_details": object{"cached_tokens": u.CacheRead}}
|
||||||
|
case Responses:
|
||||||
|
return object{"input_tokens": u.Input, "output_tokens": u.Output, "total_tokens": u.Input + u.Output, "input_tokens_details": object{"cached_tokens": u.CacheRead}}
|
||||||
|
default:
|
||||||
|
return object{"input_tokens": u.Input - u.CacheRead - u.CacheWrite, "output_tokens": u.Output, "cache_read_input_tokens": u.CacheRead, "cache_creation_input_tokens": u.CacheWrite}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type part struct{ typ, text, id, name, args, url string }
|
||||||
|
type message struct {
|
||||||
|
role string
|
||||||
|
parts []part
|
||||||
|
}
|
||||||
|
type request struct {
|
||||||
|
messages []message
|
||||||
|
tools []object
|
||||||
|
options object
|
||||||
|
choice any
|
||||||
|
}
|
||||||
|
|
||||||
|
func textPart(s string) part { return part{typ: "text", text: s} }
|
||||||
|
|
||||||
|
func ConvertRequest(from, to Kind, body []byte, model string) ([]byte, error) {
|
||||||
|
if !valid(from) || !valid(to) {
|
||||||
|
return nil, fmt.Errorf("unknown protocol")
|
||||||
|
}
|
||||||
|
m, err := decode(body)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if from == to {
|
||||||
|
if model != "" {
|
||||||
|
m["model"] = model
|
||||||
|
}
|
||||||
|
return json.Marshal(m)
|
||||||
|
}
|
||||||
|
r, err := readRequest(from, m)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
out, err := writeRequest(to, r)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if model == "" {
|
||||||
|
model = str(m["model"])
|
||||||
|
}
|
||||||
|
out["model"] = model
|
||||||
|
return json.Marshal(out)
|
||||||
|
}
|
||||||
|
func readRequest(k Kind, m object) (request, error) {
|
||||||
|
r := request{options: object{}}
|
||||||
|
var err error
|
||||||
|
common := []string{"model", "stream", "temperature", "top_p", "tools", "tool_choice"}
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
err = fields(m, append(common, "messages", "max_tokens", "max_completion_tokens", "stop", "stream_options", "parallel_tool_calls", "n")...)
|
||||||
|
case Responses:
|
||||||
|
err = fields(m, append(common, "input", "instructions", "max_output_tokens", "parallel_tool_calls", "store")...)
|
||||||
|
case Anthropic:
|
||||||
|
err = fields(m, append(common, "messages", "system", "max_tokens", "stop_sequences")...)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return r, err
|
||||||
|
}
|
||||||
|
if n, ok := m["n"]; ok && num(n) != 1 {
|
||||||
|
return r, fmt.Errorf("multiple choices unsupported")
|
||||||
|
}
|
||||||
|
if v, ok := m["store"]; ok && v != false {
|
||||||
|
return r, fmt.Errorf("stored Responses unsupported")
|
||||||
|
}
|
||||||
|
if v, ok := m["stream_options"]; ok {
|
||||||
|
o := obj(v)
|
||||||
|
if o == nil {
|
||||||
|
return r, fmt.Errorf("invalid stream_options")
|
||||||
|
}
|
||||||
|
if e := fields(o, "include_usage"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, key := range []string{"stream", "temperature", "top_p", "parallel_tool_calls"} {
|
||||||
|
if v, ok := m[key]; ok {
|
||||||
|
r.options[key] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, key := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens"} {
|
||||||
|
if v, ok := m[key]; ok {
|
||||||
|
if _, exists := r.options["max"]; exists {
|
||||||
|
return r, fmt.Errorf("conflicting token limits")
|
||||||
|
}
|
||||||
|
r.options["max"] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, key := range []string{"stop", "stop_sequences"} {
|
||||||
|
if v, ok := m[key]; ok {
|
||||||
|
if s, ok := v.(string); ok {
|
||||||
|
v = []any{s}
|
||||||
|
}
|
||||||
|
if _, ok := v.([]any); !ok {
|
||||||
|
return r, fmt.Errorf("invalid stop sequences")
|
||||||
|
}
|
||||||
|
r.options["stop"] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k == Responses {
|
||||||
|
if s, ok := m["instructions"]; ok {
|
||||||
|
v, ok := s.(string)
|
||||||
|
if !ok {
|
||||||
|
return r, fmt.Errorf("invalid instructions")
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, message{"system", []part{textPart(v)}})
|
||||||
|
}
|
||||||
|
if s, ok := m["input"].(string); ok {
|
||||||
|
r.messages = append(r.messages, message{"user", []part{textPart(s)}})
|
||||||
|
} else {
|
||||||
|
items, ok := m["input"].([]any)
|
||||||
|
if !ok {
|
||||||
|
return r, fmt.Errorf("invalid input")
|
||||||
|
}
|
||||||
|
for _, v := range items {
|
||||||
|
x := obj(v)
|
||||||
|
typ := str(x["type"])
|
||||||
|
switch typ {
|
||||||
|
case "function_call":
|
||||||
|
if e := fields(x, "type", "id", "call_id", "name", "arguments", "status"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
p, e := callPart(x, "call_id", "arguments")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, message{"assistant", []part{p}})
|
||||||
|
case "function_call_output":
|
||||||
|
if e := fields(x, "type", "id", "call_id", "output", "status"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
id, e := required(x, "call_id")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
s, ok := x["output"].(string)
|
||||||
|
if !ok {
|
||||||
|
return r, fmt.Errorf("non-text function output unsupported")
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, message{"tool", []part{{typ: "result", id: id, text: s}}})
|
||||||
|
case "", "message":
|
||||||
|
mm, e := readMessage(k, x)
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, mm)
|
||||||
|
default:
|
||||||
|
return r, fmt.Errorf("unsupported input item %q", typ)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if k == Anthropic {
|
||||||
|
if s, ok := m["system"]; ok {
|
||||||
|
ps, e := readContent(k, s)
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
for _, p := range ps {
|
||||||
|
if p.typ != "text" {
|
||||||
|
return r, fmt.Errorf("non-text system")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, message{"system", ps})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ms, ok := m["messages"].([]any)
|
||||||
|
if !ok {
|
||||||
|
return r, fmt.Errorf("invalid messages")
|
||||||
|
}
|
||||||
|
for _, v := range ms {
|
||||||
|
mm, e := readMessage(k, obj(v))
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.messages = append(r.messages, mm)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if v, ok := m["tools"]; ok {
|
||||||
|
ts, ok := v.([]any)
|
||||||
|
if !ok {
|
||||||
|
return r, fmt.Errorf("invalid tools")
|
||||||
|
}
|
||||||
|
for _, v := range ts {
|
||||||
|
t := obj(v)
|
||||||
|
var f object
|
||||||
|
switch k {
|
||||||
|
case Chat:
|
||||||
|
if e := fields(t, "type", "function"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if str(t["type"]) != "function" {
|
||||||
|
return r, fmt.Errorf("unsupported tool")
|
||||||
|
}
|
||||||
|
f = obj(t["function"])
|
||||||
|
case Responses:
|
||||||
|
if str(t["type"]) != "function" {
|
||||||
|
return r, fmt.Errorf("unsupported tool")
|
||||||
|
}
|
||||||
|
f = object{}
|
||||||
|
for key, val := range t {
|
||||||
|
if key != "type" {
|
||||||
|
f[key] = val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case Anthropic:
|
||||||
|
f = object{}
|
||||||
|
for key, val := range t {
|
||||||
|
if key == "input_schema" {
|
||||||
|
key = "parameters"
|
||||||
|
}
|
||||||
|
f[key] = val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e := fields(f, "name", "description", "parameters", "strict"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if _, e := required(f, "name"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if obj(f["parameters"]) == nil {
|
||||||
|
return r, fmt.Errorf("missing tool schema")
|
||||||
|
}
|
||||||
|
r.tools = append(r.tools, f)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if v, ok := m["tool_choice"]; ok {
|
||||||
|
if s, ok := v.(string); ok {
|
||||||
|
if s != "auto" && s != "none" && s != "required" {
|
||||||
|
return r, fmt.Errorf("unsupported tool_choice")
|
||||||
|
}
|
||||||
|
r.choice = s
|
||||||
|
} else {
|
||||||
|
x := obj(v)
|
||||||
|
if k == Chat {
|
||||||
|
if e := fields(x, "type", "function"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if str(x["type"]) != "function" {
|
||||||
|
return r, fmt.Errorf("unsupported tool_choice")
|
||||||
|
}
|
||||||
|
f := obj(x["function"])
|
||||||
|
if e := fields(f, "name"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
s, e := required(f, "name")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.choice = object{"name": s}
|
||||||
|
} else if k == Responses {
|
||||||
|
if e := fields(x, "type", "name"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
if str(x["type"]) != "function" {
|
||||||
|
return r, fmt.Errorf("unsupported tool_choice")
|
||||||
|
}
|
||||||
|
s, e := required(x, "name")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.choice = object{"name": s}
|
||||||
|
} else {
|
||||||
|
if e := fields(x, "type", "name", "disable_parallel_tool_use"); e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "auto", "none":
|
||||||
|
r.choice = x["type"]
|
||||||
|
case "any":
|
||||||
|
r.choice = "required"
|
||||||
|
case "tool":
|
||||||
|
s, e := required(x, "name")
|
||||||
|
if e != nil {
|
||||||
|
return r, e
|
||||||
|
}
|
||||||
|
r.choice = object{"name": s}
|
||||||
|
default:
|
||||||
|
return r, fmt.Errorf("unsupported tool_choice")
|
||||||
|
}
|
||||||
|
if b, ok := x["disable_parallel_tool_use"].(bool); ok {
|
||||||
|
r.options["parallel_tool_calls"] = !b
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return r, nil
|
||||||
|
}
|
||||||
|
func callPart(x object, idKey, argsKey string) (part, error) {
|
||||||
|
id, e := required(x, idKey)
|
||||||
|
if e != nil {
|
||||||
|
return part{}, e
|
||||||
|
}
|
||||||
|
name, e := required(x, "name")
|
||||||
|
if e != nil {
|
||||||
|
return part{}, e
|
||||||
|
}
|
||||||
|
s, ok := x[argsKey].(string)
|
||||||
|
if !ok {
|
||||||
|
return part{}, fmt.Errorf("invalid tool arguments")
|
||||||
|
}
|
||||||
|
if _, e := decode([]byte(s)); e != nil {
|
||||||
|
return part{}, fmt.Errorf("tool arguments must be JSON object: %w", e)
|
||||||
|
}
|
||||||
|
return part{typ: "call", id: id, name: name, args: s}, nil
|
||||||
|
}
|
||||||
|
func readMessage(k Kind, x object) (message, error) {
|
||||||
|
m := message{role: str(x["role"])}
|
||||||
|
if x == nil {
|
||||||
|
return m, fmt.Errorf("invalid message")
|
||||||
|
}
|
||||||
|
if e := fields(x, "role", "content", "type", "id", "status", "tool_calls", "tool_call_id", "refusal"); e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
if x["refusal"] != nil {
|
||||||
|
return m, fmt.Errorf("refusal content cannot be preserved")
|
||||||
|
}
|
||||||
|
switch m.role {
|
||||||
|
case "user", "assistant", "system", "developer", "tool":
|
||||||
|
default:
|
||||||
|
return m, fmt.Errorf("unsupported role %q", m.role)
|
||||||
|
}
|
||||||
|
if m.role == "tool" {
|
||||||
|
if k != Chat {
|
||||||
|
return m, fmt.Errorf("invalid tool role")
|
||||||
|
}
|
||||||
|
id, e := required(x, "tool_call_id")
|
||||||
|
if e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
s, ok := x["content"].(string)
|
||||||
|
if !ok {
|
||||||
|
return m, fmt.Errorf("non-text tool result unsupported")
|
||||||
|
}
|
||||||
|
m.parts = []part{{typ: "result", id: id, text: s}}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
if c := x["content"]; c != nil {
|
||||||
|
ps, e := readContent(k, c)
|
||||||
|
if e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
m.parts = ps
|
||||||
|
}
|
||||||
|
if tc, ok := x["tool_calls"]; ok {
|
||||||
|
if k != Chat || m.role != "assistant" {
|
||||||
|
return m, fmt.Errorf("invalid tool_calls")
|
||||||
|
}
|
||||||
|
ts, ok := tc.([]any)
|
||||||
|
if !ok {
|
||||||
|
return m, fmt.Errorf("invalid tool_calls")
|
||||||
|
}
|
||||||
|
for _, v := range ts {
|
||||||
|
t := obj(v)
|
||||||
|
if e := fields(t, "id", "type", "function"); e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
if str(t["type"]) != "function" {
|
||||||
|
return m, fmt.Errorf("unsupported tool call")
|
||||||
|
}
|
||||||
|
f := obj(t["function"])
|
||||||
|
if e := fields(f, "name", "arguments"); e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
p, e := callPart(object{"id": t["id"], "name": f["name"], "arguments": f["arguments"]}, "id", "arguments")
|
||||||
|
if e != nil {
|
||||||
|
return m, e
|
||||||
|
}
|
||||||
|
m.parts = append(m.parts, p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
func readContent(k Kind, v any) ([]part, error) {
|
||||||
|
if s, ok := v.(string); ok {
|
||||||
|
return []part{textPart(s)}, nil
|
||||||
|
}
|
||||||
|
a, ok := v.([]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid content")
|
||||||
|
}
|
||||||
|
var ps []part
|
||||||
|
for _, v := range a {
|
||||||
|
x := obj(v)
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "text", "input_text", "output_text":
|
||||||
|
if e := fields(x, "type", "text", "annotations"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if an, ok := x["annotations"]; ok && len(arr(an)) != 0 {
|
||||||
|
return nil, fmt.Errorf("annotations unsupported")
|
||||||
|
}
|
||||||
|
s, ok := x["text"].(string)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid text")
|
||||||
|
}
|
||||||
|
ps = append(ps, textPart(s))
|
||||||
|
case "image_url":
|
||||||
|
if k != Chat {
|
||||||
|
return nil, fmt.Errorf("invalid image block")
|
||||||
|
}
|
||||||
|
if e := fields(x, "type", "image_url"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
im := obj(x["image_url"])
|
||||||
|
if e := fields(im, "url", "detail"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if d := str(im["detail"]); d != "" && d != "auto" {
|
||||||
|
return nil, fmt.Errorf("image detail cannot be preserved")
|
||||||
|
}
|
||||||
|
s, e := required(im, "url")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
ps = append(ps, part{typ: "image", url: s})
|
||||||
|
case "input_image":
|
||||||
|
if k != Responses {
|
||||||
|
return nil, fmt.Errorf("invalid image block")
|
||||||
|
}
|
||||||
|
if e := fields(x, "type", "image_url", "detail"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if d := str(x["detail"]); d != "" && d != "auto" {
|
||||||
|
return nil, fmt.Errorf("image detail cannot be preserved")
|
||||||
|
}
|
||||||
|
s, e := required(x, "image_url")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
ps = append(ps, part{typ: "image", url: s})
|
||||||
|
case "image":
|
||||||
|
if k != Anthropic {
|
||||||
|
return nil, fmt.Errorf("invalid image block")
|
||||||
|
}
|
||||||
|
if e := fields(x, "type", "source"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
src := obj(x["source"])
|
||||||
|
var url string
|
||||||
|
switch str(src["type"]) {
|
||||||
|
case "url":
|
||||||
|
if e := fields(src, "type", "url"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
url = str(src["url"])
|
||||||
|
case "base64":
|
||||||
|
if e := fields(src, "type", "media_type", "data"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if str(src["media_type"]) == "" || str(src["data"]) == "" {
|
||||||
|
return nil, fmt.Errorf("invalid image source")
|
||||||
|
}
|
||||||
|
url = "data:" + str(src["media_type"]) + ";base64," + str(src["data"])
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unsupported image source")
|
||||||
|
}
|
||||||
|
if url == "" {
|
||||||
|
return nil, fmt.Errorf("empty image URL")
|
||||||
|
}
|
||||||
|
ps = append(ps, part{typ: "image", url: url})
|
||||||
|
case "tool_use":
|
||||||
|
if k != Anthropic {
|
||||||
|
return nil, fmt.Errorf("invalid tool block")
|
||||||
|
}
|
||||||
|
if e := fields(x, "type", "id", "name", "input"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
b, e := json.Marshal(x["input"])
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
p, e := callPart(object{"id": x["id"], "name": x["name"], "arguments": string(b)}, "id", "arguments")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
ps = append(ps, p)
|
||||||
|
case "tool_result":
|
||||||
|
if k != Anthropic {
|
||||||
|
return nil, fmt.Errorf("invalid result block")
|
||||||
|
}
|
||||||
|
if e := fields(x, "type", "tool_use_id", "content", "is_error"); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if x["is_error"] == true {
|
||||||
|
return nil, fmt.Errorf("tool error flag cannot be preserved")
|
||||||
|
}
|
||||||
|
id, e := required(x, "tool_use_id")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
s, ok := x["content"].(string)
|
||||||
|
if !ok {
|
||||||
|
pp, e := readContent(k, x["content"])
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
for _, p := range pp {
|
||||||
|
if p.typ != "text" {
|
||||||
|
return nil, fmt.Errorf("non-text tool result unsupported")
|
||||||
|
}
|
||||||
|
s += p.text
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ps = append(ps, part{typ: "result", id: id, text: s})
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unsupported content type %q", str(x["type"]))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ps, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,267 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRequestMatrix(t *testing.T) {
|
||||||
|
chat := []byte(`{"model":"old","max_tokens":100,"stream":true,"messages":[{"role":"system","content":"help"},{"role":"user","content":[{"type":"text","text":"look"},{"type":"image_url","image_url":{"url":"data:image/png;base64,YQ=="}}]},{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"read","arguments":"{\"path\":\"a\"}"}}]},{"role":"tool","tool_call_id":"c1","content":"ok"},{"role":"user","content":"continue"}],"tools":[{"type":"function","function":{"name":"read","parameters":{"type":"object","properties":{"path":{"type":"string"}}}}}],"tool_choice":"auto"}`)
|
||||||
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
src := chat
|
||||||
|
var e error
|
||||||
|
if from != Chat {
|
||||||
|
src, e = ConvertRequest(Chat, from, chat, "new")
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
||||||
|
b, e := ConvertRequest(from, to, src, "target")
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
m, e := decode(b)
|
||||||
|
if e != nil || m["model"] != "target" {
|
||||||
|
t.Fatalf("%s %v", b, e)
|
||||||
|
}
|
||||||
|
r, e := readRequest(to, m)
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
var calls, results, images int
|
||||||
|
for _, m := range r.messages {
|
||||||
|
for _, p := range m.parts {
|
||||||
|
switch p.typ {
|
||||||
|
case "call":
|
||||||
|
calls++
|
||||||
|
if p.id != "c1" {
|
||||||
|
t.Fatal(p)
|
||||||
|
}
|
||||||
|
case "result":
|
||||||
|
results++
|
||||||
|
case "image":
|
||||||
|
images++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if calls != 1 || results != 1 || images != 1 {
|
||||||
|
t.Fatalf("lost content: %s", b)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStrictRejection(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
k Kind
|
||||||
|
b string
|
||||||
|
to Kind
|
||||||
|
}{
|
||||||
|
{Chat, `{"messages":[],"response_format":{"type":"json_object"}}`, Responses},
|
||||||
|
{Chat, `{"messages":[]}`, Anthropic},
|
||||||
|
{Responses, `{"input":[],"previous_response_id":"r"}`, Chat},
|
||||||
|
{Anthropic, `{"max_tokens":10,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x"}]}]}`, Chat},
|
||||||
|
{Chat, `{"messages":[],"n":2}`, Responses},
|
||||||
|
{Chat, `{"messages":[],"stop":["x"]}`, Responses},
|
||||||
|
{Chat, `{"max_tokens":10,"messages":[],"tools":[{"type":"function","function":{"name":"x","parameters":{},"strict":true}}]}`, Anthropic},
|
||||||
|
}
|
||||||
|
for _, c := range cases {
|
||||||
|
if b, e := ConvertRequest(c.k, c.to, []byte(c.b), "m"); e == nil {
|
||||||
|
t.Fatalf("accepted %s => %s", c.b, b)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestResponseMatrix(t *testing.T) {
|
||||||
|
chat := []byte(`{"id":"r","model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"working","tool_calls":[{"id":"c","type":"function","function":{"name":"read","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":3}}}`)
|
||||||
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
src, e := ConvertResponse(Chat, from, chat)
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
b, e := ConvertResponse(from, to, src)
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(from, to, e)
|
||||||
|
}
|
||||||
|
m, _ := decode(b)
|
||||||
|
r, e := readResponse(to, m)
|
||||||
|
if e != nil || len(r.parts) != 2 || r.finish != "tool" {
|
||||||
|
t.Fatalf("%s %v", b, e)
|
||||||
|
}
|
||||||
|
u := ExtractUsage(to, b)
|
||||||
|
if u.Input != 12 || u.Output != 4 || u.CacheRead != 3 || !u.Known {
|
||||||
|
t.Fatal(u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type flushBuffer struct {
|
||||||
|
bytes.Buffer
|
||||||
|
flushes int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *flushBuffer) Flush() { b.flushes++ }
|
||||||
|
func chatStream() string {
|
||||||
|
return strings.Join([]string{
|
||||||
|
`data: {"id":"r","model":"m","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"}}]}`,
|
||||||
|
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"read","arguments":"{\"x\":"}},{"index":1,"id":"c2","type":"function","function":{"name":"write","arguments":"{"}}]}}]}`,
|
||||||
|
`data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"}"}},{"index":0,"function":{"arguments":"1}"}}]}}]}`,
|
||||||
|
`data: {"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}`,
|
||||||
|
`data: {"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":7}}`,
|
||||||
|
`data: [DONE]`,
|
||||||
|
}, "\n\n") + "\n\n"
|
||||||
|
}
|
||||||
|
func TestStreamMatrix(t *testing.T) {
|
||||||
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
src := chatStream()
|
||||||
|
if from != Chat {
|
||||||
|
var b bytes.Buffer
|
||||||
|
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, from, "m"); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
src = b.String()
|
||||||
|
}
|
||||||
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
t.Run(string(from)+"/"+string(to), func(t *testing.T) {
|
||||||
|
var b flushBuffer
|
||||||
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), from, to, "m")
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
if u.Input != 10 || u.Output != 7 || !u.Known {
|
||||||
|
t.Fatal(u)
|
||||||
|
}
|
||||||
|
if b.flushes < 3 {
|
||||||
|
t.Fatal("not flushed")
|
||||||
|
}
|
||||||
|
if from == to && b.String() != src {
|
||||||
|
t.Fatal("passthrough changed bytes")
|
||||||
|
}
|
||||||
|
if to == Responses {
|
||||||
|
for _, event := range []string{"response.output_item.added", "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.completed"} {
|
||||||
|
if !strings.Contains(b.String(), event) {
|
||||||
|
t.Fatal("missing", event)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if to == Anthropic && !strings.Contains(b.String(), "message_stop") {
|
||||||
|
t.Fatal("missing stop")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestUsageCumulative(t *testing.T) {
|
||||||
|
src := "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":5,\"output_tokens\":1,\"cache_read_input_tokens\":8,\"cache_creation_input_tokens\":2}}}\n\n" + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":9}}\n\n"
|
||||||
|
var b bytes.Buffer
|
||||||
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Anthropic, Anthropic, "")
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
if u.Input != 15 || u.Output != 9 || u.CacheRead != 8 || u.CacheWrite != 2 {
|
||||||
|
t.Fatal(u)
|
||||||
|
}
|
||||||
|
raw, e := decode(u.Raw)
|
||||||
|
if e != nil || num(raw["input_tokens"]) != 5 || num(raw["output_tokens"]) != 9 || num(raw["cache_read_input_tokens"]) != 8 {
|
||||||
|
t.Fatal(string(u.Raw), e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestRealtimeBeforeEOF(t *testing.T) {
|
||||||
|
r, w := io.Pipe()
|
||||||
|
defer r.Close()
|
||||||
|
defer w.Close()
|
||||||
|
done := make(chan error, 1)
|
||||||
|
out := &signalWriter{signal: make(chan struct{}, 1)}
|
||||||
|
go func() { _, e := Stream(context.Background(), out, r, Chat, Responses, "m"); done <- e }()
|
||||||
|
first, rest, _ := strings.Cut(chatStream(), "\n\n")
|
||||||
|
if _, e := io.WriteString(w, first+"\n\n"); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-out.signal:
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("first delta buffered until EOF")
|
||||||
|
}
|
||||||
|
if _, e := io.WriteString(w, rest); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case e := <-done:
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("terminal marker waits for EOF")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type signalWriter struct{ signal chan struct{} }
|
||||||
|
|
||||||
|
func (w *signalWriter) Write(p []byte) (int, error) {
|
||||||
|
select {
|
||||||
|
case w.signal <- struct{}{}:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
return len(p), nil
|
||||||
|
}
|
||||||
|
func TestUnknownUsageAndUniqueID(t *testing.T) {
|
||||||
|
src := "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"a\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
|
||||||
|
var previous string
|
||||||
|
for i := 0; i < 2; i++ {
|
||||||
|
var b bytes.Buffer
|
||||||
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Anthropic, "m")
|
||||||
|
if e != nil || u.Known {
|
||||||
|
t.Fatal(u, e)
|
||||||
|
}
|
||||||
|
line := strings.Split(b.String(), "\n")[1]
|
||||||
|
m, e := decode([]byte(strings.TrimPrefix(line, "data: ")))
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
msg := obj(m["message"])
|
||||||
|
if obj(msg["usage"]) == nil {
|
||||||
|
t.Fatal("missing initial usage")
|
||||||
|
}
|
||||||
|
id := str(msg["id"])
|
||||||
|
if id == "" || id == previous {
|
||||||
|
t.Fatal("nonunique id")
|
||||||
|
}
|
||||||
|
previous = id
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStreamErrors(t *testing.T) {
|
||||||
|
for _, src := range []string{"data: nope\n\n", "data: {\"error\":{\"message\":\"bad\"}}\n\n", strings.Replace(chatStream(), "data: [DONE]\n\n", "", 1), strings.Replace(chatStream(), `"arguments":"1}"`, `"arguments":"oops}"`, 1)} {
|
||||||
|
var b bytes.Buffer
|
||||||
|
if _, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Responses, "m"); e == nil {
|
||||||
|
t.Fatal("accepted broken stream", src)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
if _, e := Stream(ctx, io.Discard, strings.NewReader(chatStream()), Chat, Chat, ""); !errors.Is(e, context.Canceled) {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
if _, e := Stream(context.Background(), badWriter{}, strings.NewReader(chatStream()), Chat, Anthropic, ""); !errors.Is(e, io.ErrClosedPipe) {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type badWriter struct{}
|
||||||
|
|
||||||
|
func (badWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe }
|
||||||
|
func TestSSEMultilineAndCRLF(t *testing.T) {
|
||||||
|
src := ": comment\r\n\r\ndata: {\r\ndata: \"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3}}\r\n\r\ndata: [DONE]\r\n\r\n"
|
||||||
|
var b flushBuffer
|
||||||
|
u, e := Stream(context.Background(), &b, strings.NewReader(src), Chat, Chat, "")
|
||||||
|
if e != nil || u.Input != 2 || u.Output != 3 || b.String() != src {
|
||||||
|
t.Fatal(u, e, b.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,352 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Recognized Anthropic thinking blocks are omitted only from converted output.
|
||||||
|
func anthropicThinking(x object) (bool, error) {
|
||||||
|
var names []string
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "thinking":
|
||||||
|
names = []string{"thinking", "signature"}
|
||||||
|
case "redacted_thinking":
|
||||||
|
names = []string{"data"}
|
||||||
|
default:
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if err := fields(x, append([]string{"type"}, names...)...); err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
for _, name := range names {
|
||||||
|
if _, ok := x[name].(string); !ok {
|
||||||
|
return true, fmt.Errorf("invalid thinking %s", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func anthropicOutput(v any) (any, error) {
|
||||||
|
content, ok := v.([]any)
|
||||||
|
if !ok {
|
||||||
|
return nil, fmt.Errorf("invalid Anthropic content")
|
||||||
|
}
|
||||||
|
clean := []any{}
|
||||||
|
for _, block := range content {
|
||||||
|
thinking, err := anthropicThinking(obj(block))
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if !thinking {
|
||||||
|
clean = append(clean, block)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return clean, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type thinkingBlock struct {
|
||||||
|
kind string
|
||||||
|
closed bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamState) thinkingEvent(typ string, m object) (bool, error) {
|
||||||
|
if typ == "message_delta" || typ == "message_stop" {
|
||||||
|
for _, b := range s.thinking {
|
||||||
|
if !b.closed {
|
||||||
|
return true, fmt.Errorf("unfinished thinking block")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if typ != "content_block_start" && typ != "content_block_delta" && typ != "content_block_stop" {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
index, err := responseIndex(m, "index")
|
||||||
|
if err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
b := s.thinking[index]
|
||||||
|
if typ == "content_block_start" {
|
||||||
|
if b != nil || s.source[index] != nil {
|
||||||
|
return true, fmt.Errorf("duplicate content index")
|
||||||
|
}
|
||||||
|
x := obj(m["content_block"])
|
||||||
|
thinking, err := anthropicThinking(x)
|
||||||
|
if err != nil || !thinking {
|
||||||
|
return thinking, err
|
||||||
|
}
|
||||||
|
if !s.started || s.finished {
|
||||||
|
return true, fmt.Errorf("thinking outside active message")
|
||||||
|
}
|
||||||
|
if s.thinking == nil {
|
||||||
|
s.thinking = map[int]*thinkingBlock{}
|
||||||
|
}
|
||||||
|
s.thinking[index] = &thinkingBlock{kind: str(x["type"])}
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
if b == nil {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if b.closed || s.finished {
|
||||||
|
return true, fmt.Errorf("event after thinking block closed")
|
||||||
|
}
|
||||||
|
if typ == "content_block_stop" {
|
||||||
|
b.closed = true
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
d := obj(m["delta"])
|
||||||
|
field := ""
|
||||||
|
switch str(d["type"]) {
|
||||||
|
case "thinking_delta":
|
||||||
|
field = "thinking"
|
||||||
|
case "signature_delta":
|
||||||
|
field = "signature"
|
||||||
|
}
|
||||||
|
if b.kind != "thinking" || field == "" {
|
||||||
|
return true, fmt.Errorf("invalid thinking delta")
|
||||||
|
}
|
||||||
|
if err := fields(d, "type", field); err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
if _, ok := d[field].(string); !ok {
|
||||||
|
return true, fmt.Errorf("invalid thinking delta value")
|
||||||
|
}
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// chatOutput omits output-only reasoning without exposing it as visible text.
|
||||||
|
func chatOutput(x object) (object, error) {
|
||||||
|
if v := x["reasoning_content"]; v != nil {
|
||||||
|
if _, ok := v.(string); !ok {
|
||||||
|
return nil, fmt.Errorf("reasoning_content must be a string or null")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
m := object{}
|
||||||
|
for k, v := range x {
|
||||||
|
if k != "reasoning_content" {
|
||||||
|
m[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Some Chat providers include streaming-style indices in complete tool calls.
|
||||||
|
func chatResponseMessage(x object) (object, error) {
|
||||||
|
m, err := chatOutput(x)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if calls, ok := x["tool_calls"].([]any); ok {
|
||||||
|
clean := make([]any, len(calls))
|
||||||
|
for i, v := range calls {
|
||||||
|
call := obj(v)
|
||||||
|
if _, exists := call["index"]; exists {
|
||||||
|
index, err := responseIndex(call, "index")
|
||||||
|
if err != nil || index != i {
|
||||||
|
return nil, fmt.Errorf("invalid nonstream tool index")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
c := object{}
|
||||||
|
for k, v := range call {
|
||||||
|
if k != "index" {
|
||||||
|
c[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
clean[i] = c
|
||||||
|
}
|
||||||
|
m["tool_calls"] = clean
|
||||||
|
}
|
||||||
|
return m, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase describes output metadata, not message content or request semantics.
|
||||||
|
func responseMessage(x object) object {
|
||||||
|
m := object{}
|
||||||
|
for k, v := range x {
|
||||||
|
if k != "phase" {
|
||||||
|
m[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if content, ok := x["content"].([]any); ok {
|
||||||
|
clean := make([]any, len(content))
|
||||||
|
for i, v := range content {
|
||||||
|
p := obj(v)
|
||||||
|
if logs, ok := p["logprobs"].([]any); ok && len(logs) == 0 && str(p["type"]) == "output_text" {
|
||||||
|
q := object{}
|
||||||
|
for k, v := range p {
|
||||||
|
if k != "logprobs" {
|
||||||
|
q[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
clean[i] = q
|
||||||
|
} else {
|
||||||
|
clean[i] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
m["content"] = clean
|
||||||
|
}
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
|
||||||
|
func validateReasoning(x object) error {
|
||||||
|
if err := fields(x, "type", "id", "status", "content", "summary", "encrypted_content"); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
for field, typ := range map[string]string{"content": "reasoning_text", "summary": "summary_text"} {
|
||||||
|
if v, exists := x[field]; exists && v != nil {
|
||||||
|
if _, ok := v.([]any); !ok {
|
||||||
|
return fmt.Errorf("reasoning %s must be an array", field)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, v := range arr(x[field]) {
|
||||||
|
p := obj(v)
|
||||||
|
if str(p["type"]) != typ {
|
||||||
|
return fmt.Errorf("unsupported reasoning %s content", field)
|
||||||
|
}
|
||||||
|
if err := fields(p, "type", "text"); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type reasoningItem struct {
|
||||||
|
id string
|
||||||
|
closed bool
|
||||||
|
parts map[string]bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func responseIndex(m object, field string) (int, error) {
|
||||||
|
n, ok := m[field].(json.Number)
|
||||||
|
i, err := n.Int64()
|
||||||
|
if !ok || err != nil || i < 0 || int64(int(i)) != i {
|
||||||
|
return 0, fmt.Errorf("invalid %s", field)
|
||||||
|
}
|
||||||
|
return int(i), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamState) reasoningEvent(typ string, m object) (bool, error) {
|
||||||
|
if typ == "response.completed" || typ == "response.incomplete" {
|
||||||
|
seen := 0
|
||||||
|
for i, v := range arr(obj(m["response"])["output"]) {
|
||||||
|
x := obj(v)
|
||||||
|
r := s.reasoning[i]
|
||||||
|
if str(x["type"]) == "reasoning" {
|
||||||
|
if r == nil || !r.closed || r.id != str(x["id"]) {
|
||||||
|
return true, fmt.Errorf("final reasoning item disagrees with stream")
|
||||||
|
}
|
||||||
|
seen++
|
||||||
|
} else if r != nil {
|
||||||
|
return true, fmt.Errorf("final reasoning index changed type")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if seen != len(s.reasoning) {
|
||||||
|
return true, fmt.Errorf("final response missing reasoning item")
|
||||||
|
}
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
known := false
|
||||||
|
switch typ {
|
||||||
|
case "response.reasoning_text.delta", "response.reasoning_text.done",
|
||||||
|
"response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
|
||||||
|
"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
|
||||||
|
known = true
|
||||||
|
case "response.output_item.added", "response.output_item.done",
|
||||||
|
"response.content_part.added", "response.content_part.done",
|
||||||
|
"response.output_text.delta", "response.output_text.done",
|
||||||
|
"response.function_call_arguments.delta", "response.function_call_arguments.done":
|
||||||
|
default:
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
key, err := responseIndex(m, "output_index")
|
||||||
|
if err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
r := s.reasoning[key]
|
||||||
|
x := obj(m["item"])
|
||||||
|
if typ == "response.output_item.added" {
|
||||||
|
if r != nil || s.source[key] != nil {
|
||||||
|
return true, fmt.Errorf("duplicate output index %d", key)
|
||||||
|
}
|
||||||
|
if str(x["type"]) != "reasoning" {
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if err := validateReasoning(x); err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
if str(x["id"]) == "" {
|
||||||
|
return true, fmt.Errorf("missing reasoning item id")
|
||||||
|
}
|
||||||
|
if s.reasoning == nil {
|
||||||
|
s.reasoning = map[int]*reasoningItem{}
|
||||||
|
}
|
||||||
|
s.reasoning[key] = &reasoningItem{id: str(x["id"]), parts: map[string]bool{}}
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
if r == nil {
|
||||||
|
if known || str(x["type"]) == "reasoning" {
|
||||||
|
return true, fmt.Errorf("reasoning event outside reasoning item")
|
||||||
|
}
|
||||||
|
return false, nil
|
||||||
|
}
|
||||||
|
if r.closed {
|
||||||
|
return true, fmt.Errorf("event after reasoning item closed")
|
||||||
|
}
|
||||||
|
if typ == "response.output_item.done" {
|
||||||
|
if str(x["type"]) != "reasoning" || str(x["id"]) != r.id {
|
||||||
|
return true, fmt.Errorf("reasoning item identity mismatch")
|
||||||
|
}
|
||||||
|
if err := validateReasoning(x); err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
for _, closed := range r.parts {
|
||||||
|
if !closed {
|
||||||
|
return true, fmt.Errorf("reasoning item has open parts")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
r.closed = true
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
if str(m["item_id"]) != r.id {
|
||||||
|
return true, fmt.Errorf("reasoning event item identity mismatch")
|
||||||
|
}
|
||||||
|
field, partType := "content_index", "reasoning_text"
|
||||||
|
switch typ {
|
||||||
|
case "response.reasoning_summary_text.delta", "response.reasoning_summary_text.done",
|
||||||
|
"response.reasoning_summary_part.added", "response.reasoning_summary_part.done":
|
||||||
|
field, partType = "summary_index", "summary_text"
|
||||||
|
case "response.reasoning_text.delta", "response.reasoning_text.done",
|
||||||
|
"response.content_part.added", "response.content_part.done":
|
||||||
|
default:
|
||||||
|
return true, fmt.Errorf("non-reasoning event on reasoning item")
|
||||||
|
}
|
||||||
|
i, err := responseIndex(m, field)
|
||||||
|
if err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
partKey := fmt.Sprintf("%s:%d", field, i)
|
||||||
|
closed, exists := r.parts[partKey]
|
||||||
|
switch typ {
|
||||||
|
case "response.content_part.added", "response.reasoning_summary_part.added",
|
||||||
|
"response.content_part.done", "response.reasoning_summary_part.done":
|
||||||
|
p := obj(m["part"])
|
||||||
|
if str(p["type"]) != partType {
|
||||||
|
return true, fmt.Errorf("unsupported reasoning part")
|
||||||
|
}
|
||||||
|
if err := fields(p, "type", "text"); err != nil {
|
||||||
|
return true, err
|
||||||
|
}
|
||||||
|
added := typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added"
|
||||||
|
if (added && exists) || (!added && (!exists || closed)) {
|
||||||
|
return true, fmt.Errorf("invalid reasoning part lifecycle")
|
||||||
|
}
|
||||||
|
r.parts[partKey] = !added
|
||||||
|
default:
|
||||||
|
if !exists || closed {
|
||||||
|
return true, fmt.Errorf("reasoning text outside active part")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
const reasoningJSON = `{"type":"reasoning","id":"r","content":[{"type":"reasoning_text","text":"SECRET"}],"summary":[{"type":"summary_text","text":"SUMMARY"}],"encrypted_content":"CIPHER"}`
|
||||||
|
const visibleJSON = `{"type":"message","id":"m","role":"assistant","phase":"final_answer","content":[{"type":"output_text","text":"OK","logprobs":[]}]}`
|
||||||
|
|
||||||
|
func reasoningSSE(events ...string) string {
|
||||||
|
sep := string([]byte{10, 10})
|
||||||
|
return "data: " + strings.Join(events, sep+"data: ") + sep
|
||||||
|
}
|
||||||
|
func TestReasoningConversions(t *testing.T) {
|
||||||
|
for _, tool := range []bool{false, true} {
|
||||||
|
for _, incomplete := range []bool{false, true} {
|
||||||
|
visible, initial, delta := visibleJSON, `{"type":"message","content":[]}`, `{"type":"response.output_text.delta","output_index":1,"delta":"OK"}`
|
||||||
|
if tool {
|
||||||
|
visible = `{"type":"function_call","call_id":"c","name":"lookup","arguments":"{}"}`
|
||||||
|
initial = `{"type":"function_call","call_id":"c","name":"lookup","arguments":""}`
|
||||||
|
delta = `{"type":"response.function_call_arguments.delta","output_index":1,"delta":"{}"}`
|
||||||
|
}
|
||||||
|
status, details := "completed", ""
|
||||||
|
if incomplete {
|
||||||
|
status = "incomplete"
|
||||||
|
details = `,"incomplete_details":{"reason":"max_output_tokens"}`
|
||||||
|
}
|
||||||
|
body := `{"id":"res","model":"m","status":"` + status + `","output":[` + reasoningJSON + `,` + visible + `],"usage":{"input_tokens":9,"output_tokens":7,"output_tokens_details":{"reasoning_tokens":5}}` + details + `}`
|
||||||
|
src := reasoningSSE(`{"type":"response.output_item.added","output_index":0,"item":{"type":"reasoning","id":"r"}}`,
|
||||||
|
`{"type":"response.content_part.added","output_index":0,"item_id":"r","content_index":0,"part":{"type":"reasoning_text","text":""}}`,
|
||||||
|
`{"type":"response.reasoning_text.delta","output_index":0,"item_id":"r","content_index":0,"delta":"SECRET"}`,
|
||||||
|
`{"type":"response.reasoning_text.done","output_index":0,"item_id":"r","content_index":0,"text":"SECRET"}`,
|
||||||
|
`{"type":"response.content_part.done","output_index":0,"item_id":"r","content_index":0,"part":{"type":"reasoning_text","text":"SECRET"}}`,
|
||||||
|
`{"type":"response.reasoning_summary_part.added","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":""}}`,
|
||||||
|
`{"type":"response.reasoning_summary_text.delta","output_index":0,"item_id":"r","summary_index":0,"delta":"SUMMARY"}`,
|
||||||
|
`{"type":"response.reasoning_summary_text.done","output_index":0,"item_id":"r","summary_index":0,"text":"SUMMARY"}`,
|
||||||
|
`{"type":"response.reasoning_summary_part.done","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":"SUMMARY"}}`,
|
||||||
|
`{"type":"response.output_item.done","output_index":0,"item":`+reasoningJSON+`}`,
|
||||||
|
`{"type":"response.output_item.added","output_index":1,"item":`+initial+`}`, delta,
|
||||||
|
`{"type":"response.output_item.done","output_index":1,"item":`+visible+`}`,
|
||||||
|
`{"type":"response.`+status+`","response":`+body+`}`)
|
||||||
|
for _, target := range []Kind{Chat, Anthropic, Responses} {
|
||||||
|
out, err := ConvertResponse(Responses, target, []byte(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var b bytes.Buffer
|
||||||
|
u, err := Stream(context.Background(), &b, strings.NewReader(src), Responses, target, "m")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !u.Known || u.Input != 9 || u.Output != 7 || !bytes.Contains(u.Raw, []byte("reasoning_tokens")) {
|
||||||
|
t.Fatalf("usage: %+v", u)
|
||||||
|
}
|
||||||
|
if target == Responses {
|
||||||
|
if string(out) != body || b.String() != src {
|
||||||
|
t.Fatal("changed same-protocol bytes")
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, result := range []string{string(out), b.String()} {
|
||||||
|
for _, secret := range []string{"SECRET", "SUMMARY", "CIPHER", "signature", "thinking", "phase"} {
|
||||||
|
if strings.Contains(result, secret) {
|
||||||
|
t.Fatal("leaked " + secret)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
want := "OK"
|
||||||
|
if tool {
|
||||||
|
want = "lookup"
|
||||||
|
}
|
||||||
|
if !strings.Contains(result, want) {
|
||||||
|
t.Fatal("lost visible content")
|
||||||
|
}
|
||||||
|
finish := "stop"
|
||||||
|
if target == Anthropic {
|
||||||
|
finish = "end_turn"
|
||||||
|
}
|
||||||
|
if tool {
|
||||||
|
finish = "tool_calls"
|
||||||
|
if target == Anthropic {
|
||||||
|
finish = "tool_use"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if incomplete {
|
||||||
|
finish = "length"
|
||||||
|
if target == Anthropic {
|
||||||
|
finish = "max_tokens"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !strings.Contains(result, finish) {
|
||||||
|
t.Fatal("lost finish " + finish)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if target == Chat && !strings.Contains(b.String(), "[DONE]") {
|
||||||
|
t.Fatal("missing terminal")
|
||||||
|
}
|
||||||
|
if target == Anthropic && (!strings.Contains(b.String(), "message_stop") || !strings.Contains(b.String(), `"index":0`)) {
|
||||||
|
t.Fatal("terminal/index")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReasoningRejectsInvalidStreams(t *testing.T) {
|
||||||
|
start := `{"type":"response.output_item.added","output_index":0,"item":{"type":"reasoning","id":"r"}}`
|
||||||
|
part := `{"type":"response.reasoning_summary_part.added","output_index":0,"item_id":"r","summary_index":0,"part":{"type":"summary_text","text":""}}`
|
||||||
|
done := `{"type":"response.output_item.done","output_index":0,"item":` + reasoningJSON + `}`
|
||||||
|
for name, events := range map[string][]string{
|
||||||
|
"duplicate": {start, start}, "part before item": {part}, "open part": {start, part, done}, "closed item": {start, done, part},
|
||||||
|
"wrong id": {start, strings.Replace(part, `"item_id":"r"`, `"item_id":"x"`, 1)},
|
||||||
|
"unknown": {start, `{"type":"response.reasoning_unknown.delta","output_index":0}`},
|
||||||
|
"text collision": {start, `{"type":"response.output_text.delta","output_index":0,"item_id":"r","delta":"lost"}`},
|
||||||
|
"tool collision": {start, `{"type":"response.function_call_arguments.delta","output_index":0,"item_id":"r","delta":"{}"}`},
|
||||||
|
"unknown part": {start, strings.Replace(part, "summary_text", "output_text", 1)},
|
||||||
|
"negative": {strings.Replace(start, `"output_index":0`, `"output_index":-1`, 1)},
|
||||||
|
"fractional": {strings.Replace(start, `"output_index":0`, `"output_index":0.5`, 1)},
|
||||||
|
"final missing": {start, done, `{"type":"response.completed","response":{"status":"completed","output":[]}}`},
|
||||||
|
} {
|
||||||
|
t.Run(name, func(t *testing.T) {
|
||||||
|
var b bytes.Buffer
|
||||||
|
_, err := Stream(context.Background(), &b, strings.NewReader(reasoningSSE(events...)), Responses, Chat, "")
|
||||||
|
if err == nil || strings.Contains(err.Error(), "truncat") {
|
||||||
|
t.Fatalf("expected validation error: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
for _, bad := range []string{strings.Replace(reasoningJSON, "reasoning_text", "output_text", 1), strings.Replace(reasoningJSON, "summary_text", "unknown", 1), `{"type":"reasoning","summary":{}}`, strings.Replace(visibleJSON, `"logprobs":[]`, `"logprobs":[{}]`, 1), strings.Replace(visibleJSON, `"phase":`, `"unknown":`, 1)} {
|
||||||
|
if _, err := ConvertResponse(Responses, Chat, []byte(`{"status":"completed","output":[`+bad+`]}`)); err == nil {
|
||||||
|
t.Fatal("accepted invalid output")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,879 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const maxStreamBytes = 16 << 20
|
||||||
|
const maxStreamBlocks = 4096
|
||||||
|
const maxStreamUsageBytes = 16 << 10
|
||||||
|
const maxPassthroughBytes = 64 << 20
|
||||||
|
|
||||||
|
type streamBlock struct {
|
||||||
|
part
|
||||||
|
content, identity, functionName strings.Builder
|
||||||
|
index int
|
||||||
|
closed bool
|
||||||
|
emitted bool
|
||||||
|
sent int
|
||||||
|
}
|
||||||
|
type streamState struct {
|
||||||
|
thinking map[int]*thinkingBlock
|
||||||
|
ctx context.Context
|
||||||
|
w io.Writer
|
||||||
|
from, to Kind
|
||||||
|
r response
|
||||||
|
blocks []*streamBlock
|
||||||
|
source map[int]*streamBlock
|
||||||
|
reasoning map[int]*reasoningItem
|
||||||
|
started, finished, terminal bool
|
||||||
|
sequence int
|
||||||
|
retained, blockCount int
|
||||||
|
byteLimit, blockLimit int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamState) retain(n int) error {
|
||||||
|
if n > s.byteLimit-s.retained {
|
||||||
|
return fmt.Errorf("stream retained content exceeds %d bytes", s.byteLimit)
|
||||||
|
}
|
||||||
|
s.retained += n
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) reserveBlock() error {
|
||||||
|
if s.blockCount >= s.blockLimit {
|
||||||
|
return fmt.Errorf("stream exceeds %d blocks", s.blockLimit)
|
||||||
|
}
|
||||||
|
s.blockCount++
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeFlush(w io.Writer, b []byte) error {
|
||||||
|
n, e := w.Write(b)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if n != len(b) {
|
||||||
|
return io.ErrShortWrite
|
||||||
|
}
|
||||||
|
if f, ok := w.(interface{ Flush() error }); ok {
|
||||||
|
return f.Flush()
|
||||||
|
}
|
||||||
|
if f, ok := w.(interface{ Flush() }); ok {
|
||||||
|
f.Flush()
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) event(name string, m object) error {
|
||||||
|
if e := s.ctx.Err(); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if s.to == Responses {
|
||||||
|
m["sequence_number"] = s.sequence
|
||||||
|
s.sequence++
|
||||||
|
}
|
||||||
|
b, e := json.Marshal(m)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
prefix := ""
|
||||||
|
if name != "" {
|
||||||
|
prefix = "event: " + name + "\n"
|
||||||
|
}
|
||||||
|
return writeFlush(s.w, []byte(prefix+"data: "+string(b)+"\n\n"))
|
||||||
|
}
|
||||||
|
func (s *streamState) typed(name string, m object) error { m["type"] = name; return s.event(name, m) }
|
||||||
|
func (s *streamState) chat(delta object, finish any) error {
|
||||||
|
return s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{object{"index": 0, "delta": delta, "finish_reason": finish}}})
|
||||||
|
}
|
||||||
|
func (s *streamState) start() error {
|
||||||
|
if s.started {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
s.started = true
|
||||||
|
if s.r.id == "" {
|
||||||
|
s.r.id = newID()
|
||||||
|
}
|
||||||
|
switch s.to {
|
||||||
|
case Chat:
|
||||||
|
return s.chat(object{"role": "assistant", "content": ""}, nil)
|
||||||
|
case Anthropic:
|
||||||
|
m := object{"id": s.r.id, "type": "message", "role": "assistant", "model": s.r.model, "content": []any{}, "stop_reason": nil, "stop_sequence": nil}
|
||||||
|
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
||||||
|
return s.typed("message_start", object{"message": m})
|
||||||
|
case Responses:
|
||||||
|
r := object{"id": s.r.id, "object": "response", "created_at": time.Now().Unix(), "model": s.r.model, "status": "in_progress", "output": []any{}, "error": nil, "incomplete_details": nil, "usage": nil}
|
||||||
|
if e := s.typed("response.created", object{"response": r}); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
return s.typed("response.in_progress", object{"response": r})
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) itemID(b *streamBlock) string {
|
||||||
|
if b.typ == "call" {
|
||||||
|
return fmt.Sprintf("fc_%s_%d", s.r.id, b.index)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("msg_%s_%d", s.r.id, b.index)
|
||||||
|
}
|
||||||
|
func (s *streamState) item(b *streamBlock, done bool) object {
|
||||||
|
status := "in_progress"
|
||||||
|
if done {
|
||||||
|
status = "completed"
|
||||||
|
}
|
||||||
|
if b.typ == "call" {
|
||||||
|
args := ""
|
||||||
|
if done {
|
||||||
|
args = b.args
|
||||||
|
}
|
||||||
|
return object{"id": s.itemID(b), "type": "function_call", "call_id": b.id, "name": b.name, "arguments": args, "status": status}
|
||||||
|
}
|
||||||
|
content := []any{}
|
||||||
|
if done {
|
||||||
|
content = append(content, object{"type": "output_text", "text": b.text, "annotations": []any{}})
|
||||||
|
}
|
||||||
|
return object{"id": s.itemID(b), "type": "message", "role": "assistant", "status": status, "content": content}
|
||||||
|
}
|
||||||
|
func (s *streamState) open(b *streamBlock) error {
|
||||||
|
if b.emitted {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if b.typ == "call" && (b.id == "" || b.name == "") {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if e := s.start(); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
b.emitted = true
|
||||||
|
switch s.to {
|
||||||
|
case Chat:
|
||||||
|
if b.typ == "call" {
|
||||||
|
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "id": b.id, "type": "function", "function": object{"name": b.name, "arguments": ""}}}}, nil)
|
||||||
|
}
|
||||||
|
case Anthropic:
|
||||||
|
c := object{"type": "text", "text": ""}
|
||||||
|
if b.typ == "call" {
|
||||||
|
c = object{"type": "tool_use", "id": b.id, "name": b.name, "input": object{}}
|
||||||
|
}
|
||||||
|
return s.typed("content_block_start", object{"index": b.index, "content_block": c})
|
||||||
|
case Responses:
|
||||||
|
if e := s.typed("response.output_item.added", object{"output_index": b.index, "item": s.item(b, false)}); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if b.typ == "text" {
|
||||||
|
return s.typed("response.content_part.added", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": "", "annotations": []any{}}})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) toolIndex(b *streamBlock) int {
|
||||||
|
i := 0
|
||||||
|
for _, x := range s.blocks {
|
||||||
|
if x == b {
|
||||||
|
return i
|
||||||
|
}
|
||||||
|
if x.typ == "call" {
|
||||||
|
i++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return i
|
||||||
|
}
|
||||||
|
func (s *streamState) add(key int, p part) (*streamBlock, error) {
|
||||||
|
if s.finished {
|
||||||
|
return nil, fmt.Errorf("content after finish")
|
||||||
|
}
|
||||||
|
if _, ok := s.source[key]; ok {
|
||||||
|
return nil, fmt.Errorf("duplicate content index %d", key)
|
||||||
|
}
|
||||||
|
if e := s.reserveBlock(); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if e := s.retain(len(p.text) + len(p.args) + len(p.id) + len(p.name)); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
b := &streamBlock{part: p, index: len(s.blocks)}
|
||||||
|
b.content.WriteString(p.text)
|
||||||
|
b.content.WriteString(p.args)
|
||||||
|
b.identity.WriteString(p.id)
|
||||||
|
b.functionName.WriteString(p.name)
|
||||||
|
s.blocks = append(s.blocks, b)
|
||||||
|
s.source[key] = b
|
||||||
|
return b, s.open(b)
|
||||||
|
}
|
||||||
|
func (s *streamState) delta(b *streamBlock, text string) error {
|
||||||
|
if b == nil || b.closed || s.finished {
|
||||||
|
return fmt.Errorf("delta outside active content block")
|
||||||
|
}
|
||||||
|
if e := s.retain(len(text)); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
b.content.WriteString(text)
|
||||||
|
// Builder.String shares storage without copying accumulated content.
|
||||||
|
if b.typ == "text" {
|
||||||
|
b.text = b.content.String()
|
||||||
|
} else {
|
||||||
|
b.args = b.content.String()
|
||||||
|
}
|
||||||
|
if e := s.open(b); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if !b.emitted {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
whole := b.text
|
||||||
|
if b.typ == "call" {
|
||||||
|
whole = b.args
|
||||||
|
}
|
||||||
|
text = whole[b.sent:]
|
||||||
|
b.sent = len(whole)
|
||||||
|
if text == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
switch s.to {
|
||||||
|
case Chat:
|
||||||
|
if b.typ == "text" {
|
||||||
|
return s.chat(object{"content": text}, nil)
|
||||||
|
}
|
||||||
|
return s.chat(object{"tool_calls": []any{object{"index": s.toolIndex(b), "function": object{"arguments": text}}}}, nil)
|
||||||
|
case Anthropic:
|
||||||
|
d := object{"type": "text_delta", "text": text}
|
||||||
|
if b.typ == "call" {
|
||||||
|
d = object{"type": "input_json_delta", "partial_json": text}
|
||||||
|
}
|
||||||
|
return s.typed("content_block_delta", object{"index": b.index, "delta": d})
|
||||||
|
case Responses:
|
||||||
|
name := "response.output_text.delta"
|
||||||
|
m := object{"output_index": b.index, "item_id": s.itemID(b), "delta": text}
|
||||||
|
if b.typ == "call" {
|
||||||
|
name = "response.function_call_arguments.delta"
|
||||||
|
} else {
|
||||||
|
m["content_index"] = 0
|
||||||
|
}
|
||||||
|
return s.typed(name, m)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) close(b *streamBlock) error {
|
||||||
|
if b == nil {
|
||||||
|
return fmt.Errorf("unknown content block")
|
||||||
|
}
|
||||||
|
if b.closed {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if b.typ == "call" {
|
||||||
|
if b.id == "" || b.name == "" {
|
||||||
|
return fmt.Errorf("incomplete tool identity")
|
||||||
|
}
|
||||||
|
if b.args == "" {
|
||||||
|
if e := s.delta(b, "{}"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, e := decode([]byte(b.args)); e != nil {
|
||||||
|
return fmt.Errorf("invalid streamed tool arguments: %w", e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e := s.open(b); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if e := s.delta(b, ""); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
b.closed = true
|
||||||
|
switch s.to {
|
||||||
|
case Anthropic:
|
||||||
|
return s.typed("content_block_stop", object{"index": b.index})
|
||||||
|
case Responses:
|
||||||
|
m := object{"output_index": b.index, "item_id": s.itemID(b)}
|
||||||
|
name := "response.function_call_arguments.done"
|
||||||
|
if b.typ == "call" {
|
||||||
|
m["arguments"] = b.args
|
||||||
|
} else {
|
||||||
|
name = "response.output_text.done"
|
||||||
|
m["text"] = b.text
|
||||||
|
m["content_index"] = 0
|
||||||
|
}
|
||||||
|
if e := s.typed(name, m); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if b.typ == "text" {
|
||||||
|
if e := s.typed("response.content_part.done", object{"output_index": b.index, "item_id": s.itemID(b), "content_index": 0, "part": object{"type": "output_text", "text": b.text, "annotations": []any{}}}); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return s.typed("response.output_item.done", object{"output_index": b.index, "item": s.item(b, true)})
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) finish(reason string) error {
|
||||||
|
if s.finished {
|
||||||
|
return fmt.Errorf("duplicate finish")
|
||||||
|
}
|
||||||
|
if e := s.start(); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
for _, b := range s.blocks {
|
||||||
|
if e := s.close(b); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
s.r.finish = reason
|
||||||
|
s.finished = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (s *streamState) end() error {
|
||||||
|
if !s.finished {
|
||||||
|
return fmt.Errorf("stream ended without finish reason")
|
||||||
|
}
|
||||||
|
switch s.to {
|
||||||
|
case Chat:
|
||||||
|
if e := s.chat(object{}, wireFinish(Chat, s.r.finish)); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if s.r.usage.Known {
|
||||||
|
if e := s.event("", object{"id": s.r.id, "object": "chat.completion.chunk", "created": 0, "model": s.r.model, "choices": []any{}, "usage": usageJSON(Chat, s.r.usage)}); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return writeFlush(s.w, []byte("data: [DONE]\n\n"))
|
||||||
|
case Anthropic:
|
||||||
|
m := object{"delta": object{"stop_reason": wireFinish(Anthropic, s.r.finish), "stop_sequence": nil}}
|
||||||
|
if s.r.usage.Known {
|
||||||
|
m["usage"] = usageJSON(Anthropic, s.r.usage)
|
||||||
|
}
|
||||||
|
if e := s.typed("message_delta", m); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
return s.typed("message_stop", object{})
|
||||||
|
case Responses:
|
||||||
|
r := s.r
|
||||||
|
for _, b := range s.blocks {
|
||||||
|
r.parts = append(r.parts, b.part)
|
||||||
|
}
|
||||||
|
m, e := writeResponse(Responses, r)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
items := []any{}
|
||||||
|
for _, b := range s.blocks {
|
||||||
|
items = append(items, s.item(b, true))
|
||||||
|
}
|
||||||
|
m["output"] = items
|
||||||
|
name := "response.completed"
|
||||||
|
if r.finish == "length" {
|
||||||
|
name = "response.incomplete"
|
||||||
|
}
|
||||||
|
return s.typed(name, object{"response": m})
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func mergeUsage(k Kind, old Usage, m object) (Usage, error) {
|
||||||
|
if r := obj(m["response"]); r != nil {
|
||||||
|
m = r
|
||||||
|
}
|
||||||
|
if r := obj(m["message"]); r != nil {
|
||||||
|
m = r
|
||||||
|
}
|
||||||
|
u := obj(m["usage"])
|
||||||
|
if u == nil {
|
||||||
|
return old, nil
|
||||||
|
}
|
||||||
|
raw, err := json.Marshal(u)
|
||||||
|
if err != nil {
|
||||||
|
return old, err
|
||||||
|
}
|
||||||
|
if len(raw) > maxStreamUsageBytes {
|
||||||
|
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
||||||
|
}
|
||||||
|
if len(old.Raw) > 0 {
|
||||||
|
prior, e := decode(old.Raw)
|
||||||
|
if e != nil {
|
||||||
|
return old, e
|
||||||
|
}
|
||||||
|
for key, v := range u {
|
||||||
|
prior[key] = v
|
||||||
|
}
|
||||||
|
u = prior
|
||||||
|
}
|
||||||
|
n := usage(k, u)
|
||||||
|
if len(n.Raw) > maxStreamUsageBytes {
|
||||||
|
return old, fmt.Errorf("stream usage exceeds %d bytes", maxStreamUsageBytes)
|
||||||
|
}
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
func streamError(m object) error {
|
||||||
|
if m["error"] != nil {
|
||||||
|
return fmt.Errorf("upstream stream error: %v", m["error"])
|
||||||
|
}
|
||||||
|
switch str(m["type"]) {
|
||||||
|
case "error", "response.failed":
|
||||||
|
return fmt.Errorf("upstream stream error: %v", m)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamState) consume(m object, event string) error {
|
||||||
|
if e := streamError(m); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
u, err := mergeUsage(s.from, s.r.usage, m)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.r.usage = u
|
||||||
|
if s.from == Chat {
|
||||||
|
if id := str(m["id"]); id != "" && !s.started {
|
||||||
|
s.r.id = id
|
||||||
|
}
|
||||||
|
if s.r.model == "" {
|
||||||
|
s.r.model = str(m["model"])
|
||||||
|
}
|
||||||
|
cs := arr(m["choices"])
|
||||||
|
if len(cs) > 1 {
|
||||||
|
return fmt.Errorf("multiple streaming choices unsupported")
|
||||||
|
}
|
||||||
|
for _, v := range cs {
|
||||||
|
c := obj(v)
|
||||||
|
if num(c["index"]) != 0 {
|
||||||
|
return fmt.Errorf("nonzero choice index unsupported")
|
||||||
|
}
|
||||||
|
if c["logprobs"] != nil {
|
||||||
|
return fmt.Errorf("logprobs unsupported")
|
||||||
|
}
|
||||||
|
d, e := chatOutput(obj(c["delta"]))
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if e := fields(d, "role", "content", "tool_calls"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if role := str(d["role"]); role != "" && role != "assistant" {
|
||||||
|
return fmt.Errorf("invalid delta role")
|
||||||
|
}
|
||||||
|
if e := s.start(); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if text, ok := d["content"]; ok && text != nil {
|
||||||
|
t, ok := text.(string)
|
||||||
|
if !ok {
|
||||||
|
return fmt.Errorf("non-text streamed content")
|
||||||
|
}
|
||||||
|
b := s.source[-1]
|
||||||
|
if b == nil {
|
||||||
|
var e error
|
||||||
|
b, e = s.add(-1, part{typ: "text"})
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e := s.delta(b, t); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, v := range arr(d["tool_calls"]) {
|
||||||
|
t := obj(v)
|
||||||
|
if e := fields(t, "index", "id", "type", "function"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
idx, ok := t["index"]
|
||||||
|
if !ok || num(idx) < 0 {
|
||||||
|
return fmt.Errorf("invalid tool index")
|
||||||
|
}
|
||||||
|
if typ := str(t["type"]); typ != "" && typ != "function" {
|
||||||
|
return fmt.Errorf("unsupported streamed tool")
|
||||||
|
}
|
||||||
|
f := obj(t["function"])
|
||||||
|
if e := fields(f, "name", "arguments"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
key := int(num(idx))
|
||||||
|
b := s.source[key]
|
||||||
|
if b == nil {
|
||||||
|
var e error
|
||||||
|
b, e = s.add(key, part{typ: "call"})
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if b.emitted && (str(t["id"]) != "" || str(f["name"]) != "") {
|
||||||
|
return fmt.Errorf("fragmented tool identity after arguments unsupported")
|
||||||
|
}
|
||||||
|
if e := s.retain(len(str(t["id"])) + len(str(f["name"]))); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
b.identity.WriteString(str(t["id"]))
|
||||||
|
b.functionName.WriteString(str(f["name"]))
|
||||||
|
b.id = b.identity.String()
|
||||||
|
b.name = b.functionName.String()
|
||||||
|
if e := s.delta(b, str(f["arguments"])); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if reason := str(c["finish_reason"]); reason != "" {
|
||||||
|
f, e := finishReason(Chat, reason)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if e = s.finish(f); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
typ := str(m["type"])
|
||||||
|
if typ == "" {
|
||||||
|
typ = event
|
||||||
|
}
|
||||||
|
if event != "" && typ != event {
|
||||||
|
return fmt.Errorf("SSE event/type mismatch")
|
||||||
|
}
|
||||||
|
// Discarded blocks still retain lifecycle bookkeeping in separate maps.
|
||||||
|
if (s.from == Anthropic && typ == "content_block_start" && (str(obj(m["content_block"])["type"]) == "thinking" || str(obj(m["content_block"])["type"]) == "redacted_thinking")) ||
|
||||||
|
(s.from == Responses && typ == "response.output_item.added" && str(obj(m["item"])["type"]) == "reasoning") ||
|
||||||
|
(s.from == Responses && (typ == "response.content_part.added" || typ == "response.reasoning_summary_part.added") && s.reasoning[int(num(m["output_index"]))] != nil) {
|
||||||
|
if e := s.reserveBlock(); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if typ == "response.output_item.added" {
|
||||||
|
if e := s.retain(len(str(obj(m["item"])["id"]))); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if s.from == Anthropic {
|
||||||
|
if handled, err := s.thinkingEvent(typ, m); handled || err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
switch typ {
|
||||||
|
case "ping":
|
||||||
|
return nil
|
||||||
|
case "message_start":
|
||||||
|
if s.started {
|
||||||
|
return fmt.Errorf("duplicate message_start")
|
||||||
|
}
|
||||||
|
x := obj(m["message"])
|
||||||
|
s.r.id = str(x["id"])
|
||||||
|
if s.r.model == "" {
|
||||||
|
s.r.model = str(x["model"])
|
||||||
|
}
|
||||||
|
if len(arr(x["content"])) != 0 {
|
||||||
|
return fmt.Errorf("nonempty message_start content")
|
||||||
|
}
|
||||||
|
return s.start()
|
||||||
|
case "content_block_start":
|
||||||
|
x := obj(m["content_block"])
|
||||||
|
p := part{typ: "text"}
|
||||||
|
initial := ""
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "text":
|
||||||
|
if e := fields(x, "type", "text"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
initial = str(x["text"])
|
||||||
|
case "tool_use":
|
||||||
|
if e := fields(x, "type", "id", "name", "input"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
p = part{typ: "call", id: str(x["id"]), name: str(x["name"])}
|
||||||
|
if obj(x["input"]) == nil {
|
||||||
|
return fmt.Errorf("tool input must be an object")
|
||||||
|
}
|
||||||
|
if len(obj(x["input"])) > 0 {
|
||||||
|
b, e := json.Marshal(x["input"])
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
initial = string(b)
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported streamed block %q", str(x["type"]))
|
||||||
|
}
|
||||||
|
b, e := s.add(int(num(m["index"])), p)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
return s.delta(b, initial)
|
||||||
|
case "content_block_delta":
|
||||||
|
b := s.source[int(num(m["index"]))]
|
||||||
|
d := obj(m["delta"])
|
||||||
|
switch str(d["type"]) {
|
||||||
|
case "text_delta":
|
||||||
|
if b == nil || b.typ != "text" {
|
||||||
|
return fmt.Errorf("text delta block mismatch")
|
||||||
|
}
|
||||||
|
return s.delta(b, str(d["text"]))
|
||||||
|
case "input_json_delta":
|
||||||
|
if b == nil || b.typ != "call" {
|
||||||
|
return fmt.Errorf("tool delta block mismatch")
|
||||||
|
}
|
||||||
|
return s.delta(b, str(d["partial_json"]))
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported delta %q", str(d["type"]))
|
||||||
|
}
|
||||||
|
case "content_block_stop":
|
||||||
|
return s.close(s.source[int(num(m["index"]))])
|
||||||
|
case "message_delta":
|
||||||
|
d := obj(m["delta"])
|
||||||
|
if e := fields(d, "stop_reason", "stop_sequence"); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if d["stop_sequence"] != nil {
|
||||||
|
return fmt.Errorf("stop sequence cannot be preserved")
|
||||||
|
}
|
||||||
|
f, e := finishReason(Anthropic, str(d["stop_reason"]))
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
return s.finish(f)
|
||||||
|
case "message_stop":
|
||||||
|
s.terminal = true
|
||||||
|
return nil
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported Anthropic event %q", typ)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if handled, err := s.reasoningEvent(typ, m); handled || err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
switch typ {
|
||||||
|
case "response.created", "response.in_progress":
|
||||||
|
x := obj(m["response"])
|
||||||
|
if !s.started {
|
||||||
|
s.r.id = str(x["id"])
|
||||||
|
if s.r.model == "" {
|
||||||
|
s.r.model = str(x["model"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return s.start()
|
||||||
|
case "response.output_item.added":
|
||||||
|
x := obj(m["item"])
|
||||||
|
key := int(num(m["output_index"]))
|
||||||
|
p := part{typ: "text"}
|
||||||
|
switch str(x["type"]) {
|
||||||
|
case "message":
|
||||||
|
if len(arr(x["content"])) > 0 {
|
||||||
|
return fmt.Errorf("nonempty initial response content")
|
||||||
|
}
|
||||||
|
case "function_call":
|
||||||
|
if _, ok := x["arguments"].(string); !ok {
|
||||||
|
return fmt.Errorf("function arguments must be a string")
|
||||||
|
}
|
||||||
|
p = part{typ: "call", id: str(x["call_id"]), name: str(x["name"])}
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported output item %q", str(x["type"]))
|
||||||
|
}
|
||||||
|
b, e := s.add(key, p)
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if p.typ == "call" {
|
||||||
|
return s.delta(b, str(x["arguments"]))
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
case "response.content_part.added":
|
||||||
|
p := obj(m["part"])
|
||||||
|
if str(p["type"]) != "output_text" || num(m["content_index"]) != 0 {
|
||||||
|
return fmt.Errorf("unsupported response content part")
|
||||||
|
}
|
||||||
|
if len(arr(p["annotations"])) > 0 {
|
||||||
|
return fmt.Errorf("annotations unsupported")
|
||||||
|
}
|
||||||
|
return s.delta(s.source[int(num(m["output_index"]))], str(p["text"]))
|
||||||
|
case "response.output_text.delta", "response.function_call_arguments.delta":
|
||||||
|
b := s.source[int(num(m["output_index"]))]
|
||||||
|
if b == nil {
|
||||||
|
return fmt.Errorf("delta before output item")
|
||||||
|
}
|
||||||
|
if (typ == "response.output_text.delta") != (b.typ == "text") || num(m["content_index"]) != 0 {
|
||||||
|
return fmt.Errorf("response delta block mismatch")
|
||||||
|
}
|
||||||
|
return s.delta(b, str(m["delta"]))
|
||||||
|
case "response.output_text.done", "response.function_call_arguments.done":
|
||||||
|
b := s.source[int(num(m["output_index"]))]
|
||||||
|
if b == nil {
|
||||||
|
return fmt.Errorf("done before output item")
|
||||||
|
}
|
||||||
|
full := str(m["text"])
|
||||||
|
got := b.text
|
||||||
|
if b.typ == "call" {
|
||||||
|
full = str(m["arguments"])
|
||||||
|
got = b.args
|
||||||
|
}
|
||||||
|
if full != got {
|
||||||
|
return fmt.Errorf("final content disagrees with deltas")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
case "response.content_part.done":
|
||||||
|
p := obj(m["part"])
|
||||||
|
if str(p["type"]) != "output_text" || len(arr(p["annotations"])) > 0 || num(m["content_index"]) != 0 {
|
||||||
|
return fmt.Errorf("unsupported content part")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
case "response.output_item.done":
|
||||||
|
return s.close(s.source[int(num(m["output_index"]))])
|
||||||
|
case "response.completed", "response.incomplete":
|
||||||
|
r, e := readResponse(Responses, obj(m["response"]))
|
||||||
|
if e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if len(r.parts) != len(s.blocks) {
|
||||||
|
return fmt.Errorf("final response disagrees with streamed items")
|
||||||
|
}
|
||||||
|
for i, p := range r.parts {
|
||||||
|
b := s.blocks[i]
|
||||||
|
if p.typ != b.typ || p.text != b.text || p.args != b.args || p.id != b.id || p.name != b.name {
|
||||||
|
return fmt.Errorf("final response disagrees with streamed content")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e = s.finish(r.finish); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
s.terminal = true
|
||||||
|
return nil
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unsupported Responses event %q", typ)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stream does not own r. A caller must cancel or close a blocked upstream read.
|
||||||
|
func Stream(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string) (Usage, error) {
|
||||||
|
return streamWithLimits(ctx, w, r, from, to, model, maxStreamBytes, maxStreamBlocks, maxPassthroughBytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
func readSSELine(ctx context.Context, br *bufio.Reader, remaining int) (string, error) {
|
||||||
|
var line strings.Builder
|
||||||
|
for {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
fragment, err := br.ReadSlice('\n')
|
||||||
|
if len(fragment) > remaining-line.Len() {
|
||||||
|
return "", fmt.Errorf("SSE line/event exceeds size limit")
|
||||||
|
}
|
||||||
|
line.Write(fragment)
|
||||||
|
if err != bufio.ErrBufferFull {
|
||||||
|
return line.String(), err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func streamWithLimits(ctx context.Context, w io.Writer, r io.Reader, from, to Kind, model string, byteLimit, blockLimit, totalLimit int) (Usage, error) {
|
||||||
|
if !valid(from) || !valid(to) {
|
||||||
|
return Usage{}, fmt.Errorf("unknown protocol")
|
||||||
|
}
|
||||||
|
s := &streamState{ctx: ctx, w: w, from: from, to: to, r: response{model: model}, source: map[int]*streamBlock{}, byteLimit: byteLimit, blockLimit: blockLimit}
|
||||||
|
br := bufio.NewReader(r)
|
||||||
|
data := []string{}
|
||||||
|
event := ""
|
||||||
|
size := 0
|
||||||
|
total := 0
|
||||||
|
dispatch := func() error {
|
||||||
|
if len(data) == 0 {
|
||||||
|
event = ""
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
payload := strings.Join(data, "\n")
|
||||||
|
data = nil
|
||||||
|
ev := event
|
||||||
|
event = ""
|
||||||
|
size = 0
|
||||||
|
if payload == "[DONE]" {
|
||||||
|
if from != Chat {
|
||||||
|
return fmt.Errorf("unexpected DONE marker")
|
||||||
|
}
|
||||||
|
s.terminal = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
m, e := decode([]byte(payload))
|
||||||
|
if e != nil {
|
||||||
|
return fmt.Errorf("invalid SSE JSON: %w", e)
|
||||||
|
}
|
||||||
|
if e = streamError(m); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
if from == to {
|
||||||
|
u, err := mergeUsage(from, s.r.usage, m)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.r.usage = u
|
||||||
|
switch str(m["type"]) {
|
||||||
|
case "message_stop", "response.completed", "response.incomplete":
|
||||||
|
s.terminal = true
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if s.terminal {
|
||||||
|
return fmt.Errorf("data after terminal event")
|
||||||
|
}
|
||||||
|
return s.consume(m, ev)
|
||||||
|
}
|
||||||
|
for {
|
||||||
|
if e := ctx.Err(); e != nil {
|
||||||
|
return s.r.usage, e
|
||||||
|
}
|
||||||
|
line, e := readSSELine(ctx, br, byteLimit-size)
|
||||||
|
if e != nil && e != io.EOF {
|
||||||
|
return s.r.usage, e
|
||||||
|
}
|
||||||
|
if len(line) > 0 {
|
||||||
|
if from == to {
|
||||||
|
if total+len(line) > totalLimit {
|
||||||
|
return s.r.usage, fmt.Errorf("stream passthrough exceeds %d bytes", totalLimit)
|
||||||
|
}
|
||||||
|
total += len(line)
|
||||||
|
if err := writeFlush(w, []byte(line)); err != nil {
|
||||||
|
return s.r.usage, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
size += len(line)
|
||||||
|
line = strings.TrimSuffix(strings.TrimSuffix(line, "\n"), "\r")
|
||||||
|
if line == "" {
|
||||||
|
if err := dispatch(); err != nil {
|
||||||
|
return s.r.usage, err
|
||||||
|
}
|
||||||
|
size = 0
|
||||||
|
if s.terminal {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
} else if !strings.HasPrefix(line, ":") {
|
||||||
|
key, val, _ := strings.Cut(line, ":")
|
||||||
|
val = strings.TrimPrefix(val, " ")
|
||||||
|
switch key {
|
||||||
|
case "data":
|
||||||
|
data = append(data, val)
|
||||||
|
case "event":
|
||||||
|
event = val
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e != nil {
|
||||||
|
if e != io.EOF {
|
||||||
|
return s.r.usage, e
|
||||||
|
}
|
||||||
|
if err := dispatch(); err != nil {
|
||||||
|
return s.r.usage, err
|
||||||
|
}
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if e := ctx.Err(); e != nil {
|
||||||
|
return s.r.usage, e
|
||||||
|
}
|
||||||
|
if from == to {
|
||||||
|
return s.r.usage, nil
|
||||||
|
}
|
||||||
|
if !s.terminal {
|
||||||
|
return s.r.usage, fmt.Errorf("truncated SSE stream: missing terminal event")
|
||||||
|
}
|
||||||
|
if e := s.end(); e != nil {
|
||||||
|
return s.r.usage, e
|
||||||
|
}
|
||||||
|
return s.r.usage, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,164 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func limitEvent(s string) string { return "data: " + s + "\n\n" }
|
||||||
|
func TestSSEReadLimits(t *testing.T) {
|
||||||
|
for _, to := range []Kind{Chat, Anthropic} {
|
||||||
|
var out bytes.Buffer
|
||||||
|
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat("x", 10000)), Chat, to, "", 64, 8, 64)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "size limit") || out.Len() != 0 {
|
||||||
|
t.Fatalf("%s: %v", to, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, input := range []string{"12345678", "1234567\n"} {
|
||||||
|
line, err := readSSELine(context.Background(), bufio.NewReader(strings.NewReader(input)), 8)
|
||||||
|
if line != input || (err != nil && err != io.EOF) {
|
||||||
|
t.Fatalf("%q %v", line, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var out bytes.Buffer
|
||||||
|
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(":123456\n:123456\n"), Chat, Chat, "", 12, 8, 16)
|
||||||
|
if err == nil || out.String() != ":123456\n" {
|
||||||
|
t.Fatalf("%q %v", out.String(), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStreamPassthroughTotalLimit(t *testing.T) {
|
||||||
|
frame := "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":null}]}\n\n"
|
||||||
|
var out bytes.Buffer
|
||||||
|
_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat(frame, 4)), Chat, Chat, "", 512, 8, 2*len(frame))
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "passthrough") || out.Len() != 2*len(frame) {
|
||||||
|
t.Fatalf("out=%d %q %v", out.Len(), out.String(), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStreamCumulativeLimits(t *testing.T) {
|
||||||
|
for _, field := range []string{"text", "arguments", "identity"} {
|
||||||
|
delta := object{"content": strings.Repeat("x", 100)}
|
||||||
|
if field != "text" {
|
||||||
|
call := object{"index": 0}
|
||||||
|
if field == "arguments" {
|
||||||
|
call["function"] = object{"arguments": strings.Repeat("x", 100)}
|
||||||
|
} else {
|
||||||
|
call["id"] = strings.Repeat("x", 100)
|
||||||
|
}
|
||||||
|
delta = object{"tool_calls": []any{call}}
|
||||||
|
}
|
||||||
|
event := limitEvent(fmt.Sprintf("{\"choices\":[{\"index\":0,\"delta\":%s}]}", mustLimitJSON(t, delta)))
|
||||||
|
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(strings.Repeat(event, 6)), Chat, Responses, "", 512, 8, 512)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "retained content") {
|
||||||
|
t.Fatalf("%s: %v", field, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStreamBlockLimits(t *testing.T) {
|
||||||
|
for _, kind := range []string{"text", "thinking", "reasoning", "summary", "tool"} {
|
||||||
|
from := Anthropic
|
||||||
|
in := limitEvent("{\"type\":\"message_start\",\"message\":{\"content\":[]}}")
|
||||||
|
if kind == "reasoning" || kind == "summary" || kind == "tool" {
|
||||||
|
from = Responses
|
||||||
|
in = ""
|
||||||
|
}
|
||||||
|
if kind == "summary" {
|
||||||
|
in = limitEvent("{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"id\":\"r\"}}")
|
||||||
|
}
|
||||||
|
for i := 0; i < 4; i++ {
|
||||||
|
m := object{"type": "content_block_start", "index": i, "content_block": object{"type": kind, "text": ""}}
|
||||||
|
switch kind {
|
||||||
|
case "thinking":
|
||||||
|
m["content_block"] = object{"type": "thinking", "thinking": "", "signature": ""}
|
||||||
|
case "reasoning", "tool":
|
||||||
|
item := object{"type": "reasoning", "id": "r"}
|
||||||
|
if kind == "tool" {
|
||||||
|
item = object{"type": "function_call", "call_id": "c", "name": "f", "arguments": ""}
|
||||||
|
}
|
||||||
|
m = object{"type": "response.output_item.added", "output_index": i, "item": item}
|
||||||
|
case "summary":
|
||||||
|
m = object{"type": "response.reasoning_summary_part.added", "output_index": 0, "item_id": "r", "summary_index": i, "part": object{"type": "summary_text", "text": ""}}
|
||||||
|
}
|
||||||
|
in += limitEvent(mustLimitJSON(t, m))
|
||||||
|
}
|
||||||
|
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in), from, Chat, "", 1024, 3, 1024)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "3 blocks") {
|
||||||
|
t.Fatalf("%s: %v", kind, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReasoningIdentityCumulativeLimit(t *testing.T) {
|
||||||
|
var in strings.Builder
|
||||||
|
for i := 0; i < 6; i++ {
|
||||||
|
in.WriteString(limitEvent(mustLimitJSON(t, object{"type": "response.output_item.added", "output_index": i, "item": object{"type": "reasoning", "id": strings.Repeat("r", 100)}})))
|
||||||
|
}
|
||||||
|
_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in.String()), Responses, Chat, "", 512, 8, 512)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "retained content") {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStreamBufferSemantics(t *testing.T) {
|
||||||
|
var out bytes.Buffer
|
||||||
|
s := &streamState{ctx: context.Background(), w: &out, from: Chat, to: Responses, source: map[int]*streamBlock{}, byteLimit: 64, blockLimit: 4}
|
||||||
|
b, err := s.add(0, part{typ: "call"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err = s.delta(b, "{\"a\":"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if out.Len() != 0 {
|
||||||
|
t.Fatal("incomplete identity emitted")
|
||||||
|
}
|
||||||
|
b.id, b.name = "c", "f"
|
||||||
|
if err = s.delta(b, "1}"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if b.args != "{\"a\":1}" || b.sent != 7 {
|
||||||
|
t.Fatal(b.part)
|
||||||
|
}
|
||||||
|
if err = s.close(b); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
text, err := s.add(1, part{typ: "text"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, v := range []string{"ab", "cd"} {
|
||||||
|
if err = s.delta(text, v); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if text.text != "abcd" {
|
||||||
|
t.Fatal(text.text)
|
||||||
|
}
|
||||||
|
if err = s.finish("stop"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err = s.end(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(out.String(), "abcd") || !strings.Contains(out.String(), "response.completed") {
|
||||||
|
t.Fatal(out.String())
|
||||||
|
}
|
||||||
|
s.finished = false
|
||||||
|
text.closed = false
|
||||||
|
s.byteLimit = s.retained
|
||||||
|
if err = s.delta(text, "x"); err == nil || text.text != "abcd" {
|
||||||
|
t.Fatal("overflow changed content")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func mustLimitJSON(t *testing.T, m object) string {
|
||||||
|
t.Helper()
|
||||||
|
b, err := json.Marshal(m)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return string(b)
|
||||||
|
}
|
||||||
@@ -0,0 +1,89 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"reflect"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestMergeUsageSizeBoundary(t *testing.T) {
|
||||||
|
for _, kind := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
t.Run(string(kind), func(t *testing.T) {
|
||||||
|
m := object{"padding": strings.Repeat("x", maxStreamUsageBytes-len(`{"padding":""}`))}
|
||||||
|
u, err := mergeUsage(kind, Usage{}, object{"usage": m})
|
||||||
|
if err != nil || len(u.Raw) != maxStreamUsageBytes {
|
||||||
|
t.Fatalf("exact limit: size=%d err=%v", len(u.Raw), err)
|
||||||
|
}
|
||||||
|
for _, update := range []object{
|
||||||
|
{"usage": object{"padding": str(m["padding"]) + "x"}},
|
||||||
|
{"usage": object{"new_key": 1}},
|
||||||
|
} {
|
||||||
|
got, err := mergeUsage(kind, u, update)
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") || !reflect.DeepEqual(got, u) {
|
||||||
|
t.Fatalf("overflow: retained=%d err=%v", len(got.Raw), err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
got, err := mergeUsage(kind, u, object{"usage": m})
|
||||||
|
if err != nil || !reflect.DeepEqual(got, u) {
|
||||||
|
t.Fatalf("replacement at limit: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStreamUsageSizeLimits(t *testing.T) {
|
||||||
|
for _, from := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
for _, to := range []Kind{Chat, Responses, Anthropic} {
|
||||||
|
for _, cumulative := range []bool{false, true} {
|
||||||
|
t.Run(fmt.Sprintf("%s/%s/cumulative=%t", from, to, cumulative), func(t *testing.T) {
|
||||||
|
var input strings.Builder
|
||||||
|
writeUsage := func(u object) {
|
||||||
|
m := object{"usage": u}
|
||||||
|
switch from {
|
||||||
|
case Anthropic:
|
||||||
|
m["type"] = "ping"
|
||||||
|
case Responses:
|
||||||
|
m = object{"type": "response.in_progress", "response": m}
|
||||||
|
}
|
||||||
|
input.WriteString(limitEvent(mustLimitJSON(t, m)))
|
||||||
|
}
|
||||||
|
writeUsage(object{"prompt_tokens": 7, "input_tokens": 7})
|
||||||
|
if cumulative {
|
||||||
|
for i := 0; i < 32; i++ {
|
||||||
|
writeUsage(object{fmt.Sprintf("key_%d", i): strings.Repeat("x", 1024)})
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
writeUsage(object{"padding": strings.Repeat("x", maxStreamUsageBytes)})
|
||||||
|
}
|
||||||
|
u, err := Stream(context.Background(), io.Discard, strings.NewReader(input.String()), from, to, "")
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") {
|
||||||
|
t.Fatalf("expected usage limit error, got %v", err)
|
||||||
|
}
|
||||||
|
if !u.Known || u.Input != 7 || len(u.Raw) > maxStreamUsageBytes {
|
||||||
|
t.Fatalf("lost valid usage: known=%t input=%d size=%d", u.Known, u.Input, len(u.Raw))
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMergeUsagePartialUpdates(t *testing.T) {
|
||||||
|
u, err := mergeUsage(Anthropic, Usage{}, object{"message": object{"usage": object{"input_tokens": 5, "cache_read_input_tokens": 2}}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
u, err = mergeUsage(Anthropic, u, object{"usage": object{"output_tokens": 3}})
|
||||||
|
if err != nil || !u.Known || u.Input != 7 || u.Output != 3 || u.CacheRead != 2 {
|
||||||
|
t.Fatalf("partial update: %+v err=%v", u, err)
|
||||||
|
}
|
||||||
|
got, err := mergeUsage(Anthropic, u, object{})
|
||||||
|
if err != nil || !reflect.DeepEqual(got, u) {
|
||||||
|
t.Fatalf("missing usage changed totals: %+v err=%v", got, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestMalformedInitialToolArguments(t *testing.T) {
|
||||||
|
for _, bad := range []string{`null`, `[]`, `12`, `"not-an-object"`} {
|
||||||
|
src := `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"tool1","name":"fn","input":` + bad + `}}` + "\n\n"
|
||||||
|
_, err := Stream(context.Background(), io.Discard, strings.NewReader(src), Anthropic, Chat, "m")
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "tool input must be an object") {
|
||||||
|
t.Fatalf("input %s: %v", bad, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, field := range []string{``, `,"arguments":null`, `,"arguments":{}`, `,"arguments":12`} {
|
||||||
|
src := `data: {"type":"response.output_item.added","output_index":0,"item":{"type":"function_call","call_id":"call1","name":"fn"` + field + `}}` + "\n\n"
|
||||||
|
_, err := Stream(context.Background(), io.Discard, strings.NewReader(src), Responses, Chat, "m")
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "function arguments must be a string") {
|
||||||
|
t.Fatalf("field %s: %v", field, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecodeNodeLimit(t *testing.T) {
|
||||||
|
dense := `{"model":"m","pad":[` + strings.Repeat("0,", maxJSONNodes) + `0]}`
|
||||||
|
if _, err := ConvertRequest(Chat, Chat, []byte(dense), "m"); err == nil || !strings.Contains(err.Error(), "exceeds") {
|
||||||
|
t.Fatalf("dense body accepted: %v", err)
|
||||||
|
}
|
||||||
|
under := `{"model":"m","pad":[` + strings.Repeat("0,", maxJSONNodes/2) + `0]}`
|
||||||
|
if _, err := ConvertRequest(Chat, Chat, []byte(under), "m"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := ConvertRequest(Chat, Chat, []byte(`{"broken":`), "m"); err == nil || strings.Contains(err.Error(), "exceeds") {
|
||||||
|
t.Fatalf("malformed JSON misreported: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,265 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type channel struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
BaseURL string `json:"base_url"`
|
||||||
|
Protocol string `json:"protocol"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
Models map[string]string `json:"models"`
|
||||||
|
APIKey string `json:"api_key,omitempty"`
|
||||||
|
}
|
||||||
|
type apiKey struct {
|
||||||
|
Revoked bool `json:"revoked"`
|
||||||
|
ID string `json:"id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Prefix string `json:"prefix"`
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
AllChannels bool `json:"all_channels"`
|
||||||
|
Channels []string `json:"channels"`
|
||||||
|
ExpiresAt string `json:"expires_at"`
|
||||||
|
CreatedAt string `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) loadChannels() ([]channel, error) {
|
||||||
|
rows, e := s.db.Query("SELECT id,name,base_url,protocol,enabled,models,api_key FROM channels ORDER BY name")
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
out := []channel{}
|
||||||
|
for rows.Next() {
|
||||||
|
var c channel
|
||||||
|
var m string
|
||||||
|
if e = rows.Scan(&c.ID, &c.Name, &c.BaseURL, &c.Protocol, &c.Enabled, &m, &c.APIKey); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if e = json.Unmarshal([]byte(m), &c.Models); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
out = append(out, c)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
func (s *Server) channels(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == "GET" {
|
||||||
|
cs, e := s.loadChannels()
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for i := range cs {
|
||||||
|
cs[i].APIKey = ""
|
||||||
|
}
|
||||||
|
writeJSON(w, cs)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.Method != "POST" {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var c channel
|
||||||
|
if decode(w, r, &c) != nil {
|
||||||
|
errJSON(w, 400, "invalid channel")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
u, e := url.Parse(c.BaseURL)
|
||||||
|
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(c.BaseURL, "#") || c.Name == "" || strings.Contains(c.Name, "/") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||||
|
errJSON(w, 400, "invalid channel")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if c.Models == nil {
|
||||||
|
c.Models = map[string]string{}
|
||||||
|
}
|
||||||
|
for a, m := range c.Models {
|
||||||
|
if a == "" || m == "" {
|
||||||
|
errJSON(w, 400, "invalid model mapping")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
c.BaseURL = strings.TrimRight(c.BaseURL, "/")
|
||||||
|
m, _ := json.Marshal(c.Models)
|
||||||
|
if c.ID == "" {
|
||||||
|
c.ID = randomID()
|
||||||
|
_, e = s.db.Exec("INSERT INTO channels VALUES(?,?,?,?,?,?,?,?)", c.ID, c.Name, c.BaseURL, c.Protocol, c.Enabled, false, string(m), c.APIKey)
|
||||||
|
} else {
|
||||||
|
var old string
|
||||||
|
if e = s.db.QueryRow("SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old); e == sql.ErrNoRows {
|
||||||
|
errJSON(w, 404, "not found")
|
||||||
|
return
|
||||||
|
} else if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if c.APIKey == "" {
|
||||||
|
c.APIKey = old
|
||||||
|
}
|
||||||
|
_, e = s.db.Exec("UPDATE channels SET name=?,base_url=?,protocol=?,enabled=?,models=?,api_key=? WHERE id=?", c.Name, c.BaseURL, c.Protocol, c.Enabled, string(m), c.APIKey, c.ID)
|
||||||
|
}
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 400, "channel could not be saved; names must be unique")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.APIKey = ""
|
||||||
|
writeJSON(w, c)
|
||||||
|
}
|
||||||
|
func (s *Server) channelID(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != "DELETE" {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
res, e := s.db.Exec("DELETE FROM channels WHERE id=?", strings.TrimPrefix(r.URL.Path, "/admin/api/channels/"))
|
||||||
|
mutation(w, res, e)
|
||||||
|
}
|
||||||
|
func mutation(w http.ResponseWriter, res sql.Result, e error) {
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
n, e := res.RowsAffected()
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if n == 0 {
|
||||||
|
errJSON(w, 404, "not found")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, map[string]bool{"ok": true})
|
||||||
|
}
|
||||||
|
func (s *Server) keys(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method == "GET" {
|
||||||
|
rows, e := s.db.Query("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at,revoked FROM keys ORDER BY created_at DESC")
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
out := []apiKey{}
|
||||||
|
for rows.Next() {
|
||||||
|
var k apiKey
|
||||||
|
var ch string
|
||||||
|
if rows.Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt, &k.Revoked) != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if json.Unmarshal([]byte(ch), &k.Channels) != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
out = append(out, k)
|
||||||
|
}
|
||||||
|
if rows.Err() != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, out)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.Method != "POST" {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var k apiKey
|
||||||
|
if decode(w, r, &k) != nil || k.Name == "" {
|
||||||
|
errJSON(w, 400, "invalid key")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if k.ExpiresAt != "" {
|
||||||
|
t, e := time.Parse(time.RFC3339, k.ExpiresAt)
|
||||||
|
if e != nil || !t.After(time.Now()) {
|
||||||
|
errJSON(w, 400, "invalid expiry")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
k.ExpiresAt = t.UTC().Format(time.RFC3339Nano)
|
||||||
|
}
|
||||||
|
if k.Channels == nil {
|
||||||
|
k.Channels = []string{}
|
||||||
|
}
|
||||||
|
for _, id := range k.Channels {
|
||||||
|
var n int
|
||||||
|
if s.db.QueryRow("SELECT count(*) FROM channels WHERE id=?", id).Scan(&n) != nil || n == 0 {
|
||||||
|
errJSON(w, 400, "invalid channel permission")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
raw := "sk-" + randomID()
|
||||||
|
ch, _ := json.Marshal(k.Channels)
|
||||||
|
_, e := s.db.Exec("INSERT INTO keys(id,name,prefix,hash,enabled,all_channels,channels,expires_at,created_at) VALUES(?,?,?,?,?,?,?,?,?)", randomID(), k.Name, raw[:11], hash(raw), true, k.AllChannels, string(ch), k.ExpiresAt, time.Now().UTC().Format(time.RFC3339Nano))
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, map[string]string{"key": raw})
|
||||||
|
}
|
||||||
|
func (s *Server) keyID(w http.ResponseWriter, r *http.Request) {
|
||||||
|
id := strings.TrimPrefix(r.URL.Path, "/admin/api/keys/")
|
||||||
|
if r.Method == "DELETE" {
|
||||||
|
res, e := s.db.Exec("UPDATE keys SET enabled=0,revoked=1 WHERE id=?", id)
|
||||||
|
mutation(w, res, e)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.Method == "PATCH" {
|
||||||
|
var x struct {
|
||||||
|
Enabled *bool `json:"enabled"`
|
||||||
|
}
|
||||||
|
if decode(w, r, &x) != nil || x.Enabled == nil {
|
||||||
|
errJSON(w, 400, "invalid request")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
res, e := s.db.Exec("UPDATE keys SET enabled=? WHERE id=? AND revoked=0", *x.Enabled, id)
|
||||||
|
mutation(w, res, e)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
}
|
||||||
|
func (s *Server) authenticate(r *http.Request) (apiKey, error) {
|
||||||
|
var k apiKey
|
||||||
|
raw := r.Header.Get("x-api-key")
|
||||||
|
if a := r.Header.Get("Authorization"); a != "" {
|
||||||
|
parts := strings.Fields(a)
|
||||||
|
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
|
||||||
|
return k, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
if raw != "" && raw != parts[1] {
|
||||||
|
return k, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
raw = parts[1]
|
||||||
|
}
|
||||||
|
if raw == "" {
|
||||||
|
return k, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
var ch string
|
||||||
|
e := s.db.QueryRow("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at FROM keys WHERE hash=? AND enabled=1 AND revoked=0", hash(raw)).Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt)
|
||||||
|
if e != nil {
|
||||||
|
return k, e
|
||||||
|
}
|
||||||
|
if k.ExpiresAt != "" {
|
||||||
|
t, e := time.Parse(time.RFC3339Nano, k.ExpiresAt)
|
||||||
|
if e != nil || !time.Now().Before(t) {
|
||||||
|
return apiKey{}, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
}
|
||||||
|
e = json.Unmarshal([]byte(ch), &k.Channels)
|
||||||
|
return k, e
|
||||||
|
}
|
||||||
|
func (k apiKey) allows(id string) bool {
|
||||||
|
if k.AllChannels {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
for _, v := range k.Channels {
|
||||||
|
if v == id {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
@@ -0,0 +1,136 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/url"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
discoveryTimeout = 30 * time.Second
|
||||||
|
discoveryMaxBody = 1 << 20
|
||||||
|
discoveryMaxPages = 20
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *Server) discoverModels(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != http.MethodPost {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var c channel
|
||||||
|
if decode(w, r, &c) != nil {
|
||||||
|
errJSON(w, 400, "invalid channel")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
u, e := url.Parse(c.BaseURL)
|
||||||
|
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || strings.Contains(c.BaseURL, "#") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||||
|
errJSON(w, 400, "invalid channel")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if c.ID != "" {
|
||||||
|
var old string
|
||||||
|
e = s.db.QueryRowContext(r.Context(), "SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old)
|
||||||
|
if e == sql.ErrNoRows {
|
||||||
|
errJSON(w, 404, "not found")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if c.APIKey == "" {
|
||||||
|
c.APIKey = old
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithTimeout(r.Context(), discoveryTimeout)
|
||||||
|
defer cancel()
|
||||||
|
models, e := s.fetchModels(ctx, c)
|
||||||
|
if e != nil {
|
||||||
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||||
|
errJSON(w, 504, "model discovery timed out")
|
||||||
|
} else {
|
||||||
|
errJSON(w, 502, "model discovery failed")
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, map[string][]string{"models": models})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) fetchModels(ctx context.Context, c channel) ([]string, error) {
|
||||||
|
// Share the bounded transport, but never forward credentials across redirects.
|
||||||
|
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||||
|
endpoint := strings.TrimRight(c.BaseURL, "/") + "/models"
|
||||||
|
if c.Protocol == "anthropic" {
|
||||||
|
endpoint = strings.TrimRight(c.BaseURL, "/") + "/v1/models"
|
||||||
|
}
|
||||||
|
seen := map[string]bool{}
|
||||||
|
cursors := map[string]bool{}
|
||||||
|
after := ""
|
||||||
|
for page := 0; page < discoveryMaxPages; page++ {
|
||||||
|
u := endpoint
|
||||||
|
if after != "" {
|
||||||
|
u += "?" + url.Values{"after_id": {after}}.Encode()
|
||||||
|
}
|
||||||
|
req, e := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
req.Header.Set("Accept", "application/json")
|
||||||
|
if c.Protocol == "anthropic" {
|
||||||
|
req.Header.Set("x-api-key", c.APIKey)
|
||||||
|
req.Header.Set("anthropic-version", "2023-06-01")
|
||||||
|
} else {
|
||||||
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||||
|
}
|
||||||
|
res, e := client.Do(req)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if res.StatusCode != http.StatusOK {
|
||||||
|
res.Body.Close()
|
||||||
|
return nil, errors.New("upstream status")
|
||||||
|
}
|
||||||
|
body, e := io.ReadAll(io.LimitReader(res.Body, discoveryMaxBody+1))
|
||||||
|
res.Body.Close()
|
||||||
|
if e != nil || len(body) > discoveryMaxBody {
|
||||||
|
return nil, errors.New("upstream body")
|
||||||
|
}
|
||||||
|
var result struct {
|
||||||
|
Data []struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
} `json:"data"`
|
||||||
|
HasMore bool `json:"has_more"`
|
||||||
|
LastID string `json:"last_id"`
|
||||||
|
}
|
||||||
|
if json.Unmarshal(body, &result) != nil || result.Data == nil {
|
||||||
|
return nil, errors.New("upstream models")
|
||||||
|
}
|
||||||
|
for _, model := range result.Data {
|
||||||
|
if strings.TrimSpace(model.ID) == "" {
|
||||||
|
return nil, errors.New("upstream model id")
|
||||||
|
}
|
||||||
|
seen[model.ID] = true
|
||||||
|
}
|
||||||
|
if c.Protocol != "anthropic" || !result.HasMore {
|
||||||
|
models := make([]string, 0, len(seen))
|
||||||
|
for id := range seen {
|
||||||
|
models = append(models, id)
|
||||||
|
}
|
||||||
|
sort.Strings(models)
|
||||||
|
return models, nil
|
||||||
|
}
|
||||||
|
if result.LastID == "" || cursors[result.LastID] {
|
||||||
|
return nil, errors.New("upstream pagination")
|
||||||
|
}
|
||||||
|
cursors[result.LastID] = true
|
||||||
|
after = result.LastID
|
||||||
|
}
|
||||||
|
return nil, errors.New("upstream page limit")
|
||||||
|
}
|
||||||
@@ -0,0 +1,273 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"reflect"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const discoveryPath = "/admin/api/channels/discover-models"
|
||||||
|
|
||||||
|
func discover(t *testing.T, s *Server, cookie *http.Cookie, c channel) *httptest.ResponseRecorder {
|
||||||
|
t.Helper()
|
||||||
|
b, e := json.Marshal(c)
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
return call(s, "POST", discoveryPath, string(b), cookie)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiscoveryProtocolsAndReadOnly(t *testing.T) {
|
||||||
|
for _, proto := range []string{"openai-chat", "openai-responses", "anthropic"} {
|
||||||
|
t.Run(proto, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
calls := 0
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
calls++
|
||||||
|
if r.Method != "GET" || r.URL.Path != "/v1/models" {
|
||||||
|
t.Error(r.Method, r.URL)
|
||||||
|
}
|
||||||
|
if proto == "anthropic" {
|
||||||
|
if r.Header.Get("x-api-key") != "secret" || r.Header.Get("anthropic-version") != "2023-06-01" || r.Header.Get("Authorization") != "" {
|
||||||
|
t.Error("wrong anthropic headers")
|
||||||
|
}
|
||||||
|
if r.URL.Query().Get("after_id") == "" {
|
||||||
|
fmt.Fprint(w, `{"data":[{"id":"z"}],"has_more":true,"last_id":"cursor/&?"}`)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.URL.Query().Get("after_id") != "cursor/&?" {
|
||||||
|
t.Error(r.URL)
|
||||||
|
}
|
||||||
|
} else if r.Header.Get("Authorization") != "Bearer secret" || r.Header.Get("x-api-key") != "" {
|
||||||
|
t.Error("wrong openai headers")
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"a"}]}`)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
draft := channel{BaseURL: up.URL + "/v1/", Protocol: proto, APIKey: "secret"}
|
||||||
|
if proto == "anthropic" {
|
||||||
|
draft.BaseURL = up.URL + "/"
|
||||||
|
}
|
||||||
|
w := discover(t, s, cookie, draft)
|
||||||
|
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":["a","z"]}` {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
wantCalls := 1
|
||||||
|
if proto == "anthropic" {
|
||||||
|
wantCalls = 2
|
||||||
|
}
|
||||||
|
if calls != wantCalls {
|
||||||
|
t.Fatal(calls)
|
||||||
|
}
|
||||||
|
cs, e := s.loadChannels()
|
||||||
|
if e != nil || len(cs) != 0 {
|
||||||
|
t.Fatal(cs, e)
|
||||||
|
}
|
||||||
|
draft.Name = "saved"
|
||||||
|
draft.Models = map[string]string{"alias": "a"}
|
||||||
|
b, _ := json.Marshal(draft)
|
||||||
|
w = call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var saved channel
|
||||||
|
if e := json.Unmarshal(w.Body.Bytes(), &saved); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
before, _ := s.loadChannels()
|
||||||
|
draft.ID, draft.APIKey, draft.Name, draft.Models = saved.ID, "", "unsaved", nil
|
||||||
|
if w = discover(t, s, cookie, draft); w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
after, _ := s.loadChannels()
|
||||||
|
if !reflect.DeepEqual(before, after) {
|
||||||
|
t.Fatal("discovery changed channel")
|
||||||
|
}
|
||||||
|
replacement := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if proto == "anthropic" {
|
||||||
|
if r.Header.Get("x-api-key") != "replacement" {
|
||||||
|
t.Error("draft key not used")
|
||||||
|
}
|
||||||
|
} else if r.Header.Get("Authorization") != "Bearer replacement" {
|
||||||
|
t.Error("draft key not used")
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, `{"data":[]}`)
|
||||||
|
}))
|
||||||
|
defer replacement.Close()
|
||||||
|
draft.BaseURL, draft.APIKey = replacement.URL, "replacement"
|
||||||
|
w = discover(t, s, cookie, draft)
|
||||||
|
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":[]}` {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
after, _ = s.loadChannels()
|
||||||
|
if !reflect.DeepEqual(before, after) {
|
||||||
|
t.Fatal("draft replacement was saved")
|
||||||
|
}
|
||||||
|
var logs int
|
||||||
|
if e := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&logs); e != nil || logs != 0 {
|
||||||
|
t.Fatal(logs, e)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiscoveryAuthAndValidation(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
if w := call(s, "POST", discoveryPath, "{}", nil); w.Code != 401 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||||
|
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader("{}"))
|
||||||
|
r.AddCookie(cookie)
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
r.Header.Del(header)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 403 {
|
||||||
|
t.Fatal(header, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, method := range []string{"GET", "DELETE", "PUT"} {
|
||||||
|
if w := call(s, method, discoveryPath, "{}", cookie); w.Code != 405 {
|
||||||
|
t.Fatal(method, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, base := range []string{"", "ftp://example.com", "http://", "http://user:[email protected]", "http://example.com?x=1", "http://example.com?", "http://example.com#"} {
|
||||||
|
if w := discover(t, s, cookie, channel{BaseURL: base, Protocol: "openai-chat"}); w.Code != 400 {
|
||||||
|
t.Fatal(base, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if w := discover(t, s, cookie, channel{BaseURL: "http://example.com", Protocol: "unknown"}); w.Code != 400 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
if w := discover(t, s, cookie, channel{ID: "missing", BaseURL: "http://example.com", Protocol: "anthropic"}); w.Code != 404 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
for _, body := range []string{"{", "{} {}", strings.Repeat(" ", maxBody+1)} {
|
||||||
|
if w := call(s, "POST", discoveryPath, body, cookie); w.Code != 400 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiscoveryUpstreamErrors(t *testing.T) {
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name, body string
|
||||||
|
status, calls int
|
||||||
|
}{
|
||||||
|
{"status", "secret upstream details", 401, 1},
|
||||||
|
{"json", "secret", 200, 1},
|
||||||
|
{"missing data", "{}", 200, 1},
|
||||||
|
{"null", `{"data":null}`, 200, 1},
|
||||||
|
{"empty id", `{"data":[{"id":""}]}`, 200, 1},
|
||||||
|
{"trailing", `{"data":[]}{}`, 200, 1},
|
||||||
|
{"oversized", strings.Repeat("x", discoveryMaxBody+1), 200, 1},
|
||||||
|
{"missing cursor", `{"data":[],"has_more":true}`, 200, 1},
|
||||||
|
{"repeated cursor", `{"data":[],"has_more":true,"last_id":"same"}`, 200, 2},
|
||||||
|
{"page limit", "", 200, discoveryMaxPages},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
calls := 0
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
calls++
|
||||||
|
w.WriteHeader(tc.status)
|
||||||
|
if tc.name == "page limit" {
|
||||||
|
fmt.Fprintf(w, `{"data":[],"has_more":true,"last_id":"%d"}`, calls)
|
||||||
|
} else {
|
||||||
|
fmt.Fprint(w, tc.body)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: "anthropic", APIKey: "secret"})
|
||||||
|
if w.Code != 502 || strings.TrimSpace(w.Body.String()) != `{"error":"model discovery failed"}` || calls != tc.calls {
|
||||||
|
t.Fatal(w.Code, w.Body.String(), calls)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiscoveryRedirectAndDeadline(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("redirect followed") }))
|
||||||
|
defer target.Close()
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||||
|
defer up.Close()
|
||||||
|
for _, proto := range []string{"openai-chat", "anthropic"} {
|
||||||
|
if w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: proto, APIKey: "secret"}); w.Code != 502 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
|
||||||
|
defer slow.Close()
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
b, _ := json.Marshal(channel{BaseURL: slow.URL, Protocol: "openai-chat"})
|
||||||
|
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader(string(b))).WithContext(ctx)
|
||||||
|
r.AddCookie(cookie)
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 504 {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelsListSelectedAliasesAndPermissions(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
var allowed string
|
||||||
|
for _, name := range []string{"selected", "forbidden", "disabled"} {
|
||||||
|
b, _ := json.Marshal(channel{Name: name, BaseURL: "http://example.com", Protocol: "openai-chat", Enabled: name != "disabled", Models: map[string]string{"z": "vendor/real", "alias": "vendor/real"}})
|
||||||
|
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var c channel
|
||||||
|
json.Unmarshal(w.Body.Bytes(), &c)
|
||||||
|
if name != "forbidden" {
|
||||||
|
if allowed != "" {
|
||||||
|
allowed += ","
|
||||||
|
}
|
||||||
|
allowed += fmt.Sprintf("%q", c.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
w := call(s, "POST", "/admin/api/keys", `{"name":"client","channels":[`+allowed+`]}`, cookie)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var key map[string]string
|
||||||
|
json.Unmarshal(w.Body.Bytes(), &key)
|
||||||
|
r := httptest.NewRequest("GET", "/v1/models", nil)
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
var result struct{ Data []struct{ ID string } }
|
||||||
|
if e := json.Unmarshal(w.Body.Bytes(), &result); e != nil || w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, e)
|
||||||
|
}
|
||||||
|
if len(result.Data) != 2 || result.Data[0].ID != "selected/alias" || result.Data[1].ID != "selected/z" {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
for _, model := range []string{"forbidden/alias", "disabled/alias"} {
|
||||||
|
r = httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[]}`, model)))
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 403 {
|
||||||
|
t.Fatal(model, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,443 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type integrationKey struct{ raw, id string }
|
||||||
|
|
||||||
|
type integrationStreamRecorder struct {
|
||||||
|
*httptest.ResponseRecorder
|
||||||
|
flushed chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *integrationStreamRecorder) Flush() {
|
||||||
|
w.ResponseRecorder.Flush()
|
||||||
|
select {
|
||||||
|
case w.flushed <- struct{}{}:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationAdmin(t *testing.T, s *Server, cookie *http.Cookie, method, path string, body any, out any) {
|
||||||
|
t.Helper()
|
||||||
|
b, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w := call(s, method, path, string(b), cookie)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Fatalf("%s %s: %d %s", method, path, w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if out != nil {
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), out); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationChannel(t *testing.T, s *Server, cookie *http.Cookie, name, kind, base string) channel {
|
||||||
|
t.Helper()
|
||||||
|
if kind != "anthropic" {
|
||||||
|
base += "/v1"
|
||||||
|
}
|
||||||
|
var ch channel
|
||||||
|
integrationAdmin(t, s, cookie, "POST", "/admin/api/channels", channel{
|
||||||
|
Name: name, Protocol: kind, BaseURL: base, Enabled: true,
|
||||||
|
APIKey: "upstream-secret", Models: map[string]string{"gpt-5": "gpt-4", "gpt-4": "abc/gpt-4"},
|
||||||
|
}, &ch)
|
||||||
|
if ch.ID == "" || ch.APIKey != "" {
|
||||||
|
t.Fatalf("invalid channel response: %+v", ch)
|
||||||
|
}
|
||||||
|
return ch
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationCreateKey(t *testing.T, s *Server, cookie *http.Cookie, name string, ch channel) integrationKey {
|
||||||
|
t.Helper()
|
||||||
|
var created struct {
|
||||||
|
Key string `json:"key"`
|
||||||
|
}
|
||||||
|
integrationAdmin(t, s, cookie, "POST", "/admin/api/keys", apiKey{Name: name, Channels: []string{ch.ID}}, &created)
|
||||||
|
var keys []apiKey
|
||||||
|
integrationAdmin(t, s, cookie, "GET", "/admin/api/keys", nil, &keys)
|
||||||
|
for _, k := range keys {
|
||||||
|
if k.Name == name && created.Key != "" {
|
||||||
|
return integrationKey{created.Key, k.ID}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.Fatal("created key missing from management API")
|
||||||
|
return integrationKey{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationRequest(kind, model string, stream bool) string {
|
||||||
|
if kind == "openai-responses" {
|
||||||
|
return fmt.Sprintf(`{"model":%q,"input":"hello","max_output_tokens":32,"stream":%t}`, model, stream)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":32,"stream":%t}`, model, stream)
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationPath(kind string) string {
|
||||||
|
return map[string]string{"openai-chat": "/v1/chat/completions", "openai-responses": "/v1/responses", "anthropic": "/v1/messages"}[kind]
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationProxy(s *Server, key integrationKey, kind, model string, stream bool) *httptest.ResponseRecorder {
|
||||||
|
r := httptest.NewRequest("POST", integrationPath(kind), strings.NewReader(integrationRequest(kind, model, stream)))
|
||||||
|
if kind == "anthropic" {
|
||||||
|
r.Header.Set("x-api-key", key.raw)
|
||||||
|
} else {
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||||
|
}
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
return w
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationResponse(kind string) string {
|
||||||
|
switch kind {
|
||||||
|
case "openai-chat":
|
||||||
|
return `{"id":"chat-test","object":"chat.completion","model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}`
|
||||||
|
case "openai-responses":
|
||||||
|
return `{"id":"resp-test","object":"response","model":"gpt-4","status":"completed","output":[{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}`
|
||||||
|
default:
|
||||||
|
return `{"id":"msg-test","type":"message","role":"assistant","model":"gpt-4","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":6,"output_tokens":3,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}`
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationSSE(kind string) string {
|
||||||
|
switch kind {
|
||||||
|
case "openai-chat":
|
||||||
|
return "data: " + `{"id":"chat-test","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}` + "\n\ndata: " + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}` + "\n\ndata: [DONE]\n\n"
|
||||||
|
case "openai-responses":
|
||||||
|
return "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp-test","status":"in_progress"}}` + "\n\n" +
|
||||||
|
"event: response.output_item.added\ndata: " + `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}` + "\n\n" +
|
||||||
|
"event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n" +
|
||||||
|
"event: response.completed\ndata: " + `{"type":"response.completed","response":` + integrationResponse(kind) + "}\n\n"
|
||||||
|
default:
|
||||||
|
return "event: message_start\ndata: " + `{"type":"message_start","message":{"id":"msg-test","role":"assistant","content":[],"usage":{"input_tokens":6,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}}` + "\n\n" +
|
||||||
|
"event: content_block_start\ndata: " + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
|
||||||
|
"event: content_block_delta\ndata: " + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}` + "\n\n" +
|
||||||
|
"event: content_block_stop\ndata: " + `{"type":"content_block_stop","index":0}` + "\n\n" +
|
||||||
|
"event: message_delta\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n" +
|
||||||
|
"event: message_stop\ndata: " + `{"type":"message_stop"}` + "\n\n"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationLogs(t *testing.T, s *Server, c *http.Cookie, key integrationKey) []requestLog {
|
||||||
|
t.Helper()
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var logs []requestLog
|
||||||
|
integrationAdmin(t, s, c, "GET", "/admin/api/logs?key_id="+key.id, nil, &logs)
|
||||||
|
for _, l := range logs {
|
||||||
|
if l.KeyID != key.id {
|
||||||
|
t.Fatalf("cross-key log leak: %+v", l)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return logs
|
||||||
|
}
|
||||||
|
|
||||||
|
func integrationCount(t *testing.T, name string, got *int64, want int64) {
|
||||||
|
t.Helper()
|
||||||
|
if got == nil || *got != want {
|
||||||
|
t.Fatalf("%s: got %v, want %d", name, got, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIntegrationProtocolMatrix(t *testing.T) {
|
||||||
|
kinds := []string{"openai-chat", "openai-responses", "anthropic"}
|
||||||
|
for _, upKind := range kinds {
|
||||||
|
for _, downKind := range kinds {
|
||||||
|
t.Run(downKind+"_to_"+upKind, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
var calls atomic.Int64
|
||||||
|
var expectedModel atomic.Value
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
calls.Add(1)
|
||||||
|
if r.URL.Path != integrationPath(upKind) {
|
||||||
|
t.Errorf("upstream path: %s", r.URL.Path)
|
||||||
|
}
|
||||||
|
if upKind == "anthropic" {
|
||||||
|
if r.Header.Get("x-api-key") != "upstream-secret" || r.Header.Get("anthropic-version") == "" {
|
||||||
|
t.Error("missing Anthropic credentials/version")
|
||||||
|
}
|
||||||
|
} else if r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
||||||
|
t.Error("wrong upstream credentials")
|
||||||
|
}
|
||||||
|
var body struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Stream bool `json:"stream"`
|
||||||
|
}
|
||||||
|
b, err := io.ReadAll(r.Body)
|
||||||
|
if err != nil || json.Unmarshal(b, &body) != nil {
|
||||||
|
t.Error("invalid upstream request")
|
||||||
|
w.WriteHeader(400)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if body.Model != expectedModel.Load() || !strings.Contains(string(b), "hello") {
|
||||||
|
t.Errorf("upstream request: %s; want model %s", b, expectedModel.Load())
|
||||||
|
}
|
||||||
|
if body.Stream {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
for _, frame := range strings.Split(integrationSSE(upKind), "\n\n") {
|
||||||
|
if frame != "" {
|
||||||
|
fmt.Fprint(w, frame+"\n\n")
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
fmt.Fprint(w, integrationResponse(upKind))
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
abc := integrationChannel(t, s, cookie, "abc", upKind, up.URL)
|
||||||
|
other := integrationChannel(t, s, cookie, "other", upKind, up.URL)
|
||||||
|
keys := []integrationKey{integrationCreateKey(t, s, cookie, "first", abc), integrationCreateKey(t, s, cookie, "second", other)}
|
||||||
|
for i, key := range keys {
|
||||||
|
name, forbidden := "abc", "other"
|
||||||
|
if i == 1 {
|
||||||
|
name, forbidden = forbidden, name
|
||||||
|
}
|
||||||
|
for _, stream := range []bool{false, true} {
|
||||||
|
for _, alias := range []string{"gpt-5", "gpt-4"} {
|
||||||
|
expectedModel.Store("gpt-4")
|
||||||
|
if alias == "gpt-4" {
|
||||||
|
expectedModel.Store("abc/gpt-4")
|
||||||
|
}
|
||||||
|
w := integrationProxy(s, key, downKind, name+"/"+alias, stream)
|
||||||
|
if w.Code != 200 || !strings.Contains(w.Body.String(), "hello") {
|
||||||
|
t.Fatalf("stream=%t model=%s: %d %s", stream, alias, w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if stream {
|
||||||
|
terminal := map[string]string{"openai-chat": "data: [DONE]", "openai-responses": "event: response.completed", "anthropic": "event: message_stop"}[downKind]
|
||||||
|
if !w.Flushed || !strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(w.Body.String(), terminal) {
|
||||||
|
t.Fatalf("incomplete SSE: %s", w.Body.String())
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
var body map[string]any
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
field := map[string]string{"openai-chat": "choices", "openai-responses": "output", "anthropic": "content"}[downKind]
|
||||||
|
if body[field] == nil {
|
||||||
|
t.Fatalf("wrong downstream shape: %s", w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
before := calls.Load()
|
||||||
|
if w := integrationProxy(s, key, downKind, forbidden+"/gpt-5", false); w.Code != 403 {
|
||||||
|
t.Fatalf("permission: %d %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if calls.Load() != before {
|
||||||
|
t.Fatal("forbidden request reached upstream")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if calls.Load() != 8 {
|
||||||
|
t.Fatalf("upstream calls: %d", calls.Load())
|
||||||
|
}
|
||||||
|
for _, key := range keys {
|
||||||
|
logs := integrationLogs(t, s, cookie, key)
|
||||||
|
if len(logs) != 5 {
|
||||||
|
t.Fatalf("logs: %+v", logs)
|
||||||
|
}
|
||||||
|
for _, l := range logs {
|
||||||
|
if l.Status == 403 {
|
||||||
|
if l.UsageKnown || l.Input != nil {
|
||||||
|
t.Fatal("denied request has usage")
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
actual := "gpt-4"
|
||||||
|
if strings.HasSuffix(l.RequestModel, "/gpt-4") {
|
||||||
|
actual = "abc/gpt-4"
|
||||||
|
}
|
||||||
|
if l.Status != 200 || l.Error != "" || !l.UsageKnown || l.ActualModel != actual || l.Downstream != downKind || l.Upstream != upKind {
|
||||||
|
t.Fatalf("wrong log: %+v", l)
|
||||||
|
}
|
||||||
|
integrationCount(t, "input", l.Input, 12)
|
||||||
|
integrationCount(t, "output", l.Output, 3)
|
||||||
|
integrationCount(t, "cache read", l.CacheRead, 4)
|
||||||
|
write := int64(0)
|
||||||
|
if upKind == "anthropic" {
|
||||||
|
write = 2
|
||||||
|
}
|
||||||
|
integrationCount(t, "cache write", l.CacheWrite, write)
|
||||||
|
}
|
||||||
|
var stats struct {
|
||||||
|
Summary summary `json:"summary"`
|
||||||
|
ByKey []struct {
|
||||||
|
KeyID string `json:"key_id"`
|
||||||
|
} `json:"by_key"`
|
||||||
|
}
|
||||||
|
integrationAdmin(t, s, cookie, "GET", "/admin/api/stats?key_id="+key.id, nil, &stats)
|
||||||
|
if stats.Summary.Requests != 5 || stats.Summary.Success != 4 || stats.Summary.Failed != 1 || stats.Summary.Unknown != 1 || len(stats.ByKey) != 1 || stats.ByKey[0].KeyID != key.id {
|
||||||
|
t.Fatalf("wrong isolated stats: %+v", stats)
|
||||||
|
}
|
||||||
|
integrationCount(t, "total input", stats.Summary.Input, 48)
|
||||||
|
integrationCount(t, "total output", stats.Summary.Output, 12)
|
||||||
|
integrationCount(t, "total cache read", stats.Summary.CacheRead, 16)
|
||||||
|
write := int64(0)
|
||||||
|
if upKind == "anthropic" {
|
||||||
|
write = 8
|
||||||
|
}
|
||||||
|
integrationCount(t, "total cache write", stats.Summary.CacheWrite, write)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIntegrationStreamFailures(t *testing.T) {
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name, body string
|
||||||
|
status int
|
||||||
|
}{
|
||||||
|
{"upstream_http_error", `{"error":"private upstream detail"}`, 429},
|
||||||
|
{"error_event", "event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"private upstream detail\"}}\n\n", 200},
|
||||||
|
{"malformed_event", "data: {invalid}\n\n", 200},
|
||||||
|
{"truncated_stream", "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n", 200},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(tc.status)
|
||||||
|
fmt.Fprint(w, tc.body)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||||
|
key := integrationCreateKey(t, s, cookie, "errors", ch)
|
||||||
|
w := integrationProxy(s, key, "openai-responses", "abc/gpt-5", true)
|
||||||
|
if w.Code != tc.status {
|
||||||
|
t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if strings.Contains(w.Body.String(), "response.completed") || strings.Contains(w.Body.String(), "private upstream detail") {
|
||||||
|
t.Fatalf("false success or leaked upstream detail: %s", w.Body.String())
|
||||||
|
}
|
||||||
|
logs := integrationLogs(t, s, cookie, key)
|
||||||
|
if len(logs) != 1 {
|
||||||
|
t.Fatalf("logs: %+v", logs)
|
||||||
|
}
|
||||||
|
l := logs[0]
|
||||||
|
if l.Status < 400 || l.Error == "" || l.UsageKnown || l.Input != nil || l.Output != nil {
|
||||||
|
t.Fatalf("failure not recorded correctly: %+v", l)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAnthropicHeaderForwarding(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
var mu sync.Mutex
|
||||||
|
got := map[string]string{}
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
mu.Lock()
|
||||||
|
got["version"], got["beta"] = r.Header.Get("anthropic-version"), r.Header.Get("anthropic-beta")
|
||||||
|
mu.Unlock()
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
fmt.Fprint(w, integrationResponse("anthropic"))
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "anthropic", up.URL)
|
||||||
|
key := integrationCreateKey(t, s, cookie, "beta", ch)
|
||||||
|
request := func(version, beta string) {
|
||||||
|
r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(integrationRequest("anthropic", "abc/gpt-5", false)))
|
||||||
|
r.Header.Set("x-api-key", key.raw)
|
||||||
|
if version != "" {
|
||||||
|
r.Header.Set("anthropic-version", version)
|
||||||
|
}
|
||||||
|
if beta != "" {
|
||||||
|
r.Header.Set("anthropic-beta", beta)
|
||||||
|
}
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatalf("status %d: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
request("2023-06-01", "context-1m-2025-08-07,interleaved-thinking-2025-05-14")
|
||||||
|
mu.Lock()
|
||||||
|
if got["version"] != "2023-06-01" || got["beta"] != "context-1m-2025-08-07,interleaved-thinking-2025-05-14" {
|
||||||
|
t.Fatalf("valid headers dropped: %v", got)
|
||||||
|
}
|
||||||
|
mu.Unlock()
|
||||||
|
request("bad version; injected", `weird"value`)
|
||||||
|
mu.Lock()
|
||||||
|
if got["version"] != "2023-06-01" || got["beta"] != "" {
|
||||||
|
t.Fatalf("unsafe headers forwarded: %v", got)
|
||||||
|
}
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestIntegrationCancellationClosesUpstream(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
started, stopped := make(chan struct{}), make(chan struct{})
|
||||||
|
release := make(chan struct{})
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
io.Copy(io.Discard, r.Body)
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n")
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
close(started)
|
||||||
|
select {
|
||||||
|
case <-r.Context().Done():
|
||||||
|
close(stopped)
|
||||||
|
case <-release:
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
defer close(release)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||||
|
key := integrationCreateKey(t, s, cookie, "cancel", ch)
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
r := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(integrationRequest("openai-responses", "abc/gpt-5", true))).WithContext(ctx)
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||||
|
done := make(chan struct{})
|
||||||
|
w := &integrationStreamRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{}, 1)}
|
||||||
|
go func() { s.ServeHTTP(w, r); close(done) }()
|
||||||
|
select {
|
||||||
|
case <-started:
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("upstream did not start")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-w.flushed:
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("downstream did not receive streaming output")
|
||||||
|
}
|
||||||
|
cancel()
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("proxy did not stop on cancellation")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-stopped:
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("upstream request was not cancelled")
|
||||||
|
}
|
||||||
|
logs := integrationLogs(t, s, cookie, key)
|
||||||
|
if len(logs) != 1 || logs[0].Status < 400 || logs[0].Error == "" || logs[0].UsageKnown || logs[0].Input != nil {
|
||||||
|
t.Fatalf("cancel log: %+v", logs)
|
||||||
|
}
|
||||||
|
if strings.Contains(w.Body.String(), "response.completed") {
|
||||||
|
t.Fatal("cancelled response marked completed")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,158 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"log"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
maxLogFieldBytes = 512
|
||||||
|
maxLogErrorBytes = 4096
|
||||||
|
maxLogUsageBytes = 16384
|
||||||
|
logCleanupBatch = 1000
|
||||||
|
)
|
||||||
|
|
||||||
|
func (s *Server) SetLogRetention(days int) {
|
||||||
|
if days < 0 {
|
||||||
|
days = 0
|
||||||
|
}
|
||||||
|
s.logRetentionDays.Store(int64(days))
|
||||||
|
}
|
||||||
|
|
||||||
|
func truncateLogField(value string, limit int) string {
|
||||||
|
if len(value) <= limit {
|
||||||
|
return value
|
||||||
|
}
|
||||||
|
for limit > 0 && !utf8.RuneStart(value[limit]) {
|
||||||
|
limit--
|
||||||
|
}
|
||||||
|
return strings.Clone(value[:limit])
|
||||||
|
}
|
||||||
|
|
||||||
|
func boundedLog(l requestLog) requestLog {
|
||||||
|
for _, field := range []*string{&l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream} {
|
||||||
|
*field = truncateLogField(*field, maxLogFieldBytes)
|
||||||
|
}
|
||||||
|
l.Error = truncateLogField(l.Error, maxLogErrorBytes)
|
||||||
|
if len(l.Usage.Raw) > maxLogUsageBytes {
|
||||||
|
// Keep valid JSON without retaining an attacker-controlled oversized object.
|
||||||
|
l.Usage.Raw = []byte(`{"truncated":true}`)
|
||||||
|
}
|
||||||
|
l.RawUsage = nil
|
||||||
|
return l
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) purgeExpiredLogs(ctx context.Context, days int) error {
|
||||||
|
if days <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
cutoff := time.Now().UTC().AddDate(0, 0, -days).Truncate(time.Second)
|
||||||
|
_, err := s.db.ExecContext(ctx, "DELETE FROM logs WHERE id IN (SELECT id FROM logs WHERE created_at<? AND julianday(created_at)<julianday(?) LIMIT ?)", cutoff.Format("2006-01-02T15:04:05"), cutoff.Format(time.RFC3339), logCleanupBatch)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
type logEntry struct {
|
||||||
|
value requestLog
|
||||||
|
barrier chan error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) record(l requestLog) {
|
||||||
|
s.logMu.RLock()
|
||||||
|
defer s.logMu.RUnlock()
|
||||||
|
if s.logsClosed {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
l = boundedLog(l)
|
||||||
|
l.Usage.Raw = append([]byte(nil), l.Usage.Raw...)
|
||||||
|
s.logQueue <- logEntry{value: l}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) flushLogs() error {
|
||||||
|
s.logMu.RLock()
|
||||||
|
if s.logsClosed {
|
||||||
|
s.logMu.RUnlock()
|
||||||
|
<-s.logDone
|
||||||
|
return s.logErr
|
||||||
|
}
|
||||||
|
done := make(chan error, 1)
|
||||||
|
s.logQueue <- logEntry{barrier: done}
|
||||||
|
s.logMu.RUnlock()
|
||||||
|
return <-done
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) writeLogs() {
|
||||||
|
defer close(s.logDone)
|
||||||
|
ticker := time.NewTicker(100 * time.Millisecond)
|
||||||
|
defer ticker.Stop()
|
||||||
|
cleanup := time.NewTicker(time.Minute)
|
||||||
|
defer cleanup.Stop()
|
||||||
|
batch := make([]requestLog, 0, 100)
|
||||||
|
flush := func() {
|
||||||
|
if len(batch) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for attempt := 0; ; attempt++ {
|
||||||
|
if s.persistLogs(batch) == nil {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if attempt == 3 {
|
||||||
|
s.logErr = errors.New("request log persistence failed after 3 retries")
|
||||||
|
log.Printf("ai-router: request log persistence failed; batch_size=%d retries=3", len(batch))
|
||||||
|
break
|
||||||
|
}
|
||||||
|
log.Printf("ai-router: retrying request log persistence; batch_size=%d retry=%d", len(batch), attempt+1)
|
||||||
|
time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
|
||||||
|
}
|
||||||
|
clear(batch)
|
||||||
|
batch = batch[:0]
|
||||||
|
}
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case entry, ok := <-s.logQueue:
|
||||||
|
if !ok {
|
||||||
|
flush()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if entry.barrier != nil {
|
||||||
|
flush()
|
||||||
|
entry.barrier <- s.logErr
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
batch = append(batch, entry.value)
|
||||||
|
if len(batch) == cap(batch) {
|
||||||
|
flush()
|
||||||
|
}
|
||||||
|
case <-ticker.C:
|
||||||
|
flush()
|
||||||
|
case <-cleanup.C:
|
||||||
|
if days := s.logRetentionDays.Load(); days > 0 {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||||
|
err := s.purgeExpiredLogs(ctx, int(days))
|
||||||
|
cancel()
|
||||||
|
if err != nil {
|
||||||
|
log.Print("ai-router: request log retention cleanup failed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) persistLogs(batch []requestLog) error {
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
tx, err := s.db.BeginTx(ctx, nil)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer tx.Rollback()
|
||||||
|
for _, l := range batch {
|
||||||
|
if err := insertLog(ctx, tx, l); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return tx.Commit()
|
||||||
|
}
|
||||||
@@ -0,0 +1,248 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func waitLogCondition(t *testing.T, condition func() bool) {
|
||||||
|
t.Helper()
|
||||||
|
deadline := time.Now().Add(3 * time.Second)
|
||||||
|
for !condition() {
|
||||||
|
if time.Now().After(deadline) {
|
||||||
|
t.Fatal("condition timed out")
|
||||||
|
}
|
||||||
|
time.Sleep(time.Millisecond)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func logCount(t *testing.T, db *sql.DB) int {
|
||||||
|
t.Helper()
|
||||||
|
var n int
|
||||||
|
if err := db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&n); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLogAsyncAndRawCopy(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
conn, err := s.db.Conn(context.Background())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
raw := []byte(`{"input_tokens":12}`)
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
s.record(requestLog{Usage: protocol.Usage{Raw: raw}})
|
||||||
|
close(done)
|
||||||
|
}()
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("record waited for database")
|
||||||
|
}
|
||||||
|
raw[0] = 'x'
|
||||||
|
conn.Close()
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var got string
|
||||||
|
if err := s.db.QueryRow("SELECT raw_usage FROM logs").Scan(&got); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got != `{"input_tokens":12}` {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLogTimerAndBatchDrain(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
s.record(requestLog{})
|
||||||
|
waitLogCondition(t, func() bool { return logCount(t, s.db) == 1 })
|
||||||
|
for i := 0; i < 257; i++ {
|
||||||
|
s.record(requestLog{})
|
||||||
|
}
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != 258 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLogBackpressureAndConcurrentClose(t *testing.T) {
|
||||||
|
path := filepath.Join(t.TempDir(), "logs.db")
|
||||||
|
s, err := New(path, "password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { s.Close() })
|
||||||
|
conn, err := s.db.Conn(context.Background())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
s.record(requestLog{})
|
||||||
|
waitLogCondition(t, func() bool { return s.db.Stats().WaitCount > 0 })
|
||||||
|
if cap(s.logQueue) != 1024 {
|
||||||
|
t.Fatal(cap(s.logQueue))
|
||||||
|
}
|
||||||
|
for i := 0; i < 1024; i++ {
|
||||||
|
s.record(requestLog{})
|
||||||
|
}
|
||||||
|
enqueued := make(chan struct{})
|
||||||
|
go func() { s.record(requestLog{}); close(enqueued) }()
|
||||||
|
select {
|
||||||
|
case <-enqueued:
|
||||||
|
t.Fatal("full queue did not apply backpressure")
|
||||||
|
case <-time.After(20 * time.Millisecond):
|
||||||
|
}
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i := 0; i < 8; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
if err := s.Close(); err != nil {
|
||||||
|
t.Error(err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
conn.Close()
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() { wg.Wait(); close(done) }()
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
case <-time.After(5 * time.Second):
|
||||||
|
t.Fatal("Close deadlocked")
|
||||||
|
}
|
||||||
|
<-enqueued
|
||||||
|
s.record(requestLog{})
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
db, err := sql.Open("sqlite", path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
if n := logCount(t, db); n != 1026 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLogPersistenceErrorIsStickyAndRedacted(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
var output bytes.Buffer
|
||||||
|
previous := log.Writer()
|
||||||
|
log.SetOutput(&output)
|
||||||
|
defer log.SetOutput(previous)
|
||||||
|
if _, err := s.db.Exec(`CREATE TRIGGER reject_log BEFORE INSERT ON logs WHEN NEW.key_name='secret' BEGIN SELECT RAISE(FAIL,'private database secret'); END`); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
s.record(requestLog{KeyName: "ok"})
|
||||||
|
s.record(requestLog{KeyName: "secret"})
|
||||||
|
start := time.Now()
|
||||||
|
err := s.flushLogs()
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("missing persistence error")
|
||||||
|
}
|
||||||
|
if time.Since(start) > time.Second {
|
||||||
|
t.Fatal("unbounded retries")
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != 0 {
|
||||||
|
t.Fatalf("partial transaction: %d", n)
|
||||||
|
}
|
||||||
|
if strings.Contains(output.String(), "secret") || strings.Contains(err.Error(), "secret") {
|
||||||
|
t.Fatal("error leaked metadata")
|
||||||
|
}
|
||||||
|
if n := strings.Count(output.String(), "retrying"); n != 3 {
|
||||||
|
t.Fatalf("retries: %d", n)
|
||||||
|
}
|
||||||
|
if _, err := s.db.Exec("DROP TRIGGER reject_log"); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
s.record(requestLog{})
|
||||||
|
if err := s.flushLogs(); err == nil {
|
||||||
|
t.Fatal("error was not sticky")
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != 1 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
if err := s.Close(); err == nil {
|
||||||
|
t.Fatal("Close lost persistence error")
|
||||||
|
}
|
||||||
|
if err := s.Close(); err == nil {
|
||||||
|
t.Fatal("second Close lost error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCloseWaitsForActiveRequest(t *testing.T) {
|
||||||
|
path := filepath.Join(t.TempDir(), "logs.db")
|
||||||
|
s, err := New(path, "test-password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { s.Close() })
|
||||||
|
started, release := make(chan struct{}), make(chan struct{})
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
close(started)
|
||||||
|
<-release
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.Write([]byte(integrationResponse("openai-chat")))
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
defer func() {
|
||||||
|
select {
|
||||||
|
case <-release:
|
||||||
|
default:
|
||||||
|
close(release)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||||
|
key := integrationCreateKey(t, s, cookie, "active", ch)
|
||||||
|
requestDone := make(chan struct{})
|
||||||
|
go func() { integrationProxy(s, key, "openai-chat", "abc/gpt-5", false); close(requestDone) }()
|
||||||
|
<-started
|
||||||
|
closed := make(chan error, 1)
|
||||||
|
go func() { closed <- s.Close() }()
|
||||||
|
waitLogCondition(t, func() bool { s.lifecycleMu.Lock(); defer s.lifecycleMu.Unlock(); return s.closing })
|
||||||
|
if w := call(s, "GET", "/v1/models", "", nil); w.Code != 503 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-closed:
|
||||||
|
t.Fatal("Close skipped active request")
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
close(release)
|
||||||
|
<-requestDone
|
||||||
|
select {
|
||||||
|
case err := <-closed:
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("Close did not drain")
|
||||||
|
}
|
||||||
|
db, err := sql.Open("sqlite", path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
if n := logCount(t, db); n != 1 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,229 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
const modelTestTimeout = 30 * time.Second
|
||||||
|
const modelTestMaxBody = 1 << 20
|
||||||
|
|
||||||
|
type modelTestResult struct {
|
||||||
|
OK bool `json:"ok"`
|
||||||
|
DurationMS int64 `json:"duration_ms"`
|
||||||
|
Status int `json:"status"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// testModel is an administrative diagnostic, excluded from call logs and usage
|
||||||
|
// statistics. It creates no API key; upstream providers may still bill the probe.
|
||||||
|
func (s *Server) testModel(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != http.MethodPost {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var input struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Prompt *string `json:"prompt"`
|
||||||
|
}
|
||||||
|
if decode(w, r, &input) != nil {
|
||||||
|
errJSON(w, 400, "invalid request")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
name, alias, ok := strings.Cut(input.Model, "/")
|
||||||
|
prompt := "Reply with OK."
|
||||||
|
if input.Prompt != nil {
|
||||||
|
prompt = *input.Prompt
|
||||||
|
}
|
||||||
|
if !ok || name == "" || alias == "" || strings.TrimSpace(prompt) == "" || utf8.RuneCountInString(prompt) > 2000 {
|
||||||
|
errJSON(w, 400, "invalid model or prompt")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout)
|
||||||
|
defer cancel()
|
||||||
|
var c channel
|
||||||
|
var models string
|
||||||
|
e := s.db.QueryRowContext(ctx, "SELECT base_url,protocol,enabled,models,api_key FROM channels WHERE name=?", name).Scan(&c.BaseURL, &c.Protocol, &c.Enabled, &models, &c.APIKey)
|
||||||
|
if e == sql.ErrNoRows {
|
||||||
|
errJSON(w, 403, "model or channel unavailable")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if e != nil || json.Unmarshal([]byte(models), &c.Models) != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
actual, mapped := c.Models[alias]
|
||||||
|
if !c.Enabled || !mapped || actual == "" {
|
||||||
|
errJSON(w, 403, "model or channel unavailable")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, s.probeModel(ctx, c, actual, prompt))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) probeModel(ctx context.Context, c channel, actual, prompt string) (result modelTestResult) {
|
||||||
|
start := time.Now()
|
||||||
|
defer func() {
|
||||||
|
result.DurationMS = time.Since(start).Milliseconds()
|
||||||
|
if ctx.Err() != nil {
|
||||||
|
result.OK = false
|
||||||
|
result.Error = "model test canceled"
|
||||||
|
if ctx.Err() == context.DeadlineExceeded {
|
||||||
|
result.Error = "model test timed out"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
result.Error = "upstream request failed"
|
||||||
|
up := protocol.Kind(c.Protocol)
|
||||||
|
body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"role": "user", "content": prompt}}, "stream": false, "max_completion_tokens": 128})
|
||||||
|
payload, e := protocol.ConvertRequest(protocol.Chat, up, body, actual)
|
||||||
|
if e != nil {
|
||||||
|
result.Error = "invalid upstream configuration"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||||
|
req, e := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+suffix, bytes.NewReader(payload))
|
||||||
|
if e != nil {
|
||||||
|
result.Error = "invalid upstream configuration"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
if up == protocol.Anthropic {
|
||||||
|
req.Header.Set("x-api-key", c.APIKey)
|
||||||
|
req.Header.Set("anthropic-version", "2023-06-01")
|
||||||
|
} else {
|
||||||
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||||
|
}
|
||||||
|
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||||
|
res, e := client.Do(req)
|
||||||
|
if e != nil {
|
||||||
|
if errors.Is(e, context.DeadlineExceeded) {
|
||||||
|
result.Error = "model test timed out"
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer res.Body.Close()
|
||||||
|
result.Status = res.StatusCode
|
||||||
|
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||||
|
result.Error = "upstream rejected request"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
body, e = io.ReadAll(io.LimitReader(res.Body, modelTestMaxBody+1))
|
||||||
|
if len(body) > modelTestMaxBody {
|
||||||
|
result.Error = "upstream response too large"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if e != nil || !validModelTestResponse(up, body) {
|
||||||
|
result.Error = "invalid upstream response"
|
||||||
|
return
|
||||||
|
}
|
||||||
|
result.OK, result.Error = true, ""
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func validModelTestResponse(kind protocol.Kind, body []byte) bool {
|
||||||
|
var r struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Object string `json:"object"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
Error json.RawMessage `json:"error"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
StopReason string `json:"stop_reason"`
|
||||||
|
Incomplete struct {
|
||||||
|
Reason string `json:"reason"`
|
||||||
|
} `json:"incomplete_details"`
|
||||||
|
Output []json.RawMessage `json:"output"`
|
||||||
|
Content []json.RawMessage `json:"content"`
|
||||||
|
Choices []struct {
|
||||||
|
Finish string `json:"finish_reason"`
|
||||||
|
Message struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content *string `json:"content"`
|
||||||
|
} `json:"message"`
|
||||||
|
} `json:"choices"`
|
||||||
|
}
|
||||||
|
if json.Unmarshal(body, &r) != nil || r.ID == "" || (len(r.Error) != 0 && string(r.Error) != "null") {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
switch kind {
|
||||||
|
case protocol.Chat:
|
||||||
|
if r.Object != "chat.completion" || len(r.Choices) == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, c := range r.Choices {
|
||||||
|
if c.Message.Role != "assistant" || (c.Finish != "stop" && c.Finish != "length") || (c.Message.Content == nil && c.Finish != "length") {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
case protocol.Responses:
|
||||||
|
if r.Object != "response" || r.Output == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
exhausted := r.Status == "incomplete" && r.Incomplete.Reason == "max_output_tokens"
|
||||||
|
if !exhausted && (r.Status != "completed" || len(r.Output) == 0) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, raw := range r.Output {
|
||||||
|
var item struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content []json.RawMessage `json:"content"`
|
||||||
|
Summary []json.RawMessage `json:"summary"`
|
||||||
|
}
|
||||||
|
if json.Unmarshal(raw, &item) != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if item.Type == "reasoning" && item.Summary != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if item.Type != "message" || item.Role != "assistant" || !validProbeBlocks(item.Content, "output_text") {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
case protocol.Anthropic:
|
||||||
|
if r.Type != "message" || r.Role != "assistant" || r.Content == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if r.StopReason == "max_tokens" && len(r.Content) == 0 {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return (r.StopReason == "end_turn" || r.StopReason == "max_tokens") && validProbeBlocks(r.Content, "text")
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func validProbeBlocks(blocks []json.RawMessage, typ string) bool {
|
||||||
|
if len(blocks) == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, raw := range blocks {
|
||||||
|
var b struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text *string `json:"text"`
|
||||||
|
Thinking *string `json:"thinking"`
|
||||||
|
Signature *string `json:"signature"`
|
||||||
|
Data *string `json:"data"`
|
||||||
|
}
|
||||||
|
if json.Unmarshal(raw, &b) != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if typ == "text" && ((b.Type == "thinking" && b.Thinking != nil && b.Signature != nil) || (b.Type == "redacted_thinking" && b.Data != nil)) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if b.Type != typ || b.Text == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
@@ -0,0 +1,267 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const probePath = "/admin/api/models/test"
|
||||||
|
const probeChat = `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":"OK"},"finish_reason":"stop"}]}`
|
||||||
|
|
||||||
|
func saveProbeChannel(t *testing.T, s *Server, cookie *http.Cookie, url, kind string, enabled bool) {
|
||||||
|
t.Helper()
|
||||||
|
b, _ := json.Marshal(channel{Name: "saved", BaseURL: url, Protocol: kind, Enabled: enabled, Models: map[string]string{"alias": "actual/model"}, APIKey: "secret-key"})
|
||||||
|
if w := call(s, "POST", "/admin/api/channels", string(b), cookie); w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func probeResult(t *testing.T, w *httptest.ResponseRecorder) modelTestResult {
|
||||||
|
t.Helper()
|
||||||
|
var result modelTestResult
|
||||||
|
if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &result) != nil {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if result.DurationMS < 0 || strings.Contains(w.Body.String(), "secret-key") {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelProbeProtocols(t *testing.T) {
|
||||||
|
for _, tc := range []struct{ kind, path, limit, response string }{
|
||||||
|
{"openai-chat", "/chat/completions", "max_completion_tokens", probeChat},
|
||||||
|
{"openai-responses", "/responses", "max_output_tokens", `{"id":"r","object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}]}`},
|
||||||
|
{"anthropic", "/v1/messages", "max_tokens", `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}]}`},
|
||||||
|
} {
|
||||||
|
t.Run(tc.kind, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
expectedPrompt := "Reply with OK."
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
b, _ := io.ReadAll(r.Body)
|
||||||
|
var obj map[string]any
|
||||||
|
if json.Unmarshal(b, &obj) != nil || obj["model"] != "actual/model" || obj[tc.limit] != float64(128) || obj["stream"] != false || !strings.Contains(string(b), expectedPrompt) {
|
||||||
|
t.Error(string(b))
|
||||||
|
}
|
||||||
|
if tc.kind == "openai-chat" && obj["max_tokens"] != nil {
|
||||||
|
t.Error("legacy token limit")
|
||||||
|
}
|
||||||
|
if r.Method != "POST" || r.URL.Path != tc.path {
|
||||||
|
t.Error(r.Method, r.URL.Path)
|
||||||
|
}
|
||||||
|
if tc.kind == "anthropic" {
|
||||||
|
if r.Header.Get("x-api-key") != "secret-key" || r.Header.Get("anthropic-version") != "2023-06-01" {
|
||||||
|
t.Error("headers")
|
||||||
|
}
|
||||||
|
} else if r.Header.Get("Authorization") != "Bearer secret-key" {
|
||||||
|
t.Error("authorization")
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, tc.response)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
saveProbeChannel(t, s, cookie, up.URL, tc.kind, true)
|
||||||
|
for _, prompt := range []string{"", strings.Repeat("中", 2000)} {
|
||||||
|
input := map[string]string{"model": "saved/alias"}
|
||||||
|
if prompt != "" {
|
||||||
|
expectedPrompt = prompt
|
||||||
|
input["prompt"] = prompt
|
||||||
|
}
|
||||||
|
b, _ := json.Marshal(input)
|
||||||
|
result := probeResult(t, call(s, "POST", probePath, string(b), cookie))
|
||||||
|
if !result.OK || result.Status != 200 || result.Error != "" {
|
||||||
|
t.Fatal(result)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, table := range []string{"keys", "logs"} {
|
||||||
|
var n int
|
||||||
|
if e := s.db.QueryRow("SELECT count(*) FROM " + table).Scan(&n); e != nil || n != 0 {
|
||||||
|
t.Fatal(table, n, e)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(s.logQueue) != 0 {
|
||||||
|
t.Fatal("queued formal log")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelProbeGuardAndInput(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
for _, prompt := range []string{"", " ", "\t\r\n", "\u3000\u00a0"} {
|
||||||
|
body, _ := json.Marshal(map[string]string{"model": "saved/alias", "prompt": prompt})
|
||||||
|
if w := call(s, "POST", probePath, string(body), cookie); w.Code != 400 {
|
||||||
|
t.Fatalf("blank prompt %q: status %d", prompt, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("unexpected upstream request") }))
|
||||||
|
defer up.Close()
|
||||||
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", false)
|
||||||
|
if w := call(s, "POST", probePath, `{"model":"saved/alias"}`, nil); w.Code != 401 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||||
|
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`))
|
||||||
|
r.AddCookie(cookie)
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
r.Header.Del(header)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 403 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, body := range []string{`{"model":"saved/alias"}`, `{"model":"missing/alias"}`} {
|
||||||
|
if w := call(s, "POST", probePath, body, cookie); w.Code != 403 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if _, e := s.db.Exec("UPDATE channels SET enabled=1"); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
if w := call(s, "POST", probePath, `{"model":"saved/unmapped"}`, cookie); w.Code != 403 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
for _, body := range []string{`{}`, `null`, `{"model":"alias"}`, `{"model":"/alias"}`, `{"model":"saved/"}`, `{"model":"saved/alias","prompt":2}`, `{"model":"saved/alias"} {}`, `{"model":"saved/alias","prompt":"` + strings.Repeat("中", 2001) + `"}`} {
|
||||||
|
if w := call(s, "POST", probePath, body, cookie); w.Code != 400 {
|
||||||
|
t.Fatal(body, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if w := call(s, "GET", probePath, "", cookie); w.Code != 405 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelProbeFailures(t *testing.T) {
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name, body string
|
||||||
|
status int
|
||||||
|
error string
|
||||||
|
}{
|
||||||
|
{"rejected", "secret-key upstream body", 401, "upstream rejected request"},
|
||||||
|
{"redirect", "secret-key", 302, "upstream rejected request"},
|
||||||
|
{"html", "<html>secret-key</html>", 200, "invalid upstream response"},
|
||||||
|
{"empty", "{}", 200, "invalid upstream response"},
|
||||||
|
{"oversize", strings.Repeat("x", modelTestMaxBody+1), 200, "upstream response too large"},
|
||||||
|
{"limit", probeChat + strings.Repeat(" ", modelTestMaxBody-len(probeChat)), 200, ""},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Path != "/chat/completions" {
|
||||||
|
t.Error("followed redirect")
|
||||||
|
}
|
||||||
|
w.Header().Set("Location", "/redirected")
|
||||||
|
w.WriteHeader(tc.status)
|
||||||
|
fmt.Fprint(w, tc.body)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||||
|
result := probeResult(t, call(s, "POST", probePath, `{"model":"saved/alias"}`, cookie))
|
||||||
|
if result.Status != tc.status || result.Error != tc.error || result.OK != (tc.error == "") {
|
||||||
|
t.Fatal(result)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestModelProbeResponseValidation(t *testing.T) {
|
||||||
|
if !validModelTestResponse(protocol.Anthropic, []byte(`{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[{"type":"thinking","thinking":"reasoning","signature":"sig"}]}`)) {
|
||||||
|
t.Fatal("valid thinking-only response rejected")
|
||||||
|
}
|
||||||
|
for _, tc := range []struct {
|
||||||
|
kind protocol.Kind
|
||||||
|
body string
|
||||||
|
ok bool
|
||||||
|
}{
|
||||||
|
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":null},"finish_reason":"length"}]}`, true},
|
||||||
|
{protocol.Responses, `{"id":"r","object":"response","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"output":[{"type":"reasoning","summary":[]}]}`, true},
|
||||||
|
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[]}`, true},
|
||||||
|
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{}]}`, false},
|
||||||
|
{protocol.Responses, `{"id":"r","object":"response","status":"failed","output":[]}`, false},
|
||||||
|
{protocol.Responses, `{"id":"r","object":"response","status":"completed","output":[{}]}`, false},
|
||||||
|
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{}]}`, false},
|
||||||
|
} {
|
||||||
|
if got := validModelTestResponse(tc.kind, []byte(tc.body)); got != tc.ok {
|
||||||
|
t.Error(tc.kind, tc.body, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, kind := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||||
|
for _, body := range []string{"{}", "null", "[]", probeChat + "{}", `{"error":{"message":"secret-key"}}`} {
|
||||||
|
if validModelTestResponse(kind, []byte(body)) {
|
||||||
|
t.Error(kind, body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type probeRoundTripper func(*http.Request) (*http.Response, error)
|
||||||
|
|
||||||
|
func (f probeRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||||
|
|
||||||
|
func TestModelProbeTimeoutAndCancellation(t *testing.T) {
|
||||||
|
for _, cancelRequest := range []bool{false, true} {
|
||||||
|
t.Run(fmt.Sprint(cancelRequest), func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||||
|
defer up.Close()
|
||||||
|
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
s.client.Transport = probeRoundTripper(func(r *http.Request) (*http.Response, error) {
|
||||||
|
deadline, ok := r.Context().Deadline()
|
||||||
|
if !ok || time.Until(deadline) > 30*time.Second || time.Until(deadline) < 29*time.Second {
|
||||||
|
t.Error("missing 30s deadline", deadline)
|
||||||
|
}
|
||||||
|
if cancelRequest {
|
||||||
|
cancel()
|
||||||
|
<-r.Context().Done()
|
||||||
|
return nil, r.Context().Err()
|
||||||
|
}
|
||||||
|
return nil, context.DeadlineExceeded
|
||||||
|
})
|
||||||
|
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`)).WithContext(ctx)
|
||||||
|
r.AddCookie(cookie)
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
result := probeResult(t, w)
|
||||||
|
if result.OK || result.Status != 0 || result.Error == "" {
|
||||||
|
t.Fatal(result)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
s := testServer(t)
|
||||||
|
canceled := make(chan struct{})
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(200)
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
<-r.Context().Done()
|
||||||
|
close(canceled)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
result := s.probeModel(ctx, channel{BaseURL: up.URL, Protocol: "openai-chat"}, "actual", "test")
|
||||||
|
if result.OK || result.Status != 200 || result.Error != "model test timed out" {
|
||||||
|
t.Fatal(result)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-canceled:
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("upstream was not canceled")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,203 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"runtime"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Explicitly opt in: ROUTER_PERF=1 go test ./internal/server -run TestPerformanceHTTP -v -count=1.
|
||||||
|
func TestPerformanceHTTP(t *testing.T) {
|
||||||
|
if os.Getenv("ROUTER_PERF") != "1" {
|
||||||
|
t.Skip("opt-in load test")
|
||||||
|
}
|
||||||
|
for _, stream := range []bool{false, true} {
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
io.Copy(io.Discard, r.Body)
|
||||||
|
r.Body.Close()
|
||||||
|
if stream {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
for _, f := range strings.Split(integrationSSE("openai-chat"), "\n\n") {
|
||||||
|
if f != "" {
|
||||||
|
fmt.Fprint(w, f+"\n\n")
|
||||||
|
w.(http.Flusher).Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
fmt.Fprint(w, integrationResponse("openai-chat"))
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
for _, mode := range []string{"direct", "memory", "disk"} {
|
||||||
|
for _, down := range []string{"openai-chat", "anthropic"} {
|
||||||
|
if mode == "direct" && down != "openai-chat" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, concurrency := range []int{1, 8, 32} {
|
||||||
|
t.Run(fmt.Sprintf("stream=%t/%s/%s/c%d", stream, mode, down, concurrency), func(t *testing.T) {
|
||||||
|
target := up.URL + integrationPath("openai-chat")
|
||||||
|
raw := "test-key"
|
||||||
|
var s *Server
|
||||||
|
var router *httptest.Server
|
||||||
|
if mode != "direct" {
|
||||||
|
path := ":memory:"
|
||||||
|
if mode == "disk" {
|
||||||
|
path = filepath.Join(t.TempDir(), "router.db")
|
||||||
|
}
|
||||||
|
var err error
|
||||||
|
s, err = New(path, "test-password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
if err := s.SetLimits(1000000, 64, 128); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||||
|
key := integrationCreateKey(t, s, cookie, "perf", ch)
|
||||||
|
raw = key.raw
|
||||||
|
router = httptest.NewServer(s)
|
||||||
|
defer router.Close()
|
||||||
|
target = router.URL + integrationPath(down)
|
||||||
|
}
|
||||||
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||||
|
tr.MaxIdleConnsPerHost = 64
|
||||||
|
defer tr.CloseIdleConnections()
|
||||||
|
client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
|
||||||
|
body := integrationRequest(down, "abc/gpt-5", stream)
|
||||||
|
request := func() (time.Duration, time.Duration, bool) {
|
||||||
|
start := time.Now()
|
||||||
|
r, _ := http.NewRequest("POST", target, strings.NewReader(body))
|
||||||
|
r.Header.Set("Authorization", "Bearer "+raw)
|
||||||
|
r.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, err := client.Do(r)
|
||||||
|
if err != nil {
|
||||||
|
return time.Since(start), 0, false
|
||||||
|
}
|
||||||
|
first := make([]byte, 1)
|
||||||
|
_, err = resp.Body.Read(first)
|
||||||
|
ttfb := time.Since(start)
|
||||||
|
_, e := io.Copy(io.Discard, resp.Body)
|
||||||
|
resp.Body.Close()
|
||||||
|
return time.Since(start), ttfb, resp.StatusCode == 200 && (err == nil || err == io.EOF) && e == nil
|
||||||
|
}
|
||||||
|
for i := 0; i < 20; i++ {
|
||||||
|
if _, _, ok := request(); !ok {
|
||||||
|
t.Fatal("warmup failed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
const n = 400
|
||||||
|
dur := make([]float64, n)
|
||||||
|
first := make([]float64, n)
|
||||||
|
var index atomic.Int64
|
||||||
|
var failures atomic.Int64
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
runtime.GC()
|
||||||
|
var before, after runtime.MemStats
|
||||||
|
runtime.ReadMemStats(&before)
|
||||||
|
start := time.Now()
|
||||||
|
for c := 0; c < concurrency; c++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
for {
|
||||||
|
i := int(index.Add(1)) - 1
|
||||||
|
if i >= n {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
d, f, ok := request()
|
||||||
|
dur[i] = float64(d) / 1e6
|
||||||
|
first[i] = float64(f) / 1e6
|
||||||
|
if !ok {
|
||||||
|
failures.Add(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
elapsed := time.Since(start)
|
||||||
|
runtime.ReadMemStats(&after)
|
||||||
|
sort.Float64s(dur)
|
||||||
|
sort.Float64s(first)
|
||||||
|
var total float64
|
||||||
|
for _, d := range dur {
|
||||||
|
total += d
|
||||||
|
}
|
||||||
|
drained := elapsed
|
||||||
|
if s != nil {
|
||||||
|
router.Close()
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
drained = time.Since(start)
|
||||||
|
var count int
|
||||||
|
if err := s.db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&count); err != nil || count != n+20 {
|
||||||
|
t.Fatalf("log count %d err %v", count, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.Logf("RESULT rps=%.1f drained_rps=%.1f mean_ms=%.3f p50_ms=%.3f p95_ms=%.3f p99_ms=%.3f first_p50_ms=%.3f first_p95_ms=%.3f alloc_B_req=%d errors=%d", float64(n)/elapsed.Seconds(), float64(n)/drained.Seconds(), total/n, dur[n/2], dur[n*95/100], dur[n*99/100], first[n/2], first[n*95/100], (after.TotalAlloc-before.TotalAlloc)/n, failures.Load())
|
||||||
|
if failures.Load() != 0 {
|
||||||
|
t.Fail()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
up.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkPerformanceConversion(b *testing.B) {
|
||||||
|
for _, size := range []int{1024, 65536} {
|
||||||
|
for _, to := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||||
|
b.Run(fmt.Sprintf("%d/%s", size, to), func(b *testing.B) {
|
||||||
|
body := []byte(`{"model":"abc/gpt-5","max_tokens":1024,"messages":[{"role":"user","content":"` + strings.Repeat("x", size) + `"}]}`)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
if _, err := protocol.ConvertRequest(protocol.Chat, to, body, "gpt-4"); err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkPerformanceStats(b *testing.B) {
|
||||||
|
for _, n := range []int{1000, 10000, 100000} {
|
||||||
|
b.Run(fmt.Sprint(n), func(b *testing.B) {
|
||||||
|
s, err := New(filepath.Join(b.TempDir(), "stats.db"), "test-password", nil)
|
||||||
|
if err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
_, err = s.db.Exec(`WITH RECURSIVE seq(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM seq WHERE x<?) INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) SELECT '2026-10-20T10:00:00Z','key','perf','abc','abc/gpt-5','gpt-4','openai-chat','openai-chat',200,10,0,1000,100,100,0,1,'','{}' FROM seq`, n)
|
||||||
|
if err != nil {
|
||||||
|
b.Fatal(err)
|
||||||
|
}
|
||||||
|
r := httptest.NewRequest("GET", "/admin/api/stats", nil)
|
||||||
|
b.ReportAllocs()
|
||||||
|
b.ResetTimer()
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.stats(w, r)
|
||||||
|
if w.Code != 200 {
|
||||||
|
b.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,273 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func resolve(cs []channel, k apiKey, model string) (channel, string, error) {
|
||||||
|
name, alias, ok := strings.Cut(model, "/")
|
||||||
|
if !ok || alias == "" {
|
||||||
|
return channel{}, "", errors.New("model must be channel/alias")
|
||||||
|
}
|
||||||
|
for _, c := range cs {
|
||||||
|
if c.Name != name {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !c.Enabled || !k.allows(c.ID) {
|
||||||
|
return channel{}, "", errors.New("channel unavailable")
|
||||||
|
}
|
||||||
|
if actual, ok := c.Models[alias]; ok {
|
||||||
|
return c, actual, nil
|
||||||
|
}
|
||||||
|
return channel{}, "", errors.New("model unavailable")
|
||||||
|
}
|
||||||
|
return channel{}, "", errors.New("channel unavailable")
|
||||||
|
}
|
||||||
|
func (s *Server) proxy(w http.ResponseWriter, r *http.Request) {
|
||||||
|
kinds := map[string]protocol.Kind{"/v1/chat/completions": protocol.Chat, "/v1/responses": protocol.Responses, "/v1/messages": protocol.Anthropic}
|
||||||
|
down, ok := kinds[r.URL.Path]
|
||||||
|
models := r.URL.Path == "/v1/models"
|
||||||
|
if !ok && !models {
|
||||||
|
errJSON(w, 404, "not found")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if (models && r.Method != "GET") || (!models && r.Method != "POST") {
|
||||||
|
errJSON(w, 405, "method not allowed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ip := s.clientIP(r)
|
||||||
|
if !s.authAllowed(ip) {
|
||||||
|
w.Header().Set("Retry-After", "60")
|
||||||
|
errJSON(w, 429, "request limit exceeded")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
k, e := s.authenticate(r)
|
||||||
|
if e != nil {
|
||||||
|
s.authFail(ip)
|
||||||
|
errJSON(w, 401, "invalid API key")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
release, retry := s.limiter.acquire(k.ID, time.Now())
|
||||||
|
if release == nil {
|
||||||
|
w.Header().Set("Retry-After", retry)
|
||||||
|
errJSON(w, 429, "request limit exceeded")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
start := time.Now()
|
||||||
|
l := requestLog{CreatedAt: start.UTC().Format(time.RFC3339Nano), KeyID: k.ID, KeyName: k.Name, Downstream: string(down), Status: 500}
|
||||||
|
// Release the concurrency slot before the log write, which can block on a
|
||||||
|
// full queue; rate-limited and unauthenticated requests stay unlogged.
|
||||||
|
defer func() {
|
||||||
|
release()
|
||||||
|
if !models {
|
||||||
|
l.DurationMS = time.Since(start).Milliseconds()
|
||||||
|
s.record(l)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
fail := func(status int, msg string) { l.Status = status; l.Error = msg; errJSON(w, status, msg) }
|
||||||
|
cs, e := s.loadChannels()
|
||||||
|
if e != nil {
|
||||||
|
fail(500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if models {
|
||||||
|
ids := []string{}
|
||||||
|
for _, c := range cs {
|
||||||
|
if c.Enabled && k.allows(c.ID) {
|
||||||
|
for alias := range c.Models {
|
||||||
|
ids = append(ids, c.Name+"/"+alias)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
sort.Strings(ids)
|
||||||
|
out := []any{}
|
||||||
|
for _, id := range ids {
|
||||||
|
out = append(out, map[string]string{"id": id, "object": "model", "owned_by": "ai-router"})
|
||||||
|
}
|
||||||
|
l.Status = 200
|
||||||
|
writeJSON(w, map[string]any{"object": "list", "data": out})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var b []byte
|
||||||
|
e = readBody(w, r, maxBody, bodyReadTimeout, func() error {
|
||||||
|
var err error
|
||||||
|
b, err = io.ReadAll(r.Body)
|
||||||
|
return err
|
||||||
|
})
|
||||||
|
if e != nil {
|
||||||
|
var tooBig *http.MaxBytesError
|
||||||
|
if errors.As(e, &tooBig) {
|
||||||
|
fail(413, "request body too large")
|
||||||
|
} else {
|
||||||
|
fail(400, "request body unreadable")
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var req struct {
|
||||||
|
Model string `json:"model"`
|
||||||
|
Stream bool `json:"stream"`
|
||||||
|
}
|
||||||
|
if json.Unmarshal(b, &req) != nil || req.Model == "" {
|
||||||
|
fail(400, "invalid request")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
l.RequestModel = req.Model
|
||||||
|
l.Stream = req.Stream
|
||||||
|
c, actual, e := resolve(cs, k, req.Model)
|
||||||
|
if e != nil {
|
||||||
|
fail(403, "model or channel unavailable")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
l.Channel = c.Name
|
||||||
|
l.ActualModel = actual
|
||||||
|
l.Upstream = c.Protocol
|
||||||
|
up := protocol.Kind(c.Protocol)
|
||||||
|
payload, e := protocol.ConvertRequest(down, up, b, actual)
|
||||||
|
if e != nil {
|
||||||
|
fail(400, "unsupported request conversion")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if req.Stream && up == protocol.Chat {
|
||||||
|
var obj map[string]json.RawMessage
|
||||||
|
if json.Unmarshal(payload, &obj) == nil {
|
||||||
|
options := map[string]json.RawMessage{}
|
||||||
|
if raw, ok := obj["stream_options"]; ok {
|
||||||
|
_ = json.Unmarshal(raw, &options)
|
||||||
|
}
|
||||||
|
if options == nil {
|
||||||
|
options = map[string]json.RawMessage{}
|
||||||
|
}
|
||||||
|
options["include_usage"] = json.RawMessage("true")
|
||||||
|
obj["stream_options"], _ = json.Marshal(options)
|
||||||
|
payload, _ = json.Marshal(obj)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||||
|
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||||
|
defer cancel()
|
||||||
|
request, e := http.NewRequestWithContext(ctx, "POST", c.BaseURL+suffix, bytes.NewReader(payload))
|
||||||
|
if e != nil {
|
||||||
|
fail(502, "invalid upstream configuration")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
request.Header.Set("Content-Type", "application/json")
|
||||||
|
if up == protocol.Anthropic {
|
||||||
|
request.Header.Set("x-api-key", c.APIKey)
|
||||||
|
request.Header.Set("anthropic-version", "2023-06-01")
|
||||||
|
if down == up {
|
||||||
|
if v := r.Header.Get("anthropic-version"); safeForwardedHeader(v, 32) {
|
||||||
|
request.Header.Set("anthropic-version", v)
|
||||||
|
}
|
||||||
|
if v := r.Header.Get("anthropic-beta"); safeForwardedHeader(v, 256) {
|
||||||
|
request.Header.Set("anthropic-beta", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
request.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||||
|
}
|
||||||
|
response, e := s.client.Do(request)
|
||||||
|
if e != nil {
|
||||||
|
fail(502, "upstream request failed")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer response.Body.Close()
|
||||||
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||||
|
status := response.StatusCode
|
||||||
|
if status < 400 {
|
||||||
|
status = 502
|
||||||
|
}
|
||||||
|
fail(status, "upstream rejected request")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if req.Stream {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.Header().Set("Cache-Control", "no-cache")
|
||||||
|
w.Header().Set("X-Accel-Buffering", "no")
|
||||||
|
w.WriteHeader(response.StatusCode)
|
||||||
|
l.Status = response.StatusCode
|
||||||
|
streamWriter := &deadlineWriter{ResponseWriter: w}
|
||||||
|
u, e := protocol.Stream(ctx, streamWriter, response.Body, up, down, req.Model)
|
||||||
|
l.Usage = u
|
||||||
|
if e != nil {
|
||||||
|
l.Status = 502
|
||||||
|
l.Error = "upstream stream interrupted"
|
||||||
|
l.Usage.Known = false
|
||||||
|
if ctx.Err() == nil {
|
||||||
|
event := "data: {\"error\":{\"type\":\"upstream_error\",\"message\":\"upstream stream interrupted\"}}\n\n"
|
||||||
|
if down != protocol.Chat {
|
||||||
|
event = "event: error\n" + event
|
||||||
|
}
|
||||||
|
_, _ = io.WriteString(streamWriter, event)
|
||||||
|
if f, ok := w.(http.Flusher); ok {
|
||||||
|
f.Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
b, e = io.ReadAll(io.LimitReader(response.Body, maxBody+1))
|
||||||
|
if e != nil || len(b) > maxBody {
|
||||||
|
fail(502, "invalid upstream response")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
l.Usage = protocol.ExtractUsage(up, b)
|
||||||
|
out, e := protocol.ConvertResponse(up, down, b)
|
||||||
|
if e != nil {
|
||||||
|
fail(502, "unsupported response conversion")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
l.Status = response.StatusCode
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(response.StatusCode)
|
||||||
|
if _, e = w.Write(out); e != nil {
|
||||||
|
l.Status = 499
|
||||||
|
l.Error = "client disconnected"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type deadlineWriter struct{ http.ResponseWriter }
|
||||||
|
|
||||||
|
// safeForwardedHeader reports whether v is a short printable token list that
|
||||||
|
// is safe to forward upstream; anything else is dropped so clients cannot
|
||||||
|
// smuggle arbitrary features or exotic header content past the channel
|
||||||
|
// configuration.
|
||||||
|
func safeForwardedHeader(v string, max int) bool {
|
||||||
|
if len(v) == 0 || len(v) > max {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for i := 0; i < len(v); i++ {
|
||||||
|
c := v[i]
|
||||||
|
switch {
|
||||||
|
case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9':
|
||||||
|
case c == '-' || c == '_' || c == '.' || c == ',' || c == ' ':
|
||||||
|
default:
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *deadlineWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
|
||||||
|
func (w *deadlineWriter) WriteHeader(status int) {
|
||||||
|
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||||
|
w.ResponseWriter.WriteHeader(status)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *deadlineWriter) Write(b []byte) (int, error) {
|
||||||
|
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||||
|
return w.ResponseWriter.Write(b)
|
||||||
|
}
|
||||||
|
func (w *deadlineWriter) FlushError() error {
|
||||||
|
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||||
|
return http.NewResponseController(w.ResponseWriter).Flush()
|
||||||
|
}
|
||||||
|
func (w *deadlineWriter) Flush() { _ = w.FlushError() }
|
||||||
@@ -0,0 +1,203 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/netip"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const bodyReadTimeout = 30 * time.Second
|
||||||
|
const maxLimitEntries = 10000
|
||||||
|
|
||||||
|
// Login rate limits: failures only, per IP and globally; successful logins
|
||||||
|
// never consume or hit these buckets.
|
||||||
|
const (
|
||||||
|
loginFailuresPerMinute = 10
|
||||||
|
loginGlobalFailuresPerMinute = 300
|
||||||
|
)
|
||||||
|
|
||||||
|
// authFailuresPerMinute bounds how many failed /v1 authentications one client
|
||||||
|
// IP may attempt per minute before further attempts are rejected before the
|
||||||
|
// database lookup.
|
||||||
|
const authFailuresPerMinute = 60
|
||||||
|
|
||||||
|
func pruneFailures(m map[string][]time.Time, now time.Time) {
|
||||||
|
for k, ts := range m {
|
||||||
|
if len(ts) == 0 || now.Sub(ts[len(ts)-1]) > time.Minute {
|
||||||
|
delete(m, k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// trimRecent drops entries older than one minute while reusing the backing
|
||||||
|
// array, so a long-lived failure window cannot grow without bound.
|
||||||
|
func trimRecent(ts []time.Time, now time.Time) []time.Time {
|
||||||
|
i := 0
|
||||||
|
for i < len(ts) && now.Sub(ts[i]) >= time.Minute {
|
||||||
|
i++
|
||||||
|
}
|
||||||
|
if i > 0 {
|
||||||
|
ts = append(ts[:0], ts[i:]...)
|
||||||
|
}
|
||||||
|
return ts
|
||||||
|
}
|
||||||
|
|
||||||
|
// authAllowed reports whether the client IP may attempt authentication. Only
|
||||||
|
// failures are counted, so legitimate traffic is unaffected.
|
||||||
|
func (s *Server) authAllowed(ip string) bool {
|
||||||
|
now := time.Now()
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
if now.Sub(s.authFailuresSeen) >= time.Minute {
|
||||||
|
pruneFailures(s.authFailures, now)
|
||||||
|
s.authFailuresSeen = now
|
||||||
|
}
|
||||||
|
recent := 0
|
||||||
|
for _, t := range s.authFailures[ip] {
|
||||||
|
if now.Sub(t) < time.Minute {
|
||||||
|
recent++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return recent < authFailuresPerMinute && (recent > 0 || len(s.authFailures) < maxLimitEntries)
|
||||||
|
}
|
||||||
|
|
||||||
|
// authFail records one failed authentication from ip.
|
||||||
|
func (s *Server) authFail(ip string) {
|
||||||
|
now := time.Now()
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
if _, ok := s.authFailures[ip]; !ok && len(s.authFailures) >= maxLimitEntries {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.authFailures[ip] = append(s.authFailures[ip], now)
|
||||||
|
}
|
||||||
|
|
||||||
|
func readBody(w http.ResponseWriter, r *http.Request, limit int64, timeout time.Duration, read func() error) error {
|
||||||
|
controller := http.NewResponseController(w)
|
||||||
|
_ = controller.SetReadDeadline(time.Now().Add(timeout))
|
||||||
|
defer controller.SetReadDeadline(time.Time{})
|
||||||
|
r.Body = http.MaxBytesReader(w, r.Body, limit)
|
||||||
|
defer r.Body.Close()
|
||||||
|
return read()
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetTrustedProxies must be called before serving requests.
|
||||||
|
func (s *Server) SetTrustedProxies(cidrs []string) error {
|
||||||
|
prefixes := make([]netip.Prefix, 0, len(cidrs))
|
||||||
|
for _, cidr := range cidrs {
|
||||||
|
p, err := netip.ParsePrefix(cidr)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
|
||||||
|
}
|
||||||
|
if p.Bits() == 0 {
|
||||||
|
return fmt.Errorf("trusted proxy CIDR must not trust all addresses")
|
||||||
|
}
|
||||||
|
prefixes = append(prefixes, p)
|
||||||
|
}
|
||||||
|
s.trustedProxies = prefixes
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) trusted(ip netip.Addr) bool {
|
||||||
|
for _, p := range s.trustedProxies {
|
||||||
|
if p.Contains(ip) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) clientIP(r *http.Request) string {
|
||||||
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||||
|
if err != nil {
|
||||||
|
host = r.RemoteAddr
|
||||||
|
}
|
||||||
|
peer, err := netip.ParseAddr(host)
|
||||||
|
if err != nil {
|
||||||
|
return host
|
||||||
|
}
|
||||||
|
peer = peer.Unmap()
|
||||||
|
if !s.trusted(peer) {
|
||||||
|
return peer.String()
|
||||||
|
}
|
||||||
|
parts := strings.Split(strings.Join(r.Header.Values("X-Forwarded-For"), ","), ",")
|
||||||
|
ips := make([]netip.Addr, len(parts))
|
||||||
|
for i, part := range parts {
|
||||||
|
ip, err := netip.ParseAddr(strings.TrimSpace(part))
|
||||||
|
if err != nil || ip.Zone() != "" {
|
||||||
|
return peer.String()
|
||||||
|
}
|
||||||
|
ips[i] = ip.Unmap()
|
||||||
|
}
|
||||||
|
for i := len(ips) - 1; i >= 0; i-- {
|
||||||
|
if !s.trusted(ips[i]) {
|
||||||
|
return ips[i].String()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return peer.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
type limitEntry struct {
|
||||||
|
window time.Time
|
||||||
|
requests, active int
|
||||||
|
}
|
||||||
|
type requestLimiter struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
rpm, perKey, global, active int
|
||||||
|
entries map[string]*limitEntry
|
||||||
|
nextCleanup time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func newRequestLimiter(rpm, perKey, global int) *requestLimiter {
|
||||||
|
return &requestLimiter{rpm: rpm, perKey: perKey, global: global, entries: make(map[string]*limitEntry)}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetLimits must be called before serving requests.
|
||||||
|
func (s *Server) SetLimits(rpm, perKey, global int) error {
|
||||||
|
if rpm <= 0 || perKey <= 0 || global <= 0 {
|
||||||
|
return fmt.Errorf("limits must be positive integers")
|
||||||
|
}
|
||||||
|
s.limiter = newRequestLimiter(rpm, perKey, global)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (l *requestLimiter) acquire(key string, now time.Time) (func(), string) {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
if !now.Before(l.nextCleanup) {
|
||||||
|
for k, e := range l.entries {
|
||||||
|
if e.active == 0 && now.Sub(e.window) >= time.Minute {
|
||||||
|
delete(l.entries, k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
l.nextCleanup = now.Add(time.Minute)
|
||||||
|
}
|
||||||
|
e := l.entries[key]
|
||||||
|
if e == nil {
|
||||||
|
if len(l.entries) >= maxLimitEntries {
|
||||||
|
return nil, "60"
|
||||||
|
}
|
||||||
|
e = &limitEntry{window: now}
|
||||||
|
l.entries[key] = e
|
||||||
|
}
|
||||||
|
if now.Sub(e.window) >= time.Minute {
|
||||||
|
e.window = now
|
||||||
|
e.requests = 0
|
||||||
|
}
|
||||||
|
if e.requests >= l.rpm {
|
||||||
|
seconds := int((time.Minute - now.Sub(e.window) + time.Second - 1) / time.Second)
|
||||||
|
return nil, strconv.Itoa(seconds)
|
||||||
|
}
|
||||||
|
if e.active >= l.perKey || l.active >= l.global {
|
||||||
|
return nil, "1"
|
||||||
|
}
|
||||||
|
e.requests++
|
||||||
|
e.active++
|
||||||
|
l.active++
|
||||||
|
var once sync.Once
|
||||||
|
return func() { once.Do(func() { l.mu.Lock(); defer l.mu.Unlock(); e.active--; l.active-- }) }, ""
|
||||||
|
}
|
||||||
@@ -0,0 +1,161 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSecurityStatsSQLSemantics(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||||
|
for _, l := range []requestLog{
|
||||||
|
{CreatedAt: now, KeyID: "a", KeyName: "old", Status: 200, Usage: protocol.Usage{Known: true, Input: 8}},
|
||||||
|
{CreatedAt: now, KeyID: "a", KeyName: "new", Status: 200, Error: "failed", Usage: protocol.Usage{Raw: []byte("{}"), Input: 100}},
|
||||||
|
{CreatedAt: now, KeyID: "b", Status: 400},
|
||||||
|
{CreatedAt: now, KeyID: "c", Status: 204, Usage: protocol.Usage{Known: true}},
|
||||||
|
} {
|
||||||
|
s.record(l)
|
||||||
|
}
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, tc := range []struct {
|
||||||
|
query string
|
||||||
|
requests, success, unknown int64
|
||||||
|
input *int64
|
||||||
|
}{
|
||||||
|
{"", 4, 2, 2, new(int64(8))},
|
||||||
|
{"?key_id=b", 1, 0, 1, nil},
|
||||||
|
{"?key_id=c", 1, 1, 0, new(int64(0))},
|
||||||
|
{"?key_id=missing", 0, 0, 0, nil},
|
||||||
|
} {
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats"+tc.query, nil))
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var got struct {
|
||||||
|
Summary summary
|
||||||
|
ByKey []struct{ Name string } `json:"by_key"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if got.Summary.Requests != tc.requests || got.Summary.Success != tc.success || got.Summary.Unknown != tc.unknown {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
if (got.Summary.Input == nil) != (tc.input == nil) || tc.input != nil && *got.Summary.Input != *tc.input {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
if tc.query == "" && got.ByKey[0].Name != "new" {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSecurityLogFieldsBounded(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
huge := strings.Repeat("界", 10000)
|
||||||
|
s.record(requestLog{RequestModel: huge, ActualModel: huge, Channel: huge, KeyID: huge, KeyName: huge, Upstream: huge, Downstream: huge, Error: huge, Usage: protocol.Usage{Known: true, Input: 12, Raw: []byte(huge)}})
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var model, actual, channel, key, name, up, down, message, raw string
|
||||||
|
var input int64
|
||||||
|
if err := s.db.QueryRow("SELECT request_model,actual_model,channel,key_id,key_name,upstream,downstream,error,raw_usage,input_tokens FROM logs").Scan(&model, &actual, &channel, &key, &name, &up, &down, &message, &raw, &input); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, field := range []string{model, actual, channel, key, name, up, down} {
|
||||||
|
if len(field) > maxLogFieldBytes || !utf8.ValidString(field) {
|
||||||
|
t.Fatal("unbounded or invalid field")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(message) > maxLogErrorBytes || !utf8.ValidString(message) || len(raw) > maxLogUsageBytes || !json.Valid([]byte(raw)) || input != 12 {
|
||||||
|
t.Fatal("invalid bounded log")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSecurityLogRetentionOptInAndBatch(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
old := time.Now().Add(-91 * 24 * time.Hour).Format(time.RFC3339Nano)
|
||||||
|
for i := 0; i < logCleanupBatch+3; i++ {
|
||||||
|
s.record(requestLog{CreatedAt: old})
|
||||||
|
}
|
||||||
|
s.record(requestLog{CreatedAt: time.Now().Format(time.RFC3339Nano)})
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if s.logRetentionDays.Load() != 0 {
|
||||||
|
t.Fatal("retention enabled by default")
|
||||||
|
}
|
||||||
|
if err := s.purgeExpiredLogs(context.Background(), 0); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != logCleanupBatch+4 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
s.SetLogRetention(90)
|
||||||
|
if s.logRetentionDays.Load() != 90 {
|
||||||
|
t.Fatal("retention not set")
|
||||||
|
}
|
||||||
|
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != 4 {
|
||||||
|
t.Fatal("unbounded purge", n)
|
||||||
|
}
|
||||||
|
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if n := logCount(t, s.db); n != 1 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
s.SetLogRetention(0)
|
||||||
|
if s.logRetentionDays.Load() != 0 {
|
||||||
|
t.Fatal("retention not disabled")
|
||||||
|
}
|
||||||
|
s.SetLogRetention(90)
|
||||||
|
if err := s.Close(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStatsGroupLimitAndCancellation(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
tx, err := s.db.Begin()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for i := 0; i < 10001; i++ {
|
||||||
|
if err := insertLog(context.Background(), tx, requestLog{CreatedAt: "2026-01-01T00:00:00Z", KeyID: fmt.Sprint(i)}); err != nil {
|
||||||
|
tx.Rollback()
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := tx.Commit(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil))
|
||||||
|
if w.Code != 422 {
|
||||||
|
t.Fatal("unbounded group result", w.Code)
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cancel()
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil).WithContext(ctx))
|
||||||
|
if w.Code == 200 {
|
||||||
|
t.Fatal("cancelled query succeeded")
|
||||||
|
}
|
||||||
|
var day string
|
||||||
|
if err := s.db.QueryRow("SELECT router_log_day(NULL, 'UTC')").Scan(&day); err == nil {
|
||||||
|
t.Fatal("invalid timestamp accepted")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,359 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestClientIPTrust(t *testing.T) {
|
||||||
|
s := &Server{}
|
||||||
|
if err := s.SetTrustedProxies([]string{"10.0.0.0/8", "::1/128"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := s.SetTrustedProxies([]string{"bad"}); err == nil {
|
||||||
|
t.Fatal("invalid CIDR accepted")
|
||||||
|
}
|
||||||
|
for _, tt := range []struct{ peer, xff, want string }{
|
||||||
|
{"192.0.2.1:42", "198.51.100.1", "192.0.2.1"},
|
||||||
|
{"10.0.0.1:42", "198.51.100.1, 10.0.0.2", "198.51.100.1"},
|
||||||
|
{"10.0.0.1:42", "198.51.100.1, 192.0.2.2", "192.0.2.2"},
|
||||||
|
{"10.0.0.1:42", "bad, 192.0.2.2", "10.0.0.1"},
|
||||||
|
{"10.0.0.1:42", "", "10.0.0.1"},
|
||||||
|
{"10.0.0.1:42", "10.0.0.2", "10.0.0.1"},
|
||||||
|
{"[::1]:42", "2001:db8::1", "2001:db8::1"},
|
||||||
|
} {
|
||||||
|
r := httptest.NewRequest("POST", "/", nil)
|
||||||
|
r.RemoteAddr = tt.peer
|
||||||
|
r.Header.Set("X-Forwarded-For", tt.xff)
|
||||||
|
r.Header.Set("X-Real-IP", "203.0.113.1")
|
||||||
|
if got := s.clientIP(r); got != tt.want {
|
||||||
|
t.Errorf("%+v: got %s", tt, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoginProxyIsolation(t *testing.T) {
|
||||||
|
s, err := New(":memory:", "password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
login := func(peer, forwarded string) int {
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"wrong"}`))
|
||||||
|
r.RemoteAddr = peer
|
||||||
|
r.Header.Set("X-Forwarded-For", forwarded)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.login(w, r)
|
||||||
|
return w.Code
|
||||||
|
}
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.1"); got != 401 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.1"); got != 429 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.2"); got != 401 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
login("192.0.2.3:1", fmt.Sprintf("198.51.100.%d", i))
|
||||||
|
}
|
||||||
|
if got := login("192.0.2.3:1", "203.0.113.1"); got != 429 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBodyDeadline(t *testing.T) {
|
||||||
|
result := make(chan error, 1)
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||||
|
result <- err
|
||||||
|
w.WriteHeader(400)
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 20\r\n\r\nx")
|
||||||
|
select {
|
||||||
|
case err := <-result:
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("slow body accepted")
|
||||||
|
}
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Fatal("body deadline did not fire")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBodyDeadlineCleared(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "read", 400)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
select {
|
||||||
|
case <-r.Context().Done():
|
||||||
|
http.Error(w, "canceled", 500)
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, "data: ok\n\n")
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
conn.SetDeadline(time.Now().Add(2 * time.Second))
|
||||||
|
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\n\r\nx")
|
||||||
|
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != 200 {
|
||||||
|
t.Fatal(resp.StatusCode)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoginBodyLimit(t *testing.T) {
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+strings.Repeat("x", 16<<10)+`"}`))
|
||||||
|
var v any
|
||||||
|
if decode(httptest.NewRecorder(), r, &v) == nil {
|
||||||
|
t.Fatal("oversize login accepted")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLimiterReleaseAndCleanup(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
l := newRequestLimiter(2, 1, 2)
|
||||||
|
a, _ := l.acquire("a", now)
|
||||||
|
if a == nil {
|
||||||
|
t.Fatal("denied")
|
||||||
|
}
|
||||||
|
if release, _ := l.acquire("a", now); release != nil {
|
||||||
|
t.Fatal("per-key concurrency bypass")
|
||||||
|
}
|
||||||
|
b, _ := l.acquire("b", now)
|
||||||
|
if b == nil {
|
||||||
|
t.Fatal("denied")
|
||||||
|
}
|
||||||
|
if release, _ := l.acquire("c", now); release != nil {
|
||||||
|
t.Fatal("global concurrency bypass")
|
||||||
|
}
|
||||||
|
a()
|
||||||
|
a()
|
||||||
|
b()
|
||||||
|
a, _ = l.acquire("a", now)
|
||||||
|
if a == nil {
|
||||||
|
t.Fatal("not released")
|
||||||
|
}
|
||||||
|
a()
|
||||||
|
if release, retry := l.acquire("a", now); release != nil || retry != "60" {
|
||||||
|
t.Fatal("rpm bypass")
|
||||||
|
}
|
||||||
|
c, _ := l.acquire("c", now.Add(time.Minute))
|
||||||
|
if c == nil {
|
||||||
|
t.Fatal("window not reset")
|
||||||
|
}
|
||||||
|
c()
|
||||||
|
if len(l.entries) != 1 {
|
||||||
|
t.Fatalf("stale entries: %d", len(l.entries))
|
||||||
|
}
|
||||||
|
for i := 0; i < maxLimitEntries; i++ {
|
||||||
|
l.entries[fmt.Sprint(i)] = &limitEntry{window: now.Add(time.Minute)}
|
||||||
|
}
|
||||||
|
if release, _ := l.acquire("overflow", now.Add(time.Minute)); release != nil {
|
||||||
|
t.Fatal("map cap bypass")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLimiterConcurrent(t *testing.T) {
|
||||||
|
l := newRequestLimiter(10000, 8, 64)
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
if release, _ := l.acquire("key", time.Now()); release != nil {
|
||||||
|
release()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
if l.active != 0 || l.entries["key"].active != 0 {
|
||||||
|
t.Fatal("permit leak")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProxyLimitRejectsBeforeBodyAndLogging(t *testing.T) {
|
||||||
|
s, err := New(":memory:", "test-password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
if err := s.SetLimits(1, 1, 1); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||||
|
key := integrationCreateKey(t, s, cookie, "limited", ch)
|
||||||
|
request := func(body io.Reader) *httptest.ResponseRecorder {
|
||||||
|
r := httptest.NewRequest("POST", "/v1/chat/completions", body)
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
return w
|
||||||
|
}
|
||||||
|
if w := request(strings.NewReader("{}")); w.Code != 400 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
if s.limiter.active != 0 {
|
||||||
|
t.Fatal("permit leaked after invalid request")
|
||||||
|
}
|
||||||
|
body := &observedReader{}
|
||||||
|
w := request(body)
|
||||||
|
if w.Code != 429 || w.Header().Get("Retry-After") == "" || body.read {
|
||||||
|
t.Fatalf("status=%d read=%t", w.Code, body.read)
|
||||||
|
}
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
var count int
|
||||||
|
if err := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&count); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if count != 1 {
|
||||||
|
t.Fatalf("rejected request logged: %d logs", count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type observedReader struct{ read bool }
|
||||||
|
|
||||||
|
func (r *observedReader) Read(p []byte) (int, error) { r.read = true; return 0, io.EOF }
|
||||||
|
|
||||||
|
func TestLoginFailuresOnlyAndGlobalCap(t *testing.T) {
|
||||||
|
s, err := New(":memory:", "password", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer s.Close()
|
||||||
|
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
login := func(peer, forwarded, password string) int {
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+password+`"}`))
|
||||||
|
r.RemoteAddr = peer
|
||||||
|
if forwarded != "" {
|
||||||
|
r.Header.Set("X-Forwarded-For", forwarded)
|
||||||
|
}
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.login(w, r)
|
||||||
|
return w.Code
|
||||||
|
}
|
||||||
|
for i := 0; i < 15; i++ {
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.1", "password"); got != 200 {
|
||||||
|
t.Fatalf("correct login %d rejected: %d", i, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 401 {
|
||||||
|
t.Fatalf("failure %d: %d", i, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 429 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
// The per-IP lockout is uniform: it must not confirm candidate passwords.
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.2", "password"); got != 429 {
|
||||||
|
t.Fatal("lockout leaks password validity:", got)
|
||||||
|
}
|
||||||
|
// A successful login clears its IP bucket, so fat-fingers do not accumulate.
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.3", "password"); got != 200 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||||
|
t.Fatalf("post-clear failure %d: %d", i, got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 429 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
for i := 0; i < loginGlobalFailuresPerMinute; i++ {
|
||||||
|
login("10.0.0.1:1", fmt.Sprintf("198.18.%d.%d", i/256, i%256), "wrong")
|
||||||
|
}
|
||||||
|
if got := login("10.0.0.1:1", "203.0.113.1", "wrong"); got != 429 {
|
||||||
|
t.Fatal(got)
|
||||||
|
}
|
||||||
|
// The global cap fails closed too: aggregate brute force locks every
|
||||||
|
// login out without revealing candidate-password validity.
|
||||||
|
if got := login("10.0.0.1:1", "203.0.113.1", "password"); got != 429 {
|
||||||
|
t.Fatal("global lockout leaks password validity:", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAuthFailureGate(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||||
|
key := integrationCreateKey(t, s, cookie, "gated", ch)
|
||||||
|
request := func(auth string) *httptest.ResponseRecorder {
|
||||||
|
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(integrationRequest("openai-chat", "abc/gpt-5", false)))
|
||||||
|
r.Header.Set("Authorization", "Bearer "+auth)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
return w
|
||||||
|
}
|
||||||
|
if w := request(key.raw); w.Code != 502 {
|
||||||
|
t.Fatal("valid key setup:", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
for i := 0; i < authFailuresPerMinute; i++ {
|
||||||
|
if w := request("wrong"); w.Code != 401 {
|
||||||
|
t.Fatalf("%d: %d", i, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if w := request("wrong"); w.Code != 429 || w.Header().Get("Retry-After") == "" {
|
||||||
|
t.Fatal(w.Code, w.Header())
|
||||||
|
}
|
||||||
|
// The gate applies before the key lookup: a valid key from the same
|
||||||
|
// (attacker-shared) IP is throttled the same way.
|
||||||
|
if w := request(key.raw); w.Code != 429 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProxyBodyTooLarge(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||||
|
key := integrationCreateKey(t, s, cookie, "limit", ch)
|
||||||
|
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"abc/gpt-5","messages":[],"pad":"`+strings.Repeat("x", maxBody)+`"}`))
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 413 || !strings.Contains(w.Body.String(), "too large") {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,327 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/sha256"
|
||||||
|
"crypto/subtle"
|
||||||
|
"database/sql"
|
||||||
|
"encoding/hex"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"io/fs"
|
||||||
|
_ "modernc.org/sqlite"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"net/netip"
|
||||||
|
"net/url"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Server struct {
|
||||||
|
logRetentionDays atomic.Int64
|
||||||
|
trustedProxies []netip.Prefix
|
||||||
|
limiter *requestLimiter
|
||||||
|
db *sql.DB
|
||||||
|
password [32]byte
|
||||||
|
assets fs.FS
|
||||||
|
mu sync.Mutex
|
||||||
|
sessions map[string]time.Time
|
||||||
|
attempts map[string][]time.Time
|
||||||
|
authFailures map[string][]time.Time
|
||||||
|
authFailuresSeen time.Time
|
||||||
|
loginFailures []time.Time
|
||||||
|
client *http.Client
|
||||||
|
publicOrigin string
|
||||||
|
lifecycleMu sync.Mutex
|
||||||
|
closing bool
|
||||||
|
active sync.WaitGroup
|
||||||
|
closeOnce sync.Once
|
||||||
|
closeErr error
|
||||||
|
logMu sync.RWMutex
|
||||||
|
logsClosed bool
|
||||||
|
logQueue chan logEntry
|
||||||
|
logDone chan struct{}
|
||||||
|
logErr error
|
||||||
|
}
|
||||||
|
|
||||||
|
const maxBody = 16 << 20
|
||||||
|
|
||||||
|
func New(p, password string, a fs.FS) (*Server, error) {
|
||||||
|
if password == "" {
|
||||||
|
return nil, errors.New("empty admin password")
|
||||||
|
}
|
||||||
|
if p != ":memory:" {
|
||||||
|
if e := os.MkdirAll(filepath.Dir(p), 0700); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
d, e := sql.Open("sqlite", p)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
d.SetMaxOpenConns(1)
|
||||||
|
schema := `PRAGMA busy_timeout=5000; CREATE TABLE IF NOT EXISTS channels(id TEXT PRIMARY KEY,name TEXT NOT NULL UNIQUE,base_url TEXT NOT NULL,protocol TEXT NOT NULL,enabled INTEGER NOT NULL,passthrough INTEGER NOT NULL,models TEXT NOT NULL,api_key TEXT NOT NULL); CREATE TABLE IF NOT EXISTS keys(id TEXT PRIMARY KEY,name TEXT NOT NULL,prefix TEXT NOT NULL,hash TEXT NOT NULL UNIQUE,enabled INTEGER NOT NULL,all_channels INTEGER NOT NULL,channels TEXT NOT NULL,expires_at TEXT NOT NULL,created_at TEXT NOT NULL,revoked INTEGER NOT NULL DEFAULT 0); CREATE TABLE IF NOT EXISTS logs(id INTEGER PRIMARY KEY AUTOINCREMENT,created_at TEXT NOT NULL,key_id TEXT NOT NULL,key_name TEXT NOT NULL,channel TEXT NOT NULL,request_model TEXT NOT NULL,actual_model TEXT NOT NULL,downstream TEXT NOT NULL,upstream TEXT NOT NULL,status INTEGER NOT NULL,duration_ms INTEGER NOT NULL,stream INTEGER NOT NULL,input_tokens INTEGER,output_tokens INTEGER,cache_read_tokens INTEGER,cache_write_tokens INTEGER,usage_known INTEGER NOT NULL,error TEXT NOT NULL,raw_usage TEXT NOT NULL); CREATE INDEX IF NOT EXISTS logs_time ON logs(created_at); CREATE INDEX IF NOT EXISTS logs_key ON logs(key_id);`
|
||||||
|
if _, e = d.Exec(schema); e != nil {
|
||||||
|
d.Close()
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if p != ":memory:" {
|
||||||
|
if e = os.Chmod(p, 0600); e != nil {
|
||||||
|
d.Close()
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||||
|
// Upstreams must not be able to turn compressed streams into proxy-side
|
||||||
|
// decompression work; the router asks for identity encoding.
|
||||||
|
tr.DisableCompression = true
|
||||||
|
tr.ResponseHeaderTimeout = 60 * time.Second
|
||||||
|
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||||
|
c, e := (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, addr)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
return &idleConn{Conn: c}, nil
|
||||||
|
}
|
||||||
|
s := &Server{db: d, password: sha256.Sum256([]byte(password)), assets: a, sessions: map[string]time.Time{}, attempts: map[string][]time.Time{}, authFailures: map[string][]time.Time{}, client: &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, logQueue: make(chan logEntry, 1024), logDone: make(chan struct{})}
|
||||||
|
go s.writeLogs()
|
||||||
|
s.limiter = newRequestLimiter(120, 8, 64)
|
||||||
|
return s, nil
|
||||||
|
}
|
||||||
|
func (s *Server) Close() error {
|
||||||
|
s.closeOnce.Do(func() {
|
||||||
|
s.lifecycleMu.Lock()
|
||||||
|
s.closing = true
|
||||||
|
s.lifecycleMu.Unlock()
|
||||||
|
s.active.Wait()
|
||||||
|
s.logMu.Lock()
|
||||||
|
s.logsClosed = true
|
||||||
|
close(s.logQueue)
|
||||||
|
s.logMu.Unlock()
|
||||||
|
<-s.logDone
|
||||||
|
s.client.CloseIdleConnections()
|
||||||
|
s.closeErr = errors.Join(s.logErr, s.db.Close())
|
||||||
|
})
|
||||||
|
return s.closeErr
|
||||||
|
}
|
||||||
|
func writeJSON(w http.ResponseWriter, v any) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
json.NewEncoder(w).Encode(v)
|
||||||
|
}
|
||||||
|
func errJSON(w http.ResponseWriter, n int, m string) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(n)
|
||||||
|
json.NewEncoder(w).Encode(map[string]string{"error": m})
|
||||||
|
}
|
||||||
|
func decode(w http.ResponseWriter, r *http.Request, v any) error {
|
||||||
|
limit := int64(maxBody)
|
||||||
|
if r.URL.Path == "/admin/api/login" {
|
||||||
|
limit = 16 << 10
|
||||||
|
}
|
||||||
|
return readBody(w, r, limit, bodyReadTimeout, func() error {
|
||||||
|
d := json.NewDecoder(r.Body)
|
||||||
|
if e := d.Decode(v); e != nil {
|
||||||
|
return e
|
||||||
|
}
|
||||||
|
var extra any
|
||||||
|
if d.Decode(&extra) != io.EOF {
|
||||||
|
return errors.New("trailing JSON")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
}
|
||||||
|
func randomID() string {
|
||||||
|
b := make([]byte, 32)
|
||||||
|
if _, e := rand.Read(b); e != nil {
|
||||||
|
panic(e)
|
||||||
|
}
|
||||||
|
return hex.EncodeToString(b)
|
||||||
|
}
|
||||||
|
func hash(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
|
||||||
|
func (s *Server) session(r *http.Request) bool {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
c, e := r.Cookie("ai_router_session")
|
||||||
|
if e != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
t, ok := s.sessions[c.Value]
|
||||||
|
return ok && time.Now().Before(t)
|
||||||
|
}
|
||||||
|
func (s *Server) guard(r *http.Request) bool {
|
||||||
|
if r.Header.Get("X-Requested-With") != "ai-router" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
o, e := url.Parse(r.Header.Get("Origin"))
|
||||||
|
if s.publicOrigin != "" {
|
||||||
|
// Origins are scheme/host case-insensitive; a mixed-case deployment
|
||||||
|
// value must not lock the admin out.
|
||||||
|
return e == nil && strings.EqualFold(r.Header.Get("Origin"), s.publicOrigin)
|
||||||
|
}
|
||||||
|
scheme := "http"
|
||||||
|
if r.TLS != nil {
|
||||||
|
scheme = "https"
|
||||||
|
}
|
||||||
|
return e == nil && o.Scheme == scheme && strings.EqualFold(o.Host, r.Host) && o.User == nil && o.RawQuery == "" && o.Fragment == "" && o.Path == ""
|
||||||
|
}
|
||||||
|
func (s *Server) secureCookie() bool {
|
||||||
|
return len(s.publicOrigin) >= 8 && strings.EqualFold(s.publicOrigin[:8], "https://")
|
||||||
|
}
|
||||||
|
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w = &deadlineWriter{ResponseWriter: w}
|
||||||
|
_ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(bodyReadTimeout))
|
||||||
|
s.lifecycleMu.Lock()
|
||||||
|
if s.closing {
|
||||||
|
s.lifecycleMu.Unlock()
|
||||||
|
errJSON(w, http.StatusServiceUnavailable, "server shutting down")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.active.Add(1)
|
||||||
|
s.lifecycleMu.Unlock()
|
||||||
|
defer s.active.Done()
|
||||||
|
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||||
|
w.Header().Set("Referrer-Policy", "no-referrer")
|
||||||
|
w.Header().Set("X-Frame-Options", "DENY")
|
||||||
|
w.Header().Set("Content-Security-Policy", "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'none'; connect-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:")
|
||||||
|
if strings.HasPrefix(r.URL.Path, "/admin/api/") || r.URL.Path == "/admin/api" {
|
||||||
|
w.Header().Set("Cache-Control", "no-store")
|
||||||
|
s.admin(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if strings.HasPrefix(r.URL.Path, "/v1/") {
|
||||||
|
s.proxy(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if strings.HasPrefix(r.URL.Path, "/api/") {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if s.assets == nil || (r.Method != "GET" && r.Method != "HEAD") {
|
||||||
|
http.NotFound(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
p := strings.TrimPrefix(r.URL.Path, "/")
|
||||||
|
if p == "" {
|
||||||
|
p = "index.html"
|
||||||
|
}
|
||||||
|
if st, e := fs.Stat(s.assets, p); e != nil || st.IsDir() {
|
||||||
|
p = "index.html"
|
||||||
|
}
|
||||||
|
http.ServeFileFS(w, r, s.assets, p)
|
||||||
|
}
|
||||||
|
func (s *Server) admin(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Method != "GET" && !s.guard(r) {
|
||||||
|
errJSON(w, 403, "forbidden")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.URL.Path == "/admin/api/login" && r.Method == "POST" {
|
||||||
|
s.login(w, r)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.URL.Path == "/admin/api/session" && r.Method == "GET" {
|
||||||
|
writeJSON(w, map[string]bool{"authenticated": s.session(r)})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !s.session(r) {
|
||||||
|
errJSON(w, 401, "unauthorized")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if r.URL.Path == "/admin/api/logout" && r.Method == "POST" {
|
||||||
|
c, _ := r.Cookie("ai_router_session")
|
||||||
|
s.mu.Lock()
|
||||||
|
delete(s.sessions, c.Value)
|
||||||
|
s.mu.Unlock()
|
||||||
|
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", MaxAge: -1, Path: "/", HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||||
|
writeJSON(w, map[string]bool{"authenticated": false})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case r.URL.Path == "/admin/api/models/test":
|
||||||
|
s.testModel(w, r)
|
||||||
|
case r.URL.Path == "/admin/api/channels":
|
||||||
|
s.channels(w, r)
|
||||||
|
case r.URL.Path == "/admin/api/channels/discover-models":
|
||||||
|
s.discoverModels(w, r)
|
||||||
|
case strings.HasPrefix(r.URL.Path, "/admin/api/channels/"):
|
||||||
|
s.channelID(w, r)
|
||||||
|
case r.URL.Path == "/admin/api/keys":
|
||||||
|
s.keys(w, r)
|
||||||
|
case strings.HasPrefix(r.URL.Path, "/admin/api/keys/"):
|
||||||
|
s.keyID(w, r)
|
||||||
|
case r.URL.Path == "/admin/api/logs" && r.Method == "GET":
|
||||||
|
s.logs(w, r)
|
||||||
|
case r.URL.Path == "/admin/api/stats" && r.Method == "GET":
|
||||||
|
s.stats(w, r)
|
||||||
|
default:
|
||||||
|
errJSON(w, 404, "not found")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// login gates on failure counts BEFORE verifying the password. While an IP is
|
||||||
|
// locked — or the global failure budget is spent — every attempt, including a
|
||||||
|
// correct password, gets the same 429 without a comparison, so responses never
|
||||||
|
// reveal whether a candidate password is right and brute force stays capped
|
||||||
|
// at loginFailuresPerMinute per IP and loginGlobalFailuresPerMinute overall.
|
||||||
|
// Successful logins are not failures: their gate slot is released by clearing
|
||||||
|
// the IP bucket, and only wrong passwords feed the global window.
|
||||||
|
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
|
||||||
|
ip := s.clientIP(r)
|
||||||
|
now := time.Now()
|
||||||
|
var x struct {
|
||||||
|
Password string `json:"password"`
|
||||||
|
}
|
||||||
|
if decode(w, r, &x) != nil {
|
||||||
|
errJSON(w, 400, "invalid request")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.mu.Lock()
|
||||||
|
pruneFailures(s.attempts, now)
|
||||||
|
s.loginFailures = trimRecent(s.loginFailures, now)
|
||||||
|
recent := []time.Time{}
|
||||||
|
for _, t := range s.attempts[ip] {
|
||||||
|
if now.Sub(t) < time.Minute {
|
||||||
|
recent = append(recent, t)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
blocked := len(recent) >= loginFailuresPerMinute || len(s.loginFailures) >= loginGlobalFailuresPerMinute || (len(recent) == 0 && len(s.attempts) >= 10000)
|
||||||
|
if !blocked {
|
||||||
|
// Reserve the gate slot atomically with the check so a concurrent
|
||||||
|
// burst cannot slip past the limit; a success erases the bucket below.
|
||||||
|
s.attempts[ip] = append(recent, now)
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
if blocked {
|
||||||
|
w.Header().Set("Retry-After", "60")
|
||||||
|
errJSON(w, 429, "too many attempts")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
h := sha256.Sum256([]byte(x.Password))
|
||||||
|
if subtle.ConstantTimeCompare(h[:], s.password[:]) != 1 {
|
||||||
|
s.mu.Lock()
|
||||||
|
s.loginFailures = append(s.loginFailures, now)
|
||||||
|
s.mu.Unlock()
|
||||||
|
errJSON(w, 401, "invalid credentials")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
v := randomID()
|
||||||
|
s.mu.Lock()
|
||||||
|
delete(s.attempts, ip)
|
||||||
|
for k, t := range s.sessions {
|
||||||
|
if now.After(t) {
|
||||||
|
delete(s.sessions, k)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if c, e := r.Cookie("ai_router_session"); e == nil {
|
||||||
|
delete(s.sessions, c.Value)
|
||||||
|
}
|
||||||
|
s.sessions[v] = now.Add(24 * time.Hour)
|
||||||
|
s.mu.Unlock()
|
||||||
|
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", Value: v, Path: "/", MaxAge: 86400, HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||||
|
writeJSON(w, map[string]bool{"authenticated": true})
|
||||||
|
}
|
||||||
@@ -0,0 +1,340 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"testing/fstest"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func testServer(t *testing.T) *Server {
|
||||||
|
t.Helper()
|
||||||
|
s, e := New(":memory:", "test-password", fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("SPA")}})
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { s.Close() })
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
func call(s *Server, m, p, b string, c *http.Cookie) *httptest.ResponseRecorder {
|
||||||
|
r := httptest.NewRequest(m, p, strings.NewReader(b))
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
if c != nil {
|
||||||
|
r.AddCookie(c)
|
||||||
|
}
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
return w
|
||||||
|
}
|
||||||
|
func loginCookie(t *testing.T, s *Server) *http.Cookie {
|
||||||
|
t.Helper()
|
||||||
|
w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
return w.Result().Cookies()[0]
|
||||||
|
}
|
||||||
|
func TestAuthentication(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
for _, origin := range []string{"", "http://evil.test", "https://example.com"} {
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||||
|
r.Header.Set("Origin", origin)
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 403 {
|
||||||
|
t.Fatal(origin, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
c := loginCookie(t, s)
|
||||||
|
if !c.HttpOnly || c.SameSite != http.SameSiteStrictMode {
|
||||||
|
t.Fatal(c)
|
||||||
|
}
|
||||||
|
if w := call(s, "GET", "/admin/api/channels", "", nil); w.Code != 401 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
if w := call(s, "POST", "/admin/api/logout", "", c); w.Code != 200 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
if w := call(s, "GET", "/admin/api/channels", "", c); w.Code != 401 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
if w := call(s, "GET", "/admin/api/session", "", nil); w.Code != 200 || w.Header().Get("Content-Security-Policy") == "" || w.Header().Get("X-Frame-Options") != "DENY" || w.Header().Get("Referrer-Policy") != "no-referrer" {
|
||||||
|
t.Fatal("security headers missing", w.Code, w.Header())
|
||||||
|
}
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil)
|
||||||
|
}
|
||||||
|
if w := call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil); w.Code != 429 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
// While locked, a correct password gets the same 429: the lockout must
|
||||||
|
// not become an oracle that confirms candidate passwords.
|
||||||
|
if w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil); w.Code != 429 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
// The lockout is per IP; a fresh address still logs in.
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||||
|
r.Header.Set("Origin", "http://example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
r.RemoteAddr = "192.0.2.9:1"
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal("fresh IP rejected:", w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestMappingPermissions(t *testing.T) {
|
||||||
|
cs := []channel{{ID: "stable", Name: "renamed", Enabled: true, Models: map[string]string{"alias": "vendor/real", "vendor/real": "wrong"}}}
|
||||||
|
k := apiKey{Channels: []string{"stable"}}
|
||||||
|
_, actual, e := resolve(cs, k, "renamed/alias")
|
||||||
|
if e != nil || actual != "vendor/real" {
|
||||||
|
t.Fatal(actual, e)
|
||||||
|
}
|
||||||
|
if _, _, e = resolve(cs, apiKey{Channels: []string{"renamed"}}, "renamed/alias"); e == nil {
|
||||||
|
t.Fatal("name granted access")
|
||||||
|
}
|
||||||
|
if _, _, e = resolve(cs, k, "renamed/vendor/unknown"); e == nil {
|
||||||
|
t.Fatal("unknown allowed")
|
||||||
|
}
|
||||||
|
cs[0].Enabled = false
|
||||||
|
if _, _, e = resolve(cs, k, "renamed/alias"); e == nil {
|
||||||
|
t.Fatal("disabled allowed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestLegacyPassthroughIgnored(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
t.Error("unmapped model reached upstream")
|
||||||
|
w.WriteHeader(http.StatusBadGateway)
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
body := map[string]any{
|
||||||
|
"name": "legacy", "base_url": up.URL, "protocol": "openai-chat",
|
||||||
|
"enabled": true, "passthrough": true, "models": map[string]string{"alias": "vendor/real"},
|
||||||
|
}
|
||||||
|
save := func() channel {
|
||||||
|
t.Helper()
|
||||||
|
b, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||||
|
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var c channel
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &c); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
c := save()
|
||||||
|
var stored bool
|
||||||
|
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
||||||
|
t.Fatal("new channel enabled passthrough", stored, err)
|
||||||
|
}
|
||||||
|
body["id"] = c.ID
|
||||||
|
save()
|
||||||
|
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
||||||
|
t.Fatal("update enabled passthrough", stored, err)
|
||||||
|
}
|
||||||
|
if _, err := s.db.Exec("UPDATE channels SET passthrough=1 WHERE id=?", c.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
cs, err := s.loadChannels()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, actual, err := resolve(cs, apiKey{AllChannels: true}, "legacy/alias"); err != nil || actual != "vendor/real" {
|
||||||
|
t.Fatal(actual, err)
|
||||||
|
}
|
||||||
|
w := call(s, "GET", "/admin/api/channels", "", cookie)
|
||||||
|
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
w = call(s, "POST", "/admin/api/keys", `{"name":"client","all_channels":true}`, cookie)
|
||||||
|
var key map[string]string
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &key); err != nil || w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, w.Body.String(), err)
|
||||||
|
}
|
||||||
|
for _, path := range []string{"/v1/chat/completions", "/v1/responses", "/v1/messages"} {
|
||||||
|
for _, model := range []string{"legacy/unknown", "legacy/vendor/real"} {
|
||||||
|
t.Run(path+"/"+model, func(t *testing.T) {
|
||||||
|
r := httptest.NewRequest("POST", path, strings.NewReader(`{"model":"`+model+`","messages":[]}`))
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != http.StatusForbidden {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
save()
|
||||||
|
cs, err = s.loadChannels()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, _, err := resolve(cs, apiKey{AllChannels: true}, "legacy/unknown"); err == nil {
|
||||||
|
t.Fatal("updated legacy channel allowed unknown model")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestKeyRevocationAndProxy(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
c := loginCookie(t, s)
|
||||||
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
||||||
|
t.Error(r.URL.Path, r.Header)
|
||||||
|
}
|
||||||
|
var b map[string]any
|
||||||
|
json.NewDecoder(r.Body).Decode(&b)
|
||||||
|
if b["model"] != "vendor/real" || b["custom"] != "keep" {
|
||||||
|
t.Error(b)
|
||||||
|
}
|
||||||
|
writeJSON(w, map[string]any{"choices": []any{}, "usage": map[string]int{"prompt_tokens": 7, "completion_tokens": 3}})
|
||||||
|
}))
|
||||||
|
defer up.Close()
|
||||||
|
b, _ := json.Marshal(channel{Name: "demo", BaseURL: up.URL + "/v1", Protocol: "openai-chat", Enabled: true, Models: map[string]string{"alias": "vendor/real"}, APIKey: "upstream-secret"})
|
||||||
|
w := call(s, "POST", "/admin/api/channels", string(b), c)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var ch channel
|
||||||
|
json.Unmarshal(w.Body.Bytes(), &ch)
|
||||||
|
if strings.Contains(w.Body.String(), "upstream-secret") {
|
||||||
|
t.Fatal("secret leak")
|
||||||
|
}
|
||||||
|
w = call(s, "POST", "/admin/api/keys", `{"name":"client","channels":["`+ch.ID+`"]}`, c)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var key map[string]string
|
||||||
|
json.Unmarshal(w.Body.Bytes(), &key)
|
||||||
|
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"demo/alias","messages":[],"custom":"keep"}`))
|
||||||
|
r.Header.Set("x-api-key", key["key"])
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var id, stored string
|
||||||
|
if e := s.db.QueryRow("SELECT id,hash FROM keys").Scan(&id, &stored); e != nil || stored == key["key"] {
|
||||||
|
t.Fatal(e, stored)
|
||||||
|
}
|
||||||
|
w = call(s, "DELETE", "/admin/api/keys/"+id, "", c)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
w = call(s, "PATCH", "/admin/api/keys/"+id, `{"enabled":true}`, c)
|
||||||
|
if w.Code != 404 {
|
||||||
|
t.Fatal("revived revoked key", w.Code)
|
||||||
|
}
|
||||||
|
r = httptest.NewRequest("GET", "/v1/models", nil)
|
||||||
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 401 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
var n int
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
s.db.QueryRow("SELECT count(*) FROM logs").Scan(&n)
|
||||||
|
if n != 1 {
|
||||||
|
t.Fatal(n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func TestStatsUnknownAndFilters(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
c := loginCookie(t, s)
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||||
|
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 200, Usage: protocol.Usage{Input: 8, Output: 2, Known: true}})
|
||||||
|
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 400})
|
||||||
|
s.record(requestLog{CreatedAt: now, KeyID: "b", Status: 200})
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w := call(s, "GET", "/admin/api/stats?key_id=a", "", c)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
var out struct {
|
||||||
|
Summary summary `json:"summary"`
|
||||||
|
}
|
||||||
|
if e := json.Unmarshal(w.Body.Bytes(), &out); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
if out.Summary.Requests != 2 || out.Summary.Success != 1 || out.Summary.Unknown != 1 || out.Summary.Input == nil || *out.Summary.Input != 8 {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
w = call(s, "GET", "/admin/api/logs?key_id=b", "", c)
|
||||||
|
if !strings.Contains(w.Body.String(), `"input_tokens":null`) {
|
||||||
|
t.Fatal(w.Body.String())
|
||||||
|
}
|
||||||
|
w = call(s, "GET", "/admin/api/stats?from=invalid", "", c)
|
||||||
|
if w.Code != 400 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
w = call(s, "GET", "/admin/api/unknown", "", c)
|
||||||
|
if w.Code != 404 {
|
||||||
|
t.Fatal(w.Code)
|
||||||
|
}
|
||||||
|
w = call(s, "GET", "/some/spa/path", "", nil)
|
||||||
|
if w.Code != 200 || w.Body.String() != "SPA" {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPublicOriginCaseInsensitive(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
if e := s.SetPublicOrigin("https://Router.Example.com"); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||||
|
r.Header.Set("Origin", "https://router.example.com")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
||||||
|
t.Fatal(w.Code, w.Header())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPublicOriginAndRedirect(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
if e := s.SetPublicOrigin("https://router.example"); e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
||||||
|
r.Header.Set("Origin", "https://router.example")
|
||||||
|
r.Header.Set("X-Requested-With", "ai-router")
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
s.ServeHTTP(w, r)
|
||||||
|
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
||||||
|
t.Fatal(w.Code, w.Header())
|
||||||
|
}
|
||||||
|
reached := false
|
||||||
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true }))
|
||||||
|
defer target.Close()
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||||
|
defer upstream.Close()
|
||||||
|
req, _ := http.NewRequest("POST", upstream.URL, strings.NewReader("body"))
|
||||||
|
req.Header.Set("Authorization", "Bearer private")
|
||||||
|
res, e := s.client.Do(req)
|
||||||
|
if e != nil {
|
||||||
|
t.Fatal(e)
|
||||||
|
}
|
||||||
|
res.Body.Close()
|
||||||
|
if reached || res.StatusCode != 307 {
|
||||||
|
t.Fatal("redirect followed")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,306 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"database/sql/driver"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
_ "time/tzdata"
|
||||||
|
|
||||||
|
"modernc.org/sqlite"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
var locationMu sync.Mutex
|
||||||
|
var cachedZone string
|
||||||
|
var cachedLocation *time.Location
|
||||||
|
if err := sqlite.RegisterDeterministicScalarFunction("router_log_day", 2, func(_ *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
|
||||||
|
stampText, ok := args[0].(string)
|
||||||
|
zone, zoneOK := args[1].(string)
|
||||||
|
if !ok || !zoneOK {
|
||||||
|
return nil, fmt.Errorf("invalid log timestamp or timezone")
|
||||||
|
}
|
||||||
|
stamp, err := time.Parse(time.RFC3339Nano, stampText)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
locationMu.Lock()
|
||||||
|
defer locationMu.Unlock()
|
||||||
|
if cachedLocation == nil || cachedZone != zone {
|
||||||
|
loc, err := time.LoadLocation(zone)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
cachedZone, cachedLocation = zone, loc
|
||||||
|
}
|
||||||
|
return stamp.In(cachedLocation).Format("2006-01-02"), nil
|
||||||
|
}); err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type requestLog struct {
|
||||||
|
ID int64 `json:"id"`
|
||||||
|
CreatedAt string `json:"created_at"`
|
||||||
|
KeyID string `json:"key_id"`
|
||||||
|
KeyName string `json:"key_name"`
|
||||||
|
Channel string `json:"channel"`
|
||||||
|
RequestModel string `json:"request_model"`
|
||||||
|
ActualModel string `json:"actual_model"`
|
||||||
|
Downstream string `json:"downstream"`
|
||||||
|
Upstream string `json:"upstream"`
|
||||||
|
Status int `json:"status"`
|
||||||
|
DurationMS int64 `json:"duration_ms"`
|
||||||
|
Stream bool `json:"stream"`
|
||||||
|
Input *int64 `json:"input_tokens"`
|
||||||
|
Output *int64 `json:"output_tokens"`
|
||||||
|
CacheRead *int64 `json:"cache_read_tokens"`
|
||||||
|
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||||
|
UsageKnown bool `json:"usage_known"`
|
||||||
|
Error string `json:"error"`
|
||||||
|
RawUsage json.RawMessage `json:"raw_usage"`
|
||||||
|
Usage protocol.Usage `json:"-"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertLog(ctx context.Context, tx *sql.Tx, l requestLog) error {
|
||||||
|
l = boundedLog(l)
|
||||||
|
var in, out, cr, cw any
|
||||||
|
if l.Usage.Known || len(l.Usage.Raw) > 0 {
|
||||||
|
in = l.Usage.Input
|
||||||
|
out = l.Usage.Output
|
||||||
|
cr = l.Usage.CacheRead
|
||||||
|
cw = l.Usage.CacheWrite
|
||||||
|
}
|
||||||
|
_, err := tx.ExecContext(ctx, "INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", l.CreatedAt, l.KeyID, l.KeyName, l.Channel, l.RequestModel, l.ActualModel, l.Downstream, l.Upstream, l.Status, l.DurationMS, l.Stream, in, out, cr, cw, l.Usage.Known, l.Error, string(l.Usage.Raw))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
func filter(r *http.Request) (string, []any, error) {
|
||||||
|
parts := []string{"1=1"}
|
||||||
|
args := []any{}
|
||||||
|
for _, name := range []string{"channel", "model"} {
|
||||||
|
if v := r.URL.Query().Get(name); v != "" {
|
||||||
|
col := name
|
||||||
|
if name == "model" {
|
||||||
|
col = "actual_model"
|
||||||
|
}
|
||||||
|
parts = append(parts, col+"=?")
|
||||||
|
args = append(args, v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if k := r.URL.Query().Get("key_id"); k != "" {
|
||||||
|
parts = append(parts, "key_id=?")
|
||||||
|
args = append(args, k)
|
||||||
|
}
|
||||||
|
var from, to time.Time
|
||||||
|
for _, n := range []string{"from", "to"} {
|
||||||
|
if v := r.URL.Query().Get(n); v != "" {
|
||||||
|
t, e := time.Parse(time.RFC3339Nano, v)
|
||||||
|
if e != nil {
|
||||||
|
return "", nil, e
|
||||||
|
}
|
||||||
|
if n == "from" {
|
||||||
|
from = t
|
||||||
|
parts = append(parts, "julianday(created_at)>=julianday(?)")
|
||||||
|
} else {
|
||||||
|
to = t
|
||||||
|
parts = append(parts, "julianday(created_at)<=julianday(?)")
|
||||||
|
}
|
||||||
|
args = append(args, t.UTC().Format(time.RFC3339Nano))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !from.IsZero() && !to.IsZero() && from.After(to) {
|
||||||
|
return "", nil, sql.ErrNoRows
|
||||||
|
}
|
||||||
|
return strings.Join(parts, " AND "), args, nil
|
||||||
|
}
|
||||||
|
func (s *Server) readLogs(r *http.Request, limited bool) ([]requestLog, error) {
|
||||||
|
where, args, e := filter(r)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
q := "SELECT id,created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage FROM logs WHERE " + where + " ORDER BY id DESC"
|
||||||
|
if limited {
|
||||||
|
q += " LIMIT 200"
|
||||||
|
}
|
||||||
|
rows, e := s.db.Query(q, args...)
|
||||||
|
if e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
out := []requestLog{}
|
||||||
|
for rows.Next() {
|
||||||
|
var l requestLog
|
||||||
|
var raw string
|
||||||
|
if e = rows.Scan(&l.ID, &l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream, &l.Status, &l.DurationMS, &l.Stream, &l.Input, &l.Output, &l.CacheRead, &l.CacheWrite, &l.UsageKnown, &l.Error, &raw); e != nil {
|
||||||
|
return nil, e
|
||||||
|
}
|
||||||
|
if json.Valid([]byte(raw)) {
|
||||||
|
l.RawUsage = json.RawMessage(raw)
|
||||||
|
}
|
||||||
|
out = append(out, l)
|
||||||
|
}
|
||||||
|
return out, rows.Err()
|
||||||
|
}
|
||||||
|
func (s *Server) logs(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if _, _, e := filter(r); e != nil {
|
||||||
|
errJSON(w, 400, "invalid time filter")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ls, e := s.readLogs(r, true)
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, ls)
|
||||||
|
}
|
||||||
|
|
||||||
|
type summary struct {
|
||||||
|
Requests int64 `json:"requests"`
|
||||||
|
Success int64 `json:"success"`
|
||||||
|
Failed int64 `json:"failed"`
|
||||||
|
Input *int64 `json:"input_tokens"`
|
||||||
|
Output *int64 `json:"output_tokens"`
|
||||||
|
CacheRead *int64 `json:"cache_read_tokens"`
|
||||||
|
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||||
|
Unknown int64 `json:"unknown_usage"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *summary) add(l requestLog) {
|
||||||
|
a.Requests++
|
||||||
|
if l.Status >= 200 && l.Status < 300 && l.Error == "" {
|
||||||
|
a.Success++
|
||||||
|
} else {
|
||||||
|
a.Failed++
|
||||||
|
}
|
||||||
|
if !l.UsageKnown {
|
||||||
|
a.Unknown++
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if a.Input == nil {
|
||||||
|
a.Input = new(int64)
|
||||||
|
a.Output = new(int64)
|
||||||
|
a.CacheRead = new(int64)
|
||||||
|
a.CacheWrite = new(int64)
|
||||||
|
}
|
||||||
|
if l.Input != nil {
|
||||||
|
*a.Input += *l.Input
|
||||||
|
}
|
||||||
|
if l.Output != nil {
|
||||||
|
*a.Output += *l.Output
|
||||||
|
}
|
||||||
|
if l.CacheRead != nil {
|
||||||
|
*a.CacheRead += *l.CacheRead
|
||||||
|
}
|
||||||
|
if l.CacheWrite != nil {
|
||||||
|
*a.CacheWrite += *l.CacheWrite
|
||||||
|
}
|
||||||
|
}
|
||||||
|
func (s *Server) stats(w http.ResponseWriter, r *http.Request) {
|
||||||
|
zone := r.URL.Query().Get("timezone")
|
||||||
|
if zone == "" {
|
||||||
|
zone = "UTC"
|
||||||
|
}
|
||||||
|
_, err := time.LoadLocation(zone)
|
||||||
|
if err != nil || zone == "Local" {
|
||||||
|
errJSON(w, 400, "invalid timezone")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
where, args, e := filter(r)
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 400, "invalid time filter")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
const aggregates = `COUNT(*), COALESCE(SUM(status>=200 AND status<300 AND error=''),0),
|
||||||
|
COALESCE(SUM(NOT(status>=200 AND status<300 AND error='')),0),
|
||||||
|
SUM(CASE WHEN usage_known THEN COALESCE(input_tokens,0) END),
|
||||||
|
SUM(CASE WHEN usage_known THEN COALESCE(output_tokens,0) END),
|
||||||
|
SUM(CASE WHEN usage_known THEN COALESCE(cache_read_tokens,0) END),
|
||||||
|
SUM(CASE WHEN usage_known THEN COALESCE(cache_write_tokens,0) END),
|
||||||
|
COALESCE(SUM(NOT usage_known),0), MAX(id)`
|
||||||
|
// NOT MATERIALIZED avoids copying the full log payload into a temporary table.
|
||||||
|
query := `WITH filtered AS NOT MATERIALIZED (SELECT * FROM logs WHERE ` + where + `)
|
||||||
|
SELECT 'total','','',` + aggregates + ` FROM filtered
|
||||||
|
UNION ALL SELECT 'key',key_id,key_name,` + aggregates + ` FROM filtered GROUP BY key_id
|
||||||
|
UNION ALL SELECT 'day',router_log_day(created_at,?), '',` + aggregates + ` FROM filtered GROUP BY router_log_day(created_at,?)
|
||||||
|
UNION ALL SELECT 'model',channel,actual_model,` + aggregates + ` FROM filtered GROUP BY channel,actual_model`
|
||||||
|
args = append(args, zone, zone)
|
||||||
|
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
rows, e := s.db.QueryContext(ctx, query, args...)
|
||||||
|
if e != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
type byKey struct {
|
||||||
|
KeyID string `json:"key_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
summary
|
||||||
|
}
|
||||||
|
type daily struct {
|
||||||
|
Day string `json:"day"`
|
||||||
|
summary
|
||||||
|
}
|
||||||
|
type byModel struct {
|
||||||
|
Channel string `json:"channel"`
|
||||||
|
Model string `json:"model"`
|
||||||
|
summary
|
||||||
|
}
|
||||||
|
total := summary{}
|
||||||
|
ks := map[string]*byKey{}
|
||||||
|
ds := map[string]*daily{}
|
||||||
|
ms := map[[2]string]*byModel{}
|
||||||
|
groups := 0
|
||||||
|
for rows.Next() {
|
||||||
|
groups++
|
||||||
|
if groups > 10000 {
|
||||||
|
errJSON(w, 422, "too many statistics groups; narrow the time range")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
var kind, a, b string
|
||||||
|
var v summary
|
||||||
|
var latest sql.NullInt64
|
||||||
|
if err := rows.Scan(&kind, &a, &b, &v.Requests, &v.Success, &v.Failed, &v.Input, &v.Output, &v.CacheRead, &v.CacheWrite, &v.Unknown, &latest); err != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
switch kind {
|
||||||
|
case "total":
|
||||||
|
total = v
|
||||||
|
case "key":
|
||||||
|
// SQLite's sole MAX selects key_name from the latest ID, not timestamp.
|
||||||
|
ks[a] = &byKey{KeyID: a, Name: b, summary: v}
|
||||||
|
case "day":
|
||||||
|
ds[a] = &daily{Day: a, summary: v}
|
||||||
|
case "model":
|
||||||
|
ms[[2]string{a, b}] = &byModel{Channel: a, Model: b, summary: v}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := rows.Err(); err != nil {
|
||||||
|
errJSON(w, 500, "database error")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
keys := []*byKey{}
|
||||||
|
days := []*daily{}
|
||||||
|
mods := []*byModel{}
|
||||||
|
for _, v := range ks {
|
||||||
|
keys = append(keys, v)
|
||||||
|
}
|
||||||
|
for _, v := range ds {
|
||||||
|
days = append(days, v)
|
||||||
|
}
|
||||||
|
for _, v := range ms {
|
||||||
|
mods = append(mods, v)
|
||||||
|
}
|
||||||
|
sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
|
||||||
|
sort.Slice(days, func(i, j int) bool { return days[i].Day < days[j].Day })
|
||||||
|
sort.Slice(mods, func(i, j int) bool { return mods[i].Channel+mods[i].Model < mods[j].Channel+mods[j].Model })
|
||||||
|
writeJSON(w, map[string]any{"summary": total, "by_key": keys, "daily": days, "by_model": mods, "timezone": zone})
|
||||||
|
}
|
||||||
@@ -0,0 +1,74 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/protocol"
|
||||||
|
"encoding/json"
|
||||||
|
"net/url"
|
||||||
|
"reflect"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestStatsBrowserTimezone(t *testing.T) {
|
||||||
|
for _, tc := range []struct {
|
||||||
|
name, zone string
|
||||||
|
stamps []string
|
||||||
|
want map[string]int
|
||||||
|
}{
|
||||||
|
{"Shanghai midnight", "Asia/Shanghai", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 1, "2026-10-21": 1}},
|
||||||
|
{"UTC default", "", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 2}},
|
||||||
|
{"New York DST spring", "America/New_York", []string{"2025-03-09T04:59:59Z", "2025-03-09T05:00:00Z", "2025-03-10T03:59:59Z", "2025-03-10T04:00:00Z"}, map[string]int{"2025-03-08": 1, "2025-03-09": 2, "2025-03-10": 1}},
|
||||||
|
{"New York DST fall", "America/New_York", []string{"2025-11-02T03:59:59Z", "2025-11-02T04:00:00Z", "2025-11-03T04:59:59Z", "2025-11-03T05:00:00Z"}, map[string]int{"2025-11-01": 1, "2025-11-02": 2, "2025-11-03": 1}},
|
||||||
|
} {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
cookie := loginCookie(t, s)
|
||||||
|
for _, stamp := range tc.stamps {
|
||||||
|
s.record(requestLog{CreatedAt: stamp, KeyID: "one", Status: 200, Usage: protocol.Usage{Known: true, Input: 10}})
|
||||||
|
}
|
||||||
|
s.record(requestLog{CreatedAt: tc.stamps[0], KeyID: "other", Status: 200})
|
||||||
|
if err := s.flushLogs(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w := call(s, "GET", "/admin/api/stats?key_id=one&timezone="+url.QueryEscape(tc.zone), "", cookie)
|
||||||
|
if w.Code != 200 {
|
||||||
|
t.Fatal(w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var result struct {
|
||||||
|
Timezone string `json:"timezone"`
|
||||||
|
Summary summary `json:"summary"`
|
||||||
|
Daily []struct {
|
||||||
|
Day string `json:"day"`
|
||||||
|
Requests int `json:"requests"`
|
||||||
|
} `json:"daily"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got := map[string]int{}
|
||||||
|
for _, d := range result.Daily {
|
||||||
|
got[d.Day] = d.Requests
|
||||||
|
}
|
||||||
|
if !reflect.DeepEqual(got, tc.want) {
|
||||||
|
t.Fatalf("got %v want %v", got, tc.want)
|
||||||
|
}
|
||||||
|
if result.Summary.Requests != int64(len(tc.stamps)) || result.Summary.Input == nil || *result.Summary.Input != int64(10*len(tc.stamps)) {
|
||||||
|
t.Fatal(result.Summary)
|
||||||
|
}
|
||||||
|
var stored string
|
||||||
|
if err := s.db.QueryRow("SELECT created_at FROM logs ORDER BY id LIMIT 1").Scan(&stored); err != nil || stored != tc.stamps[0] {
|
||||||
|
t.Fatal("UTC storage changed", stored, err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestStatsRejectsInvalidTimezone(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
c := loginCookie(t, s)
|
||||||
|
for _, zone := range []string{"Local", "not/a-zone", "../../etc/passwd"} {
|
||||||
|
w := call(s, "GET", "/admin/api/stats?timezone="+url.QueryEscape(zone), "", c)
|
||||||
|
if w.Code != 400 {
|
||||||
|
t.Fatal(zone, w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,33 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"net"
|
||||||
|
"net/url"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type idleConn struct{ net.Conn }
|
||||||
|
|
||||||
|
func (c *idleConn) Read(b []byte) (int, error) {
|
||||||
|
if e := c.SetReadDeadline(time.Now().Add(2 * time.Minute)); e != nil {
|
||||||
|
return 0, e
|
||||||
|
}
|
||||||
|
return c.Conn.Read(b)
|
||||||
|
}
|
||||||
|
func (c *idleConn) Write(b []byte) (int, error) {
|
||||||
|
if e := c.SetWriteDeadline(time.Now().Add(30 * time.Second)); e != nil {
|
||||||
|
return 0, e
|
||||||
|
}
|
||||||
|
return c.Conn.Write(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetPublicOrigin must be called before serving requests; forwarded headers are not trusted.
|
||||||
|
func (s *Server) SetPublicOrigin(origin string) error {
|
||||||
|
u, e := url.Parse(origin)
|
||||||
|
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" {
|
||||||
|
return errors.New("public origin must be an http(s) origin")
|
||||||
|
}
|
||||||
|
s.publicOrigin = origin
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,51 @@
|
|||||||
|
package server
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type observedDeadlineWriter struct {
|
||||||
|
*httptest.ResponseRecorder
|
||||||
|
writes, reads int
|
||||||
|
deadline time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *observedDeadlineWriter) SetWriteDeadline(d time.Time) error {
|
||||||
|
w.writes++
|
||||||
|
w.deadline = d
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (w *observedDeadlineWriter) SetReadDeadline(d time.Time) error { w.reads++; return nil }
|
||||||
|
|
||||||
|
func TestResponseWriteDeadlines(t *testing.T) {
|
||||||
|
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||||
|
w := &deadlineWriter{ResponseWriter: raw}
|
||||||
|
w.WriteHeader(200)
|
||||||
|
if raw.writes != 1 || !raw.deadline.After(time.Now()) {
|
||||||
|
t.Fatal("headers lack deadline")
|
||||||
|
}
|
||||||
|
if _, err := w.Write([]byte("body")); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := w.FlushError(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if raw.writes != 3 {
|
||||||
|
t.Fatal("write/flush must refresh deadline", raw.writes)
|
||||||
|
}
|
||||||
|
if err := http.NewResponseController(w).SetReadDeadline(time.Time{}); err != nil || raw.reads != 1 {
|
||||||
|
t.Fatal("wrapper must preserve read deadline control", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestNonStreamResponseHasWriteDeadline(t *testing.T) {
|
||||||
|
s := testServer(t)
|
||||||
|
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||||
|
s.ServeHTTP(raw, httptest.NewRequest("GET", "/admin/api/session", nil))
|
||||||
|
if raw.Code != 200 || raw.writes == 0 || !raw.deadline.After(time.Now()) {
|
||||||
|
t.Fatal("nonstream response lacks write deadline")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,70 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"ai-router/internal/config"
|
||||||
|
"ai-router/internal/server"
|
||||||
|
"context"
|
||||||
|
"embed"
|
||||||
|
"flag"
|
||||||
|
"io/fs"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"os/signal"
|
||||||
|
"syscall"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
//go:embed all:web/dist
|
||||||
|
var frontend embed.FS
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
path := flag.String("config", "config.yaml", "YAML configuration path")
|
||||||
|
flag.Parse()
|
||||||
|
cfg, err := config.Load(*path)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
assets, err := fs.Sub(frontend, "web/dist")
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
app, err := server.New(cfg.Storage.Path, cfg.Admin.Password, assets)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
if err := app.Close(); err != nil {
|
||||||
|
log.Print("AI Router shutdown: failed to persist all request logs or close database")
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
if cfg.Server.PublicOrigin != "" {
|
||||||
|
if err := app.SetPublicOrigin(cfg.Server.PublicOrigin); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := app.SetTrustedProxies(cfg.Server.TrustedProxies); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := app.SetLimits(cfg.Limits.RequestsPerMinute, cfg.Limits.PerKeyConcurrency, cfg.Limits.GlobalConcurrency); err != nil {
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
app.SetLogRetention(cfg.Storage.LogRetentionDays)
|
||||||
|
srv := &http.Server{Addr: cfg.Server.Listen, Handler: app, ReadTimeout: 30 * time.Second, ReadHeaderTimeout: 10 * time.Second, IdleTimeout: 90 * time.Second, MaxHeaderBytes: 1 << 20}
|
||||||
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
||||||
|
defer stop()
|
||||||
|
done := make(chan error, 1)
|
||||||
|
go func() { log.Printf("AI Router listening on %s", cfg.Server.Listen); done <- srv.ListenAndServe() }()
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
if err != http.ErrServerClosed {
|
||||||
|
log.Printf("server: %v", err)
|
||||||
|
}
|
||||||
|
case <-ctx.Done():
|
||||||
|
shutdown, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
if err := srv.Shutdown(shutdown); err != nil {
|
||||||
|
_ = srv.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
node_modules/
|
||||||
|
*.tsbuildinfo
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
<!doctype html>
|
||||||
|
<html lang="zh-CN"><head><meta charset="UTF-8"/><meta name="viewport" content="width=device-width, initial-scale=1.0"/><meta name="theme-color" content="#181a34"/><title>AI Router · 管理台</title><script src="/theme-init.js"></script></head><body><div id="root"></div><script type="module" src="/src/main.tsx"></script></body></html>
|
||||||
Generated
+1883
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
|||||||
|
{
|
||||||
|
"name": "ai-router-admin",
|
||||||
|
"version": "1.0.0",
|
||||||
|
"private": true,
|
||||||
|
"type": "module",
|
||||||
|
"scripts": {
|
||||||
|
"dev": "vite --host 127.0.0.1",
|
||||||
|
"build": "tsc -b && vite build",
|
||||||
|
"preview": "vite preview --host 127.0.0.1"
|
||||||
|
},
|
||||||
|
"dependencies": {
|
||||||
|
"react": "^19.0.0",
|
||||||
|
"react-dom": "^19.0.0"
|
||||||
|
},
|
||||||
|
"devDependencies": {
|
||||||
|
"@types/react": "^19.0.0",
|
||||||
|
"@types/react-dom": "^19.0.0",
|
||||||
|
"@vitejs/plugin-react": "^4.3.4",
|
||||||
|
"typescript": "~5.7.3",
|
||||||
|
"vite": "^6.1.0"
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
(function () {
|
||||||
|
try {
|
||||||
|
var t = localStorage.getItem('theme')
|
||||||
|
if (t !== 'light' && t !== 'dark') t = matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light'
|
||||||
|
document.documentElement.dataset.theme = t
|
||||||
|
} catch (e) {}
|
||||||
|
})()
|
||||||
+158
@@ -0,0 +1,158 @@
|
|||||||
|
import { useEffect, useState } from 'react'
|
||||||
|
import { api, APIError, type Channel, type Key, type Log, type Stats } from './api'
|
||||||
|
import { emptyFilter, type Action, type Filter, type Page } from './common'
|
||||||
|
import { applyTheme, initialTheme, type Theme } from './theme'
|
||||||
|
import { Filters } from './components/Filters'
|
||||||
|
import { Shell } from './components/Shell'
|
||||||
|
import { Feedback } from './components/ui'
|
||||||
|
import { ChannelsPage } from './pages/Channels'
|
||||||
|
import { KeysPage } from './pages/Keys'
|
||||||
|
import { Login } from './pages/Login'
|
||||||
|
import { LogsPage } from './pages/Logs'
|
||||||
|
import { ModelsPage } from './pages/Models'
|
||||||
|
import { Overview } from './pages/Overview'
|
||||||
|
|
||||||
|
export default function App() {
|
||||||
|
const timezone = Intl.DateTimeFormat().resolvedOptions().timeZone
|
||||||
|
const [theme, setTheme] = useState<Theme>(initialTheme)
|
||||||
|
const [auth, setAuth] = useState<boolean | null>(null)
|
||||||
|
const [page, setPage] = useState<Page>('概览')
|
||||||
|
const [error, setError] = useState('')
|
||||||
|
const [notice, setNotice] = useState('')
|
||||||
|
const [busy, setBusy] = useState(false)
|
||||||
|
const [loading, setLoading] = useState(false)
|
||||||
|
const [rev, setRev] = useState(0)
|
||||||
|
const [channels, setChannels] = useState<Channel[]>([])
|
||||||
|
const [keys, setKeys] = useState<Key[]>([])
|
||||||
|
const [stats, setStats] = useState<Stats | null>(null)
|
||||||
|
const [logs, setLogs] = useState<Log[]>([])
|
||||||
|
const [filter, setFilter] = useState<Filter>(emptyFilter)
|
||||||
|
const [query, setQuery] = useState('')
|
||||||
|
|
||||||
|
useEffect(() => applyTheme(theme), [theme])
|
||||||
|
|
||||||
|
function fail(e: unknown) {
|
||||||
|
setError(e instanceof Error ? e.message : '操作失败，请重试。')
|
||||||
|
// 会话失效后回到登录页，各页面的弹窗状态随组件卸载自动清理
|
||||||
|
if (e instanceof APIError && e.status === 401) setAuth(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
const action: Action = fn => {
|
||||||
|
setBusy(true)
|
||||||
|
setError('')
|
||||||
|
setNotice('')
|
||||||
|
void fn().catch(fail).finally(() => setBusy(false))
|
||||||
|
}
|
||||||
|
|
||||||
|
async function loadSession() {
|
||||||
|
const session = await api<{ authenticated: boolean }>('/session')
|
||||||
|
setAuth(session.authenticated)
|
||||||
|
}
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
loadSession().catch(e => {
|
||||||
|
setAuth(false)
|
||||||
|
if (!(e instanceof APIError && e.status === 401)) fail(e)
|
||||||
|
})
|
||||||
|
}, [])
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
if (!auth) return
|
||||||
|
let alive = true
|
||||||
|
setLoading(true)
|
||||||
|
setError('')
|
||||||
|
setStats(null)
|
||||||
|
setLogs([])
|
||||||
|
const statsParams = new URLSearchParams(query)
|
||||||
|
statsParams.set('timezone', timezone || 'UTC')
|
||||||
|
Promise.all([
|
||||||
|
api<Channel[]>('/channels'),
|
||||||
|
api<Key[]>('/keys'),
|
||||||
|
page === '概览' || page === 'Keys' ? api<Stats>('/stats?' + statsParams) : Promise.resolve(null),
|
||||||
|
page === '记录' ? api<Log[]>('/logs' + query) : Promise.resolve([] as Log[]),
|
||||||
|
])
|
||||||
|
.then(([c, k, s, l]) => { if (alive) { setChannels(c ?? []); setKeys(k ?? []); setStats(s); setLogs(l ?? []) } })
|
||||||
|
.catch(e => { if (alive) fail(e) })
|
||||||
|
.finally(() => { if (alive) setLoading(false) })
|
||||||
|
return () => { alive = false }
|
||||||
|
}, [auth, page, rev, query, timezone])
|
||||||
|
|
||||||
|
const refresh = () => setRev(v => v + 1)
|
||||||
|
|
||||||
|
/** 从 Keys 页跳转：概览页按该 Key 筛选统计。 */
|
||||||
|
function viewKeyStats(k: Key) {
|
||||||
|
setFilter({ ...filter, key_id: k.id })
|
||||||
|
const p = new URLSearchParams(query)
|
||||||
|
p.set('key_id', k.id)
|
||||||
|
setQuery('?' + p)
|
||||||
|
setPage('概览')
|
||||||
|
}
|
||||||
|
|
||||||
|
if (auth === null) return <div className="login"><p className="muted">正在检查登录状态…</p></div>
|
||||||
|
|
||||||
|
if (!auth) {
|
||||||
|
return <Login busy={busy} error={error} notice={notice} action={action} onSuccess={loadSession} />
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<Shell
|
||||||
|
page={page}
|
||||||
|
theme={theme}
|
||||||
|
busy={busy}
|
||||||
|
loading={loading}
|
||||||
|
onNavigate={p => { setPage(p); setNotice('') }}
|
||||||
|
onRefresh={refresh}
|
||||||
|
onThemeChange={setTheme}
|
||||||
|
onLogout={() => action(async () => { await api('/logout', 'POST'); setAuth(false) })}
|
||||||
|
>
|
||||||
|
<Feedback error={error} notice={notice} />
|
||||||
|
{(page === '概览' || page === '记录') && (
|
||||||
|
<Filters
|
||||||
|
keys={keys}
|
||||||
|
filter={filter}
|
||||||
|
loading={loading}
|
||||||
|
timezone={timezone}
|
||||||
|
onChange={setFilter}
|
||||||
|
onQuery={q => { setQuery(q); setError('') }}
|
||||||
|
onRefresh={refresh}
|
||||||
|
onError={fail}
|
||||||
|
/>
|
||||||
|
)}
|
||||||
|
{loading && <div className="panel empty" role="status">正在加载数据…</div>}
|
||||||
|
{/* 加载时隐藏而非卸载页面：保留弹窗与表单等本地状态（如创建后展示密钥的弹窗） */}
|
||||||
|
<div style={loading ? { display: 'none' } : undefined}>
|
||||||
|
{page === '概览' && <Overview stats={stats} channels={channels} />}
|
||||||
|
{page === '模型' && <ModelsPage channels={channels} />}
|
||||||
|
{page === '渠道' && (
|
||||||
|
<ChannelsPage
|
||||||
|
channels={channels}
|
||||||
|
busy={busy}
|
||||||
|
error={error}
|
||||||
|
notice={notice}
|
||||||
|
action={action}
|
||||||
|
refresh={refresh}
|
||||||
|
notify={setNotice}
|
||||||
|
clearError={() => setError('')}
|
||||||
|
/>
|
||||||
|
)}
|
||||||
|
{page === 'Keys' && (
|
||||||
|
<KeysPage
|
||||||
|
keys={keys}
|
||||||
|
stats={stats}
|
||||||
|
channels={channels}
|
||||||
|
timezone={timezone}
|
||||||
|
busy={busy}
|
||||||
|
error={error}
|
||||||
|
notice={notice}
|
||||||
|
action={action}
|
||||||
|
refresh={refresh}
|
||||||
|
notify={setNotice}
|
||||||
|
clearError={() => setError('')}
|
||||||
|
onViewStats={viewKeyStats}
|
||||||
|
/>
|
||||||
|
)}
|
||||||
|
{page === '记录' && <LogsPage logs={logs} channels={channels} />}
|
||||||
|
</div>
|
||||||
|
</Shell>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
export type Channel = { id: string; name: string; base_url: string; protocol: string; enabled: boolean; models: Record<string, string> }
|
||||||
|
export type Key = { id: string; name: string; prefix: string; enabled: boolean; revoked?: boolean; all_channels: boolean; channels: string[]; expires_at: string; created_at: string }
|
||||||
|
export type Usage = { requests: number; success: number; failed: number; input_tokens: number; output_tokens: number; cache_read_tokens: number; cache_write_tokens: number; unknown_usage: number }
|
||||||
|
export type Stats = { summary: Usage; by_key: (Usage & { key_id: string; name: string })[]; daily: (Usage & { day: string })[]; by_model: (Usage & { channel: string; model: string })[] }
|
||||||
|
export type Log = { id: string; created_at: string; key_id: string; key_name: string; channel: string; request_model: string; actual_model: string; downstream: string; upstream: string; status: number; duration_ms: number; stream: boolean; input_tokens: number; output_tokens: number; cache_read_tokens: number; cache_write_tokens: number; usage_known: boolean; error: string }
|
||||||
|
export class APIError extends Error { constructor(message: string, public status: number) { super(message) } }
|
||||||
|
export async function api<T>(path: string, method = 'GET', body?: unknown, signal?: AbortSignal): Promise<T> {
|
||||||
|
const res = await fetch('/admin/api' + path, { method, signal, credentials: 'same-origin', headers: method === 'GET' ? {} : { 'Content-Type': 'application/json', 'X-Requested-With': 'ai-router' }, body: body === undefined ? undefined : JSON.stringify(body) })
|
||||||
|
const text = await res.text()
|
||||||
|
let data
|
||||||
|
try { data = text ? JSON.parse(text) : undefined } catch { if (res.ok) throw new Error('服务器返回了无法识别的数据，请稍后重试。') }
|
||||||
|
if (!res.ok) throw new APIError(typeof data?.error === 'string' ? data.error : typeof data?.message === 'string' ? data.message : res.status === 401 ? '登录已失效，请重新登录。' : '请求失败（' + res.status + '），请稍后重试。', res.status)
|
||||||
|
return data as T
|
||||||
|
}
|
||||||
@@ -0,0 +1,35 @@
|
|||||||
|
import type { Channel, Usage } from './api'
|
||||||
|
|
||||||
|
export type Page = '概览' | '渠道' | '模型' | 'Keys' | '记录'
|
||||||
|
|
||||||
|
export const PAGES: Page[] = ['概览', '渠道', '模型', 'Keys', '记录']
|
||||||
|
|
||||||
|
export const PAGE_DESC: Record<Page, string> = {
|
||||||
|
概览: '查看请求趋势与模型用量。',
|
||||||
|
渠道: '配置上游服务与模型映射。',
|
||||||
|
模型: '汇总各渠道选定的展示模型；/v1/models 仅返回当前 Key 有权访问的启用渠道模型。',
|
||||||
|
Keys: '按渠道授权，安全管理每一个访问凭证。',
|
||||||
|
记录: '查看最近请求明细，快速定位问题。',
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface Filter {
|
||||||
|
key_id: string
|
||||||
|
from: string
|
||||||
|
to: string
|
||||||
|
}
|
||||||
|
|
||||||
|
export const emptyFilter: Filter = { key_id: '', from: '', to: '' }
|
||||||
|
|
||||||
|
/** 包装异步操作：统一 busy 状态与错误处理，见 App。 */
|
||||||
|
export type Action = (fn: () => Promise<void>) => void
|
||||||
|
|
||||||
|
export const num = (v: number | null | undefined) => (v == null ? '—' : v.toLocaleString('zh-CN'))
|
||||||
|
|
||||||
|
export function channelName(channels: Channel[], id: string) {
|
||||||
|
return channels.find(c => c.id === id)?.name || id
|
||||||
|
}
|
||||||
|
|
||||||
|
export const usageHeads = ['请求数', '输入 tokens', '输出 tokens', '缓存读取', '缓存写入']
|
||||||
|
|
||||||
|
export const usageCells = (u: Usage) =>
|
||||||
|
[u.requests, u.input_tokens, u.output_tokens, u.cache_read_tokens, u.cache_write_tokens].map(num)
|
||||||
@@ -0,0 +1,64 @@
|
|||||||
|
import type { FormEvent } from 'react'
|
||||||
|
import type { Key } from '../api'
|
||||||
|
import { emptyFilter, type Filter } from '../common'
|
||||||
|
import { filterParams, quickRange } from '../time'
|
||||||
|
import { IconGlobe } from './icons'
|
||||||
|
|
||||||
|
interface FiltersProps {
|
||||||
|
keys: Key[]
|
||||||
|
filter: Filter
|
||||||
|
loading: boolean
|
||||||
|
timezone: string
|
||||||
|
onChange: (f: Filter) => void
|
||||||
|
/** 提交新的查询串（含前导 ?，空串表示无筛选），同时清除已有错误。 */
|
||||||
|
onQuery: (q: string) => void
|
||||||
|
onRefresh: () => void
|
||||||
|
onError: (e: unknown) => void
|
||||||
|
}
|
||||||
|
|
||||||
|
export function Filters({ keys, filter, loading, timezone, onChange, onQuery, onRefresh, onError }: FiltersProps) {
|
||||||
|
function submit(e: FormEvent) {
|
||||||
|
e.preventDefault()
|
||||||
|
try {
|
||||||
|
onQuery('?' + filterParams(filter))
|
||||||
|
onRefresh()
|
||||||
|
} catch (err) {
|
||||||
|
onError(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function quick(days: number) {
|
||||||
|
const range = quickRange(days)
|
||||||
|
onChange({ ...filter, from: range.fromInput, to: range.toInput })
|
||||||
|
const p = new URLSearchParams({ from: range.from, to: range.to })
|
||||||
|
if (filter.key_id) p.set('key_id', filter.key_id)
|
||||||
|
onQuery('?' + p)
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<form className="filters panel" onSubmit={submit}>
|
||||||
|
<label>访问密钥
|
||||||
|
<select value={filter.key_id} onChange={e => onChange({ ...filter, key_id: e.target.value })}>
|
||||||
|
<option value="">全部 Keys</option>
|
||||||
|
{keys.map(k => <option key={k.id} value={k.id}>{k.name} · {k.prefix}</option>)}
|
||||||
|
</select>
|
||||||
|
</label>
|
||||||
|
<label>开始时间
|
||||||
|
<input type="datetime-local" value={filter.from} onChange={e => onChange({ ...filter, from: e.target.value })} />
|
||||||
|
</label>
|
||||||
|
<label>结束时间
|
||||||
|
<input type="datetime-local" value={filter.to} onChange={e => onChange({ ...filter, to: e.target.value })} />
|
||||||
|
</label>
|
||||||
|
<div className="filters-actions">
|
||||||
|
<button className="primary" disabled={loading}>查询</button>
|
||||||
|
<button type="button" onClick={() => { onChange(emptyFilter); onQuery('') }}>重置</button>
|
||||||
|
</div>
|
||||||
|
<div className="quick">
|
||||||
|
{[1, 7, 30].map(days => (
|
||||||
|
<button type="button" key={days} onClick={() => quick(days)}>{days === 1 ? '今天' : `近 ${days} 天`}</button>
|
||||||
|
))}
|
||||||
|
<small className="tz-hint" title="时间筛选与每日汇总均按浏览器本地时区解释"><IconGlobe />{timezone}</small>
|
||||||
|
</div>
|
||||||
|
</form>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
import type { ReactNode } from 'react'
|
||||||
|
import { PAGES, PAGE_DESC, type Page } from '../common'
|
||||||
|
import type { Theme } from '../theme'
|
||||||
|
import { IconChannel, IconKey, IconLog, IconLogout, IconModel, IconMonitor, IconMoon, IconOverview, IconRefresh, IconRoute, IconSun } from './icons'
|
||||||
|
|
||||||
|
const PAGE_ICONS: Record<Page, ReactNode> = {
|
||||||
|
概览: <IconOverview />,
|
||||||
|
渠道: <IconChannel />,
|
||||||
|
模型: <IconModel />,
|
||||||
|
Keys: <IconKey />,
|
||||||
|
记录: <IconLog />,
|
||||||
|
}
|
||||||
|
|
||||||
|
const THEMES: { value: Theme; label: string; icon: ReactNode }[] = [
|
||||||
|
{ value: 'system', label: '跟随系统', icon: <IconMonitor /> },
|
||||||
|
{ value: 'light', label: '浅色模式', icon: <IconSun /> },
|
||||||
|
{ value: 'dark', label: '深色模式', icon: <IconMoon /> },
|
||||||
|
]
|
||||||
|
|
||||||
|
interface ShellProps {
|
||||||
|
page: Page
|
||||||
|
theme: Theme
|
||||||
|
busy: boolean
|
||||||
|
loading: boolean
|
||||||
|
onNavigate: (p: Page) => void
|
||||||
|
onRefresh: () => void
|
||||||
|
onThemeChange: (theme: Theme) => void
|
||||||
|
onLogout: () => void
|
||||||
|
children: ReactNode
|
||||||
|
}
|
||||||
|
|
||||||
|
export function Shell({ page, theme, busy, loading, onNavigate, onRefresh, onThemeChange, onLogout, children }: ShellProps) {
|
||||||
|
return (
|
||||||
|
<div className="shell">
|
||||||
|
<aside>
|
||||||
|
<div className="brand">
|
||||||
|
<span className="mark"><IconRoute width={22} height={22} strokeWidth={2} /></span>
|
||||||
|
<div>AI Router<small>模型路由管理台</small></div>
|
||||||
|
</div>
|
||||||
|
<nav aria-label="主导航">
|
||||||
|
{PAGES.map(p => (
|
||||||
|
<button key={p} aria-current={page === p ? 'page' : undefined} className={page === p ? 'active' : ''} onClick={() => onNavigate(p)}>
|
||||||
|
{PAGE_ICONS[p]}{p}
|
||||||
|
</button>
|
||||||
|
))}
|
||||||
|
</nav>
|
||||||
|
<div className="aside-bottom">
|
||||||
|
<div className="theme-seg" role="group" aria-label="主题模式">
|
||||||
|
{THEMES.map(t => (
|
||||||
|
<button
|
||||||
|
key={t.value}
|
||||||
|
type="button"
|
||||||
|
title={t.label}
|
||||||
|
aria-label={t.label}
|
||||||
|
aria-pressed={theme === t.value}
|
||||||
|
className={theme === t.value ? 'active' : ''}
|
||||||
|
onClick={() => onThemeChange(t.value)}
|
||||||
|
>{t.icon}</button>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
<button type="button" disabled={busy} onClick={onLogout}><IconLogout />退出登录</button>
|
||||||
|
</div>
|
||||||
|
</aside>
|
||||||
|
<main>
|
||||||
|
<header>
|
||||||
|
<div>
|
||||||
|
<h1>{page === 'Keys' ? '访问密钥' : page}</h1>
|
||||||
|
<p className="muted">{PAGE_DESC[page]}</p>
|
||||||
|
</div>
|
||||||
|
<button type="button" disabled={loading || busy} onClick={onRefresh}><IconRefresh />刷新</button>
|
||||||
|
</header>
|
||||||
|
{children}
|
||||||
|
</main>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,68 @@
|
|||||||
|
import { useLayoutEffect, useRef, useState } from 'react'
|
||||||
|
import type { Usage } from '../api'
|
||||||
|
import { num } from '../common'
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 每日用量堆叠柱状图（纯 CSS 实现，无第三方依赖）。
|
||||||
|
* 按 README 的口径，缓存读/写是输入 tokens 的细分项，不能再与输入相加，
|
||||||
|
* 因此图表只堆叠 输入 + 输出 构成真实总量；缓存明细见下方表格。
|
||||||
|
* X 轴标签数量随容器宽度自适应，窄屏（手机）自动减少，避免重叠。
|
||||||
|
*/
|
||||||
|
export function UsageChart({ data }: { data: (Usage & { day: string })[] }) {
|
||||||
|
const plotRef = useRef<HTMLDivElement>(null)
|
||||||
|
const [plotWidth, setPlotWidth] = useState(720)
|
||||||
|
|
||||||
|
useLayoutEffect(() => {
|
||||||
|
const el = plotRef.current
|
||||||
|
if (!el) return
|
||||||
|
const ro = new ResizeObserver(entries => setPlotWidth(entries[0].contentRect.width))
|
||||||
|
ro.observe(el)
|
||||||
|
return () => ro.disconnect()
|
||||||
|
}, [])
|
||||||
|
|
||||||
|
if (!data.length) return <div className="chart-empty">所选时间范围内暂无用量数据。</div>
|
||||||
|
|
||||||
|
const totals = data.map(d => d.input_tokens + d.output_tokens)
|
||||||
|
const max = Math.max(...totals, 1)
|
||||||
|
// max 较小时四舍五入会产生重复刻度，去重保持纵轴整洁
|
||||||
|
const ticks = [...new Set([1, 0.75, 0.5, 0.25].map(r => Math.round(max * r)))]
|
||||||
|
// 天数较多时抽稀 X 轴标签；每个标签约需 52px，按容器实际宽度定上限
|
||||||
|
const maxLabels = Math.max(3, Math.min(12, Math.floor(plotWidth / 52)))
|
||||||
|
const labelStep = Math.ceil(data.length / maxLabels)
|
||||||
|
|
||||||
|
return (
|
||||||
|
<div className="chart">
|
||||||
|
<div className="chart-plot" ref={plotRef} role="img" aria-label="每日 tokens 用量柱状图">
|
||||||
|
{ticks.map((t, i) => (
|
||||||
|
<div className="chart-grid" key={i} style={{ bottom: `${(t / max) * 100}%` }}>
|
||||||
|
<span>{compact(t)}</span>
|
||||||
|
</div>
|
||||||
|
))}
|
||||||
|
<div className="chart-bars">
|
||||||
|
{data.map((d, i) => {
|
||||||
|
const total = totals[i]
|
||||||
|
return (
|
||||||
|
<div className="chart-col" key={d.day} title={`${d.day}\n输入 ${num(d.input_tokens)} · 输出 ${num(d.output_tokens)}\n请求 ${num(d.requests)} 次`}>
|
||||||
|
<div className="chart-bar">
|
||||||
|
<i className="seg-output" style={{ height: `${(d.output_tokens / max) * 100}%` }} />
|
||||||
|
<i className="seg-input" style={{ height: `${(d.input_tokens / max) * 100}%` }} />
|
||||||
|
</div>
|
||||||
|
{(i % labelStep === 0 || i === data.length - 1) && <span className="chart-x">{d.day.slice(5)}</span>}
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
})}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<div className="chart-legend">
|
||||||
|
<span><i className="seg-input" />输入 tokens</span>
|
||||||
|
<span><i className="seg-output" />输出 tokens</span>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
function compact(v: number): string {
|
||||||
|
if (v >= 1e6) return (v / 1e6).toFixed(1).replace(/\.0$/, '') + 'M'
|
||||||
|
if (v >= 1e3) return (v / 1e3).toFixed(1).replace(/\.0$/, '') + 'k'
|
||||||
|
return String(v)
|
||||||
|
}
|
||||||
@@ -0,0 +1,209 @@
|
|||||||
|
import type { SVGProps } from 'react'
|
||||||
|
|
||||||
|
type IconProps = SVGProps<SVGSVGElement>
|
||||||
|
|
||||||
|
function base(props: IconProps): IconProps {
|
||||||
|
return {
|
||||||
|
width: 16,
|
||||||
|
height: 16,
|
||||||
|
viewBox: '0 0 24 24',
|
||||||
|
fill: 'none',
|
||||||
|
stroke: 'currentColor',
|
||||||
|
strokeWidth: 1.8,
|
||||||
|
strokeLinecap: 'round',
|
||||||
|
strokeLinejoin: 'round',
|
||||||
|
'aria-hidden': true,
|
||||||
|
...props,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/** 路由/分流图标，用作品牌标识。 */
|
||||||
|
export const IconRoute = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M4 7h10.5l-3-3" />
|
||||||
|
<path d="M20 17H9.5l3 3" />
|
||||||
|
<circle cx="4" cy="7" r="2" />
|
||||||
|
<circle cx="20" cy="17" r="2" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconOverview = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M3 3v18h18" />
|
||||||
|
<path d="M8 16v-5" />
|
||||||
|
<path d="M13 16V7" />
|
||||||
|
<path d="M18 16v-8" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconChannel = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<rect x="3" y="4" width="18" height="7" rx="2" />
|
||||||
|
<rect x="3" y="13" width="18" height="7" rx="2" />
|
||||||
|
<path d="M7 7.5h.01" />
|
||||||
|
<path d="M7 16.5h.01" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconKey = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="8" cy="15" r="4" />
|
||||||
|
<path d="m10.8 12.2 8.2-8.2" />
|
||||||
|
<path d="m15 6 3 3" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconLog = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M8 6h13" />
|
||||||
|
<path d="M8 12h13" />
|
||||||
|
<path d="M8 18h13" />
|
||||||
|
<path d="M3.5 6h.01" />
|
||||||
|
<path d="M3.5 12h.01" />
|
||||||
|
<path d="M3.5 18h.01" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconRefresh = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M21 12a9 9 0 1 1-2.64-6.36" />
|
||||||
|
<path d="M21 3v6h-6" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconPlus = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M12 5v14" />
|
||||||
|
<path d="M5 12h14" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconSun = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="12" cy="12" r="4" />
|
||||||
|
<path d="M12 2v2" />
|
||||||
|
<path d="M12 20v2" />
|
||||||
|
<path d="m4.93 4.93 1.41 1.41" />
|
||||||
|
<path d="m17.66 17.66 1.41 1.41" />
|
||||||
|
<path d="M2 12h2" />
|
||||||
|
<path d="M20 12h2" />
|
||||||
|
<path d="m4.93 19.07 1.41-1.41" />
|
||||||
|
<path d="m17.66 6.34 1.41-1.41" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconMoon = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M21 12.8A9 9 0 1 1 11.2 3a7 7 0 0 0 9.8 9.8Z" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconMonitor = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<rect x="3" y="4" width="18" height="12" rx="2" />
|
||||||
|
<path d="M9 20h6" />
|
||||||
|
<path d="M12 16v4" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconLogout = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M9 21H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h4" />
|
||||||
|
<path d="m16 17 5-5-5-5" />
|
||||||
|
<path d="M21 12H9" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconClose = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M18 6 6 18" />
|
||||||
|
<path d="m6 6 12 12" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconGlobe = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="12" cy="12" r="9" />
|
||||||
|
<path d="M3 12h18" />
|
||||||
|
<path d="M12 3a13.7 13.7 0 0 1 0 18" />
|
||||||
|
<path d="M12 3a13.7 13.7 0 0 0 0 18" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconModel = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M12 2 3 7v10l9 5 9-5V7z" />
|
||||||
|
<path d="M12 22V12" />
|
||||||
|
<path d="m3 7 9 5 9-5" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconCopy = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<rect x="9" y="9" width="12" height="12" rx="2" />
|
||||||
|
<path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconCheck = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M20 6 9 17l-5-5" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconActivity = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M3 12h4l3 8 4-16 3 8h4" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconCheckCircle = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="12" cy="12" r="9" />
|
||||||
|
<path d="m8.5 12.5 2.5 2.5 4.5-5" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconXCircle = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="12" cy="12" r="9" />
|
||||||
|
<path d="m9 9 6 6" />
|
||||||
|
<path d="m15 9-6 6" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconHelpCircle = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<circle cx="12" cy="12" r="9" />
|
||||||
|
<path d="M9.2 9a2.9 2.9 0 0 1 5.6 1c0 1.9-2.8 2.3-2.8 4" />
|
||||||
|
<path d="M12 17.5h.01" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconArrowDown = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M12 4v14" />
|
||||||
|
<path d="m6 12 6 6 6-6" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconArrowUp = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M12 20V6" />
|
||||||
|
<path d="m6 12 6-6 6 6" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconZap = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<path d="M13 2 3 14h9l-1 8 10-12h-9l1-8z" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
|
|
||||||
|
export const IconDatabase = (p: IconProps) => (
|
||||||
|
<svg {...base(p)}>
|
||||||
|
<ellipse cx="12" cy="5" rx="8" ry="3" />
|
||||||
|
<path d="M4 5v14c0 1.66 3.58 3 8 3s8-1.34 8-3V5" />
|
||||||
|
<path d="M4 12c0 1.66 3.58 3 8 3s8-1.34 8-3" />
|
||||||
|
</svg>
|
||||||
|
)
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
import type { ReactNode } from 'react'
|
||||||
|
import { IconClose } from './icons'
|
||||||
|
|
||||||
|
export function Table({ heads, rows }: { heads: string[]; rows: ReactNode[][] }) {
|
||||||
|
return (
|
||||||
|
<div className="table-wrap">
|
||||||
|
<table>
|
||||||
|
<thead>
|
||||||
|
<tr>{heads.map(h => <th key={h}>{h}</th>)}</tr>
|
||||||
|
</thead>
|
||||||
|
<tbody>
|
||||||
|
{rows.length
|
||||||
|
? rows.map((r, i) => <tr key={i}>{r.map((v, j) => <td key={j}>{v}</td>)}</tr>)
|
||||||
|
: <tr><td className="empty" colSpan={heads.length}>暂无数据</td></tr>}
|
||||||
|
</tbody>
|
||||||
|
</table>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
export type BadgeVariant = 'good' | 'warn' | 'danger' | 'neutral'
|
||||||
|
|
||||||
|
export function Badge({ variant = 'neutral', children }: { variant?: BadgeVariant; children: ReactNode }) {
|
||||||
|
return <span className={`badge ${variant}`}><i aria-hidden="true" />{children}</span>
|
||||||
|
}
|
||||||
|
|
||||||
|
/** 统一的错误 / 成功反馈条。 */
|
||||||
|
export function Feedback({ error, notice }: { error: string; notice: string }) {
|
||||||
|
return (
|
||||||
|
<>
|
||||||
|
{error && <div className="alert error" role="alert">{error}</div>}
|
||||||
|
{notice && <div className="alert success" role="status">{notice}</div>}
|
||||||
|
</>
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 模态弹窗。传入 onClose 时标题栏显示关闭按钮；
|
||||||
|
* 刻意不响应遮罩点击，避免误触丢失表单内容（尤其是只展示一次的密钥）。
|
||||||
|
*/
|
||||||
|
export function Dialog({ id, title, onClose, children }: { id: string; title: string; onClose?: () => void; children: ReactNode }) {
|
||||||
|
return (
|
||||||
|
<div className="overlay">
|
||||||
|
<section className="dialog" role="dialog" aria-modal="true" aria-labelledby={id}>
|
||||||
|
<header className="dialog-head">
|
||||||
|
<h2 id={id}>{title}</h2>
|
||||||
|
{onClose && <button type="button" className="icon-btn" aria-label="关闭" onClick={onClose}><IconClose /></button>}
|
||||||
|
</header>
|
||||||
|
{children}
|
||||||
|
</section>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,5 @@
|
|||||||
|
import React from 'react'
|
||||||
|
import ReactDOM from 'react-dom/client'
|
||||||
|
import App from './App'
|
||||||
|
import './style.css'
|
||||||
|
ReactDOM.createRoot(document.getElementById('root')!).render(<React.StrictMode><App /></React.StrictMode>)
|
||||||
@@ -0,0 +1,218 @@
|
|||||||
|
import { useState, type FormEvent } from 'react'
|
||||||
|
import { api, type Channel } from '../api'
|
||||||
|
import type { Action } from '../common'
|
||||||
|
import { IconPlus } from '../components/icons'
|
||||||
|
import { Badge, Dialog, Feedback, Table } from '../components/ui'
|
||||||
|
|
||||||
|
const NEW_CHANNEL: Channel = { id: '', name: '', base_url: '', protocol: 'openai-chat', enabled: true, models: {} }
|
||||||
|
const KNOWN_PROTOCOLS = ['openai-chat', 'openai-responses', 'anthropic']
|
||||||
|
type ModelDraft = { alias: string; real: string; selected: boolean }
|
||||||
|
|
||||||
|
interface ChannelsProps {
|
||||||
|
channels: Channel[]
|
||||||
|
busy: boolean
|
||||||
|
error: string
|
||||||
|
notice: string
|
||||||
|
action: Action
|
||||||
|
refresh: () => void
|
||||||
|
notify: (msg: string) => void
|
||||||
|
clearError: () => void
|
||||||
|
}
|
||||||
|
|
||||||
|
export function ChannelsPage({ channels, busy, error, notice, action, refresh, notify, clearError }: ChannelsProps) {
|
||||||
|
const [draft, setDraft] = useState<Channel | null>(null)
|
||||||
|
const [secret, setSecret] = useState('')
|
||||||
|
const [modelRows, setModelRows] = useState<ModelDraft[]>([])
|
||||||
|
const [fetching, setFetching] = useState(false)
|
||||||
|
const [modelQuery, setModelQuery] = useState('')
|
||||||
|
const [candidates, setCandidates] = useState<string[] | null>(null)
|
||||||
|
const [chosen, setChosen] = useState<Set<string>>(new Set())
|
||||||
|
const [candidateQuery, setCandidateQuery] = useState('')
|
||||||
|
const existingModels = new Set(modelRows.map(row => row.real.trim()))
|
||||||
|
const visibleCandidates = (candidates ?? []).filter(real => real.toLowerCase().includes(candidateQuery.toLowerCase()))
|
||||||
|
|
||||||
|
function edit(c: Channel) {
|
||||||
|
setDraft({ ...c })
|
||||||
|
setSecret('')
|
||||||
|
setModelRows(Object.entries(c.models ?? {}).map(([alias, real]) => ({ alias, real, selected: true })))
|
||||||
|
setModelQuery('')
|
||||||
|
clearError()
|
||||||
|
}
|
||||||
|
|
||||||
|
function close() {
|
||||||
|
if (busy) return
|
||||||
|
setDraft(null)
|
||||||
|
setSecret('')
|
||||||
|
clearError()
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseMapping(): Record<string, string> {
|
||||||
|
const models: Record<string, string> = Object.create(null)
|
||||||
|
for (const row of modelRows.filter(row => row.selected)) {
|
||||||
|
const alias = row.alias.trim()
|
||||||
|
const real = row.real.trim()
|
||||||
|
if (!alias || !real) throw new Error('已展示模型的展示名称和上游模型不能为空。')
|
||||||
|
if (Object.hasOwn(models, alias)) throw new Error('展示名称重复：' + alias)
|
||||||
|
models[alias] = real
|
||||||
|
}
|
||||||
|
return models
|
||||||
|
}
|
||||||
|
|
||||||
|
function discover() {
|
||||||
|
if (!draft) return
|
||||||
|
action(async () => {
|
||||||
|
setFetching(true)
|
||||||
|
try {
|
||||||
|
const result = await api<{ models: string[] }>('/channels/discover-models', 'POST', { ...draft, api_key: secret })
|
||||||
|
setCandidates([...new Set(result.models)])
|
||||||
|
setChosen(new Set())
|
||||||
|
setCandidateQuery('')
|
||||||
|
} finally {
|
||||||
|
setFetching(false)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function updateModel(index: number, patch: Partial<ModelDraft>) {
|
||||||
|
setModelRows(rows => rows.map((row, i) => i === index ? { ...row, ...patch } : row))
|
||||||
|
}
|
||||||
|
|
||||||
|
function submit(e: FormEvent) {
|
||||||
|
e.preventDefault()
|
||||||
|
if (!draft) return
|
||||||
|
action(async () => {
|
||||||
|
const models = parseMapping()
|
||||||
|
await api('/channels', 'POST', { ...draft, models, api_key: secret })
|
||||||
|
setDraft(null)
|
||||||
|
setSecret('')
|
||||||
|
refresh()
|
||||||
|
notify('渠道已保存。')
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function addChosen() {
|
||||||
|
const additions = (candidates ?? []).filter(real => chosen.has(real) && !existingModels.has(real))
|
||||||
|
setModelRows(rows => [...rows, ...additions.map(real => ({ real, alias: real, selected: true }))])
|
||||||
|
setModelQuery('')
|
||||||
|
setCandidates(null)
|
||||||
|
notify('已加入 ' + additions.length + ' 个模型，请保存渠道使配置生效。')
|
||||||
|
}
|
||||||
|
|
||||||
|
function remove(c: Channel) {
|
||||||
|
if (!confirm(`确定删除渠道“${c.name}”？此操作无法撤销。`)) return
|
||||||
|
action(async () => {
|
||||||
|
await api('/channels/' + encodeURIComponent(c.id), 'DELETE')
|
||||||
|
refresh()
|
||||||
|
notify('渠道已删除。')
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<section className="panel">
|
||||||
|
<div className="section-head">
|
||||||
|
<h2>渠道列表 <small>{channels.length}</small></h2>
|
||||||
|
<button className="primary" onClick={() => edit(NEW_CHANNEL)}><IconPlus />新建渠道</button>
|
||||||
|
</div>
|
||||||
|
<Table
|
||||||
|
heads={['名称 / 地址', '协议', '路由策略', '状态', '操作']}
|
||||||
|
rows={channels.map(c => [
|
||||||
|
<><strong>{c.name}</strong><small className="sub">{c.base_url}</small></>,
|
||||||
|
c.protocol,
|
||||||
|
<>严格映射<small className="sub">{Object.keys(c.models ?? {}).length} 个已展示模型</small></>,
|
||||||
|
<Badge variant={c.enabled ? 'good' : 'neutral'}>{c.enabled ? '启用' : '停用'}</Badge>,
|
||||||
|
<div className="actions">
|
||||||
|
<button onClick={() => edit(c)}>编辑</button>
|
||||||
|
<button className="danger" disabled={busy} onClick={() => remove(c)}>删除</button>
|
||||||
|
</div>,
|
||||||
|
])}
|
||||||
|
/>
|
||||||
|
{draft && candidates === null && (
|
||||||
|
<Dialog id="channel-title" title={draft.id ? '编辑渠道' : '新建渠道'} onClose={close}>
|
||||||
|
<Feedback error={error} notice={notice} />
|
||||||
|
<form onSubmit={submit}>
|
||||||
|
<fieldset disabled={busy}>
|
||||||
|
<label>渠道名称
|
||||||
|
<input autoFocus required value={draft.name} onChange={e => setDraft({ ...draft, name: e.target.value })} />
|
||||||
|
</label>
|
||||||
|
<label>上游 Base URL
|
||||||
|
<input type="url" required value={draft.base_url} onChange={e => setDraft({ ...draft, base_url: e.target.value })} />
|
||||||
|
<small>{draft.protocol === 'anthropic' ? '例如 https://api.anthropic.com，不带 /v1；程序追加 /v1/messages。' : '例如 https://api.openai.com/v1；程序追加 /chat/completions 或 /responses。'}</small>
|
||||||
|
</label>
|
||||||
|
<label>协议
|
||||||
|
<select value={draft.protocol} onChange={e => setDraft({ ...draft, protocol: e.target.value })}>
|
||||||
|
<option value="openai-chat">OpenAI Chat</option>
|
||||||
|
<option value="openai-responses">OpenAI Responses</option>
|
||||||
|
<option value="anthropic">Anthropic</option>
|
||||||
|
{!KNOWN_PROTOCOLS.includes(draft.protocol) && <option value={draft.protocol}>{draft.protocol}</option>}
|
||||||
|
</select>
|
||||||
|
</label>
|
||||||
|
<label>上游 API Key
|
||||||
|
<input type="password" autoComplete="new-password" value={secret} onChange={e => setSecret(e.target.value)} placeholder={draft.id ? '留空保留原密钥' : '输入上游密钥（可选）'} />
|
||||||
|
</label>
|
||||||
|
<section className="channel-models" aria-label="渠道模型">
|
||||||
|
<div className="section-head">
|
||||||
|
<h3>展示给下游的模型 <small>{modelRows.filter(row => row.selected).length}</small></h3>
|
||||||
|
<button type="button" onClick={discover} disabled={!draft.base_url.trim()}>{fetching ? '获取中…' : '获取模型列表'}</button>
|
||||||
|
</div>
|
||||||
|
<p className="muted">获取后在弹窗中选择要加入的模型，可编辑展示名称，保存渠道后生效；上游不支持获取时可手动添加。</p>
|
||||||
|
<input aria-label="搜索渠道模型" placeholder="搜索上游模型或展示名称" value={modelQuery} onChange={e => setModelQuery(e.target.value)} />
|
||||||
|
<div className="channel-model-list">
|
||||||
|
{modelRows.map((row, index) => ({ row, index })).filter(({ row }) => `${row.real} ${row.alias}`.toLowerCase().includes(modelQuery.toLowerCase())).map(({ row, index }) => (
|
||||||
|
<div className="channel-model-row" key={index}>
|
||||||
|
<label className="check"><input type="checkbox" checked={row.selected} onChange={e => updateModel(index, { selected: e.target.checked })} />展示</label>
|
||||||
|
<label>上游模型<input aria-label={`上游模型 ${index + 1}`} value={row.real} onChange={e => updateModel(index, { real: e.target.value })} spellCheck={false} /></label>
|
||||||
|
<label>展示名称<input aria-label={`展示名称 ${index + 1}`} value={row.alias} onChange={e => updateModel(index, { alias: e.target.value })} spellCheck={false} /></label>
|
||||||
|
<button type="button" aria-label={`移除模型 ${index + 1}`} onClick={() => setModelRows(rows => rows.filter((_, i) => i !== index))}>移除</button>
|
||||||
|
</div>
|
||||||
|
))}
|
||||||
|
{modelRows.length === 0 && <p className="muted">尚未选择模型，请获取列表或手动添加。</p>}
|
||||||
|
</div>
|
||||||
|
<button type="button" onClick={() => { setModelQuery(''); setModelRows(rows => [...rows, { alias: '', real: '', selected: true }]) }}>手动添加模型</button>
|
||||||
|
<p className="muted">下游调用格式：{draft.name || '渠道名'}/展示名称。展示名称在同一渠道内不可重复。</p>
|
||||||
|
</section>
|
||||||
|
<label className="check">
|
||||||
|
<input type="checkbox" checked={draft.enabled} onChange={e => setDraft({ ...draft, enabled: e.target.checked })} />启用渠道
|
||||||
|
</label>
|
||||||
|
<p className="muted">只允许调用已勾选并保存的模型，未映射模型不可调用。</p>
|
||||||
|
<div className="dialog-actions">
|
||||||
|
<button type="button" onClick={close}>取消</button>
|
||||||
|
<button className="primary">{busy && !fetching ? '保存中…' : '保存渠道'}</button>
|
||||||
|
</div>
|
||||||
|
</fieldset>
|
||||||
|
</form>
|
||||||
|
</Dialog>
|
||||||
|
)}
|
||||||
|
{draft && candidates !== null && (
|
||||||
|
<Dialog id="model-picker-title" title="选择要加入的模型" onClose={() => setCandidates(null)}>
|
||||||
|
<p className="muted">获取到 {candidates.length} 个模型，已选择 {chosen.size} 个。已有模型不会重复加入，也不会覆盖展示名称。</p>
|
||||||
|
<label>搜索上游模型<input autoFocus value={candidateQuery} onChange={e => setCandidateQuery(e.target.value)} placeholder="输入模型名称" /></label>
|
||||||
|
<div className="actions">
|
||||||
|
<button type="button" onClick={() => setChosen(current => new Set([...current, ...visibleCandidates.filter(real => !existingModels.has(real))]))}>全选搜索结果</button>
|
||||||
|
<button type="button" onClick={() => setChosen(new Set())}>清空选择</button>
|
||||||
|
</div>
|
||||||
|
<div className="channel-model-list">
|
||||||
|
{visibleCandidates.map(real => (
|
||||||
|
<label className="check" key={real} style={{ overflowWrap: 'anywhere' }}>
|
||||||
|
<input type="checkbox" disabled={existingModels.has(real)} checked={existingModels.has(real) || chosen.has(real)} onChange={e => {
|
||||||
|
const checked = e.target.checked
|
||||||
|
setChosen(current => {
|
||||||
|
const next = new Set(current)
|
||||||
|
if (checked) next.add(real)
|
||||||
|
else next.delete(real)
|
||||||
|
return next
|
||||||
|
})
|
||||||
|
}} />
|
||||||
|
{real}{existingModels.has(real) && <small>已在列表中</small>}
|
||||||
|
</label>
|
||||||
|
))}
|
||||||
|
{visibleCandidates.length === 0 && <p className="muted">{candidates.length === 0 ? '上游未返回可选模型，可返回手动添加。' : '没有匹配的模型。'}</p>}
|
||||||
|
</div>
|
||||||
|
<div className="dialog-actions">
|
||||||
|
<button type="button" onClick={() => setCandidates(null)}>取消</button>
|
||||||
|
<button type="button" className="primary" disabled={chosen.size === 0} onClick={addChosen}>加入所选模型（{chosen.size}）</button>
|
||||||
|
</div>
|
||||||
|
</Dialog>
|
||||||
|
)}
|
||||||
|
</section>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,166 @@
|
|||||||
|
import { useState, type FormEvent } from 'react'
|
||||||
|
import { api, type Channel, type Key, type Stats } from '../api'
|
||||||
|
import { channelName, num, type Action } from '../common'
|
||||||
|
import { formatTime, keyExpiry } from '../time'
|
||||||
|
import { IconPlus } from '../components/icons'
|
||||||
|
import { Badge, Dialog, Feedback, Table } from '../components/ui'
|
||||||
|
|
||||||
|
interface KeyForm {
|
||||||
|
name: string
|
||||||
|
all_channels: boolean
|
||||||
|
channels: string[]
|
||||||
|
expires_at: string
|
||||||
|
}
|
||||||
|
|
||||||
|
const EMPTY_FORM: KeyForm = { name: '', all_channels: false, channels: [], expires_at: '' }
|
||||||
|
|
||||||
|
interface KeysProps {
|
||||||
|
keys: Key[]
|
||||||
|
stats: Stats | null
|
||||||
|
channels: Channel[]
|
||||||
|
timezone: string
|
||||||
|
busy: boolean
|
||||||
|
error: string
|
||||||
|
notice: string
|
||||||
|
action: Action
|
||||||
|
refresh: () => void
|
||||||
|
notify: (msg: string) => void
|
||||||
|
clearError: () => void
|
||||||
|
/** 跳转到概览页并按该 Key 筛选统计。 */
|
||||||
|
onViewStats: (k: Key) => void
|
||||||
|
}
|
||||||
|
|
||||||
|
export function KeysPage({ keys, stats, channels, timezone, busy, error, notice, action, refresh, notify, clearError, onViewStats }: KeysProps) {
|
||||||
|
const [creating, setCreating] = useState(false)
|
||||||
|
const [keyForm, setKeyForm] = useState<KeyForm>(EMPTY_FORM)
|
||||||
|
const [newKey, setNewKey] = useState('')
|
||||||
|
|
||||||
|
const time = formatTime
|
||||||
|
|
||||||
|
function openCreate() {
|
||||||
|
setCreating(true)
|
||||||
|
setKeyForm(EMPTY_FORM)
|
||||||
|
}
|
||||||
|
|
||||||
|
function closeCreate() {
|
||||||
|
setCreating(false)
|
||||||
|
clearError()
|
||||||
|
}
|
||||||
|
|
||||||
|
function submit(e: FormEvent) {
|
||||||
|
e.preventDefault()
|
||||||
|
action(async () => {
|
||||||
|
if (!keyForm.all_channels && !keyForm.channels.length) throw new Error('请选择至少一个渠道，或授权全部渠道。')
|
||||||
|
const expires_at = keyExpiry(keyForm.expires_at)
|
||||||
|
const r = await api<{ key: string }>('/keys', 'POST', { ...keyForm, channels: keyForm.all_channels ? [] : keyForm.channels, expires_at })
|
||||||
|
setNewKey(r.key)
|
||||||
|
setCreating(false)
|
||||||
|
refresh()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function toggle(k: Key) {
|
||||||
|
action(async () => {
|
||||||
|
await api('/keys/' + encodeURIComponent(k.id), 'PATCH', { enabled: !k.enabled })
|
||||||
|
refresh()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function revoke(k: Key) {
|
||||||
|
if (!confirm(`确定撤销 Key“${k.name}”？历史记录将保留。`)) return
|
||||||
|
action(async () => {
|
||||||
|
await api('/keys/' + encodeURIComponent(k.id), 'DELETE')
|
||||||
|
refresh()
|
||||||
|
notify('Key 已撤销，历史记录已保留。')
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function copyKey() {
|
||||||
|
action(async () => {
|
||||||
|
if (!navigator.clipboard) throw new Error('无法自动复制，请选中密钥手动复制。')
|
||||||
|
await navigator.clipboard.writeText(newKey)
|
||||||
|
notify('密钥已复制。')
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
function toggleChannel(id: string, checked: boolean) {
|
||||||
|
setKeyForm({
|
||||||
|
...keyForm,
|
||||||
|
channels: checked ? [...keyForm.channels, id] : keyForm.channels.filter(c => c !== id),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<section className="panel">
|
||||||
|
<div className="section-head">
|
||||||
|
<h2>访问密钥 <small>{keys.length}</small></h2>
|
||||||
|
<button className="primary" onClick={openCreate}><IconPlus />创建 Key</button>
|
||||||
|
</div>
|
||||||
|
<Table
|
||||||
|
heads={['名称 / 前缀', '渠道权限', '状态', '请求 / tokens', '有效期 / 创建时间', '操作']}
|
||||||
|
rows={keys.map(k => {
|
||||||
|
const u = stats?.by_key?.find(u => u.key_id === k.id)
|
||||||
|
return [
|
||||||
|
<><button className="link" onClick={() => onViewStats(k)}>{k.name}</button><code className="sub">{k.prefix}…</code></>,
|
||||||
|
k.all_channels ? '全部渠道' : k.channels?.map(id => channelName(channels, id)).join('、') || '无渠道权限',
|
||||||
|
<Badge variant={k.revoked ? 'danger' : k.enabled ? 'good' : 'neutral'}>{k.revoked ? '已撤销' : k.enabled ? '启用' : '停用'}</Badge>,
|
||||||
|
<>{num(u?.requests ?? 0)} 次<small className="sub">输入 {num(u?.input_tokens ?? 0)} / 输出 {num(u?.output_tokens ?? 0)}</small></>,
|
||||||
|
<>{time(k.expires_at) === '—' ? '永不过期' : time(k.expires_at)}<small className="sub">创建于 {time(k.created_at)}</small></>,
|
||||||
|
<div className="actions">
|
||||||
|
<button disabled={busy || k.revoked} onClick={() => toggle(k)}>{k.enabled ? '停用' : '启用'}</button>
|
||||||
|
<button className="danger" disabled={busy} onClick={() => revoke(k)}>撤销</button>
|
||||||
|
</div>,
|
||||||
|
]
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
{creating && (
|
||||||
|
<Dialog id="key-title" title="创建访问密钥" onClose={closeCreate}>
|
||||||
|
<Feedback error={error} notice={notice} />
|
||||||
|
<form onSubmit={submit}>
|
||||||
|
<fieldset disabled={busy}>
|
||||||
|
<label>名称
|
||||||
|
<input autoFocus required value={keyForm.name} onChange={e => setKeyForm({ ...keyForm, name: e.target.value })} />
|
||||||
|
</label>
|
||||||
|
<label>有效期（{timezone}，可选）
|
||||||
|
<input type="datetime-local" value={keyForm.expires_at} onChange={e => setKeyForm({ ...keyForm, expires_at: e.target.value })} />
|
||||||
|
<small>留空表示永不过期。</small>
|
||||||
|
</label>
|
||||||
|
<label className="check">
|
||||||
|
<input type="checkbox" checked={keyForm.all_channels} onChange={e => setKeyForm({ ...keyForm, all_channels: e.target.checked })} />允许全部渠道（包含未来新增渠道）
|
||||||
|
</label>
|
||||||
|
{!keyForm.all_channels && (
|
||||||
|
<div className="channel-options">
|
||||||
|
<p>选择允许访问的渠道</p>
|
||||||
|
{!channels.length && <p className="muted">暂无渠道，请先创建渠道。</p>}
|
||||||
|
{channels.map(c => (
|
||||||
|
<label className="check" key={c.id}>
|
||||||
|
<input type="checkbox" checked={keyForm.channels.includes(c.id)} onChange={e => toggleChannel(c.id, e.target.checked)} />
|
||||||
|
{c.name}{!c.enabled && '（停用）'}
|
||||||
|
</label>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
)}
|
||||||
|
<div className="dialog-actions">
|
||||||
|
<button type="button" onClick={closeCreate}>取消</button>
|
||||||
|
<button className="primary">{busy ? '创建中…' : '创建 Key'}</button>
|
||||||
|
</div>
|
||||||
|
</fieldset>
|
||||||
|
</form>
|
||||||
|
</Dialog>
|
||||||
|
)}
|
||||||
|
{newKey && (
|
||||||
|
<Dialog id="secret-title" title="Key 已创建">
|
||||||
|
<div className="alert warning">完整密钥仅展示这一次。请立即复制并妥善保存，关闭后无法再次查看。</div>
|
||||||
|
<label>完整密钥
|
||||||
|
<textarea autoFocus readOnly rows={3} value={newKey} onFocus={e => e.target.select()} />
|
||||||
|
</label>
|
||||||
|
<Feedback error={error} notice={notice} />
|
||||||
|
<div className="dialog-actions">
|
||||||
|
<button onClick={() => { setNewKey(''); notify('') }}>我已保存，关闭</button>
|
||||||
|
<button className="primary" onClick={copyKey}>复制密钥</button>
|
||||||
|
</div>
|
||||||
|
</Dialog>
|
||||||
|
)}
|
||||||
|
</section>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,42 @@
|
|||||||
|
import { useState } from 'react'
|
||||||
|
import { api } from '../api'
|
||||||
|
import type { Action } from '../common'
|
||||||
|
import { IconRoute } from '../components/icons'
|
||||||
|
import { Feedback } from '../components/ui'
|
||||||
|
|
||||||
|
interface LoginProps {
|
||||||
|
busy: boolean
|
||||||
|
error: string
|
||||||
|
notice: string
|
||||||
|
action: Action
|
||||||
|
/** 登录成功后回调，用于重新加载会话状态。 */
|
||||||
|
onSuccess: () => Promise<void>
|
||||||
|
}
|
||||||
|
|
||||||
|
export function Login({ busy, error, notice, action, onSuccess }: LoginProps) {
|
||||||
|
const [password, setPassword] = useState('')
|
||||||
|
|
||||||
|
function submit(e: React.FormEvent) {
|
||||||
|
e.preventDefault()
|
||||||
|
action(async () => {
|
||||||
|
await api('/login', 'POST', { password })
|
||||||
|
setPassword('')
|
||||||
|
await onSuccess()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<div className="login">
|
||||||
|
<form className="login-card" onSubmit={submit}>
|
||||||
|
<div className="mark"><IconRoute width={22} height={22} strokeWidth={2} /></div>
|
||||||
|
<h1>AI Router</h1>
|
||||||
|
<p className="muted">登录管理台，管理模型与访问权限。</p>
|
||||||
|
<Feedback error={error} notice={notice} />
|
||||||
|
<label>管理密码
|
||||||
|
<input type="password" required autoComplete="current-password" autoFocus value={password} onChange={e => setPassword(e.target.value)} />
|
||||||
|
</label>
|
||||||
|
<button className="primary" disabled={busy}>{busy ? '登录中…' : '登录管理台'}</button>
|
||||||
|
</form>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
import type { Channel, Log } from '../api'
|
||||||
|
import { channelName, num } from '../common'
|
||||||
|
import { formatTime } from '../time'
|
||||||
|
import { Badge, Table } from '../components/ui'
|
||||||
|
|
||||||
|
export function LogsPage({ logs, channels }: { logs: Log[]; channels: Channel[] }) {
|
||||||
|
return (
|
||||||
|
<section className="panel">
|
||||||
|
<div className="section-head">
|
||||||
|
<h2>请求记录</h2>
|
||||||
|
<small>最多 200 条，可通过时间缩小范围</small>
|
||||||
|
</div>
|
||||||
|
<Table
|
||||||
|
heads={['时间 / Key', '渠道 / 模型', '协议 / 模式', '状态 / 耗时', '用量', '详情']}
|
||||||
|
rows={logs.map(l => [
|
||||||
|
<><span className="nowrap">{formatTime(l.created_at)}</span><small className="sub">{l.key_name || l.key_id}</small></>,
|
||||||
|
<>{channelName(channels, l.channel)}<small className="sub">{l.request_model} → {l.actual_model || '—'}</small></>,
|
||||||
|
<>{l.downstream} → {l.upstream}<small className="sub">{l.stream ? '流式' : '非流式'}</small></>,
|
||||||
|
<><Badge variant={l.status >= 200 && l.status < 400 ? 'good' : 'danger'}>{l.status || '失败'}</Badge><small className="sub">{num(l.duration_ms)} ms</small></>,
|
||||||
|
l.usage_known
|
||||||
|
? <>输入 {num(l.input_tokens)} / 输出 {num(l.output_tokens)}<small className="sub">缓存读 {num(l.cache_read_tokens)} / 写 {num(l.cache_write_tokens)}</small></>
|
||||||
|
: '用量未知',
|
||||||
|
<details>
|
||||||
|
<summary>展开</summary>
|
||||||
|
<small className="sub">ID: {l.id}</small>
|
||||||
|
<p>{l.error || '无错误信息'}</p>
|
||||||
|
</details>,
|
||||||
|
])}
|
||||||
|
/>
|
||||||
|
</section>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
import { useEffect, useRef, useState } from 'react'
|
||||||
|
import { api, APIError, type Channel } from '../api'
|
||||||
|
import { IconCheck, IconCopy } from '../components/icons'
|
||||||
|
import { Badge, Dialog, Table } from '../components/ui'
|
||||||
|
|
||||||
|
type TestResult = { ok: boolean; duration_ms: number; status: number; error?: string }
|
||||||
|
type TestState = { state: 'running' | 'done' | 'cancelled'; result?: TestResult; error?: string }
|
||||||
|
const DEFAULT_PROMPT = 'Reply with OK.'
|
||||||
|
const testKey = (m: ModelRow) => JSON.stringify([m.channel.id, m.id, m.real, m.channel.base_url, m.channel.protocol])
|
||||||
|
|
||||||
|
interface ModelRow {
|
||||||
|
id: string
|
||||||
|
alias: string
|
||||||
|
real: string
|
||||||
|
channel: Channel
|
||||||
|
}
|
||||||
|
|
||||||
|
/** 汇总所有渠道的可调用模型：下游调用的模型格式为「渠道名/别名」。 */
|
||||||
|
export function ModelsPage({ channels }: { channels: Channel[] }) {
|
||||||
|
const [results, setResults] = useState<Record<string, TestState>>({})
|
||||||
|
const [manual, setManual] = useState<ModelRow | null>(null)
|
||||||
|
const [prompt, setPrompt] = useState(DEFAULT_PROMPT)
|
||||||
|
const [running, setRunning] = useState(false)
|
||||||
|
const [progress, setProgress] = useState({ done: 0, total: 0 })
|
||||||
|
const controller = useRef<AbortController | null>(null)
|
||||||
|
useEffect(() => () => { controller.current?.abort() }, [])
|
||||||
|
const rows: ModelRow[] = channels
|
||||||
|
.flatMap(c => Object.entries(c.models ?? {}).map(([alias, real]) => ({ id: `${c.name}/${alias}`, alias, real, channel: c })))
|
||||||
|
.sort((a, b) => a.id.localeCompare(b.id))
|
||||||
|
const available = rows.filter(m => m.channel.enabled)
|
||||||
|
|
||||||
|
async function runTests(models: ModelRow[], text: string) {
|
||||||
|
if (controller.current || models.length === 0) return
|
||||||
|
const ctrl = new AbortController()
|
||||||
|
controller.current = ctrl
|
||||||
|
setRunning(true)
|
||||||
|
setResults(old => {
|
||||||
|
const next = { ...old }
|
||||||
|
for (const model of models) delete next[testKey(model)]
|
||||||
|
return next
|
||||||
|
})
|
||||||
|
setProgress({ done: 0, total: models.length })
|
||||||
|
try {
|
||||||
|
for (const model of models) {
|
||||||
|
if (ctrl.signal.aborted) break
|
||||||
|
const key = testKey(model)
|
||||||
|
setResults(old => ({ ...old, [key]: { state: 'running' } }))
|
||||||
|
let stop = false
|
||||||
|
try {
|
||||||
|
const result = await api<TestResult>('/models/test', 'POST', { model: model.id, prompt: text }, ctrl.signal)
|
||||||
|
setResults(old => ({ ...old, [key]: { state: 'done', result } }))
|
||||||
|
} catch (e) {
|
||||||
|
setResults(old => ({ ...old, [key]: ctrl.signal.aborted ? { state: 'cancelled' } : { state: 'done', error: e instanceof Error ? e.message : '测试失败' } }))
|
||||||
|
stop = ctrl.signal.aborted || (e instanceof APIError && (e.status === 401 || e.status === 403))
|
||||||
|
}
|
||||||
|
if (!ctrl.signal.aborted) setProgress(p => ({ ...p, done: p.done + 1 }))
|
||||||
|
if (stop) break
|
||||||
|
}
|
||||||
|
} finally {
|
||||||
|
controller.current = null
|
||||||
|
setRunning(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function testAll() {
|
||||||
|
if (!confirm(`将依次测试 ${available.length} 个启用模型，每个模型发送一次简短请求，可能产生 Token 费用。是否继续？`)) return
|
||||||
|
void runTests(available, DEFAULT_PROMPT)
|
||||||
|
}
|
||||||
|
|
||||||
|
return (
|
||||||
|
<section className="panel">
|
||||||
|
<div className="section-head models-head">
|
||||||
|
<h2>已配置展示模型 <small>{rows.length}</small></h2>
|
||||||
|
<small>仅汇总渠道中已勾选并保存的模型；下游列表按渠道状态和 Key 权限过滤</small>
|
||||||
|
<div className="actions">
|
||||||
|
<button className="primary" disabled={running || available.length === 0} onClick={testAll}>一键测试全部</button>
|
||||||
|
{running && <button onClick={() => controller.current?.abort()}>停止测试</button>}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
<p className="muted panel-foot">测试会实际调用上游并可能产生费用，停用渠道不参与测试。结果仅保留在当前页面，不计入正式调用统计，不代表下游 Key 权限验证。{progress.total > 0 && <span role="status"> {running ? '测试中' : '本轮已结束'}：{progress.done}/{progress.total}</span>}</p>
|
||||||
|
<Table
|
||||||
|
heads={['下游模型名 / 展示名称', '上游真实模型', '所属渠道', '状态', '测试结果', '操作']}
|
||||||
|
rows={rows.map(m => [
|
||||||
|
<span className="model-id"><code>{m.id}</code><CopyButton text={m.id} /></span>,
|
||||||
|
<code>{m.real}</code>,
|
||||||
|
<>{m.channel.name}<small className="sub">{m.channel.protocol}</small></>,
|
||||||
|
<Badge variant={m.channel.enabled ? 'good' : 'neutral'}>{m.channel.enabled ? '可用' : '渠道停用'}</Badge>,
|
||||||
|
<TestStatus value={results[testKey(m)]} />,
|
||||||
|
<button disabled={running || !m.channel.enabled} onClick={() => { setManual(m); setPrompt(DEFAULT_PROMPT) }}>手动测试</button>,
|
||||||
|
])}
|
||||||
|
/>
|
||||||
|
{manual && (
|
||||||
|
<Dialog id="model-test-title" title={`测试 ${manual.id}`} onClose={running ? undefined : () => setManual(null)}>
|
||||||
|
<form onSubmit={e => { e.preventDefault(); void runTests([manual], prompt.trim()) }}>
|
||||||
|
<label>测试提示词<textarea required maxLength={2000} rows={3} value={prompt} disabled={running} onChange={e => setPrompt(e.target.value)} /></label>
|
||||||
|
<p className="muted">发送一次非流式请求，限制输出 Token，最长等待 30 秒。可能产生上游费用；仅检查接口响应，不评价回答质量。</p>
|
||||||
|
<TestStatus value={results[testKey(manual)]} />
|
||||||
|
<div className="dialog-actions">
|
||||||
|
{running ? <button type="button" onClick={() => controller.current?.abort()}>停止测试</button> : <button type="button" onClick={() => setManual(null)}>关闭</button>}
|
||||||
|
<button className="primary" disabled={running || !prompt.trim()}>{running ? '测试中…' : '开始测试'}</button>
|
||||||
|
</div>
|
||||||
|
</form>
|
||||||
|
</Dialog>
|
||||||
|
)}
|
||||||
|
</section>
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
function TestStatus({ value }: { value?: TestState }) {
|
||||||
|
if (!value) return <span className="muted">未测试</span>
|
||||||
|
if (value.state === 'running') return <span role="status">测试中…</span>
|
||||||
|
if (value.state === 'cancelled') return <span className="muted">已停止（已发送的请求仍可能计费）</span>
|
||||||
|
const result = value.result
|
||||||
|
return <div role="status">
|
||||||
|
<Badge variant={result?.ok ? 'good' : 'danger'}>{result?.ok ? '成功' : '失败'}</Badge>
|
||||||
|
{result && <small className="sub">{result.duration_ms} ms{result.status > 0 ? ` · HTTP ${result.status}` : ''}</small>}
|
||||||
|
{(value.error || result?.error) && <small className="sub">{value.error || result?.error}</small>}
|
||||||
|
</div>
|
||||||
|
}
|
||||||
|
|
||||||
|
function CopyButton({ text }: { text: string }) {
|
||||||
|
const [copied, setCopied] = useState(false)
|
||||||
|
async function copy() {
|
||||||
|
try {
|
||||||
|
await navigator.clipboard.writeText(text)
|
||||||
|
setCopied(true)
|
||||||
|
setTimeout(() => setCopied(false), 1500)
|
||||||
|
} catch {
|
||||||
|
// 剪贴板不可用（非安全上下文等），静默失败，用户可手动选中复制
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return (
|
||||||
|
<button type="button" className="icon-btn copy-btn" aria-label={`复制 ${text}`} title={copied ? '已复制' : '复制模型名'} onClick={copy}>
|
||||||
|
{copied ? <IconCheck /> : <IconCopy />}
|
||||||
|
</button>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
import type { ReactNode } from 'react'
|
||||||
|
import type { Channel, Stats } from '../api'
|
||||||
|
import { channelName, num, usageCells, usageHeads } from '../common'
|
||||||
|
import { Table } from '../components/ui'
|
||||||
|
import { UsageChart } from '../components/UsageChart'
|
||||||
|
import { IconActivity, IconArrowDown, IconArrowUp, IconCheckCircle, IconDatabase, IconHelpCircle, IconXCircle, IconZap } from '../components/icons'
|
||||||
|
|
||||||
|
const METRICS: { key: keyof Stats['summary']; label: string; icon: ReactNode; tone?: 'bad' }[] = [
|
||||||
|
{ key: 'requests', label: '总请求', icon: <IconActivity /> },
|
||||||
|
{ key: 'success', label: '成功请求', icon: <IconCheckCircle /> },
|
||||||
|
{ key: 'failed', label: '失败请求', icon: <IconXCircle />, tone: 'bad' },
|
||||||
|
{ key: 'unknown_usage', label: '用量未知', icon: <IconHelpCircle /> },
|
||||||
|
{ key: 'input_tokens', label: '输入 tokens', icon: <IconArrowDown /> },
|
||||||
|
{ key: 'output_tokens', label: '输出 tokens', icon: <IconArrowUp /> },
|
||||||
|
{ key: 'cache_read_tokens', label: '缓存读取', icon: <IconZap /> },
|
||||||
|
{ key: 'cache_write_tokens', label: '缓存写入', icon: <IconDatabase /> },
|
||||||
|
]
|
||||||
|
|
||||||
|
export function Overview({ stats, channels }: { stats: Stats | null; channels: Channel[] }) {
|
||||||
|
if (!stats) return <div className="panel empty">暂无统计数据，请刷新重试。</div>
|
||||||
|
|
||||||
|
return (
|
||||||
|
<>
|
||||||
|
<div className="metrics">
|
||||||
|
{METRICS.map(m => {
|
||||||
|
const value = stats.summary[m.key]
|
||||||
|
const alert = m.tone === 'bad' && value > 0
|
||||||
|
return (
|
||||||
|
<div className={`metric${alert ? ' alert' : ''}`} key={m.key}>
|
||||||
|
<span>{m.label}</span>
|
||||||
|
<span className="metric-icon" aria-hidden="true">{m.icon}</span>
|
||||||
|
<strong>{num(value)}</strong>
|
||||||
|
</div>
|
||||||
|
)
|
||||||
|
})}
|
||||||
|
</div>
|
||||||
|
<p className="muted">用量未知不代表零消耗；tokens 以已返回用量为准。</p>
|
||||||
|
<section className="panel">
|
||||||
|
<h2>每日用量</h2>
|
||||||
|
<UsageChart data={stats.daily ?? []} />
|
||||||
|
<Table heads={['日期', ...usageHeads]} rows={(stats.daily ?? []).map(u => [<span className="nowrap">{u.day}</span>, ...usageCells(u)])} />
|
||||||
|
</section>
|
||||||
|
<section className="panel">
|
||||||
|
<h2>按 Key 统计</h2>
|
||||||
|
<Table
|
||||||
|
heads={['Key', ...usageHeads, '成功 / 失败', '用量未知']}
|
||||||
|
rows={(stats.by_key ?? []).map(u => [<span className="cell-id" title={u.name || u.key_id}>{u.name || u.key_id}</span>, ...usageCells(u), `${num(u.success)} / ${num(u.failed)}`, num(u.unknown_usage)])}
|
||||||
|
/>
|
||||||
|
</section>
|
||||||
|
<section className="panel">
|
||||||
|
<h2>按模型统计</h2>
|
||||||
|
<Table
|
||||||
|
heads={['渠道', '模型', ...usageHeads]}
|
||||||
|
rows={(stats.by_model ?? []).map(u => [
|
||||||
|
<span className="cell-id" title={channelName(channels, u.channel)}>{channelName(channels, u.channel)}</span>,
|
||||||
|
<span className="cell-id" title={u.model}>{u.model}</span>,
|
||||||
|
...usageCells(u),
|
||||||
|
])}
|
||||||
|
/>
|
||||||
|
</section>
|
||||||
|
</>
|
||||||
|
)
|
||||||
|
}
|
||||||
@@ -0,0 +1,522 @@
|
|||||||
|
/* ===== 主题令牌 ===== */
|
||||||
|
:root {
|
||||||
|
--bg: #f2f3fa;
|
||||||
|
--surface: #ffffff;
|
||||||
|
--surface-2: #f6f6fc;
|
||||||
|
--input-bg: #ffffff;
|
||||||
|
--border: #e5e6f1;
|
||||||
|
--border-strong: #d3d5e8;
|
||||||
|
--text: #21243a;
|
||||||
|
--text-2: #5b5f7c;
|
||||||
|
--text-3: #8c8fa9;
|
||||||
|
--accent: #4f46e5;
|
||||||
|
--accent-hover: #4338ca;
|
||||||
|
--accent-2: #7c3aed;
|
||||||
|
--accent-soft: rgba(79, 70, 229, .09);
|
||||||
|
--on-accent: #ffffff;
|
||||||
|
--danger: #d92d20;
|
||||||
|
--danger-bg: #fdeceb;
|
||||||
|
--danger-text: #b42318;
|
||||||
|
--success-bg: #e4f3e9;
|
||||||
|
--success-text: #1a7f4b;
|
||||||
|
--warning-bg: #fbf2d9;
|
||||||
|
--warning-text: #a15c07;
|
||||||
|
--sidebar-bg: #181a34;
|
||||||
|
--sidebar-text: #a4a7cc;
|
||||||
|
--sidebar-hover: rgba(255, 255, 255, .06);
|
||||||
|
--sidebar-active: rgba(129, 140, 248, .16);
|
||||||
|
--sidebar-active-text: #cdd3ff;
|
||||||
|
--sidebar-border: rgba(255, 255, 255, .08);
|
||||||
|
--mark-bg: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%);
|
||||||
|
--mark-text: #ffffff;
|
||||||
|
--mark-glow: rgba(99, 102, 241, .38);
|
||||||
|
--chart-input: #6366f1;
|
||||||
|
--chart-output: #38bdf8;
|
||||||
|
--overlay: rgba(24, 26, 52, .45);
|
||||||
|
--shadow-sm: 0 1px 2px rgba(30, 27, 75, .05);
|
||||||
|
--shadow-md: 0 1px 2px rgba(30, 27, 75, .06), 0 8px 24px rgba(30, 27, 75, .08);
|
||||||
|
--shadow-lg: 0 24px 70px rgba(23, 20, 60, .24);
|
||||||
|
--primary-shadow: 0 1px 2px rgba(30, 27, 75, .16), 0 5px 16px rgba(79, 70, 229, .28);
|
||||||
|
--primary-shadow-hover: 0 2px 4px rgba(30, 27, 75, .16), 0 8px 22px rgba(79, 70, 229, .36);
|
||||||
|
--scrollbar-thumb: #c0c3de;
|
||||||
|
--scrollbar-thumb-hover: #9ea2c8;
|
||||||
|
--r-sm: 7px;
|
||||||
|
--r-md: 10px;
|
||||||
|
--r-lg: 14px;
|
||||||
|
color-scheme: light;
|
||||||
|
}
|
||||||
|
|
||||||
|
:root[data-theme=dark] {
|
||||||
|
--bg: #0c0e1b;
|
||||||
|
--surface: #161932;
|
||||||
|
--surface-2: #1d2040;
|
||||||
|
--input-bg: #11132a;
|
||||||
|
--border: #292d4e;
|
||||||
|
--border-strong: #3b4068;
|
||||||
|
--text: #e8eaf8;
|
||||||
|
--text-2: #a3a7cb;
|
||||||
|
--text-3: #72769d;
|
||||||
|
--accent: #818cf8;
|
||||||
|
--accent-hover: #a5b4fc;
|
||||||
|
--accent-2: #a78bfa;
|
||||||
|
--accent-soft: rgba(129, 140, 248, .14);
|
||||||
|
--on-accent: #1e1b4b;
|
||||||
|
--danger: #f97066;
|
||||||
|
--danger-bg: rgba(249, 112, 102, .12);
|
||||||
|
--danger-text: #fda29b;
|
||||||
|
--success-bg: rgba(74, 222, 128, .12);
|
||||||
|
--success-text: #5eeaa6;
|
||||||
|
--warning-bg: rgba(251, 191, 36, .12);
|
||||||
|
--warning-text: #fcd34d;
|
||||||
|
--sidebar-bg: #090a16;
|
||||||
|
--sidebar-text: #9094bb;
|
||||||
|
--sidebar-hover: rgba(255, 255, 255, .05);
|
||||||
|
--sidebar-active: rgba(129, 140, 248, .14);
|
||||||
|
--sidebar-active-text: #a5b4fc;
|
||||||
|
--sidebar-border: rgba(255, 255, 255, .07);
|
||||||
|
--mark-glow: rgba(129, 140, 248, .32);
|
||||||
|
--chart-input: #818cf8;
|
||||||
|
--chart-output: #38bdf8;
|
||||||
|
--overlay: rgba(2, 4, 18, .62);
|
||||||
|
--shadow-sm: 0 1px 2px rgba(0, 0, 0, .4);
|
||||||
|
--shadow-md: 0 1px 2px rgba(0, 0, 0, .4), 0 10px 28px rgba(0, 0, 0, .35);
|
||||||
|
--shadow-lg: 0 24px 70px rgba(0, 0, 0, .6);
|
||||||
|
--primary-shadow: 0 1px 2px rgba(0, 0, 0, .3), 0 5px 16px rgba(129, 140, 248, .25);
|
||||||
|
--primary-shadow-hover: 0 2px 4px rgba(0, 0, 0, .3), 0 8px 22px rgba(129, 140, 248, .34);
|
||||||
|
--scrollbar-thumb: #3b4068;
|
||||||
|
--scrollbar-thumb-hover: #50558a;
|
||||||
|
color-scheme: dark;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* ===== 基础 ===== */
|
||||||
|
:root {
|
||||||
|
font-family: Inter, -apple-system, BlinkMacSystemFont, Segoe UI, PingFang SC, Hiragino Sans GB, Microsoft YaHei, sans-serif;
|
||||||
|
font-size: 14px;
|
||||||
|
font-synthesis: none;
|
||||||
|
color: var(--text);
|
||||||
|
background: var(--bg);
|
||||||
|
scrollbar-width: thin;
|
||||||
|
scrollbar-color: var(--scrollbar-thumb) transparent;
|
||||||
|
}
|
||||||
|
|
||||||
|
* { box-sizing: border-box; }
|
||||||
|
|
||||||
|
body {
|
||||||
|
margin: 0;
|
||||||
|
background: var(--bg);
|
||||||
|
transition: background-color .25s ease, color .25s ease;
|
||||||
|
}
|
||||||
|
|
||||||
|
::selection { background: var(--accent-soft); }
|
||||||
|
|
||||||
|
::-webkit-scrollbar { width: 10px; height: 10px; }
|
||||||
|
::-webkit-scrollbar-track { background: transparent; }
|
||||||
|
::-webkit-scrollbar-thumb {
|
||||||
|
background: var(--scrollbar-thumb);
|
||||||
|
border-radius: 8px;
|
||||||
|
border: 2px solid transparent;
|
||||||
|
background-clip: content-box;
|
||||||
|
}
|
||||||
|
::-webkit-scrollbar-thumb:hover { background: var(--scrollbar-thumb-hover); background-clip: content-box; }
|
||||||
|
::-webkit-scrollbar-corner { background: transparent; }
|
||||||
|
|
||||||
|
/* ===== 按钮 ===== */
|
||||||
|
button, input, select, textarea { font: inherit; }
|
||||||
|
|
||||||
|
button {
|
||||||
|
display: inline-flex;
|
||||||
|
align-items: center;
|
||||||
|
justify-content: center;
|
||||||
|
gap: 7px;
|
||||||
|
cursor: pointer;
|
||||||
|
border: 1px solid var(--border-strong);
|
||||||
|
border-radius: var(--r-sm);
|
||||||
|
padding: 8px 14px;
|
||||||
|
background: var(--surface);
|
||||||
|
color: var(--text);
|
||||||
|
font-weight: 500;
|
||||||
|
white-space: nowrap;
|
||||||
|
transition: background-color .15s ease, border-color .15s ease, color .15s ease, box-shadow .15s ease, transform .15s ease, filter .15s ease;
|
||||||
|
}
|
||||||
|
button:hover { background: var(--surface-2); }
|
||||||
|
button:active { transform: translateY(1px); }
|
||||||
|
button:disabled { opacity: .5; cursor: not-allowed; }
|
||||||
|
button:disabled:active { transform: none; }
|
||||||
|
|
||||||
|
.primary {
|
||||||
|
background: linear-gradient(135deg, var(--accent) 0%, var(--accent-2) 100%);
|
||||||
|
border-color: transparent;
|
||||||
|
color: var(--on-accent);
|
||||||
|
font-weight: 600;
|
||||||
|
box-shadow: var(--primary-shadow);
|
||||||
|
}
|
||||||
|
.primary:hover { background: linear-gradient(135deg, var(--accent-hover) 0%, var(--accent-2) 100%); box-shadow: var(--primary-shadow-hover); }
|
||||||
|
|
||||||
|
.danger { color: var(--danger); }
|
||||||
|
.danger:hover { background: var(--danger-bg); border-color: var(--danger); }
|
||||||
|
|
||||||
|
.icon-btn { border: 0; background: transparent; color: var(--text-3); padding: 5px; border-radius: var(--r-sm); }
|
||||||
|
.icon-btn:hover { background: var(--surface-2); color: var(--text); }
|
||||||
|
|
||||||
|
.link { border: 0; padding: 0; background: none; color: var(--accent); font-weight: 600; white-space: normal; text-align: left; display: inline; }
|
||||||
|
.link:hover { text-decoration: underline; background: none; }
|
||||||
|
|
||||||
|
/* ===== 表单 ===== */
|
||||||
|
input, select, textarea {
|
||||||
|
border: 1px solid var(--border-strong);
|
||||||
|
border-radius: var(--r-sm);
|
||||||
|
padding: 9px 11px;
|
||||||
|
background: var(--input-bg);
|
||||||
|
color: var(--text);
|
||||||
|
min-width: 0;
|
||||||
|
width: 100%;
|
||||||
|
transition: border-color .15s ease, box-shadow .15s ease, background-color .25s ease;
|
||||||
|
}
|
||||||
|
input::placeholder, textarea::placeholder { color: var(--text-3); }
|
||||||
|
textarea { resize: vertical; font-family: ui-monospace, SF Mono, Cascadia Mono, Consolas, monospace; font-size: 12.5px; line-height: 1.7; }
|
||||||
|
|
||||||
|
button:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; }
|
||||||
|
input:focus, select:focus, textarea:focus { outline: none; border-color: var(--accent); box-shadow: 0 0 0 3px var(--accent-soft); }
|
||||||
|
|
||||||
|
label { display: flex; flex-direction: column; gap: 7px; font-size: 12.5px; font-weight: 600; color: var(--text-2); margin-bottom: 15px; }
|
||||||
|
label > input, label > select, label > textarea { font-weight: 400; color: var(--text); font-size: 13.5px; }
|
||||||
|
|
||||||
|
small, .muted { color: var(--text-3); font-weight: 400; font-size: 12px; line-height: 1.7; }
|
||||||
|
.sub { display: block; margin-top: 5px; overflow-wrap: anywhere; }
|
||||||
|
|
||||||
|
h1 { font-size: 25px; letter-spacing: -.5px; margin: 8px 0; font-weight: 700; }
|
||||||
|
h2 { font-size: 15.5px; margin: 0; padding: 20px 22px 14px; font-weight: 600; }
|
||||||
|
h2 small { margin-left: 8px; }
|
||||||
|
|
||||||
|
code { font-family: ui-monospace, SF Mono, Cascadia Mono, Consolas, monospace; font-size: 12px; background: var(--surface-2); border: 1px solid var(--border); padding: 1px 6px; border-radius: 5px; }
|
||||||
|
|
||||||
|
fieldset { padding: 0; border: 0; margin: 0; min-width: 0; }
|
||||||
|
|
||||||
|
.check { flex-direction: row; align-items: center; font-weight: 500; line-height: 1.7; color: var(--text); }
|
||||||
|
.check input { width: 15px; height: 15px; accent-color: var(--accent); flex-shrink: 0; }
|
||||||
|
|
||||||
|
summary { cursor: pointer; color: var(--accent); font-weight: 500; font-size: 12.5px; }
|
||||||
|
details p { white-space: pre-wrap; min-width: 160px; margin: 8px 0 0; color: var(--text-2); }
|
||||||
|
|
||||||
|
/* 行内模型名 + 复制按钮（触屏无 hover，按钮保底可见） */
|
||||||
|
.model-id { display: inline-flex; align-items: center; gap: 6px; }
|
||||||
|
.copy-btn { padding: 3px; opacity: .4; transition: opacity .12s ease, background-color .15s ease, color .15s ease; }
|
||||||
|
tr:hover .copy-btn, .copy-btn:hover, .copy-btn:focus-visible { opacity: 1; }
|
||||||
|
.copy-btn svg { width: 13px; height: 13px; }
|
||||||
|
|
||||||
|
.panel-foot { margin: 0; padding: 12px 22px 16px; border-top: 1px dashed var(--border); }
|
||||||
|
|
||||||
|
/* ===== 布局骨架 / 侧栏 ===== */
|
||||||
|
.shell { display: flex; min-height: 100vh; }
|
||||||
|
|
||||||
|
aside {
|
||||||
|
width: 230px;
|
||||||
|
flex-shrink: 0;
|
||||||
|
background: var(--sidebar-bg);
|
||||||
|
background-image: linear-gradient(180deg, color-mix(in srgb, var(--accent) 16%, var(--sidebar-bg)) 0%, var(--sidebar-bg) 460px);
|
||||||
|
color: #edf0fb;
|
||||||
|
display: flex;
|
||||||
|
flex-direction: column;
|
||||||
|
padding: 26px 16px 20px;
|
||||||
|
position: sticky;
|
||||||
|
top: 0;
|
||||||
|
height: 100vh;
|
||||||
|
transition: background-color .25s ease;
|
||||||
|
}
|
||||||
|
|
||||||
|
.brand { display: flex; align-items: center; gap: 12px; font-size: 17.5px; font-weight: 700; padding: 0 6px; margin-bottom: 36px; }
|
||||||
|
.brand small { display: block; color: var(--sidebar-text); font-size: 11px; font-weight: 400; letter-spacing: .8px; margin-top: 3px; }
|
||||||
|
|
||||||
|
.mark {
|
||||||
|
display: grid;
|
||||||
|
place-items: center;
|
||||||
|
background: var(--mark-bg);
|
||||||
|
color: var(--mark-text);
|
||||||
|
border-radius: 11px;
|
||||||
|
width: 40px;
|
||||||
|
height: 40px;
|
||||||
|
flex-shrink: 0;
|
||||||
|
box-shadow: 0 4px 14px var(--mark-glow), inset 0 1px 0 rgba(255, 255, 255, .22);
|
||||||
|
}
|
||||||
|
|
||||||
|
nav { display: grid; gap: 4px; }
|
||||||
|
nav button {
|
||||||
|
justify-content: flex-start;
|
||||||
|
border: 0;
|
||||||
|
background: transparent;
|
||||||
|
color: var(--sidebar-text);
|
||||||
|
padding: 10px 12px 10px 15px;
|
||||||
|
border-radius: 9px;
|
||||||
|
font-weight: 500;
|
||||||
|
position: relative;
|
||||||
|
}
|
||||||
|
nav button svg { opacity: .75; }
|
||||||
|
nav button:hover { background: var(--sidebar-hover); color: #fff; }
|
||||||
|
nav button.active { background: var(--sidebar-active); color: var(--sidebar-active-text); font-weight: 600; }
|
||||||
|
nav button.active::before {
|
||||||
|
content: "";
|
||||||
|
position: absolute;
|
||||||
|
left: 4px;
|
||||||
|
top: 50%;
|
||||||
|
transform: translateY(-50%);
|
||||||
|
width: 3px;
|
||||||
|
height: 17px;
|
||||||
|
border-radius: 2px;
|
||||||
|
background: var(--sidebar-active-text);
|
||||||
|
}
|
||||||
|
nav button.active svg { opacity: 1; }
|
||||||
|
|
||||||
|
.aside-bottom { margin-top: auto; padding-top: 16px; border-top: 1px solid var(--sidebar-border); display: grid; gap: 8px; }
|
||||||
|
.aside-bottom button { width: 100%; justify-content: flex-start; background: transparent; color: var(--sidebar-text); border-color: var(--sidebar-border); }
|
||||||
|
.aside-bottom button:hover { background: var(--sidebar-hover); color: #fff; }
|
||||||
|
|
||||||
|
/* 主题分段控件（侧栏深色底上的自定义三态开关） */
|
||||||
|
.theme-seg { display: flex; gap: 2px; padding: 3px; background: rgba(255, 255, 255, .06); border: 1px solid var(--sidebar-border); border-radius: 11px; }
|
||||||
|
.aside-bottom .theme-seg button { flex: 1; justify-content: center; border: 0; padding: 6px 0; border-radius: 8px; background: transparent; color: var(--sidebar-text); }
|
||||||
|
.aside-bottom .theme-seg button:hover { background: var(--sidebar-hover); color: #fff; }
|
||||||
|
.aside-bottom .theme-seg button.active { background: var(--sidebar-active); color: var(--sidebar-active-text); }
|
||||||
|
.theme-seg button svg { width: 15px; height: 15px; }
|
||||||
|
|
||||||
|
main { padding: 30px 40px 40px; min-width: 0; flex: 1; width: 100%; max-width: 1400px; margin: 0 auto; }
|
||||||
|
|
||||||
|
header { display: flex; align-items: flex-start; justify-content: space-between; gap: 15px; margin-bottom: 22px; }
|
||||||
|
header > button { margin-top: 6px; }
|
||||||
|
header h1 { margin-top: 0; }
|
||||||
|
header .muted { margin: 0; }
|
||||||
|
|
||||||
|
.tz-hint { display: inline-flex; align-items: center; gap: 5px; }
|
||||||
|
.tz-hint svg { opacity: .7; flex-shrink: 0; }
|
||||||
|
|
||||||
|
/* ===== 面板 / 表格 ===== */
|
||||||
|
.panel {
|
||||||
|
background: var(--surface);
|
||||||
|
border: 1px solid var(--border);
|
||||||
|
border-radius: var(--r-lg);
|
||||||
|
margin-bottom: 22px;
|
||||||
|
overflow: hidden;
|
||||||
|
box-shadow: var(--shadow-sm);
|
||||||
|
transition: background-color .25s ease, border-color .25s ease;
|
||||||
|
}
|
||||||
|
|
||||||
|
.section-head { display: flex; align-items: center; justify-content: space-between; gap: 15px; padding-right: 22px; flex-wrap: wrap; }
|
||||||
|
.section-head > small { padding-bottom: 6px; }
|
||||||
|
|
||||||
|
.empty { text-align: center; padding: 52px 20px; color: var(--text-3); }
|
||||||
|
|
||||||
|
/* 指标卡 */
|
||||||
|
.metrics { display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 14px; margin-bottom: 14px; }
|
||||||
|
.metric {
|
||||||
|
display: grid;
|
||||||
|
grid-template-columns: minmax(0, 1fr) auto;
|
||||||
|
align-content: start;
|
||||||
|
gap: 10px 12px;
|
||||||
|
padding: 18px 20px;
|
||||||
|
background: var(--surface);
|
||||||
|
border: 1px solid var(--border);
|
||||||
|
border-radius: var(--r-lg);
|
||||||
|
box-shadow: var(--shadow-sm);
|
||||||
|
transition: transform .18s ease, box-shadow .18s ease, background-color .25s ease, border-color .25s ease;
|
||||||
|
}
|
||||||
|
.metric:hover { transform: translateY(-2px); box-shadow: var(--shadow-md); }
|
||||||
|
.metric > span { color: var(--text-2); font-size: 12.5px; }
|
||||||
|
.metric > .metric-icon {
|
||||||
|
grid-column: 2;
|
||||||
|
grid-row: 1 / 3;
|
||||||
|
align-self: center;
|
||||||
|
display: grid;
|
||||||
|
place-items: center;
|
||||||
|
width: 36px;
|
||||||
|
height: 36px;
|
||||||
|
border-radius: 10px;
|
||||||
|
background: var(--accent-soft);
|
||||||
|
color: var(--accent);
|
||||||
|
}
|
||||||
|
.metric > .metric-icon svg { width: 18px; height: 18px; }
|
||||||
|
.metric strong { grid-column: 1; display: block; font-size: 27px; font-weight: 650; font-variant-numeric: tabular-nums; letter-spacing: -.5px; line-height: 1.15; }
|
||||||
|
.metric.alert { border-color: var(--danger); }
|
||||||
|
.metric.alert strong { color: var(--danger); }
|
||||||
|
.metric.alert > .metric-icon { background: var(--danger-bg); color: var(--danger-text); }
|
||||||
|
|
||||||
|
.table-wrap { overflow-x: auto; }
|
||||||
|
table { width: 100%; border-collapse: collapse; text-align: left; font-size: 13px; }
|
||||||
|
th { font-weight: 600; font-size: 12px; letter-spacing: .3px; color: var(--text-3); background: var(--surface-2); white-space: nowrap; transition: background-color .25s ease; }
|
||||||
|
th, td { padding: 13px 22px; border-top: 1px solid var(--border); vertical-align: top; }
|
||||||
|
thead th { border-top: 0; }
|
||||||
|
tbody tr { transition: background-color .12s ease; }
|
||||||
|
tbody tr:hover { background: var(--surface-2); }
|
||||||
|
td { max-width: 320px; overflow-wrap: anywhere; font-variant-numeric: tabular-nums; }
|
||||||
|
td strong { font-weight: 600; }
|
||||||
|
td.empty:hover { background: transparent; }
|
||||||
|
.nowrap { white-space: nowrap; }
|
||||||
|
|
||||||
|
.actions { display: flex; gap: 7px; }
|
||||||
|
.actions button { padding: 5px 10px; font-size: 12.5px; font-weight: 500; }
|
||||||
|
|
||||||
|
/* ===== 徽章 / 提示条 ===== */
|
||||||
|
.badge { display: inline-flex; align-items: center; gap: 6px; padding: 3px 10px; border-radius: 999px; white-space: nowrap; font-size: 11.5px; font-weight: 600; background: var(--surface-2); color: var(--text-2); border: 1px solid var(--border); }
|
||||||
|
.badge i { width: 6px; height: 6px; border-radius: 50%; background: currentColor; }
|
||||||
|
.badge.good { background: var(--success-bg); color: var(--success-text); border-color: transparent; }
|
||||||
|
.badge.warn { background: var(--warning-bg); color: var(--warning-text); border-color: transparent; }
|
||||||
|
.badge.danger { background: var(--danger-bg); color: var(--danger-text); border-color: transparent; }
|
||||||
|
|
||||||
|
.alert { padding: 12px 15px; border-radius: var(--r-md); margin-bottom: 16px; line-height: 1.6; overflow-wrap: anywhere; font-size: 13px; border: 1px solid color-mix(in srgb, currentColor 22%, transparent); }
|
||||||
|
.alert.error { background: var(--danger-bg); color: var(--danger-text); }
|
||||||
|
.alert.success { background: var(--success-bg); color: var(--success-text); }
|
||||||
|
.alert.warning { background: var(--warning-bg); color: var(--warning-text); }
|
||||||
|
|
||||||
|
/* ===== 弹窗 ===== */
|
||||||
|
.overlay {
|
||||||
|
position: fixed;
|
||||||
|
inset: 0;
|
||||||
|
background: var(--overlay);
|
||||||
|
-webkit-backdrop-filter: blur(3px);
|
||||||
|
backdrop-filter: blur(3px);
|
||||||
|
z-index: 20;
|
||||||
|
display: flex;
|
||||||
|
justify-content: center;
|
||||||
|
align-items: center;
|
||||||
|
padding: 20px;
|
||||||
|
animation: fade-in .15s ease;
|
||||||
|
}
|
||||||
|
.dialog {
|
||||||
|
background: var(--surface);
|
||||||
|
border: 1px solid var(--border);
|
||||||
|
border-radius: 18px;
|
||||||
|
padding: 24px 26px;
|
||||||
|
width: 560px;
|
||||||
|
max-width: 100%;
|
||||||
|
max-height: 90vh;
|
||||||
|
overflow: auto;
|
||||||
|
box-shadow: var(--shadow-lg);
|
||||||
|
animation: slide-up .18s ease;
|
||||||
|
}
|
||||||
|
.dialog-head { display: flex; align-items: center; justify-content: space-between; gap: 12px; margin-bottom: 18px; }
|
||||||
|
.dialog-head h2 { padding: 0; font-size: 18px; }
|
||||||
|
.dialog-actions { display: flex; justify-content: flex-end; gap: 10px; margin-top: 22px; }
|
||||||
|
|
||||||
|
@keyframes fade-in { 0% { opacity: 0; } }
|
||||||
|
@keyframes slide-up { 0% { opacity: 0; transform: translateY(10px) scale(.985); } }
|
||||||
|
|
||||||
|
/* ===== 渠道编辑 ===== */
|
||||||
|
.channel-models { margin: 20px 0; }
|
||||||
|
.channel-models .section-head { flex-wrap: wrap; gap: 10px; }
|
||||||
|
.channel-model-list { max-height: 320px; overflow-y: auto; margin: 12px 0; }
|
||||||
|
.channel-model-row { display: grid; grid-template-columns: auto minmax(0, 1fr) minmax(0, 1fr) auto; align-items: end; gap: 8px; padding: 10px 0; border-bottom: 1px solid var(--border); }
|
||||||
|
.channel-model-row label { margin: 0; min-width: 0; }
|
||||||
|
.channel-model-row .check { align-self: center; }
|
||||||
|
@media (max-width: 600px) {
|
||||||
|
.channel-model-row { grid-template-columns: minmax(0, 1fr) minmax(0, 1fr); }
|
||||||
|
.channel-model-row .check { grid-column: 1 / -1; }
|
||||||
|
}
|
||||||
|
|
||||||
|
.channel-options { background: var(--surface-2); border: 1px solid var(--border); padding: 13px 15px; border-radius: var(--r-md); max-height: 220px; overflow: auto; margin-bottom: 15px; transition: background-color .25s ease; }
|
||||||
|
.channel-options p { margin: 0 0 8px; font-size: 12.5px; color: var(--text-2); }
|
||||||
|
.channel-options .check { margin-bottom: 8px; }
|
||||||
|
.channel-options .check:last-child { margin-bottom: 0; }
|
||||||
|
|
||||||
|
/* ===== 筛选栏 ===== */
|
||||||
|
.filters { padding: 18px 20px 14px; display: flex; flex-wrap: wrap; align-items: flex-end; gap: 12px; }
|
||||||
|
.filters label { margin: 0; flex: 1; min-width: 170px; }
|
||||||
|
.filters-actions { display: flex; gap: 8px; }
|
||||||
|
.quick { display: flex; gap: 8px; align-items: center; width: 100%; border-top: 1px dashed var(--border); padding-top: 12px; margin-top: 2px; }
|
||||||
|
.quick button { padding: 5px 14px; font-size: 12px; border-radius: 999px; }
|
||||||
|
.quick small { margin-left: auto; }
|
||||||
|
|
||||||
|
/* ===== 图表 ===== */
|
||||||
|
.chart { padding: 0 22px 6px; }
|
||||||
|
.chart-empty { padding: 36px 0 30px; text-align: center; color: var(--text-3); font-size: 12.5px; }
|
||||||
|
.chart-plot { position: relative; height: 220px; margin-left: 42px; }
|
||||||
|
.chart-grid { position: absolute; left: 0; right: 0; border-top: 1px dashed var(--border); }
|
||||||
|
.chart-grid span { position: absolute; right: 100%; margin-right: 8px; top: -7px; font-size: 11px; color: var(--text-3); font-variant-numeric: tabular-nums; }
|
||||||
|
.chart-bars { position: absolute; left: 0; right: 0; top: 0; bottom: 26px; display: flex; gap: 2px; border-bottom: 1px solid var(--border); }
|
||||||
|
.chart-col { flex: 1; position: relative; display: flex; min-width: 0; }
|
||||||
|
.chart-bar { flex: 1; max-width: 46px; margin: 0 auto; display: flex; flex-direction: column; justify-content: flex-end; }
|
||||||
|
.chart-bar i { display: block; width: 100%; opacity: .92; transition: opacity .12s ease; }
|
||||||
|
.chart-col:hover .chart-bar i { opacity: 1; }
|
||||||
|
/* 柱顶圆角：output 在栈顶自带圆角；input 顶部圆角并通过 -4px 重叠遮住接缝 */
|
||||||
|
.chart-bar .seg-input { background: linear-gradient(180deg, color-mix(in srgb, var(--chart-input) 76%, #ffffff), var(--chart-input)); border-radius: 4px 4px 0 0; margin-top: -4px; }
|
||||||
|
.chart-bar .seg-output { background: linear-gradient(180deg, color-mix(in srgb, var(--chart-output) 76%, #ffffff), var(--chart-output)); border-radius: 4px 4px 0 0; }
|
||||||
|
.chart-x { position: absolute; top: 100%; margin-top: 8px; left: -14px; right: -14px; text-align: center; font-size: 11px; color: var(--text-3); white-space: nowrap; }
|
||||||
|
.chart-legend { display: flex; gap: 18px; padding: 34px 0 8px; font-size: 12px; color: var(--text-2); }
|
||||||
|
.chart-legend span { display: inline-flex; align-items: center; gap: 7px; }
|
||||||
|
.chart-legend i { width: 10px; height: 10px; border-radius: 3px; }
|
||||||
|
.seg-input { background: var(--chart-input); }
|
||||||
|
.seg-output { background: var(--chart-output); }
|
||||||
|
|
||||||
|
/* ===== 登录页 ===== */
|
||||||
|
.login {
|
||||||
|
min-height: 100vh;
|
||||||
|
display: grid;
|
||||||
|
place-items: center;
|
||||||
|
padding: 24px;
|
||||||
|
background:
|
||||||
|
radial-gradient(1000px 500px at 15% -10%, var(--accent-soft), transparent 60%),
|
||||||
|
radial-gradient(800px 400px at 100% 110%, var(--accent-soft), transparent 55%),
|
||||||
|
var(--bg);
|
||||||
|
transition: background-color .25s ease;
|
||||||
|
}
|
||||||
|
.login-card {
|
||||||
|
background: var(--surface);
|
||||||
|
width: 400px;
|
||||||
|
max-width: 100%;
|
||||||
|
padding: 40px 38px 34px;
|
||||||
|
border: 1px solid var(--border);
|
||||||
|
border-radius: 20px;
|
||||||
|
box-shadow: var(--shadow-md);
|
||||||
|
transition: background-color .25s ease, border-color .25s ease;
|
||||||
|
}
|
||||||
|
.login-card .mark { margin-bottom: 22px; width: 46px; height: 46px; border-radius: 13px; }
|
||||||
|
.login-card h1 { margin: 0 0 6px; font-size: 23px; }
|
||||||
|
.login-card > .muted { margin: 0 0 22px; }
|
||||||
|
.login-card .primary { width: 100%; margin-top: 6px; padding: 10px; }
|
||||||
|
|
||||||
|
/* ===== 响应式 ===== */
|
||||||
|
@media (max-width: 1100px) {
|
||||||
|
main { padding: 24px 22px 32px; }
|
||||||
|
aside { width: 198px; padding: 22px 12px 16px; }
|
||||||
|
.metrics { gap: 10px; }
|
||||||
|
.metric { padding: 15px 16px; }
|
||||||
|
.metric strong { font-size: 23px; }
|
||||||
|
th, td { padding: 11px 14px; }
|
||||||
|
.filters label { min-width: 150px; }
|
||||||
|
}
|
||||||
|
|
||||||
|
@media (max-width: 760px) {
|
||||||
|
.shell { display: block; }
|
||||||
|
aside { position: static; width: 100%; height: auto; padding: 16px; background-image: none; }
|
||||||
|
.brand { margin-bottom: 14px; font-size: 16px; }
|
||||||
|
.brand small { display: none; }
|
||||||
|
.mark { width: 34px; height: 34px; border-radius: 9px; }
|
||||||
|
nav { display: flex; gap: 6px; }
|
||||||
|
nav button { flex: 1; justify-content: center; padding: 9px 6px; }
|
||||||
|
nav button.active::before { display: none; }
|
||||||
|
.aside-bottom { margin-top: 12px; display: flex; gap: 8px; border-top: 0; padding-top: 0; }
|
||||||
|
.aside-bottom button { flex: 1; justify-content: center; padding: 7px 8px; font-size: 12.5px; }
|
||||||
|
.aside-bottom .theme-seg { flex: 1; }
|
||||||
|
main { padding: 20px 14px 28px; }
|
||||||
|
h1 { font-size: 22px; }
|
||||||
|
.metrics { grid-template-columns: repeat(2, minmax(0, 1fr)); }
|
||||||
|
.filters { padding: 14px; gap: 10px; }
|
||||||
|
.filters label { min-width: 100%; }
|
||||||
|
.filters-actions { width: 100%; }
|
||||||
|
.filters-actions button { flex: 1; }
|
||||||
|
.section-head { padding-right: 14px; }
|
||||||
|
/* 模型页头部：标题与按钮同行，说明文字独占一行沉底 */
|
||||||
|
.models-head > small { order: 3; width: 100%; }
|
||||||
|
.models-head .actions { margin-left: auto; }
|
||||||
|
/* 窄屏下保持列宽、横向滚动，避免单元格被挤压断词 */
|
||||||
|
.table-wrap table { min-width: 620px; }
|
||||||
|
/* 首列长标识（Key 名/模型名）截断省略，防止被挤成逐字竖排；全名见 title */
|
||||||
|
.cell-id { display: inline-block; max-width: 148px; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
|
||||||
|
th, td { padding: 11px 16px; }
|
||||||
|
h2 { padding: 16px 14px 12px; }
|
||||||
|
.chart { padding: 0 14px 4px; }
|
||||||
|
.chart-plot { margin-left: 34px; height: 180px; }
|
||||||
|
.dialog { padding: 20px; }
|
||||||
|
.overlay { padding: 12px; }
|
||||||
|
.login-card { padding: 30px 24px 26px; }
|
||||||
|
}
|
||||||
|
|
||||||
|
@media (prefers-reduced-motion: reduce) {
|
||||||
|
*, *::before, *::after { animation: none !important; transition: none !important; }
|
||||||
|
}
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
export type Theme = 'light' | 'dark' | 'system'
|
||||||
|
|
||||||
|
const STORAGE_KEY = 'theme'
|
||||||
|
|
||||||
|
/** 读取初始主题：优先用户保存的选择，否则跟随系统偏好。 */
|
||||||
|
export function initialTheme(): Theme {
|
||||||
|
try {
|
||||||
|
const saved = localStorage.getItem(STORAGE_KEY)
|
||||||
|
if (saved === 'light' || saved === 'dark' || saved === 'system') return saved
|
||||||
|
} catch {
|
||||||
|
// localStorage 不可用（隐私模式等），回退到系统偏好
|
||||||
|
}
|
||||||
|
return 'system'
|
||||||
|
}
|
||||||
|
|
||||||
|
/** 应用主题到 <html> 并持久化；同步移动端浏览器 chrome 颜色。 */
|
||||||
|
export function applyTheme(theme: Theme) {
|
||||||
|
const media = window.matchMedia('(prefers-color-scheme: dark)')
|
||||||
|
const update = () => {
|
||||||
|
const resolved = theme === 'system' ? (media.matches ? 'dark' : 'light') : theme
|
||||||
|
document.documentElement.dataset.theme = resolved
|
||||||
|
document.documentElement.style.colorScheme = resolved
|
||||||
|
const meta = document.querySelector('meta[name="theme-color"]')
|
||||||
|
if (meta) meta.setAttribute('content', resolved === 'dark' ? '#090a16' : '#181a34')
|
||||||
|
}
|
||||||
|
update()
|
||||||
|
if (theme === 'system') media.addEventListener('change', update)
|
||||||
|
try {
|
||||||
|
localStorage.setItem(STORAGE_KEY, theme)
|
||||||
|
} catch {
|
||||||
|
// 忽略持久化失败，主题仍在本次会话内生效
|
||||||
|
}
|
||||||
|
return () => { if (theme === 'system') media.removeEventListener('change', update) }
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
export function formatTime(value: string): string {
|
||||||
|
return !value || value.startsWith('0001-') ? '—' : new Date(value).toLocaleString('zh-CN')
|
||||||
|
}
|
||||||
|
|
||||||
|
function localInput(d: Date): string {
|
||||||
|
return new Date(d.getTime() - d.getTimezoneOffset() * 60000).toISOString().slice(0, 16)
|
||||||
|
}
|
||||||
|
|
||||||
|
export function localToUTC(value: string): string {
|
||||||
|
const d = new Date(value)
|
||||||
|
if (!/^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}$/.test(value) || !Number.isFinite(d.getTime()) || localInput(d) !== value) throw new Error('无效的本地时间。')
|
||||||
|
return d.toISOString()
|
||||||
|
}
|
||||||
|
|
||||||
|
export function quickRange(days: number, now = Date.now()) {
|
||||||
|
const to = new Date(now), from = new Date(now)
|
||||||
|
from.setHours(0, 0, 0, 0)
|
||||||
|
from.setDate(from.getDate() - days + 1)
|
||||||
|
return { from: from.toISOString(), to: to.toISOString(), fromInput: localInput(from), toInput: localInput(to) }
|
||||||
|
}
|
||||||
|
|
||||||
|
export function filterParams(filter: { key_id: string; from: string; to: string }) {
|
||||||
|
const from = filter.from ? localToUTC(filter.from) : ''
|
||||||
|
const to = filter.to ? localToUTC(filter.to) : ''
|
||||||
|
if (from && to && from > to) throw new Error('开始时间不能晚于结束时间。')
|
||||||
|
const params = new URLSearchParams()
|
||||||
|
if (filter.key_id) params.set('key_id', filter.key_id)
|
||||||
|
if (from) params.set('from', from)
|
||||||
|
if (to) params.set('to', to)
|
||||||
|
return params
|
||||||
|
}
|
||||||
|
|
||||||
|
export function keyExpiry(value: string, now = Date.now()) {
|
||||||
|
if (!value) return ''
|
||||||
|
const expires = localToUTC(value)
|
||||||
|
if (new Date(expires).getTime() <= now) throw new Error('有效期必须晚于当前时间。')
|
||||||
|
return expires
|
||||||
|
}
|
||||||
@@ -0,0 +1,4 @@
|
|||||||
|
{
|
||||||
|
"compilerOptions": { "target": "ES2022", "lib": ["ES2022", "DOM", "DOM.Iterable"], "module": "ESNext", "moduleResolution": "Bundler", "jsx": "react-jsx", "strict": true, "skipLibCheck": true, "noEmit": true, "allowImportingTsExtensions": true },
|
||||||
|
"include": ["src", "vite.config.ts"]
|
||||||
|
}
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
import { defineConfig, type ProxyOptions } from 'vite'
|
||||||
|
import react from '@vitejs/plugin-react'
|
||||||
|
const proxy: ProxyOptions = {
|
||||||
|
target: 'http://127.0.0.1:8080',
|
||||||
|
changeOrigin: true,
|
||||||
|
configure(server) {
|
||||||
|
server.on('proxyReq', request => {
|
||||||
|
if (request.getHeader('origin')) request.setHeader('origin', 'http://127.0.0.1:8080')
|
||||||
|
})
|
||||||
|
},
|
||||||
|
}
|
||||||
|
export default defineConfig({ plugins: [react()], server: { proxy: { '/admin': proxy, '/v1': proxy } } })
|
||||||
Reference in New Issue
Block a user