Compare commits

..
Author SHA1 Message Date
wangbo cfa44b5f28 fix(nginx): 移除已退役香港上游
生产香港旧节点已经离线,固定上游地址会让请求等待连接超时后才回退宁波。\n\n移除失效的 API 与 Web NodePort,并补充仅配置已部署站点入口的说明。\n\n验证:生产 nginx -t 通过;本地配置与线上 active 配置 SHA-256 一致。
2026-08-06 12:01:02 +08:00
wangbo 335b21d1f8 fix(routing): 避免路由探测耗尽数据库连接
将跨实例探测协调从持有会话级 advisory lock 改为带过期时间的数据库短租约,避免网络探测期间占用关键 PostgreSQL 连接。新增可配置的探测并发上限和真实 PostgreSQL 回归测试,确保租约生效期间连接池仍可全部获取。\n\n验证:go test ./... -count=1;go vet ./...;真实 PostgreSQL 17 集成测试;迁移安全检查。
2026-08-05 23:13:32 +08:00
wangbo 21f72da7a8 fix(routing): 修正 Worker 心跳时间查询
生产 shadow 发布证明 PostgreSQL 将未显式定型的时间参数推断为 interval,导致 Worker 列表和容量查询返回 SQLSTATE 42883。

改为在 Go 中计算心跳截止时间并以 timestamptz 参数查询,补充真实 PostgreSQL 集成回归。
2026-08-05 22:57:33 +08:00
wangbo 7786692d32 feat(routing): 引入多执行池智能调度
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口,新增 Kubernetes 和静态容量适配器,并以 shadow 模式接入生产配置。

实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌,以及提交状态不明时禁止重复分配的安全语义。

新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
2026-08-05 22:25:37 +08:00
wangbo 03c0873649 fix(metrics): 补齐媒体结果存储读取与兼容响应指标
增加对象存储 GET、结果 URL 签名和同步 Base64 限流指标,使生产轮询可直接验证零对象下载,并区分签名失败、容量等待超时和 20MiB 超限。\n\n验证:API 全量 go test、go vet、gofmt 与 git diff --check。
2026-08-05 18:55:41 +08:00
wangbo b13392ef50 fix(media): 统一图片结果 URL 化并限制同步 Base64
将上游 URL 直接持久化,内联媒体经对象存储后仅保留 URL 与内部定位元数据;异步轮询、任务详情和幂等重放统一使用零对象读取的 URL 投影,并增加 64KiB 响应门禁。

OpenAI 图片接口接受 url 与 b64_json,同步 Base64 限制为 20MiB 和每 Pod 2 并发;新增历史结果迁移清零门禁、结果指标和 API GOMEMLIMIT。

验证:API go test ./...、go vet、聚焦 race、pnpm openapi、pnpm lint/test/build、迁移安全检查与 docker compose config 均通过。
2026-08-05 18:15:06 +08:00
wangbo f9b945e4aa fix(provider): 统一响应体超时错误分类
将收到响应头后在读取 JSON 或 Veo 视频响应体阶段发生的 deadline/网络超时统一归类为 terminal timeout,保留 HTTP 状态、request ID 和 wire 证据,避免继续返回 response_read_error。

验证:go test ./... -count=1;go test ./internal/clients ./internal/runner -count=1;go vet ./internal/clients ./internal/runner;gofmt。
2026-08-05 14:35:54 +08:00
wangbo dd10a807df fix(provider): 延长媒体超时并终止超时轮转
将图像和视频的默认 HTTP/轮询超时分别提高到 20 分钟和 30 分钟。媒体任务超时后直接失败,不再重试、切换客户端或标记为 upstream_submission_unknown。OpenAI 图像端点遇到上游明确拒绝 response_format 时自动移除并缓存兼容结论。

验证:go test ./... -count=1;go vet ./...;gofmt;相关 Shell bash -n、ShellCheck 与发布脚本回归测试。
2026-08-05 13:58:02 +08:00
wangbo 9ce9053d7b fix(worker): 记住轮询间隔内的容量需求
异步视频任务在 River 轮询空档会释放本地 lease,采样瞬间可能观察不到已经触顶的领取需求,导致自适应容量长期停在初始值。本提交记录采样窗口内出现过的容量饱和,同时保持空闲 Worker 不盲目升容。\n\n移植时保留主线现有的单节点容量阶梯、请求下限和严格失败门禁,仅合入失败路径先记录并排空任务与回调、再删除模拟器资源的清理顺序,并补充脚本守卫。\n\n验证:go test ./... -count=1;go vet ./...;bash -n;ShellCheck;production-acceptance-script-test.sh。
2026-08-05 00:53:06 +08:00
wangbo ebdb96e7d7 fix(provider): 修正媒体请求转换与上游错误透传
按上游协议能力延迟处理媒体资源:OpenAI 兼容平台默认使用 multipart,显式配置后才发送 JSON URL;Gemini 官方协议使用 Files API,兼容协议使用内嵌 Base64,并同步覆盖相关媒体客户端。\n\n安全的上游 400/422 原始错误会作为下游 message 返回,同时保留结构化诊断信息和历史任务兼容。\n\n验证:API 全量无缓存测试、go vet、pnpm lint、pnpm test、pnpm build、pnpm openapi、git diff --check。
2026-08-05 00:40:42 +08:00
wangbo c79c2a7b44 fix(cluster): 兼容 CNPG 单实例预检
发布前先通过 merge patch 删除旧同步复制配置,期望清单不再写入 CRD 不接受的 null 值,确保 server-side dry-run 与正式 apply 使用相同的有效单实例状态。
2026-08-04 23:29:48 +08:00
wangbo 8ef9de0b06 fix(cluster): 收敛宁波单实例数据库拓扑
将生产 CNPG 固定为宁波单实例,并显式关闭旧同步复制配置。发布流程在拓扑收敛前校验主库位置并创建 OSS 备份,应用后验证单实例、无复制连接和主库可用。同步更新生产监控、集群验收和发布脚本回归测试。
2026-08-04 23:24:37 +08:00
wangbo 4b639df30a fix(worker): 香港站点直连 Gemini 官方平台
官方 Gemini 平台配置的共享 HTTP Proxy 会拒绝香港出口的 CONNECT 请求,导致香港 Worker 的真实 VEO 与图片任务进入 upstream_submission_unknown。

新增仅按平台 UUID 生效的代理直连白名单,并只在香港 Worker 为官方 Gemini 平台启用;宁波与其他平台继续使用原代理策略。

验证:API 全量 go test、代理路由单测、bash -n、ShellCheck、cluster release helper、manual release 与差异检查均通过。
2026-08-04 21:40:32 +08:00
wangbo 60d16b11ba fix(capacity): 忽略未就绪的 Worker 节点
香港站点保留了已离线的旧 control-plane 节点标签,容量控制器仍尝试读取其 metrics,导致 404 并阻断 readiness。容量采集现在只纳入 Ready=True 的候选节点,实际 Worker Pod 所在的就绪节点继续参与容量计算。

验证:容量控制器聚焦测试与 API 全量 go test 均通过,并覆盖同站点 NotReady 旧节点与 Ready Worker 节点并存场景。
2026-08-04 20:54:25 +08:00
wangbo a5dd7f36f5 fix(release): 补齐香港 Worker 权限与精确回滚
生产容量控制器原 Role 仅允许访问宁波 Worker,香港 Worker 启用后导致 reconciliation 返回 403。补齐香港 Deployment 及 scale 子资源权限。

同时将失败发布的 Deployment 恢复从合并式 apply 改为基于资源版本的精确 replace,并优先删除发布前不存在的站点,避免新增环境变量残留导致回滚卡死。

验证:bash -n、ShellCheck、cluster release helper、manual release 与差异检查均通过。
2026-08-04 20:49:04 +08:00
wangbo 0b2bd1d512 feat(release): 固定启用香港异步 Worker
将生产 Worker 拓扑固定为宁波 1 个、香港 1 个,并让发布、容量配置和回滚流程持续保留香港 Worker。当前仍只使用宁波 PostgreSQL,关闭横向自动扩缩容,避免依赖数据库同步副本。

验证:API 全量 go test、迁移安全检查、cluster release helper、manual release、bash -n、ShellCheck、Kustomize 渲染均通过。
2026-08-04 20:33:06 +08:00
wangbo fe8dcb40ca feat(openai): 完善 Chat 与 Responses 参数转发
原生 Chat/Responses 改为透明转发,保留标准工具结构并保护调用方显式参数。补齐 Responses 到 Chat 的兼容转换、协议路由边界、完整响应和流式事件,并同步更新 Swagger、回归测试与真实验收脚本。

验证:
- cd apps/api && env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1
- pnpm openapi
- pnpm lint
- pnpm test
- pnpm build
- gofmt -l 无输出
- git diff --check 通过

风险:
- Chat 回退无法等价表达的 Responses 原生能力现在会返回 unsupported_response_parameter
- 真实供应商 E2E 因本地没有已启用的平台模型候选而未完成
2026-08-04 19:26:48 +08:00
wangbo b2c9b4f6d9 test(runner): 补齐延迟 429 流式轮转回归
验证上游延迟返回 429 时,在首个 delta 前仍分类为可重试限流并冷却轮转到第二候选,响应不混入失败候选内容。

验证:完整 Go 单元测试通过;隔离 PostgreSQL 下新增 HTTP acceptance 场景通过;gofmt 与差异检查通过。
2026-08-04 19:26:48 +08:00
wangbo 484b05d005 fix(queue): 回收终态任务的孤儿 River Job
原因:Worker 滚动切换后,执行租约恢复会终态化提交结果不确定的任务,但失活 Worker 持有的 River Job 仍需等待一小时通用救援窗口。\n\n影响:仅在任务已终态、执行租约清空、River owner 无活跃 Worker 且超过 45 秒时,将孤儿 Job 幂等标记为 completed;活跃 owner 不受影响。\n\n验证:PostgreSQL 集成测试覆盖失活 owner、活跃 owner 与重复回收;API 全量 go test 和 go vet 通过。
2026-08-04 18:51:47 +08:00
wangbo 9bcfec890b fix(release): 默认使用 Kubernetes 生产通道
原因:旧 Compose 生产环境已退役,默认通道仍指向 Compose 会在每次发布时先命中无效数据库。\n\n影响:发布和部署脚本默认读取 Kubernetes release helper;仍可通过显式 AI_GATEWAY_DEPLOY_MODE=compose 操作遗留环境。\n\n验证:bash -n、ShellCheck 和 manual-release-test 全部通过。
2026-08-04 18:44:19 +08:00
wangbo f1402309c8 chore(release): 对齐 Kubernetes 生产基线 2026-08-04 18:31:10 +08:00
wangbo 0d9d773ce8 chore(release): 刷新队列修复发布基线
线上在镜像发布期间回退到较早的已知 manifest,原发布产物的 production base 已失效。创建新的完整 SHA 以保持镜像 Tag 不可变,并让正式发布脚本基于当前线上版本重新生成可部署 manifest。
2026-08-04 18:26:44 +08:00
wangbo f4214dd489 fix(queue): 隔离异步准入故障并持久化退避
将 X-Async 调度改为逐任务阻塞协议,区分全局容量、平台容量、用户组 FIFO、任务级异常和系统级故障,避免单个历史毒任务触发整批回滚。\n\n新增持久化退避、单任务 CAS 重选、幂等终态事务、固定标签指标和损坏快照自愈,并修复 Worker 从 preparing 直接进入 finalizing 时的自等待死锁。\n\n验证:API 全量测试、PostgreSQL 集成场景、race、go vet、govulncheck、pnpm lint/test/build、Compose 与发布脚本测试通过;pnpm audit 命中未改动的 Nx 工具链既有漏洞。
2026-08-04 18:23:35 +08:00
wangbo 44d5cf2b9d feat(gemini): 接入 Veo 官方异步生成协议
将 Gemini Veo 视频请求改为 predictLongRunning 提交与 Operation 轮询,支持任务恢复、首尾帧和参考图的官方 Base64 字段映射。

完成后由网关携带 API Key 从受信任的 Gemini 文件域下载视频,再进入现有对象存储转存链路,避免暴露鉴权地址。

验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;真实 Veo 3.1 任务生成并下载 1633544 字节视频。
2026-08-04 16:27:24 +08:00
wangbo 5c7d6ac9aa fix(runner): 修复转存后残留二进制误判
统一结果二进制探测与上传规则,避免将 thinking_bytes 等上游元数据误判为待转存媒体,同时保留显式媒体字段和签名识别。\n\n补充不含原始内容的安全诊断,并兼容对象存储前缀旧字段;修正真实 OSS 验收脚本使用的正式字段。\n\n验证:Go 全量测试、go vet、迁移安全检查、真实 Gemini 上游响应及阿里云 OSS 转存均通过。
2026-08-04 14:33:04 +08:00
wangbo 394cc6288f fix(runner): 修复图像结果转存与任务租约失效
修正非语义字段长 Base64 元数据被误判为媒体的问题,并允许按文件签名识别真实媒体后转存 OSS;同时将网关本地存储失败与上游失败分离,避免重复调用和错误降权。

异步任务复用已分配并发名额前立即刷新租约,租约已丢失时重新排队获取,避免排队耗时侵蚀租约后错误调用上游。生产环境关闭 server-main 尚未实现的进度回调,保留 Gateway 任务详情与事件查询。

验证:API 全量 go test、go vet、gofmt、Kubernetes 渲染、迁移安全检查及真实阿里云 OSS 读写/生命周期验收通过。
2026-08-04 13:48:44 +08:00
wangbo 6a64d3936a feat(storage): 完善对象存储配置与过期策略
补齐 OSS/S3 的 Endpoint、Region、Bucket、CDN、对象前缀和签名有效期配置,并为生成结果与请求素材自动维护分级生命周期规则。普通上传继续保持永久,私有资源按配置生成限时签名 URL,管理端连接测试覆盖生命周期、上传、读取和删除。\n\n新增可重复的真实 OSS 验收脚本,凭据仅从本地环境读取,接口响应继续保持脱敏。\n\n验证:Go 全量测试、迁移安全检查、pnpm lint、pnpm test、pnpm build、本地阿里云 OSS 真实上传下载删除验收。
2026-08-04 12:44:24 +08:00
wangbo fe56aa46b9 fix(errors): 区分平台限流并保留上游状态码
原因:公开错误层将平台并发限流误标为上游限流,并把多种上游 4xx 统一压成 400,影响定位和客户端处理。

影响:新增公开错误 source,平台限流使用 gateway_rate_limited,上游请求按安全分类返回对应状态;数据库与管理端继续保留原始错误码、消息和状态用于审计。

验证:Go 全量测试、pnpm test、pnpm lint、pnpm build、pnpm openapi、gofmt 和 diff 检查均通过。
2026-08-04 10:46:46 +08:00
wangbo 0f0998cbcf feat(storage): 统一二进制对象存储与公开错误
新增 Aliyun OSS 与 S3 协议、通道内重试和按优先级跨通道切换,保留 server-main 兼容与环境 OSS 内存通道。

将请求及结果中的 Base64、Data URI、Buffer、multipart 和内联二进制统一对象化,生产路径不再写入本机静态目录,历史本地资源仅保留只读兼容。

引入 PublicErrorV1 并统一 API、异步查询、兼容协议和失败回调的安全错误输出,同时补充迁移、管理端、指标、OpenAPI 与本地模拟验收。

验证:go test ./... -count=1;go vet ./...;pnpm lint;pnpm test;pnpm build;pnpm openapi;tests/ci/migrations-test.sh。
2026-08-04 08:14:39 +08:00
wangbo d129bcccbd fix(runner): 保留同平台重试的并发租约
同平台 retry_same 之前会在首个 attempt 失败时释放任务级 admission 租约,下一次 attempt 继续续租旧租约并进入 upstream_submission_unknown。

现在明确区分 attempt 自有租约与任务级 admission 租约,只在 attempt 结束时释放前者;补充回归测试覆盖同平台重试继续续租 admission 租约。

验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;gofmt;git diff --check。
2026-08-03 21:06:23 +08:00
wangbo 3af6655e86 fix(acceptance): 严格校验单节点诊断终态
修复视频容量阶梯失败后仍被标记为完成的问题,并将低并发视频任务量按槽位缩放,避免验收负载本身超过 admission 等待窗口。

增加任务终态总量校验,确保任一失败、排队或运行中任务都会使诊断返回失败;已执行 bash -n、ShellCheck -x 与 diff check。
2026-08-03 19:19:29 +08:00
wangbo eccfe568ad test(acceptance): 区分路由尝试与上游提交
额度原子抢占失败会留下 quota_race_rotated 尝试,但不会调用上游。验收报告新增 submitted 与轮转原因,并以真实提交数校验任务去重和同优先级负载分布,避免把安全轮转误判为重复提交。\n\n验证:bash -n scripts/acceptance/provider-burst.sh;shellcheck scripts/acceptance/provider-burst.sh;使用本地三 Worker 运行数据验证 25 次路由尝试中仅 24 次提交、重复远端任务和账务均为 0。
2026-08-03 18:21:50 +08:00
wangbo c8d04ca731 fix(billing): 重排任务复用钱包预留
平台额度跨窗口重排时保留任务已有的钱包预留,后续 Worker 领取复用同一幂等键,避免重复 reserve/release 事务。\n\n补齐排队任务取消后的原子结算释放,并允许任务 attempt 持久化软负载、满载原因和选择原因等有界路由快照。\n\n验证:go test ./...;go vet ./...;钱包预留复用与排队取消 PostgreSQL 集成测试。
2026-08-03 18:09:47 +08:00
wangbo 13faf1d072 test(acceptance): 校准平台突发负载门禁
将长轮询视频的 RPM 与 TPM 额度调整为可稳定触达的窗口值,并以按并发容量归一化后的近似均衡替代脆弱的精确任务数断言。\n\n门禁继续要求高优先级平台并发、RPM、TPM 跑满,低优先级发生溢出,三类额度均不越界且队列、Worker、回调和租约最终清空。\n\n验证:bash -n scripts/acceptance/provider-burst.sh;shellcheck scripts/acceptance/provider-burst.sh。
2026-08-03 17:54:04 +08:00
wangbo 8362d6d27a feat(routing): 完善平台满载避让与故障轮转
为未配置并发上限的平台增加基于运行和等待任务数的软负载,并按非满载、有效优先级、缓存亲和与负载稳定排序。\n\n平台模型 RPM、TPM 和并发额度竞争失败时只轮转候选,不触发冷却、禁用或降级;用户组额度保持不可绕过。补齐异步冷却排队恢复、满载原因、选择原因与低基数指标。\n\n验证:go test ./...;go vet ./...;PostgreSQL 原子额度、准入队列、Worker 容量回收及故障策略 HTTP acceptance。
2026-08-03 17:42:34 +08:00
chengcheng b125599354 fix(oidc): 同步认证中心用户资料到本地用户
从已验证的 OIDC Claim 提取用户名、显示名称、邮箱、手机号和头像,并覆盖单租户、多租户及平台用户的 JIT 创建与重复登录同步。\n\n保留 metadata.manualProfile 标记下的人工资料,限制字段长度且仅接收已验证联系方式与 HTTPS 头像。已通过 auth、httpapi、store 测试及临时 PostgreSQL 集成验证。
2026-08-03 15:55:50 +08:00
chengcheng 3c8d9839a4 fix(web): 避免刷新时闪现未登录页面
统一认证会话需要异步恢复,首次渲染不能把尚未确认的状态当作未登录。新增 checking、authenticated、unauthenticated 三态鉴权,在检查期间显示中性加载状态并隐藏登录相关操作。

验证:前端 156 项测试通过;TypeScript 类型检查、lint、生产构建和浏览器刷新验收通过。
2026-08-03 15:55:49 +08:00
wangbo b56769512f feat(admin): 默认过滤实时负载禁用项
实时负载默认仅展示平台和模型均启用的来源,并增加全部状态、仅看禁用筛选及重置行为测试。
2026-08-03 15:43:49 +08:00
wangbo ad8cdd525b feat(identity): 支持用户和用户组批量管理
增加原子批量启用、禁用和删除接口及管理端多选操作,目标缺失时整批回滚。\n\n拆分管理端与 API Key 权限缓存并在弹窗保存后刷新候选;补齐失效规则一键清理样式、固定右侧操作列和 OpenAPI 契约。
2026-08-03 15:43:49 +08:00
wangbo 7376d6fab6 refactor(access): 统一分层白名单权限语义
取消跨主体专属占用,按租户、用户组、用户、当前 API Key 和 scope 分层求交,并在任务落库前统一校验候选。\n\n增加旧 allow 规则归档清理迁移、脱敏审计工具和回滚运行手册,补齐主体隔离、deny 优先及列表与运行时一致性测试。
2026-08-03 15:43:49 +08:00
wangbo c9393af43a test(acceptance): 在单窗口验证厂商额度峰值
厂商额度验收此前复用了多图与超大图容量素材,3 Worker 的下载和预处理使 24 个任务跨越多个固定分钟窗口,无法精确证明 RPM 与 TPM 上限。\n\n额度场景固定使用三张正常尺寸引用图并为每个任务增加唯一 URL 变体,使请求满足视频协议且能在单窗口完成;6、9 图和超大图转换继续由独立视频容量 profile 覆盖。\n\n验证:\n- go test ./cmd/acceptance-load ./internal/acceptanceworkload -count=1\n- go vet ./cmd/acceptance-load ./internal/acceptanceworkload\n- gofmt -l cmd/acceptance-load/main.go cmd/acceptance-load/main_test.go
2026-08-03 13:04:03 +08:00
wangbo 76bb730b13 fix(routing): 将 admission 排队计入平台负载
候选路由此前只统计 River 队列和历史尝试,未统计尚未生成 River Job 的异步 admission 等待任务,突发提交会集中绑定同一高优先级平台,导致其他平台容量闲置。\n\n将等待中的 gateway_task_admissions 纳入平台模型排队负载,并通过 task_id 去重保留原有统计语义,使候选路由可以在集群 admission 阶段按并发利用率分流。\n\n验证:\n- go test ./internal/store ./internal/runner -count=1\n- go vet ./internal/store ./internal/runner\n- gofmt -l internal/store/candidates.go
2026-08-03 12:33:49 +08:00
wangbo b933d59783 fix(routing): 额度饱和时解除任务候选固定
已取得并发租约的任务继续保持候选绑定;仅当已绑定平台的 RPM 或 TPM 已满且存在可用候选时,允许任务在执行前原子迁移到下一平台,避免等待窗口重置。\n\n验证:\n- go test ./internal/runner -count=1\n- go vet ./internal/runner
2026-08-03 12:19:52 +08:00
wangbo d740e9f676 fix(acceptance): 按隔离额度执行平台降级路由
让 acceptance 候选负载读取同一 Run ID 下的并发、RPM 和 TPM 计数器,并在单个平台模型额度不足时继续尝试下一候选,生产与 canary 额度作用域保持不变。\n\n补齐本地容量控制器对宁波和香港测试 Worker 的最小 RBAC,避免初次应用生产反亲和配置时发生滚动死锁,并将突发负载对齐到完整限流窗口。\n\n验证:\n- go test ./internal/runner ./internal/store -count=1\n- go vet ./internal/runner ./internal/store\n- bash -n scripts/acceptance/local-cluster.sh scripts/acceptance/provider-burst.sh scripts/acceptance/run-local-acceptance.sh\n- shellcheck -x scripts/acceptance/local-cluster.sh scripts/acceptance/provider-burst.sh scripts/acceptance/run-local-acceptance.sh\n- kubectl apply --dry-run=client -f deploy/kubernetes/local-acceptance/capacity-controller-rbac.yaml
2026-08-03 12:06:09 +08:00
wangbo 8485f14bc4 test(acceptance): 验证厂商三类额度与优先路由
为 provider-burst 固定三平台并发、RPM、TPM 和优先级组合,使用稳定 token 用量的本地视频负载验证同级按容量分流及低优先级溢出。\n\n报告升级为 v2,强制校验三类额度峰值精确命中上限、不越界、队列回收、三 Worker 分发以及重复提交和回调安全。\n\n验证:\n- env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1\n- go vet ./...\n- bash -n scripts/acceptance/provider-burst.sh scripts/acceptance/run-local-acceptance.sh\n- shellcheck -x scripts/acceptance/provider-burst.sh scripts/acceptance/run-local-acceptance.sh
2026-08-03 11:48:54 +08:00
wangbo aa4f8532bb test(acceptance): 扩大视频容量突发样本
将单节点视频容量阶梯的任务数提高到每槽位 6 个且最低 64 个,确保自适应容量达到目标后仍有足够排队任务用于观测真实运行峰值。\n\n已执行 bash -n、ShellCheck 和生产验收脚本测试。
2026-08-03 11:06:28 +08:00
wangbo 0d33400e09 test(acceptance): 增加视频容量阶梯场景
新增 30 至 45 秒模拟上游等待的视频容量场景,使单 Worker 自适应控制在每级突发负载中获得连续占满窗口,同时保持提交、轮询、结果处理和回调全链路。\n\n未修改生产 Worker 和厂商限流逻辑。已执行完整 Go 测试、go vet、gofmt、bash -n、ShellCheck 和生产验收脚本测试。
2026-08-03 10:45:09 +08:00
wangbo 75b3ed4748 fix(acceptance): 对齐负载报告第二版契约
生产验收统一校验并合并 acceptance-load-report/v2,避免成功负载因旧版 schemaVersion 断言被误判失败。\n\n已执行 bash -n、ShellCheck 和 production acceptance 脚本测试。
2026-08-03 10:26:05 +08:00
wangbo 2636e93353 fix(acceptance): 按自适应容量执行单节点压测
修正单节点验收将 hard limit 等同空闲分配容量的问题,从冷启动容量 4 开始逐级加压,并要求实际分配容量和运行任务峰值均达到目标后才计入稳定并发。\n\n异常退出时同步清理模拟器和回调收集器资源。已执行 bash -n、ShellCheck 和 production acceptance 脚本测试。
2026-08-03 10:20:11 +08:00
wangbo d8b86bb88c fix(deploy): 合并香港容量环境清理 2026-08-03 10:11:09 +08:00
wangbo 0e292bb1d8 fix(deploy): 清理旧香港容量环境变量
单节点升级时显式删除旧 Deployment 模板中残留的香港副本环境变量,避免 ConfigMap 已清理后容量控制器仍访问退役站点。\n\n验证:bash -n;ShellCheck;cluster release helper 测试。
2026-08-03 10:10:56 +08:00
wangbo 699d004c5b feat(deploy): 合并宁波单节点生产拓扑 2026-08-03 10:00:07 +08:00
wangbo 4ff85d96e0 feat(deploy): 切换宁波单节点生产拓扑
移除已停用的香港 API、Web、Worker 和边缘 Service 渲染结果,将 Worker 与容量控制器固定到宁波单实例,并在发布成功后清理旧香港资源。\n\n容量控制器忽略最大副本为零的站点,集群环境允许不再声明香港主机;单节点 production acceptance 不再依赖香港 Service、Deployment 或 SSH 回退。\n\n验证:Go 全量测试与 go vet;迁移安全检查;Kustomize 单节点渲染;ShellCheck;cluster release、production acceptance 和 manual release 测试。
2026-08-03 09:59:34 +08:00
wangbo 4fa1981bb8 fix(web): 按权限隐藏管理工作台入口
基于 /api/v1/me 返回的当前用户角色控制主导航,仅向具备 power 或 manager 权限的 operator、manager、admin 展示管理工作台。所有已登录页面统一加载当前用户身份,避免停留首页时无法正确判断权限。\n\n新增角色映射与导航渲染回归测试。已验证前端 152 项测试、前端构建、pnpm lint 和 git diff --cached --check。
2026-08-03 09:38:25 +08:00
wangbo 61a74445bc chore(migration): 合并 Worker 负载迁移安全审查 2026-08-03 09:23:49 +08:00
wangbo 5e9a7888ae chore(migration): 登记 Worker 负载迁移安全审查
基于宁波生产 PostgreSQL 18 的只读预检,gateway_worker_instances 当前为 96 行、180224 字节,常量默认值新增列不会触发表重写,回填和约束验证范围可控。\n\ngateway_concurrency_leases 仅增加可空列;记录迁移校验器允许的 non-null column addition,并保留迁移文件不可变。\n\n验证:node scripts/ci-validate-migrations.mjs ac8d6dcf452ac555eb344f76caf6965c2ea42a97;./tests/ci/migrations-test.sh
2026-08-03 09:23:34 +08:00
wangbo dc2605fb39 fix(dev): 等待 PostgreSQL 容器就绪
数据库初始化脚本在容器停止时主动启动,并在执行 psql 前等待 PostgreSQL 就绪;若容器提前退出或超过等待时间则返回明确错误。

验证:bash -n、ShellCheck、git diff --check 与 pnpm db:create 通过。
2026-08-03 09:18:42 +08:00
wangbo cc97e6649c fix(access): 统一 API Key 模型权限与列表契约
将全局启用、用户组基线、API Key 专属或排除规则及 scope 按固定顺序求值,避免 Key 越过所属用户组权限,并让运行时候选与模型列表共用同一权限链。

新增 Key 级可分配模型与失效规则诊断接口、OpenAI 兼容 /v1/models 及 rich 列表迁移路径;前端权限弹窗改为按当前 Key 实时加载并支持清理失效规则。

验证:Go 全量测试与 go vet 通过;Web 22 个测试文件共 142 项通过;pnpm lint、pnpm openapi、pnpm build、Compose 配置、gofmt、ShellCheck 和 git diff --check 通过;独立 PostgreSQL 真实配置验收通过。
2026-08-03 09:17:15 +08:00
wangbo c28bf74230 feat(worker): 实现集群限流与自适应负载
保留平台模型 RPM、TPM 和并发策略语义,增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载,避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标,并扩展本地 acceptance runner,覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证:go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过;48/48 成功,无越限、重复提交、重复计费、重复回调或终态资源泄漏。
2026-08-03 00:13:46 +08:00
wangbo 9a01fd4657 fix(k3s): 修正磁盘调度 unit 启动命令
systemd 会展开 unit 命令中的百分号 specifier,导致原 oneshot 启动失败。改用 echo 写入调度器,并强制校验 unit 为 active;bash -n、ShellCheck、manual-release-test.sh 通过。
2026-08-02 10:41:16 +08:00
wangbo e20688bbe2 fix(k3s): 持久化香港控制面磁盘调度缓解
原因:香港虚拟磁盘在低吞吐下仍出现高同步写延迟,触发 etcd slow apply、Handler timeout 和 PostgreSQL API Server 连通性异常。\n\n影响:新增仅针对香港根块设备的 mq-deadline 配置脚本,应用时不重启 K3s,并提供显式 rollback 恢复 none;运行手册明确该项仅为缓解,不能替代高性能云盘。\n\n验证:新香港节点同步写探针从 13.28ms 降至 6.25ms;bash -n、ShellCheck、manual-release-test.sh 通过。
2026-08-02 10:39:55 +08:00
wangbo c55e52af11 fix(acceptance): 增强生产验收控制面容错
原因:资源前置检查中的一次瞬时 SSH 中断会直接终止 Run,失败收尾又可能因 partial 报告已存在而无法完成。\n\n影响:节点内存采集增加有界重试并区分传输故障;验收管理 GET 与 finish 支持香港本地入口安全回退,并对响应丢失后的 409 反查 Run 终态;已有 partial 报告会验证复用。\n\n验证:bash -n、ShellCheck、production-acceptance-script-test.sh、manual-release-test.sh 通过。
2026-08-02 10:13:35 +08:00
wangbo ac8d6dcf45 fix(postgres): 放宽跨地域 WAL 连接抖动窗口
将生产 WAL sender/receiver timeout 从 CNPG 隐式 5 秒显式调整为 30 秒,避免短时控制面或 Pod 网络停顿反复拆除复制流并触发异步降级。保留 preferred 同步策略,防止副本不可达时冻结全部写入。\n\n影响:仅修改生产 CNPG 参数、发布门禁和运维说明,不改变应用镜像和迁移。仍需通过同步复制、六向链路与验收硬门禁。\n\n验证:kubectl kustomize、bash -n、ShellCheck、manual-release-test。
2026-08-02 00:10:25 +08:00
wangbo a4046c2b81 perf(worker): 解耦准入调度与远端执行
原因:线上 P24 同构验收确认两台香港 Worker 的全局容量为 48,但跨地域逐行准入事务每轮只能形成 8 个活跃租约,队列最老等待超过 15 分钟。\n\n影响:新增可配置的异步准入 dispatcher 角色;生产两地 API 负责准入和过期回收,Worker 仅执行 River job。当前主库同站点 API 可低延迟填满容量,主库切换后另一地 API 通过既有数据库锁安全接管;未配置环境变量时保持原 Worker 一体化行为。\n\n验证:Go 全量测试、go vet、gofmt、真实 PostgreSQL 1000 任务双进程回归、前端 lint/test/build、Kubernetes server-side dry-run、Compose、ShellCheck、cluster/manual release tests 全部通过。
2026-08-01 23:24:30 +08:00
easyai 9c093974a1 fix(gateway): 修复 Qwen3.8 推理参数并增加上游自愈
统一 Qwen3.8 固定思考、推理档位、预算互斥与温度下限规则,并覆盖 Chat Completions 和原生 Responses 请求。

新增安全参数纠正、有限重试和并发安全的进程内 LRU 学习;补充共享行为向量、httptest 回归与真实模型验收用例。

验证:go test ./... -count=1 通过,gofmt 检查通过;真实 Qwen3.8 流式、非流式及缓存重学习验收通过。
2026-08-01 22:50:45 +08:00
wangbo 50a13de70b fix(acceptance): 允许模拟器调度到专用 Worker
原因:新香港专用 Worker 节点带有 easyai.io/worker-only NoSchedule taint,现有模拟器只有节点选择器,导致线上验收滚动更新长期 Pending。\n\n影响:协议模拟器可在排除数据库与控制面节点的前提下调度到专用 Worker;正式 API、Worker 和流量门禁不变。\n\n验证:生产 API Server server-side dry-run 通过;manual-release-test 通过;git diff --check 通过。
2026-08-01 21:41:07 +08:00
wangbo 56a5176c92 perf(worker): 持久化准入快照提升队列吞吐
原因:线上 P24 验收显示 Worker 在每轮准入前逐条恢复媒体、重算候选和查询 attempts,跨地域查询导致 48 个执行槽长期只能使用约 8 至 16 个。\n\n影响:任务首次排队时保存鉴权后的 admission scope 快照,Worker 单次批量读取并只刷新动态总容量;显式重新选路保留完整慢路径。验收模拟器改到非数据库专用 Worker 节点,运行期 85% 作为立即中止硬门禁,80% 继续作为容量认证目标。\n\n验证:Go 全量测试、go vet、govulncheck、真实 PostgreSQL 迁移与跨 Store 集成测试、ShellCheck、发布及验收脚本、OpenAPI、前端 lint/test/build、pnpm audit high、Compose 与 Kubernetes 渲染均通过。
2026-08-01 21:25:54 +08:00
wangbo c89c56ca65 perf(worker): 以有界微批次提升准入吞吐
原因:线上 P24 同构验收中,单任务同步复制提交与全局容量锁串行化,使两个 Worker 的 48 个执行槽只能维持约 10–14 个运行任务,最老等待超过 15 分钟。

影响:新增可配置的 1–32 条准入微批次,默认 8;租约、任务准入与唯一 River job 在一个有界事务内原子提交,并按确定顺序预锁任务和容量范围,避免整窗 48 条大事务和多 Dispatcher 死锁。并容忍 rebind 已被其他 Dispatcher 完成的幂等竞态。

验证:Go 全量测试、go vet、真实 PostgreSQL 跨 Store 集成测试、ShellCheck、迁移安全检查、OpenAPI、前端 lint/test/build、Compose/Kubernetes 渲染及人工发布脚本均通过。
2026-08-01 20:26:47 +08:00
wangbo 92e328a575 fix(acceptance): 隔离容量压测并稳定数据库连接
将协议模拟验收的供应商并发限额与 Worker 容量门禁分离,真实金丝雀继续保留生产限流。readyz 改用关键连接池并预热关键及 River 池,延长生产连接空闲周期,降低跨地域连接抖动。失败 Run 现在可以原子替换且失败报告记录任务数量,避免 validation 之间短暂放开正式流量。\n\n验证:Go 全量测试、go vet、PostgreSQL 集成测试、ShellCheck、迁移安全、发布脚本、pnpm lint/test/build、OpenAPI 无漂移均通过。
2026-08-01 19:29:08 +08:00
wangbo 89a0ff7e99 fix(release): 同步 API 运行 revision
生产发布滚动 API 时写入当前完整 release SHA,避免镜像 digest 已更新但运行 revision 仍指向上一版本。补充发布 helper 静态回归检查。\n\n验证:bash -n、ShellCheck、cluster release helper 测试通过。
2026-08-01 18:28:21 +08:00
wangbo edcbbd6a87 fix(worker): 拆分准入事务并收敛验收失败状态
将批量准入改为每任务独立事务,避免 P24 压测时 48 条任务共享长事务造成 transactionid 锁潮与 Worker 槽位空转。失败验收在保持 validation 的同时清理未提交任务,并记录精确压力门禁原因。统一生产 0+2 Worker 拓扑、运行时 ConfigMap 与发布配置,补充镜像预检诊断和回归测试。\n\n验证:Go 全量测试、go vet、真实 PostgreSQL 48 任务集成测试、迁移安全检查、发布脚本测试、bash -n、ShellCheck、Compose 配置均通过。
2026-08-01 18:23:00 +08:00
wangbo 215b6b607f fix(acceptance): 解除八任务准入瓶颈
线上 P24 验收证明 Worker 注册容量为 48,但异步 admission 始终只有 8 条活跃租约。将单事务调度窗口覆盖当前 2 x P32 的 64 槽,避免多 Worker 在同一 FIFO 小批次上竞争形成隐性全局上限。\n\nGemini 流式请求补充可回放 GetBody 与精确 Content-Length,使带幂等键的 POST 遇到 HTTP/2 可重试传输错误时可安全重放;同时提前部署验收回调收集器,保证旧 Run side effects 能在接管前收口。\n\n已通过 Go 全量测试、go vet、gofmt、Shell bash -n、ShellCheck、验收与人工发布脚本测试、迁移安全检查,以及独立 PostgreSQL 48 条 admission/租约/River job 原子集成测试。
2026-08-01 16:59:28 +08:00
wangbo 7623449e33 fix(acceptance): 复用集群 SSH 控制连接
线上验收在宁波 SSH 建连超时时于负载前失败。为集群脚本增加受限 ControlMaster 复用和仅建连阶段的 ConnectionAttempts,避免重放远程命令。\n\n生产快照导出增加三次只读重试、原子落盘和空 Pod 清理防护。已通过 bash -n、ShellCheck、production acceptance 脚本测试、人工发布脚本测试,并验证宁波真实 SSH 控制连接可复用。
2026-08-01 15:59:00 +08:00
wangbo 8ca3e2dbc2 fix(acceptance): 实时执行压力资源硬门禁
压力采样一旦发现节点内存、数据库连接、Pod RSS、租约或排队等待越线,立即联动停止负载,避免失败档继续堆积任务。\n\n同时保证连接池零值写入报告,并增加边界回归测试。已通过 bash -n、ShellCheck、production acceptance 脚本测试和人工发布脚本测试。
2026-08-01 15:44:57 +08:00
wangbo 3514c13b0d fix(acceptance): 按实际节点聚合多 Worker 指标
同站点存在多个物理节点时,按每个 Ready Pod 的 nodeName 映射真实 SSH 主机采集 metrics,并隔离 ssh 标准输入,避免吞掉后续 Pod 目标。所有 Worker Pod 指标按容量最小/最大、连接池最小/最大和计数器总和聚合,不再随机漏采一个实例。\n\n运行时门禁根据已启用的第四节点计算 Ready 节点数,并跳过零副本站点,同时继续逐 Pod 校验 P24/P28/P32 容量、连接池和租约。\n\n验证:bash -n、ShellCheck、双节点 metrics 映射测试、生产只读双 Worker metrics 探针、manual-release-test。
2026-08-01 15:14:39 +08:00
wangbo 53fc79691a fix(acceptance): 联动压力采样与负载停止
零副本 Worker 站点没有 metrics,应从聚合采样中跳过;有副本的站点采样失败仍保持硬失败。补充缺失字段诊断,避免只输出无上下文的采样失败。\n\n容量轮次和混合长稳压测现在持续监视压力采样进程;采样器提前退出时立即终止当前负载并走失败收口,避免在资源监控失效后继续造任务。\n\n验证:bash -n、ShellCheck、零副本站点测试、采样失败联动停止测试、manual-release-test。
2026-08-01 14:37:45 +08:00
wangbo 8749f38eb3 fix(acceptance): 完善跨版本验收状态收口
生产流量处于旧失败验收 Run 的 validation 状态时,允许新验收接管任意已验证的 Git 祖先 release,而非仅允许 release manifest 的直接 base;仍保留 Run 状态、镜像 digest、revision 与当前集群 release CAS 校验。\n\n允许 pending Run 在激活前异常时直接标记为 failed,避免验收工具在 CAS 或网络错误后遗留无法结束的 Run。新增祖先关系脚本测试及临时 PostgreSQL 集成验证。\n\n验证:Go 全量测试、go vet、临时 PostgreSQL 集成测试、bash -n、ShellCheck、manual-release-test。
2026-08-01 13:59:37 +08:00
wangbo 4d84210344 fix(release): 修正 Worker 滚动与失败传播
单站点 Worker 启用硬反亲和后,maxSurge=1 会让新 Pod 无法调度并造成滚动超时。将 Worker 改为先下线旧 Pod 再创建新 Pod,同时保持 API 的零中断滚动策略。\n\n发布辅助脚本为关键 kubectl 操作补充显式失败返回,避免函数位于条件表达式时 Bash 忽略 errexit,进而把失败发布误报为成功。新增回归测试模拟 rollout status 失败且旧 Pod 仍 Ready 的场景。\n\n验证:bash -n、ShellCheck、cluster-release-helper-test、manual-release-test、生产 API Server dry-run。
2026-08-01 13:43:33 +08:00
wangbo 80801b7bbb fix(capacity): 允许显式禁用无节点站点
当 Worker Deployment 为零副本且站点没有 eligible Worker 节点时返回零容量状态,避免 min/max=0 的禁用站点令容量控制器持续 NotReady;非零副本仍保持失败关闭。\n\n验证:容量控制器单测、Go 全量测试、go vet、gofmt、迁移安全检查和人工发布脚本测试。
2026-08-01 13:02:55 +08:00
wangbo 773d310214 fix(cluster): 强制关键副本跨节点分散
避免同站点 Worker 与容量控制器在滚动发布时集中到同一节点,并让 Kubernetes 运行态清单变更进入可追溯的 API release。\n\n验证:bash -n、ShellCheck、manual-release-test、生产集群 server-side dry-run。
2026-08-01 12:33:19 +08:00
wangbo fcc3dc3cb9 fix(acceptance): 排空后切换验收 Worker 拓扑
首次 CAS 只校验 release 与镜像,进入 validation 并排空既有任务后再快照并应用 P24 基线,随后执行副本放置与资源门禁。避免在 live 流量下从 1+1 竞态切换到 0+2,也解除验收前置副本与安全排空之间的循环依赖。\n\n验证:bash -n、ShellCheck、git diff --check。
2026-08-01 11:45:32 +08:00
wangbo c9fbf6e6f6 fix(cluster): 稳定 Worker 节点首次接入
在新增 WireGuard 对等后等待全部握手并预热双向路由,再执行严格 10 包链路门禁;K3s 安装器检测到既有 unit 时显式启动未运行的 agent,并补充 Worker 标签校验。\n\n验证:bash -n、ShellCheck、四节点六向链路复测、K3s agent Ready 与隔离状态逐项核对。
2026-08-01 11:38:19 +08:00
wangbo ed3b970178 refactor(cluster): 通用化第四个 Worker 节点接入
将原宁波专用接入脚本改为由环境变量控制节点名称、站点和 WireGuard 端口,支持香港双物理节点 Worker 池。增加候选节点启用开关和服务商网络许可门禁,未明确批准时禁止任何 WireGuard 或 K3s 写操作。同步调整生产验收、巡检、拓扑校验和运行手册。\n\n验证:\n- scripts/cluster 全部 Shell 脚本 bash -n\n- scripts/cluster 全部 ShellCheck\n- git diff --check\n- 本地凭据值扫描\n- 当前三节点 WireGuard/K3s/PostgreSQL 只读基线检查
2026-08-01 11:23:19 +08:00
wangbo e77bd17754 fix(cluster): 固化宁波 Worker 隧道端口
新宁波节点改用 UDP/443 监听,其他三台节点保留 UDP/51820,并在全量 WireGuard 引导中按节点配置 endpoint 端口。\n\n三轮六向实测均为零丢包,香港 RTT 从 73-77ms 降至 48-53ms;已通过 bash -n、ShellCheck 和差异检查。
2026-08-01 10:54:30 +08:00
wangbo 53589fe3ef fix(cluster): 以宁波专用节点替换深圳 Worker
移除深圳节点及中继拓扑,新增第二台宁波 K3s agent 的全互联 WireGuard 接入和严格 UFW 门禁。\n\nWorker Deployment 与容量控制器仅选择 easyai.io/worker=true 节点,使原宁波混部节点退出 Worker 资源预算,生产基线恢复为宁波专用节点与香港节点各一实例。\n\n已通过 Go 全量测试、go vet、gofmt、迁移安全检查、bash -n、ShellCheck、发布脚本测试和 Kubernetes 清单渲染。
2026-08-01 10:47:01 +08:00
wangbo d3b36cf63d fix(cluster): 保留深圳原有 UFW 防火墙
移除会与 UFW 冲突的 iptables-persistent 安装,改由 UFW 持久化 WireGuard、kubelet、Flannel 和 NodePort 规则。\n\n预检新增 UFW 已安装且处于 active 的硬门禁,避免包管理器切换防火墙后保留 DROP 默认策略并中断远程接入。\n\n已通过 bash -n、ShellCheck 和差异检查。
2026-08-01 09:51:03 +08:00
wangbo c63709d785 fix(cluster): 经宁波中继深圳香港链路
香港与深圳公网直连出现高丢包和高时延时,固定将两地 WireGuard 数据前缀经宁波转发,同时保留直接 peer 心跳。\n\n中继规则和转发状态持久化,接入与全量 WireGuard 引导脚本保持一致;链路验收仍使用端到端丢包和 RTT 硬门禁。\n\n已通过 bash -n、ShellCheck、深圳资源预检和差异检查。
2026-08-01 09:42:59 +08:00
wangbo 70b0ffb9ae feat(cluster): 接入深圳专用 Worker 节点
新增深圳 K3s agent 与四节点 WireGuard 全互联接入脚本,将深圳归入香港逻辑 Worker 池并用污点限制为 Worker 专用。\n\n生产验收支持宁波 0、香港逻辑池 2 的双物理节点基线,补充副本拓扑分散、容量档位、链路巡检和零副本站点校验。\n\n已通过 Go 全量测试、go vet、govulncheck、迁移安全检查、ShellCheck、发布脚本测试、前后端 lint/test/build、Compose 与 Kubernetes 渲染校验。
2026-08-01 09:38:10 +08:00
wangbo 132cda35d8 fix(acceptance): 隔离控制面抖动与租约瞬态故障
线上 P24 验收暴露出高频 kubectl exec 放大 K3s API 压力、门禁查询挤占关键连接池,以及 PostgreSQL 锁超时被误判为租约所有权丢失。

本次合并验收身份查询、在租约有效期内重试瞬态续期错误、修复人工审核残留 attempt,并增加滚动后 etcd 稳定窗口、节点直连指标和双站独立报告。

验证:Go 全量测试、go vet、聚焦 race、gofmt、迁移安全检查、bash -n、ShellCheck、manual release test。
2026-08-01 01:51:44 +08:00
wangbo be6ce7f78a fix(acceptance): 修复线上双站压测失真
将线上压力生成器拆分为宁波、香港宿主机上的 K3s 外进程,按全局序号分片并保守合并站点报告,避免操作者本机上行带宽成为容量瓶颈。\n\n为每次负载执行生成独立幂等键和 Run 级图片变体,阻止三轮验收重放旧任务或复用历史媒体缓存,并新增直接 OSS 物化一致性门禁。\n\n验证:Go 全量测试、go vet、迁移安全检查、bash -n、ShellCheck、manual-release-test 和双宿主机 linux/amd64 启动冒烟通过。
2026-08-01 01:02:07 +08:00
wangbo 9fd9c267d6 fix(acceptance): 配置化资源预热门槛
宁波节点混部 legacy 生产服务,固定 65% 预热门槛无法反映当前可用资源,也无法通过停止 staging/CI 实质改善。

预热内存门槛现在可在 50–79% 显式配置并写入报告;80%运行目标、85%硬上限和 15 分钟稳定窗口保持不变。

验证:bash -n、ShellCheck、验收报告测试。
2026-08-01 00:15:05 +08:00
wangbo 3a01a5c93e fix(acceptance): 允许无运行时差异的工具更新
验收编排修复被正确分类为 components=none,但旧 CAS 要求工具 HEAD 与线上 release 完全相等,导致无运行时变化也必须重发镜像。

现在仅允许线上 release 到工具 HEAD 之间为无迁移、无 API/Web 运行时变化的后代提交;任何运行时或迁移差异仍立即阻断。

验证:bash -n、ShellCheck、验收报告测试、release-components 分类检查。
2026-08-01 00:11:25 +08:00
wangbo a612c3cb4a fix(acceptance): 按发布隔离验收身份组
历史验收主身份仍引用固定 production-acceptance 组,严格隔离门禁因此在切换 validation 前正确阻断。

验收组改为包含 release SHA 的稳定键;历史身份保持可审计,本次分片及已有 API Key 幂等迁移到当前发布专属组,避免清理或影响正式身份。

验证:bash -n、ShellCheck、验收报告测试、线上引用关系只读审计。
2026-08-01 00:10:41 +08:00
wangbo b6ef0435ab fix(controller): 隔离无关媒体凭据校验
容量控制器只负责读取容量状态和调整 Worker 副本,按最小权限不注入 OSS 或供应商凭据;全局媒体配置校验却导致控制器在生产启动时 CrashLoop。

容量控制器角色现在跳过不会使用的直传 OSS 校验,API/Worker 仍保持原校验;发布先以零副本应用目标状态,再用精确镜像和同构配置运行 Ready 预检,成功后才扩到双副本;同时补齐迁移 Job 的 restricted PodSecurity 配置。

验证:Go 全量测试、go vet、角色回归测试、bash -n、ShellCheck、人工发布测试、生产 server-side dry-run。
2026-07-31 23:52:22 +08:00
wangbo 993fb7d8be fix(release): 兼容首次创建容量控制器
旧生产集群没有 capacity-controller,部署前精确快照此前对缺失 Deployment 直接报错,导致目标状态尚未应用就阻断发布。

快照现在允许目标 Deployment 不存在;失败恢复时,快照中原本不存在的资源会被删除,从而同时支持首次创建和精确回退。

验证:bash -n、ShellCheck、人工发布脚本测试、线上缺失控制器只读快照演练。
2026-07-31 23:38:55 +08:00
wangbo 744e4acd1f fix(cluster): 消除发布前门禁假阴性
发布前巡检在 pipefail 下直接将 K3s readyz 管给 grep -q,可能因 SIGPIPE 误报 etcd 失败;链路和 PostgreSQL 断言失败时也缺少可定位证据。

改为完整读取 readyz,输出六向链路指标,并为复制、归档、备份和配置投影增加明确失败原因;同时缩短 10 包探测间隔,不改变丢包与 RTT 门槛。

验证:bash -n、ShellCheck、生产 precutover 只读巡检。
2026-07-31 23:34:22 +08:00
wangbo 709434a256 fix(acceptance): 自举线上隔离验收身份
线上验收此前依赖人工准备管理员 Token、专属用户、API Key、参考图片和数据库连接,导致跳过本地验收后仍无法安全、可重复执行。

本提交增加短期 Manager 登录、隔离用户/钱包/API Key 幂等创建、Pod 内脱敏快照导出和合成图片真实上传;敏感材料仅保留在进程或 0600 临时文件中。

验证:bash -n、ShellCheck、验收报告测试、参考图片生成测试、git diff --check。
2026-07-31 23:25:43 +08:00
wangbo 0b681275ed fix(acceptance): 支持显式跳过本地验收
在用户明确授权时,将本地原生与 amd64 制品阶段记录为 skipped/waived,并直接使用当前生产脱敏快照进入线上模拟。修复线上报告合并未写入生产 Run ID 导致后续 promote CAS 必然失败的问题。已通过报告回归测试、bash -n、ShellCheck 和 diff 检查。
2026-07-31 23:15:22 +08:00
wangbo a95184b5b6 fix(acceptance): 阻断本地控制面漂移污染验收
本地 K3s 节点此前在 Kubelet 中登记为宿主机资源,负载可挤压 etcd 并在 server 重启后继续污染同一 Run。现在为三节点设置真实可调度预算、独立 etcd/数据卷和锁定依赖镜像,并持续核对 server 与 API/etcd 健康。\n\n每次验收生成独立 Run ID,报告使用独占或原子写入,负载错误记录具体阶段;数据库鉴权不可用返回带 Retry-After 的 503,避免基础设施故障被误报为 401。\n\n验证:Go 全量测试、go vet、gofmt、OpenAPI、bash -n、ShellCheck、报告测试和 k3d 配置解析均通过。
2026-07-31 23:07:19 +08:00
wangbo 015ff8ea6c fix(acceptance): 允许受信模拟源的媒体物化
验收任务物化最终媒体时,仅允许访问 Acceptance Run 登记的协议模拟器精确 origin;其他私网、协议、端口及带 userinfo 的地址继续由 SSRF 防护拒绝。\n\n验证:Go 全量测试、go vet、gofmt 和 git diff --check 通过。
2026-07-31 21:39:19 +08:00
wangbo add99e5421 fix(acceptance): 物化验收任务的最终媒体
仅对已认证 Acceptance Run 强制下载并持久化 URL 型结果,确保集群外压测端能验证最终视频;普通生产任务继续遵循现有文件存储策略。压测器仅对 Gateway 自有地址做双入口重写,避免向第三方媒体域名泄露验收凭据。\n\n同时修复本地集群重复使用同一快照文件时的幂等复制失败。验证:Go 全量测试、go vet、gofmt、bash -n、ShellCheck 和 git diff --check 通过。
2026-07-31 21:25:42 +08:00
wangbo 29a9b8c89f fix(acceptance): 使用可完整解码的 WebP 样本
替换截断的 WebP 验收夹具,新增全量图片像素解码回归,并保留 6K 外部水化与三参考图归一化测试,避免协议模拟样本缺陷被误判为 Worker 运行故障。\n\n验证:Go 全量测试、go vet、gofmt、git diff --check 通过。
2026-07-31 21:11:50 +08:00
wangbo 29533537ec fix(worker): 限制大图归一化峰值内存
6K 多参考图在 P24 下会并发进入高内存缩放并击穿 2 GiB Worker。新增独立可配置的图片归一化并发,生产默认 2,只约束解码、缩放和重编码阶段,不占用视频上游等待槽;缩放器改为内存稳定的近似双线性实现。\n\n验证:6K 到 6000x2400 转换、信号量串行化、Go 全量测试、gofmt、Kubernetes 客户端 dry-run。
2026-07-31 20:54:11 +08:00
wangbo f190af00be fix(acceptance): 修复快照迁移重放与派发死锁
本地同构验收在全量迁移后导入生产快照,导致本次发布的数据迁移被旧能力覆盖;增加仅允许严格本地集群标记启用的导入后迁移重放,并新增幂等 Seedance 约束校准。\n\n批量异步派发改为在事务开始按全局顺序预锁全部任务与容量作用域,同时对 PostgreSQL 死锁和序列化失败做退避重试,避免双 API 重叠批次形成环形等待。\n\n验证:Go 全量测试、gofmt、bash -n、ShellCheck、迁移安全检查。
2026-07-31 20:39:42 +08:00
wangbo bfdabd3853 fix(acceptance): 校准 Seedance 图片转换验收
按 Volces 官方输入边界补齐 Seedance 2.0 候选能力,将错误的合法 4K 转换样本替换为真实越界图片,并让协议模拟器校验物化后的 Base64 data URL。\n\n验证:Go 全量测试、迁移安全检查、gofmt。
2026-07-31 20:07:28 +08:00
wangbo 2c7ed905a9 fix(acceptance): 打通本地集群内验收链路
为本地三节点验收补齐 Gateway 内部 DNS、TLS CA 信任与端口映射,并让仅限验收命名空间的网络故障代理以所需 UID 获取 NET_ADMIN。\n\n验证:bash -n、ShellCheck、K3d 三节点实测。
2026-07-31 20:07:21 +08:00
wangbo 26f3b3fb0a fix(acceptance): 消除 bootstrap 产物竞态
bootstrap 主容器完成后 Kubernetes 不允许继续 exec,导致无法取回只存在容器文件系统中的私密运行参数。使用共享 emptyDir 与限时导出 sidecar,在主容器退出码为零后复制 0600 产物并立即删除 Pod。\n\n验证:bash -n;ShellCheck;本地 CNPG 环境 32 身份 bootstrap、sidecar 导出和 runtime Schema 校验通过。
2026-07-31 19:47:51 +08:00
wangbo 805b677e86 fix(acceptance): 明确身份序号 SQL 类型
PostgreSQL 无法从 jsonb_build_object 的多态参数推断验收身份序号类型,导致本地 32 身份 bootstrap 失败。将序号绑定参数显式转换为 integer。\n\n验证:gofmt;go test ./cmd/acceptance-bootstrap -count=1;本地 CNPG 已复现原始 SQLSTATE 42P18。
2026-07-31 19:39:59 +08:00
wangbo 3ba6026e12 fix(acceptance): 安全暂存脱敏快照
ConfigMap 投影文件是符号链接,而快照导入器必须拒绝链接输入。新增 initContainer 将只读投影复制为 emptyDir 中的 0600 普通文件,保留防链接安全门禁并让本地导入可执行。\n\n验证:bash -n;ShellCheck;本地 CNPG 双实例实际快照导入通过。
2026-07-31 19:32:28 +08:00
wangbo 38207494d7 fix(acceptance): 修复本地数据库身份标记
通过 psql 标准输入执行变量替换,并校验本地集群标识格式,确保安全标记在导入生产脱敏快照前可靠写入,避免验收编排因 -c 参数不展开变量而退出。\n\n验证:bash -n;ShellCheck;本地 CNPG 双实例实际写入与查询标记通过。
2026-07-31 19:26:25 +08:00
wangbo 0bc4c00581 fix(acceptance): 固定本地节点交换内存上限
设置 K3d 节点资源限制时同步更新 memory 与 memory-swap,避免 Docker 因既有 swap 上限较低而拒绝节点内存预算,并禁止验收节点依赖额外交换空间。\n\n验证:bash -n scripts/acceptance/local-cluster.sh;shellcheck scripts/acceptance/local-cluster.sh
2026-07-31 19:20:16 +08:00
wangbo ed3bae4ba6 fix(acceptance): 修正本地 K3d 节点资源约束
本地同构环境创建后使用 K3d 实际容器名设置宁波、香港和洛杉矶节点的 CPU 与内存限制,避免集群已创建但验收编排在组件安装前退出。\n\n验证:bash -n scripts/acceptance/local-cluster.sh;shellcheck scripts/acceptance/local-cluster.sh
2026-07-31 19:18:44 +08:00
wangbo aec61eecb1 fix(acceptance): 固定相对快照路径解析
在本地集群脚本切换到 apps/api 前,将已校验的快照路径解析为绝对路径,确保文档中的相对路径可以直接执行。\n\n验证:bash -n、ShellCheck。
2026-07-31 18:44:50 +08:00
wangbo 0b9634b74c fix(postgres): 修正跨地域 WAL 归档带宽预算
将生产 archive_timeout 从 60 秒调整为 5 分钟,避免 16 MiB WAL 段在当前宁波香港链路上持续产生高于复制吞吐的空闲流量。\n\n同步复制保持开启,双库健康时 RPO 仍为 0;单库降级时对象存储归档 RPO 保持在 5 分钟目标内。补充发布回归门禁,防止配置退回无法追平的 60 秒。\n\n验证:kubectl kustomize、bash -n、ShellCheck、manual-release-test。
2026-07-31 18:42:46 +08:00
wangbo e05922b0f4 feat(acceptance): 建立同构验收与弹性容量体系
实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护,并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
2026-07-31 18:02:24 +08:00
chengcheng 93d36d0e55 fix(identity): 统一认证启用时直接跳转登录
双上下文配置不再落入工作台本地登录页,首次点击统一入口直接发起未预选上下文的 OIDC BFF 登录。

验证:OIDC 定向 Vitest 8 项通过;Web TypeScript 检查与生产构建通过。
2026-07-31 17:33:02 +08:00
chengcheng 59f0817938 feat(identity): 收敛 AI Gateway 统一认证入口
Web 登录页合并为单一统一认证入口,未指定上下文时由认证中心在登录后决策。继续接受旧客户端显式传入 platform 或 tenant,并保持回调上下文及 tenant_hint 的绑定校验。同步更新 OpenAPI 产物。\n\n验证:Go 全量测试与 go vet 通过;Web lint、141 项 Vitest 和生产构建通过;本地 platform.admin 登录及管理工作台访问通过。
2026-07-31 15:28:12 +08:00
chengcheng c0296dbf06 fix(identity): 支持平台用户显式登录 AI Gateway
修复多租户身份配置将所有人类登录都强制解释为租户上下文的问题。Web 现在提供平台与租户两个受控入口,API 严格绑定 context_type、tid、issuer、application 和 subject,并为平台用户建立独立本地投影与可刷新会话。\n\n风险:新增会话身份列保持旧会话可读,新会话一律使用严格约束;未改变租户数据隔离和 API Key 行为。\n\n验证:Go 全量测试与 go vet 通过;PostgreSQL 平台投影、会话和安全事件集成测试通过;前端 lint、141 项测试与生产 build 通过;OpenAPI 生成和迁移安全测试通过。
2026-07-31 14:24:40 +08:00
chengcheng 88c971564a feat(web): 增加单点登录分流与本地应急入口
开启统一认证后,正常登录入口直接跳转认证中心,并将本地账号入口收敛到 /login 应急页;未开启统一认证时继续保留工作台内原有登录方式。

影响范围仅限 Web 登录路由、认证入口展示和相关测试,不修改认证接口或 OpenAPI。风险主要在运行时身份状态判断,已通过 141 项前端测试、生产构建、无缓存类型检查及浏览器分支验证。
2026-07-31 13:43:06 +08:00
wangbo cae97b6f77 perf(worker): 缩短媒体完成链路
生成媒体在启用直传 OSS 时绕过后端存储通道,异步任务持久化成功后不再在 Worker 内二次水化 Base64。\n\n将任务事件与回调 outbox 合并到同一事务,降低跨地域同步复制提交次数并消除事件已落库但回调未登记的崩溃窗口。文件通道健康遥测改为异步复制提交,避免非关键状态占用执行槽。\n\n验证:Go 全量测试、go vet、gofmt、git diff --check、相对 447e7ed701 的迁移安全检查均通过。
2026-07-31 10:15:53 +08:00
wangbo 447e7ed701 perf(worker): 提交后发送准入唤醒通知
将准入登记、派发、释放、恢复和验收清理中的 PostgreSQL NOTIFY 从同步复制业务事务移到提交后的独立异步提交事务,避免数据库对象锁串行阻塞跨城事务。通知继续作为有轮询兜底的限时提示,不改变任务状态、租约和账务的同步提交语义。

验证:Go 全量测试、go vet、迁移安全检查、gofmt、diff check 通过;新增集成回归确认 NOTIFY 对象锁存在时准入状态仍先提交可见。
2026-07-31 09:56:47 +08:00
wangbo 580f84f569 perf(acceptance): 解耦 API 媒体请求并发
生产同构验收将 API 请求体并发独立配置为 AI_GATEWAY_ACCEPTANCE_API_MEDIA_REQUEST_CONCURRENCY,默认每个 API 128;Worker 执行槽和媒体物化仍按 P24/P28/P32 阶梯变化。\n\n这避免同步 Gemini 请求在整个任务生命周期占用仅 24/28/32 个入口槽,并继续由 RSS 与节点内存门禁验证安全上限。\n\n验证:bash -n;shellcheck -x;git diff --check。
2026-07-31 09:39:23 +08:00
wangbo 387bf5842d perf(worker): 隔离异步等待登记与执行锁
为异步等待队列登记增加独立的进程内和 PostgreSQL advisory lock 域,使严格队列上限校验不再等待 Worker 执行槽事务的跨城同步提交。任务级锁保持不变,调度期间 active 与 waiting 总量守恒。\n\n新增 PostgreSQL 集成测试,验证执行 scope 锁被占用时等待登记仍可独立完成。\n\n验证:go vet ./...;go test ./... -count=1;隔离 PostgreSQL 集成测试;迁移安全检查;gofmt;git diff --check。
2026-07-31 09:29:01 +08:00
wangbo 11ec057e6a perf(worker): 消除运行态计数跨城热行阻塞
将仅用于候选负载均衡的 runtime_client_states 分配和释放写入改为事务级异步提交,避免同步复制期间持续持有单候选热行锁。\n\n运行态计数仍由本地 WAL 持久化,并在运行时恢复中按 running attempt 重建;任务、账务、租约和状态提交不受影响。新增 PostgreSQL 并发与恢复集成测试。\n\n验证:go vet ./...;go test ./... -count=1;64 并发 PostgreSQL 集成测试;迁移安全检查;gofmt;git diff --check。
2026-07-31 09:14:13 +08:00
wangbo b0dac9f8ec fix(worker): 固定等待任务的准入候选
双 Worker 会按动态运行态排序重复规划同一批 waiting 任务,导致候选绑定来回迁移并把同步事务消耗在 rebind 上。\n\n调度 waiting 任务时优先使用首次登记的候选;仅当候选不再可用或已不在合法候选集时才按现有顺序迁移。执行阶段的 admitted 候选固定逻辑保持不变。\n\n验证:新增 waiting 候选固定回归测试;Go 全量测试、go vet、gofmt 通过。
2026-07-31 08:58:57 +08:00
wangbo 1e12522f0c perf(worker): 解耦异步等待登记与跨城提交
媒体突发中的等待准入行是已同步持久化任务的可恢复派生状态,却逐条等待跨城同步提交并串行占用全局容量锁。\n\n仅对异步等待登记事务设置 synchronous_commit=off;任务实体、素材、执行准入、River job、租约、状态与账务继续同步提交。主备切换若命中极小窗口,由既有 queued-task 恢复器重建等待登记。\n\n验证:临时 PostgreSQL 18 跨 Store 分布式准入集成测试通过;Go 全量测试、go vet、gofmt 通过。
2026-07-31 08:49:37 +08:00
wangbo a34a508140 perf(worker): 公平调度跨节点异步准入锁
持续突发提交会通过 pg_try_advisory_xact_lock 反复抢占全局容量锁,使 Worker 补槽器在首批之后饥饿。\n\n跨进程锁改为 PostgreSQL 公平等待;进程内锁仍保证每个 API 或 Worker 最多一条连接参与等待,30 秒 lock_timeout 会回到原有重试路径,60 秒空闲事务超时继续处理失主会话。\n\n验证:临时 PostgreSQL 18 上 64 并发锁竞争峰值连接不超过 2;Go 全量测试、go vet、gofmt 通过。
2026-07-31 08:40:54 +08:00
wangbo 17433bf2da perf(worker): 持续批量填满异步执行槽
异步准入通知在突发提交时会合并,原调度器每次只填充一个 8 任务批次,导致全局容量长期空闲并反向拖慢入口登记。\n\n收到通知或周期唤醒后连续拉取并提交批次,直到队列为空或全局执行容量真正饱和;保留每批 8 个的短事务边界。\n\n验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;gofmt -l 无输出。
2026-07-31 08:32:39 +08:00
wangbo d5b5d4b4cd perf(worker): 缩短异步批量填槽延迟
将单次异步准入准备批次限制为 8 个任务,使双节点调度器以多个短事务批次渐进填满全局容量,避免准备 48 个媒体任务后才首次提交造成的长时间冷启动空槽。\n\n验证:Go 全量测试、go vet、runner/store race、gofmt、迁移安全检查通过。
2026-07-31 08:09:17 +08:00
wangbo 5dd7765ae3 perf(worker): 批量填充异步执行槽
跨地域同步复制下逐任务准入会为每个 River Job 支付一次事务提交 RTT,导致 P24 实际只能维持少量运行任务。调度器现在按全局容量窗口准备 FIFO 批次,在同一 PostgreSQL 事务内逐项校验队首、创建租约并插入唯一 River Job,一次提交即可填满可用槽;任一 Hook 失败时整批原子回滚。\n\n验收压力采样同时排除正在终止的 Ready Pod,避免滚动切换期间误连已移除容器。\n\n验证:真实 PostgreSQL 批量提交及整批回滚集成测试、Go 全量测试、go vet、runner/store race、gofmt、bash -n、ShellCheck、迁移安全检查通过。
2026-07-31 07:59:57 +08:00
wangbo 14d13ad3a7 fix(worker): 隔离任务创建与队列恢复窗口
高媒体并发下,任务创建后的请求恢复和候选准入可能持续数十秒;通用 River 恢复扫描会在准入记录落库前误把新任务当成崩溃遗留任务,提前创建 Job 并占满 Worker 执行槽。\n\n为无活动 River Job 的通用兜底恢复增加 5 分钟新任务保护窗。正常提交、准入调度以及已有 River Job 的 Worker 强杀恢复保持原路径;新增 PostgreSQL 集成回归,验证新任务不被抢占且超过保护窗后仍可恢复。\n\n验证:Go 全量测试、go vet、runner/store race、gofmt、迁移安全检查通过。
2026-07-31 07:46:47 +08:00
wangbo b625edcd71 fix(worker): 收敛异步重准入到调度器
准入租约失效时 Worker 仅释放执行准备并短暂让出 River 槽位,由异步调度器统一恢复 waiting 或无有效租约的任务,避免多 Worker 同时争抢全局容量锁并放大数据库同步提交等待。\n\n调度扫描覆盖到期且可运行的 scheduled/available/retryable River 任务,但跳过正在执行或仍持有有效准入租约的任务;新增 PostgreSQL/River 状态集成回归测试。\n\n验证:Go 全量测试、go vet、runner/store race、gofmt、迁移安全检查通过;专用数据库集成用例因本机未配置 AI_GATEWAY_TEST_DATABASE_URL 明确跳过。
2026-07-31 07:37:37 +08:00
wangbo d1b482c2f3 fix(worker): 容量饱和后停止重复准入扫描
双 Worker 每秒遍历全部 waiting 任务,在全局执行槽已满时仍重复探测,导致 River 执行协程长时间等待同一 worker_capacity 进程锁。

异步 FIFO 头部未准入时立即结束本轮扫描,等待租约释放通知后再继续;任务级错误仍跳过并处理后续任务。

验证:Go 全量测试、runner race、go vet、迁移安全检查通过。
2026-07-31 07:23:46 +08:00
wangbo bfd1257db0 fix(worker): 统一异步执行准入范围
Worker 执行阶段重建准入范围时遗漏 worker_capacity,导致已预留 River 槽的任务再次同步等待准入并占住执行槽。

调度、执行和候选切换统一复用 taskAdmissionScopes,确保异步任务始终携带同一全局 Worker 容量范围。

验证:Go 全量测试、runner race、go vet、迁移安全检查通过。
2026-07-31 07:15:44 +08:00
wangbo 238798e47c fix(worker): 固定异步准入候选避免执行槽空转
异步任务在调度器准入后,Worker 会再次按实时负载排序候选;排序变化会让已准入任务退回 waiting,同时继续占用 River 执行槽,导致高并发吞吐塌陷。

执行前复用已持久化的 admitted 候选并稳定置顶,候选失效时仍保留原有重选路径;增加候选固定与非准入场景单测。

验证:Go 全量测试、runner race、go vet、OpenAPI 生成一致性、迁移安全检查通过。
2026-07-31 07:06:06 +08:00
wangbo 3d9ae74b87 fix(admission): 避免高并发准入锁占满连接池
将同进程相同准入键的请求先在内存中串行化,跨节点使用 pg_try_advisory_xact_lock 非阻塞竞争并在事务外退避,避免等待 advisory lock 时长期占用 PostgreSQL 连接。\n\n新增 64 路竞争回归测试,验证被占用锁下连接池峰值仅保留持锁者和单个尝试者;256 路 Gemini Base64 端到端压力通过,advisory wait 峰值为 0。
2026-07-31 06:31:24 +08:00
wangbo 4fa44f4c41 fix(admission): 限制媒体任务准入锁等待并发
媒体请求从认证、任务创建到异步准入登记共用可配置的请求并发令牌,完成登记后立即释放。这样仍会尽早丢弃 Base64 正文,但不会让高并发请求各占一个 PostgreSQL 连接等待 advisory lock。\n\n扩展 Gemini 双角色压力测试,使用生产同档 API/Worker 连接池并记录数据库连接与 advisory lock 峰值。\n\n验证:go test ./... -count=1;go vet ./...;pnpm openapi;迁移安全检查;256 个 256 KiB 异步 Gemini 请求全部通过,连接峰值 45、锁等待峰值 14。
2026-07-31 06:05:54 +08:00
wangbo 4af86b22ee fix(worker): 将 Gemini 同步请求交由异步队列执行
将非流式 Gemini generateContent 改为 River Worker 执行,API 通过批量状态查询等待完成,避免高并发同步请求耗尽 API 数据库连接池。\n\n生成图片在持久化时保存带哈希的资产引用,响应恢复时下载并校验后重建 Base64,数据库不保存媒体原文。请求体物化增加前置内存门禁,并扩展双 API/Worker PostgreSQL 压力测试。\n\n验证:go test ./... -count=1;go vet ./...;pnpm openapi;迁移安全检查;64 与 256 请求双角色异步 Gemini 压力测试。
2026-07-31 05:49:45 +08:00
wangbo 352e23e099 perf(media): 生产媒体直传阿里云 OSS
生产同构验收确认 server_main_openapi 上传通道在媒体高并发下产生 381 个 502,并将已创建任务的 p95 拉高到 172 秒。\n\n为 request_asset 与 image_result 增加环境配置控制的阿里云 OSS 直传,保留普通上传原通道;使用 Kubernetes Secret 注入凭据,并对签名、重试、场景隔离和生产配置校验增加测试。\n\n验证:Go 全量测试、go vet、OpenAPI、迁移安全检查、Kustomize/Compose 渲染、gofmt 与 git diff --check 均通过;真实 OSS PUT/CDN 读取探针通过且对象已清理。
2026-07-31 05:20:27 +08:00
wangbo 7b2ab786a2 perf(acceptance): 流式生成唯一图片输入
避免千请求验收重复复用同一图片内容而争用同一资产记录,使每个任务走真实独立上传与物化链路。\n\nGemini 请求体改为流式 Base64 编码,降低压测客户端的大对象内存峰值;新增尺寸、哈希唯一性和请求体还原测试。\n\n验证:Go 全量测试、go vet、迁移安全检查、bash -n、ShellCheck、gofmt、git diff --check 均通过。
2026-07-31 04:55:31 +08:00
wangbo 1bb4b2ab12 fix(acceptance): 同步配置 API 媒体并发
P24/P28/P32 之前只调整 Worker 媒体并发,API 仍使用默认每实例 16,导致入口实际容量低于 Worker 执行槽。现在按档位同时设置两地 API 的媒体请求与物化并发。\n\n验证:bash -n、ShellCheck。
2026-07-31 04:40:59 +08:00
wangbo 61b4417695 fix(acceptance): 继承分片身份的验收角色
分片 API Key 已具有正确 scopes 和候选规则,但空角色无法通过基础接口权限检查。创建和重用分片身份时同步主验收用户角色,保持专属用户组与钱包隔离不变。\n\n验证:bash -n、ShellCheck。
2026-07-31 04:31:50 +08:00
wangbo 262090db7b perf(acceptance): 使用多钱包身份分片模拟并发
单一验收钱包会把预留和结算串行化,掩盖双节点 Worker 的真实吞吐。验收现在幂等准备 32 个隔离身份和钱包,按请求轮询 API Key,并在 Run 配置中登记允许的 Key/User 身份对;密钥仅经 stdin 传给集群内压测进程。\n\n仍保留真实账务、候选权限、回调、重复扣费和强杀恢复校验。\n\n验证:Go 全量测试、临时 PostgreSQL 集成测试、go vet、OpenAPI 生成、迁移安全检查、bash -n、ShellCheck。
2026-07-31 04:25:09 +08:00
wangbo 8ce120631f fix(acceptance): 仅在连接池持续满载时阻断验收
验收压力采样仍记录连接获取取消计数,但单次上下文取消不再被误判为池饱和。继续对连续六个采样周期满池、租约错误、数据库连接、节点内存和 Pod RSS 执行硬门禁。\n\n验证:bash -n、ShellCheck、真实 P24 压力报告回归。
2026-07-31 04:07:14 +08:00
wangbo add80f781e perf(wallet): 避免同钱包预留占满数据库池
千并发验收的 pg_stat_activity 显示大量事务锁等待集中在同一验收钱包,等待行锁的请求持续占用 API Pool,导致 canceled acquire 和排队增长。

在每个 Store 内增加 256 槽钱包预留分片锁,同一钱包先在进程内串行再申请数据库连接;跨 API 节点仍由 PostgreSQL 行锁保证余额与幂等正确性。

验证:Go 全量测试、临时 PostgreSQL 并发账务集成测试、gofmt 和 git diff --check 通过。
2026-07-31 03:55:22 +08:00
wangbo 36546bd5c4 perf(postgres): 回收突发连接并等待验收副作用
失败 Run 的未提交任务已取消,但退款 outbox 尚未完成时下一轮会继承 API 连接高水位和 canceled acquire 增量,污染容量验收。

新增 AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS,生产配置 30 秒回收突发空闲连接;验收中止后等待旧任务、退款、回调和数据库连接全部收敛后才允许新负载启动。

验证:Go 全量测试、gofmt、bash -n、ShellCheck、kubectl kustomize 和 git diff --check 通过。
2026-07-31 03:46:41 +08:00
wangbo 75a56b21bc perf(acceptance): 将 API 连接池提升至 64
P24 千并发在 API Pool 48 下仍出现 acquired 接近上限及 canceled acquire 增长,而 Worker Pool 仅使用 3 条连接,PostgreSQL 实际连接数仍有安全余量。

将生产 API Pool 默认值提升至 64,并允许通过 AI_GATEWAY_ACCEPTANCE_API_DATABASE_MAX_CONNS 配置单次验收;Worker 档位与 min-idle=4 保持不变。

验证:bash -n、ShellCheck、kubectl kustomize 和 git diff --check 通过。
2026-07-31 03:38:28 +08:00
wangbo 2529679b67 perf(postgres): 分离 API 池并降低连接预热
生产同构千并发请求在干净队列上复现 API Pool 31/32、idle=1 且 canceled acquire 持续增长,Worker Pool 仅使用少量连接,确认瓶颈位于 API 数据库连接预算而非 Worker 执行槽。

新增可配置的 AI_GATEWAY_DATABASE_MIN_IDLE_CONNS,生产仅预热 4 条连接;API Pool 独立固定为 48,Worker Pool 继续跟随 P24/P28/P32 的 32/36/40 档位。发布工具、Kubernetes 清单、验收门禁和文档同步更新。

验证:Go 全量测试、gofmt、bash -n、ShellCheck、kubectl kustomize 和 git diff --check 通过。
2026-07-31 03:28:43 +08:00
wangbo 2606cc8223 fix(acceptance): 清理中止验收的未提交任务
验收失败后,排队及尚未提交上游的任务会继续占用队列与 API 数据库连接池,导致下一轮容量测试被历史负载污染。

中止验收时仅取消确认未提交上游的任务,删除其未提交 attempt 与准入记录,释放租约并生成幂等退款事件;提交中或已收到上游响应的任务继续自然收敛。新建验收前也会清理历史失败 Run 的安全遗留任务,并等待其余任务终态。

验证:Go 全量测试通过;临时 PostgreSQL 集成测试覆盖排队、运行未提交、提交中与退款事件;gofmt、bash -n、ShellCheck、git diff --check 通过。
2026-07-31 03:16:38 +08:00
wangbo 1bd9e618ca fix(acceptance): 重试瞬时压力采样失败
Pod 滚动后的首次 metrics 读取可能短暂失败,原采样器会直接退出且主流程忽略后台状态,导致验收缺少资源曲线。\n\n采样器现在允许两次瞬时失败,连续三次失败才中止;主流程同时检查采样进程状态,确保任何曲线中断都会明确阻断验收。\n\n验证:bash -n、ShellCheck、git diff --check。
2026-07-31 03:03:11 +08:00
wangbo 78005aaf85 fix(acceptance): 兼容只读压测容器报告
验收模拟器 Pod 使用只读根文件系统,失败报告缓冲不能依赖 mktemp。改为在 Shell 内存中捕获并回放压测 JSON,同时保留原始退出状态。\n\n验证:bash -n、ShellCheck、远端脚本语法检查、git diff --check。
2026-07-31 02:56:11 +08:00
wangbo 489208cef9 fix(acceptance): 纳入 API 数据库池容量门禁
高并发 Gemini 验收暴露两地 API 数据库池仍固定为 16,导致连接池满载、readiness 失败且压力曲线只监测 Worker,无法及时识别真实瓶颈。\n\n将容量档位的数据库连接池同步应用到 API,发布工具和生产清单保持同一配置;压力采样及日志门禁覆盖 API 与 Worker 四个连接池,并在失败时保留验收 JSON 报告。\n\n验证:bash -n、ShellCheck、kubectl kustomize、git diff --check。
2026-07-31 02:46:10 +08:00
wangbo 7a193934ed fix(acceptance): 优先选择稳定版 Gemini 图片模型
生产同构验收默认优先选择当前稳定版 gemini-3.1-flash-image,并在回退选择中将非 preview 模型置于预览版之前,避免选中并发 5 的预览候选后与千请求八分钟门槛产生物理冲突。\n\n候选确定后的正式路由优先级与真实限流策略保持不变。\n\n验证:bash -n;ShellCheck;git diff --check。
2026-07-31 02:27:19 +08:00
wangbo e0f841e8fb fix(admission): 按候选唤醒同步队列头
同步准入释放容量后,分别唤醒每个独立平台模型队列的 FIFO 队头;只有任务实际绑定用户组时才同时约束用户组队头,避免一个已饱和候选阻塞其他候选补位。\n\n调度器单次最多处理 256 个真实队头,不唤醒整个等待队列。新增 PostgreSQL 集成测试覆盖无用户组的双候选并行唤醒,以及共享用户组下仍保持组级 FIFO。\n\n验证:gofmt;Store/Runner 聚焦测试;一次性 PostgreSQL 集成测试;git diff --check。
2026-07-31 02:18:22 +08:00
wangbo 3976cfb64d fix(acceptance): 在集群内驱动媒体压力
将验收负载二进制改为在协议模拟器 Pod 内执行,并继续通过两个公网入口均分请求,避免本机上行带宽污染 Gateway 的 8 分钟吞吐门槛。\n\nAPI Key 与 Run Token 仅通过 kubectl exec 标准输入传入进程,不写入 Pod 规格、命令参数或验收报告;响应仍由客户端流式解码并计算哈希。\n\n验证:bash -n;ShellCheck;git diff --check。
2026-07-31 02:07:50 +08:00
wangbo 66db98ec1c fix(acceptance): 避免用户组阻塞候选队列
验收用户组仅承担身份和账务隔离,不再创建非约束性的组级并发队列,避免不同生产候选之间出现 FIFO 队头阻塞。真实候选自身的并发、排队和限流策略保持不变。\n\n验证:bash -n;ShellCheck;git diff --check。
2026-07-31 01:59:09 +08:00
wangbo a2dea335e0 fix(acceptance): 隔离高并发验收用户组
为生产同构验收创建并绑定专属用户组,避免验收账号继承默认并发 5 的限流。\n\n验收脚本通过管理 API 幂等维护用户组,并在数据库侧校验唯一用户、唯一活跃 Key、并发策略和无外部引用。\n\n验证:bash -n;ShellCheck;git diff --check。
2026-07-31 01:49:12 +08:00
wangbo 2bfeb3e179 fix(acceptance): 为高并发验收启用有界排队
生产同构 Gemini 千请求在真实候选并发饱和后因 queueing disabled 返回 429,同时单台压测机建立千条 WAN TCP 产生连接层 reset。仅对 acceptance 与 acceptance_canary 覆盖为 10000 条、最长 15 分钟的有界等待队列,保留候选原始并发和 RPM/TPM 上限;压测传输强制 HTTP/2 复用连接,避免单源连接风暴。正式 production 策略不变。\n\n验证:acceptance-load 与 runner 定向测试通过,新增生产策略不变和验收队列边界单测;gofmt、bash -n、ShellCheck、git diff --check 通过。
2026-07-31 01:34:16 +08:00
wangbo 0c9960d2e6 fix(acceptance): 预授权生产候选访问规则
生产同构 Gemini 验收在入口路由修复后暴露专属 API Key 无法通过生产 Access Rule,所有已创建任务均以 no_model_candidate 失败。验收脚本现在会为选中的 Gemini 和视频候选幂等授予 platform、platform_model、base_model 三层访问权限,并在创建 Run 前反查规则完整性。\n\n验证:bash -n、ShellCheck 与 git diff --check 通过。
2026-07-31 01:18:14 +08:00
wangbo cd9c8d56ab fix(nginx): 将原生模型接口转发至 API
生产同构 Gemini 验收发现 /v1beta 请求被 Web 入口返回 405,导致请求在到达 Worker 前全部失败。补齐 /v1 和 /v1beta 原生兼容路由,沿用 API 的长请求超时、无缓冲和双上游故障切换配置。\n\n验证:使用 nginx:1.27-alpine 对完整站点配置执行 nginx -t 通过。
2026-07-31 01:03:51 +08:00
wangbo 3b67038660 fix(acceptance): 移除生产验收持久管理隧道
管理 API 改为通过短 SSH 在宁波节点调用内部 api-ningbo-edge Service,避免长期端口转发与大量集群控制连接竞争导致 SSH 握手限流。Token 和请求体以 Base64 传输且不输出。\n\n验证:内部 traffic-mode 端到端返回 200;GNU Bash 3.2;bash -n;shellcheck -x -P .;git diff --check。
2026-07-31 00:41:56 +08:00
wangbo 0e22ccb0b8 fix(acceptance): 兼容 macOS Bash 3.2
将容量报告路径中的 Bash 4 小写参数展开替换为 POSIX tr,避免验收在提交首轮负载前因 bad substitution 中断。\n\n验证:GNU Bash 3.2.57 实际执行;bash -n;shellcheck -x -P .;git diff --check。
2026-07-31 00:34:02 +08:00
wangbo 02f9071a37 fix(acceptance): 安全接管直接基线失败任务
允许新发布在 validation 状态下接管其直接 base release 的失败 Run。接管前严格核对旧 Run 状态、流量 CAS 字段和 manifest 基线,并使用旧 CAS 字段中止旧 Run,再激活新 Run。\n\n验证:bash -n;shellcheck -x -P .;git diff --check。
2026-07-31 00:28:28 +08:00
wangbo cec3d429af fix(acceptance): 支持失败验收无缝重试
在 validation 保持正式流量关闭时,允许为同一 release 与镜像 digest 创建新 Run,并仅在确认旧 Run 已失败后通过 CAS 中止旧 Run、立即激活新 Run。避免验收控制连接异常后必须先长时间切回 live 才能重试。\n\n验证:bash -n;shellcheck -x -P .;git diff --check。
2026-07-31 00:25:52 +08:00
wangbo 3053ba4925 feat(acceptance): 增加生产同构媒体压力验收模式
引入动态流量门禁、隔离验收身份与协议级 Gemini/Volces 模拟器,覆盖双站点 API、Worker、PostgreSQL、River、账务、回调和媒体物化链路。

新增 P24/P28/P32 容量阶梯、Worker 强杀恢复、真实小流量 canary、CAS 放量和失败保持 validation 的生产编排;Worker 执行槽、连接池、媒体并发和双站点副本数改为环境配置。

验证:Go 全量测试、真实 PostgreSQL 迁移集成测试、迁移安全检查、OpenAPI 生成、ShellCheck、Kustomize、gofmt 和 git diff --check。
2026-07-30 23:06:19 +08:00
wangbo f33d6d64e0 fix(queue): 回收失效任务并阻止不确定提交重放
Worker 进程心跳不能证明单个 River job goroutine 仍存活,改以任务执行租约作为回收所有权栅栏。\n\n排队任务若上一次 attempt 在租约中断时处于 submitting 或 response_received,则转入人工复核并生成 release 记录,不再重新提交上游;明确被上游拒绝的响应仍允许重试。\n\n验证:go vet ./...;env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;真实 PostgreSQL 下相关 store 与 River 集成测试通过。
2026-07-30 19:15:19 +08:00
wangbo 9e06ed6162 fix(k3s): 让数据库探针访问节点本地 API
避免 kubernetes Service 在跨地域 API endpoint 间负载,导致 CNPG liveness/readiness 探针受 WireGuard 拥塞影响。\n\n通过 Downward API 将 KUBERNETES_SERVICE_HOST 设为 Pod 所在节点名,并使用本地 K3s 6443 端口。节点名解析、TLS 主机名校验和 kubectl 服务端 dry-run 均已通过。
2026-07-30 19:00:01 +08:00
wangbo e1e18dbe58 fix(k3s): 放宽数据库隔离探针容忍窗口
避免跨节点控制面短时拥塞导致主库被默认 30 秒 liveness 窗口误判并重启。\n\n保持隔离检查开启,将 livenessProbeTimeout 调整为 300 秒,并把隔离检查连接与请求超时调整为 5 秒。已通过线上 CRD 字段核验和 kubectl 服务端 dry-run。
2026-07-30 18:55:47 +08:00
wangbo 38f87b6970 fix(queue): 防止过期执行占位阻塞与重复提交
队列恢复后,仍在运行标记中的 River job 可能保留 waiting admission,形成全局 FIFO 队头阻塞;同时旧执行在租约被接管后仍可能创建 attempt 或切换到 submitting。

新增 stale admission 自动让出机制,并用当前 execution token 对 attempt 创建和上游提交状态切换做 fencing。任务、River job 和结算状态均不在让出流程中改写。

验证:go vet ./...;env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;真实 PostgreSQL 集成测试覆盖 stale admission 让出与旧 token 提交拦截。
2026-07-30 18:25:00 +08:00
wangbo fa818e9ebf fix(queue): 重新调度终态 River 任务
等待异步准入的任务只跳过仍处于活跃状态的 River Job;missing 或 terminal Job 交由准入调度器重新创建,避免非空 river_job_id 永久阻断旧任务。\n\n通用 River 恢复排除等待准入任务,消除双恢复循环与误导日志;事务回滚继续使用独立有界上下文。\n\n验证:真实 PostgreSQL 下异步 Worker 准入验收和跨 Store 队列集成测试均通过。
2026-07-30 17:42:33 +08:00
wangbo e8893e2b8f chore(release): 重新生成可回滚发布基线
当前 K3s 工作负载已核验并登记为 3aa78d3e5c0ba83227aaf33849c098405a89ceb5。\n\n使用新的不可变 SHA 重新生成 release manifest,避免覆盖已推送镜像 Tag,并恢复部署前后的自动回滚链。
2026-07-30 17:32:18 +08:00
wangbo 6bab0f0749 fix(worker): 防止事务泄漏并恢复滞留队列
为 PostgreSQL 连接增加可配置的事务空闲与锁等待超时,并在请求取消或 Worker 退出后使用独立有界上下文回滚事务。\n\n恢复过期任务时按批次使用 SKIP LOCKED,周期重建缺失或终态 River Job;锁等待超时只在确认尚未提交上游时安全重排,避免重复执行和重复结算。\n\n验证:完整 Go 测试、go vet、生产 Kustomize 渲染、gofmt 与 git diff --check 均通过。
2026-07-30 17:27:11 +08:00
wangbo bc44af751e fix(worker): 为容量扩展预留数据库连接
原因:单实例执行容量提升到 24 后,Worker 的 24 条数据库连接会被执行任务全部占用,导致心跳、选主、健康检查和并发租约续期超时。

影响:生产 Worker 连接池提高到 32,在 24 个执行槽之外保留 8 条控制连接;PostgreSQL max_connections=200,两实例与 API 的理论连接总量仍在安全范围内。

验证:kubectl kustomize deploy/kubernetes/production、git diff --check 通过;线上两地 Worker 按 32 条连接滚动后连续 3 分钟 Ready、0 重启、0 数据库控制错误、0 租约丢失。
2026-07-30 16:17:16 +08:00
wangbo 3aa78d3e5c perf(worker): 按实例内存上限扩展异步容量
将策略推导出的集群目标与单 Worker 内存安全上限分离,避免失活实例的全部容量转移到单个 2 GiB Pod。

生产环境单实例上限和媒体物化并发设为 24;两实例正常时总容量由 16 提升至 48,后续可通过增加 Worker 节点继续扩展。新增 0093 迁移保存实例容量上限,并兼容滚动发布中的旧实例。

风险:更高媒体并发会增加 Worker 和节点内存压力,保留单实例 24 的硬边界并由持续监控观察 RSS、OOM、队列与节点余量。

验证:go test ./... -count=1;go vet ./...;真实 PostgreSQL 分配与故障转移测试;迁移安全检查;kubectl kustomize;gofmt -l;git diff --check。
2026-07-30 15:08:54 +08:00
wangbo ea58d21d03 perf(queue): 限制大媒体任务内存并消除准入惊群
将同步与异步非文本任务的准入唤醒改为按任务和全局队首推进,批量续租等待者,避免大量请求同时争抢 advisory lock 和数据库连接。

对 Base64 请求解析、素材解码、上游媒体执行和结果物化增加分层并发限制,复用请求素材并释放重复 Gemini wire 数据;生产 API 默认入口解析并发 16,媒体物化并发 8。

新增 1000 个同步 Gemini 图像编辑请求的模拟上游压力验收,10 MiB 输入和输出下全部成功,Heap 峰值增长约 2.58 GiB,并验证共享上传哈希、单次 attempt 与本地零落盘。

验证:go test ./... -count=1;go vet ./...;gofmt -l 无输出;kubectl kustomize deploy/kubernetes/production;10 MiB Gemini Base64 千任务压力测试通过。
2026-07-30 13:13:32 +08:00
wangbo fd3b6bf042 fix(worker): 增强队列启动与租约到期恢复
River 初始化遇到瞬时数据库超时时最多重试三次,并确保启动异常时先取消后台 LISTEN 和 Worker 上下文再关闭连接池,避免不就绪 Pod 卡在退出路径。Worker 维护循环每 15 秒执行幂等运行时恢复,使旧实例中断任务在执行租约到期后自动重排并继续触发孤儿 Job 救援。验证通过完整 Go 测试、go vet、竞态检查和迁移安全检查。
2026-07-29 23:55:49 +08:00
wangbo f3dd7cd262 fix(queue): 自动救援失活 Worker 遗留任务
新增基于任务状态、执行租约和 Worker 心跳的孤儿 River Job 精确识别,每 15 秒将确认失活的 Job 恢复为可重试,避免等待 River 一小时通用救援窗口。运行时恢复同步清除中断及终态任务的 admission、执行令牌和并发租约,防止状态残留污染队列指标。正常运行中的长任务仍由 running 状态和活跃心跳保护。验证通过完整 Go 测试、go vet、竞态检查、迁移安全检查和 PostgreSQL 18 集成测试。
2026-07-29 23:44:29 +08:00
wangbo 98820378b7 perf(store): 消除同步复制下的鉴权与心跳锁阻塞
将 Worker 心跳与容量分配事务设为本地异步提交,避免同步副本延迟期间长期持有全局分配锁。API Key 使用时间改为异步提交并按分钟合并,消除高并发请求对同一热行的锁排队。业务任务、钱包、结算、并发租约等关键数据仍保持同步复制。验证通过完整 Go 测试、go vet、竞态测试、迁移安全检查及 PostgreSQL 18 集成测试。
2026-07-29 23:25:54 +08:00
wangbo 91451b3c86 fix(cluster): 修复身份事务阻塞与 Worker 容量抖动
将身份协调和安全事件心跳改为 PostgreSQL 单 Leader 执行,并从独立 Worker 进程中移除身份运行时,避免多副本重复写同一状态。

为安全事件事务增加锁等待、空闲事务超时及独立回滚上下文;Worker 需连续丢失六次心跳后才判定失效,降低跨节点抖动导致的容量反复扩缩。

验证:Go 全量测试、go vet、Race 聚焦测试、PostgreSQL 18 Leader/安全事件/Worker 分配集成测试及迁移测试通过。
2026-07-29 23:04:41 +08:00
wangbo 3886048e0f fix(gemini): 移除空返图占位并暴露失败原因
Gemini 图片响应未包含可提取资源时,不再返回伪造占位图,改为非重试失败并保留安全拦截、候选结束状态和上游错误等结构化诊断。\n\n同步将诊断写入 HTTP 错误、任务结果、attempt 指标和日志,确保任务失败时不结算,并补充单元及 PostgreSQL 验收测试。\n\n验证:gofmt、go test ./... -count=1、go vet ./...、迁移安全检查、pnpm openapi。
2026-07-29 19:50:43 +08:00
wangbo 012823adff fix(cluster): 保护 API 持久队列恢复启动阶段
API 角色启动时同样会恢复持久队列状态,积压较多时可能超过 liveness 初始窗口并被误杀。为两地 API 增加最长 10 分钟 startupProbe,旧 Pod 在新 Pod Ready 前继续服务。\n\n验证:kubectl kustomize、服务端 dry-run、生产滚动更新。
2026-07-29 19:18:42 +08:00
wangbo d810881901 fix(cluster): 修正 server-main 内网端口
生产 easyai-server 实际监听 3001,原 K3s 配置指向 3000 导致内部调用和任务进度回调 connection refused。将 SERVER_MAIN_BASE_URL 与回调地址修正为 WireGuard 内网 3001。\n\n当前 server-main v3.8.3 尚未实现任务进度回调路由,后续仍需 backend 配套并重放回调。验证:两节点访问 3001、kubectl kustomize、服务端 dry-run。
2026-07-29 19:12:27 +08:00
wangbo 4371e57212 fix(cluster): 进一步收紧媒体 Worker 并发上限
全局 48 槽在持续观察中仍导致宁波 Worker OOM。按实测内存占用将生产全局异步执行上限降至 16,双节点各 8 槽,单节点接管最多 16 槽;高并发请求继续由持久队列吸收。\n\n未开启 queue_size 生产策略。验证:kubectl kustomize、服务端 dry-run。
2026-07-29 18:51:26 +08:00
wangbo 8ce1f857a4 fix(cluster): 延长 Worker 启动探针保护期
Worker 在切换后恢复积压 River 任务时可能超过原 liveness 初始窗口,导致健康端口尚未监听就被误杀。增加最长 10 分钟 startupProbe,启动完成后仍沿用原 readiness 和 liveness。\n\n验证:kubectl kustomize、服务端 dry-run、生产滚动更新。
2026-07-29 17:49:13 +08:00
wangbo ba01cbef53 fix(cluster): 限制生产异步 Worker 执行容量
切换后集中恢复积压时,单 Worker 155 槽导致数据库连接池 24/24 饱和并触发香港 Pod OOM。将集群异步执行硬上限暂定为 48,使双 Worker 正常时各分配 24 槽,与每 Pod 连接池上限对齐。\n\n未开启任何 queue_size 生产策略。验证:kubectl kustomize、服务端 dry-run。
2026-07-29 17:45:23 +08:00
wangbo 4eb394bfbd test(cluster): 支持固定成功候选执行文件验收
为跨节点真实文件验收增加可选平台模型筛选,避免同名模型的失效候选掩盖共享文件链路结果。\n\n验证:真实任务 7888708e-a870-4676-bdfc-4021853365a7 由香港 Worker 执行,输入与生成结果在两节点下载的长度和 SHA-256 一致;单 attempt、结算、扣费及回调去重断言通过。
2026-07-29 17:43:34 +08:00
wangbo de1ce2274b fix(cluster): 等待 Worker 完全退出后执行跨节点验收
修正缩容期间 Deployment 已显示零就绪但旧 Worker Pod 尚在退出时仍可能领取验收任务的竞态。真实任务提交前现在会确认宁波 Worker Pod 数量为零,并在失败路径尝试取消未完成的测试任务。\n\n验证:node --check、bash -n、ShellCheck、git diff --check。
2026-07-29 17:29:15 +08:00
wangbo a073c840ac fix(cluster): 兼容 Worker 缩容后的空就绪状态
Kubernetes 在 Deployment 缩容为 0 后可能省略 status.readyReplicas;验收脚本将空值按 0 处理,避免在提交真实任务前误判失败。
2026-07-29 17:05:59 +08:00
wangbo 39b1541fa1 build(migrations): 固化 OIDC 迁移安全复核
为已进入主干且会重建约束的 OIDC 迁移增加 SHA-256 锁定的精确规则豁免,文件内容或违规类型变化时仍会拒绝发布。

生产只读预检确认受影响表规模很小,现有 6 条身份配置均为 schema v1,OIDC 会话和撤销水位均为空;新增校验器测试覆盖允许项和校验和篡改。
2026-07-29 16:19:00 +08:00
wangbo 9e4fc7362d feat(queue): 增加非文本模型分布式准入队列
使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配,并将生产 API 与独立 Worker 角色拆分。

补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本;未默认启用任何生产 queue_size 策略。

已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收;合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
chengcheng 2e5a90731b fix(oidc): 修复登录会话落库兼容问题
部分已配对环境在 0090 首次执行后缺少 oidc_client_id,导致统一认证完成投影后无法写入 Gateway 服务端会话。新增幂等前向迁移补齐该列,并为 Token 处理及会话创建增加不泄露凭据的稳定失败分类和关联诊断。\n\n验证:OIDC Session 全量单测、HTTP 回调定向测试、迁移升级集成测试、隔离 PostgreSQL 跨仓库 OIDC E2E 和本地 Chrome 真实登录均通过。
2026-07-29 10:40:33 +08:00
chengcheng fd9bbbb508 chore(git): 合并远端主分支最新变更 2026-07-29 08:59:08 +08:00
wangbo 135fb3d5b8 fix(deploy): 延长数据库副本重同步门限
CNPG 主库在 5 秒内晋升,但跨地域 pg_rewind 需要重传约 1.27GB,原 6 分钟等待不足。将同步副本恢复等待扩展到 30 分钟,并增加退出时清理写入探针的保护。\n\n已通过 bash -n、ShellCheck 和 git diff --check。
2026-07-29 03:46:16 +08:00
wangbo ea33be5d69 fix(deploy): 强化跨节点真实任务验收门禁
修正宁波旧 Worker 在滚动退场期间仍可能领取任务的竞态,提交任务前要求仅保留唯一且明确禁用 Worker 的新 Pod。\n\n为临时验收 Key 分配 gpt-image-2 所需的最小平台访问资源,并让 multipart 参数对齐线上 2K 图像编辑请求。已通过 bash -n、ShellCheck、node --check,并完成真实跨节点任务验收。
2026-07-29 03:24:02 +08:00
wangbo a642f43cf2 fix(deploy): 校准跨节点图片验收请求
真实成功样本使用 2048x2048,验收脚本改用相同默认尺寸并允许环境覆盖。修正创建 API Key 响应中的 apiKey.id 读取,每次运行前清理本脚本遗留 Key,结束时删除新 Key。\n\n验证:bash -n、ShellCheck、node --check、git diff --check。
2026-07-29 03:13:01 +08:00
wangbo 5d5284e068 test(api): 显式启用异步执行验收 Worker
动态并发验收直接构造 Config 时需显式开启 AsyncQueueWorkerEnabled,避免零值 false 把该测试误变为仅入队模式。\n\n验证:gofmt、go test ./internal/httpapi -count=1。
2026-07-29 03:07:13 +08:00
wangbo 9a3cc582ce fix(api): 分离异步队列入队与执行开关
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED=false 时仍初始化 River schema 与控制客户端,允许 HTTP 请求持久化并入队;仅跳过执行客户端和动态容量 Worker。关闭上下文时同步清理控制客户端引用。\n\n新增数据库集成测试覆盖无执行 Worker 时成功创建 River job 且 attempted_by 为空。验证:gofmt、go test ./... -count=1、go vet ./...。
2026-07-29 03:03:14 +08:00
wangbo 0bef13fbd7 fix(deploy): 完整读取边缘验收响应
OpenAPI 响应较大,curl 通过管道交给 grep -q 时会因下游提前退出报写端关闭,造成健康服务被误判失败。改为完整读取各响应后再断言内容。\n\n验证:bash -n、ShellCheck、git diff --check。
2026-07-29 02:56:50 +08:00
wangbo 6897fb3b66 fix(deploy): 使用站点专属 NodePort
HostPort 被命名空间 PodSecurity baseline 拒绝,因此恢复无特权 Pod 配置,并为宁波、香港的 API/Web 增加只选择本站 Pod 的专属 NodePort Service。双 NGINX 和验收脚本按站点使用 31088/31089 与 31178/31179,避免主机本地访问共享 NodePort 时随机命中不可达的跨站 Pod。\n\n验证:kubectl kustomize、服务端 dry-run、bash -n、ShellCheck、无 hostPort/hostIP、git diff --check。
2026-07-29 02:40:57 +08:00
wangbo 006609569d fix(deploy): 稳定边缘到站点的固定入口
主机本地访问跨站点 NodePort 时 kube-proxy 会随机选择不可达的对端 Pod,导致约四分之一请求被拒绝。为两地应用绑定仅位于 WireGuard 地址的 HostPort 31088/31178,双 NGINX 和集群验收改用该固定入口;原 NodePort 30088/30178 继续保留。端口处于现有公网 DROP、wg0 私网放行范围。\n\n验证:kubectl kustomize、服务端 dry-run、bash -n、ShellCheck、Secret 与 digest 静态检查、git diff --check。
2026-07-29 02:36:29 +08:00
wangbo 8e387ca553 fix(deploy): 固定 API 容器运行身份
API 镜像使用名称用户 appuser,Kubelet 无法仅凭 runAsNonRoot 验证其非 root 身份,导致 CreateContainerConfigError。按照 Dockerfile 中固定 UID 10001 显式设置 runAsUser、runAsGroup 和 fsGroup,确保 EmptyDir 可写且满足 Pod 安全约束。\n\n验证:kubectl kustomize 渲染成功,两个 API Deployment 均包含 UID/GID 10001,git diff --check 通过。
2026-07-29 02:28:19 +08:00
wangbo 39a8fd2ce0 fix(deploy): 移除激活阶段外网依赖
正式激活时不再从 GitHub 下载并重装已健康运行的 cert-manager、CloudNativePG 和 Barman Operator;Operator 安装继续保留在 prepare 阶段。激活阶段只渲染、校验并启动已锁定 digest 的应用清单,避免维护窗口受外网波动影响。\n\n验证:bash -n、ShellCheck、git diff --check。
2026-07-29 02:25:43 +08:00
wangbo f12f352bf0 fix(deploy): 规范化数据库版本对比
PostgreSQL 官方与 CNPG 镜像的 server_version 展示字符串包含不同构建后缀,造成同为 18.4 的数据库被误判不一致。改用 server_version_num 进行精确版本校验。\n\n验证:旧 Docker 与已恢复 CNPG 的完整验收快照 diff 无差异,server_version_num 均为 180004。
2026-07-29 02:17:13 +08:00
wangbo 945fc83c60 fix(deploy): 修正切换快照目录表名
切换前验收快照误用了不存在的 gateway_providers,导致在数据库导出前触发安全回滚。改为实际的 model_catalog_providers。\n\n验证:完整快照 SQL 已对线上旧库只读执行,版本、扩展、表、序列和关键计数均成功生成。
2026-07-29 02:11:48 +08:00
wangbo 5d2796337b fix(deploy): 等待维护入口完成重载
NGINX 优雅重载期间旧 worker 可能短暂继续返回 200,单次状态码断言会误触发安全回滚。改为最多轮询 15 秒等待维护页返回 503。\n\n验证:bash -n、ShellCheck、git diff --check、真实维护页 503 与恢复后公网 200。
2026-07-29 02:09:45 +08:00
wangbo b7c03d348e fix(deploy): 适配宁波 NGINX 配置加载路径
宁波 NGINX 仅加载 conf.d/*.conf,原切换脚本写入 sites-enabled 后不会生效。统一将维护页、集群入口和回滚配置写入实际加载的 conf.d 入口,并保留切换前配置备份。\n\n验证:bash -n、ShellCheck、git diff --check、双节点 nginx -t 与 prepare。
2026-07-29 02:06:12 +08:00
wangbo 74cf0ed5c8 fix(deploy): 避免验收入口端口冲突
将宁波跨节点验收私有入口从被 staging Docker 占用的 18088 调整为 18089,避免 NGINX reload 在切换前失败。同步更新真实文件任务验收脚本默认地址。已通过 node --check、bash -n、ShellCheck 和 git diff --check。
2026-07-29 02:02:42 +08:00
chengcheng 021ee9ab8f fix(identity): 允许清理未收到 Manifest 的配对
0090 的 Manifest 约束要求无 issuer 的 Revision 保持 draft,原取消流程将其改为 failed,导致取消事务被数据库拒绝。现在由 cancelled Pairing 记录终态,并仅允许明确标记 pairing_cancelled 的预 Manifest 草稿完成清理。\n\n验证:IdentityPairing PostgreSQL 定向测试通过;PairingService 定向测试通过。
2026-07-28 18:42:13 +08:00
chengcheng 206e367dbf fix(identity): 兼容多租户 V2 无固定租户映射
多租户 Revision 改为按 Token tid 动态解析租户,配对和验证不再要求 default 映射;单租户 V1 继续保留固定映射门禁。同步补齐 tenantMode 共享契约与 OpenAPI 产物。\n\n验证:Web 测试 129 项通过;pnpm lint 通过;Web 生产构建通过;pnpm openapi 通过。
2026-07-28 18:02:59 +08:00
chengcheng 5c679ff13f feat(identity): 接入认证中心多租户登录
支持 Manifest V2 动态 tid 验证、Tenant Context 同步和租户内 JIT 投影,并保留 Manifest V1 与旧 Session 兼容。\n\n增加 tenantHint、租户切换、普通注册关闭及 application/principal/tenant 两级 SSF 撤销;迁移、定向安全测试和本地双租户跨仓 E2E 已通过。\n\nrelease_required=true;未执行 Release、Staging 或真实链路。
2026-07-28 17:28:35 +08:00
chengcheng 0b02e62c72 chore(git): 同步主分支发布与业务变更
# Conflicts:
#	apps/api/go.mod
#	apps/api/go.sum
2026-07-28 10:47:41 +08:00
chengcheng 9dfd1aafa5 test(identity): 禁止 Gateway 依赖认证内核语义
将 OIDC 与安全事件集成测试中的 Realm 风格 Issuer 收敛为稳定 Issuer,并增加扫描所有 Gateway 已跟踪源码类别的回归门禁,避免 Keycloak 或 Realm URL 成为消费方契约。\n\n验证:go vet ./...;go test -race ./... -count=1;go test ./... -count=1;govulncheck ./...。
2026-07-28 10:17:39 +08:00
wangbo 7c142f5960 fix(deploy): 补齐集群备份兼容与切换前门禁
阿里云 OSS 的 S3 兼容层要求 virtual-hosted addressing,且 boto3 需要使用 Signature V2。本提交为 Barman 的归档、备份、保留与恢复统一挂载 AWS 配置,并将实际 WAL 归档和近期全备设为切换硬门禁。

同时补充 CNPG/etcdutl 固定版本安装、显式主库切换、三节点逐台停机、OSS 快照下载校验与远端临时 Secret 清理。已通过 ShellCheck、Kustomize、服务端 dry-run、迁移测试、发布脚本测试、敏感信息扫描及生产 precutover 验收。
2026-07-28 06:46:09 +08:00
chengcheng 4a464cacca feat(observability): 增加 OIDC 失败关联诊断
将 OIDC Access Token 与 ID Token 校验失败映射为固定安全分类,并生成 diagnosticId 关联浏览器错误与服务端日志。日志不记录授权码、Token 或底层敏感输入,且保留 errors.Is(ErrUnauthorized) 兼容语义。

验证:go test ./... -count=1(apps/api)
2026-07-28 06:19:24 +08:00
wangbo f5b6ff72f2 fix(deploy): 无损传输证书同步 SSH 公钥
将包含空格与注释的 Ed25519 host key、公钥先 Base64 编码为单个 SSH 参数,远端解码后再次校验,并清理前次失败产生的无效 authorized_keys 单词行。

验证:Base64 完整往返、bash -n、ShellCheck,两次首次连接失败现场。
2026-07-28 05:53:43 +08:00
wangbo 5d8ed1f293 fix(deploy): 固定证书同步的香港 SSH Host Key
通过已认证的香港管理连接读取 Ed25519 host 公钥,将 WireGuard 地址 10.77.0.2 的精确键写入宁波 known_hosts,再执行证书同步;不关闭 StrictHostKeyChecking。

验证:远端 host key 格式解析、bash -n、ShellCheck,首次同步因缺少 known_hosts 的现场复现。
2026-07-28 05:53:00 +08:00
chengcheng b51ecc5eed fix(identity): 拒绝跨站浏览器会话端点
统一认证使用 HttpOnly SameSite=Strict 会话 Cookie。配对时若 Public Base 与 Web Base 不属于同一 schemeful site,浏览器回调会丢失事务 Cookie。新增基于 Public Suffix List 的同站校验,并在运行时重建时重复门禁,防止旧持久化数据绕过。

验证:go test ./... -count=1(apps/api)
2026-07-28 05:52:46 +08:00
wangbo 1c0c298267 fix(deploy): 放宽数据库 Operator 首次拉取超时
CNPG 与 Barman Cloud 的固定镜像经跨境 mirror 首次拉取可能超过 180 秒,将 rollout 门限提升到 600 秒;cert-manager 保持原 180 秒快速门禁。

验证:CNPG 已实际 Ready,Barman 证书与插件注册正常且事件显示仅处于镜像拉取。
2026-07-28 05:17:25 +08:00
wangbo 6d7b2e88ba fix(deploy): 递归验证 containerd Registry 主机配置
K3s 1.36 将 Registry 镜像端点生成在 containerd/certs.d 子目录;验收改为递归检查整个 containerd 配置树,并继续以实际拉取 pause 镜像作为最终门禁。

验证:洛杉矶节点 Ready,三个 hosts.toml 端点存在,pause 镜像拉取成功。
2026-07-28 05:07:23 +08:00
wangbo 95ddfb2a04 fix(deploy): 为 K3s 配置容器镜像端点
三节点 containerd 增加 Docker Hub、GHCR 与 Quay 的 DaoCloud 镜像端点,解决生产节点直连公共 Registry 超时导致 Pod sandbox 无法创建的问题。

新增逐节点滚动刷新脚本,每次等待节点重新 Ready 并实际拉取 pause 镜像后才继续,避免同时重启 etcd 成员。初始 K3s 安装也会预置同一 registries.yaml。

验证:三节点镜像端点均返回 Registry 认证响应,bash -n、ShellCheck。
2026-07-28 05:05:33 +08:00
wangbo 44915c1343 fix(deploy): 使用 Server-Side Apply 安装大型 CRD
cert-manager、CloudNativePG 与 Barman 清单改用 server-side apply,避免 CNPG CRD 的 last-applied annotation 超过 Kubernetes 256KiB 限制;force-conflicts 仅用于已固定 checksum 的 operator 清单。

验证:cert-manager 组件已实际滚动成功,CNPG 客户端 apply 在 annotation 门禁处复现。
2026-07-28 05:00:52 +08:00
wangbo 19ce2d4272 fix(deploy): 保留 Kustomize digest 行缩进
替换真实 release manifest 的 API/Web digest 时固定输出 Kustomization 所需的四空格缩进,避免 awk 字段重建破坏 YAML。

验证:使用已发布 971540a 完整 manifest 实际渲染 696 行清单,确认所有应用镜像均为非零 digest。
2026-07-28 04:59:43 +08:00
wangbo d342462a3d fix(deploy): 兼容 macOS Bash 3.2 下载平台清单
将 operator URL 与固定 SHA-256 从关联数组改为并行索引数组,避免 macOS Bash 3.2 将带连字符的文件名解析为算术变量。远端资源复制前仍逐项校验 checksum。

验证:bash -n、ShellCheck、Bash 3.2 索引遍历。
2026-07-28 04:59:04 +08:00
wangbo 4e60bf1c64 fix(deploy): 本地校验并分发固定 K3s 二进制
生产节点访问 GitHub Release 失败时,改由本机下载 v1.36.2+k3s1 amd64 二进制,使用固定官方 SHA-256 校验后分发到三节点,再让官方安装器跳过远端下载。

同时将运行手册顺序改为先完成公网端口隔离,再安装 K3s。

验证:bash -n、ShellCheck、官方 checksum 获取与节点失败现场。
2026-07-28 04:50:54 +08:00
chengcheng 1b48a91af0 fix(ssf): 允许同一配对恢复未完成连接
首次创建 SSF Stream 失败后,恢复管理器会因缺少尚未生成的 Audience 而将同一 Revision 误判为凭据交接冲突。现在仅对同一 owner、Transmitter、机器客户端且尚未绑定 Stream 的可恢复状态放行重试,跨 Revision 与不完整绑定配置仍保持安全拒绝。\n\n验证:gofmt、go vet ./...、env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1。
2026-07-28 04:48:15 +08:00
wangbo 85e14f2cb8 fix(deploy): 过滤 WireGuard 安装阶段标准输出
节点首次安装 wireguard-tools 时 apt 输出会混入公钥命令替换,导致格式门禁误判并安全停止。仅取远端输出最后一行作为公钥,同时保留原有长度和字符校验。

验证:bash -n、ShellCheck、远端公钥长度与权限只读检查。
2026-07-28 04:40:26 +08:00
wangbo 971540a2a4 feat(deploy): 增加三节点 K3s 高可用迁移能力
新增 WireGuard 全互联、三 server embedded-etcd K3s、CloudNativePG 双实例、Barman OSS 备份、双 NGINX、Kubernetes Secret/RBAC 与本地旧文件按严格 24 小时清理。

新增维护窗口数据迁移、digest 固定滚动发布、应用回滚、跨节点文件 E2E、节点和数据库故障演练、CNPG 恢复与 etcd 快照验收脚本;洛杉矶仅作为带 NoSchedule 污点的仲裁节点。

所有生产 Secret 只在执行时从 0600 本地环境和旧生产容器导入,仓库不保存凭据;公网入口保持人工 DNS 故障切换边界。

验证:bash -n、ShellCheck、kubectl kustomize、Node 语法检查、Secret 扫描、OSS put/head/delete 实测。
2026-07-28 04:37:31 +08:00
wangbo edf7e66941 feat(cluster): 增加跨节点任务执行与请求素材场景
新增 AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED,默认保持启用;关闭时仅停止 River 异步任务执行器,HTTP、健康检查、结算与其他后台任务继续运行。

新增 0089 迁移并同步管理端默认场景,将 request_asset 加入现有上传渠道,确保 multipart 请求素材先转换为跨节点共享 URL。

验证:gofmt、go test ./... -count=1、go vet ./...、pnpm lint、pnpm test、pnpm build、迁移测试。
2026-07-28 04:36:58 +08:00
wangbo 2c7f8c14eb chore(release): 明确生产发布授权并支持集群发布
将发布与上线授权边界调整为一次明确的生产发布指令可连续执行 publish 和 deploy,同时保留仅构建镜像时的停止边界。

为现有发布脚本增加 compose/kubernetes helper 选择和显式 SSH 私钥支持,保持完整 SHA、digest 与生产基线校验不变。

验证:bash -n、ShellCheck、manual-release-test。
2026-07-28 04:36:53 +08:00
chengcheng 9e8722dc9f fix(deps): 修复 Gateway 依赖拒绝服务风险
将 Nx 间接使用的 minimatch 收敛到 10.2.6,从而使用带长度上限的 brace-expansion 5.0.8;同时升级 DOMPurify 到 3.4.12。\n\n变更保留 Node 22 运行基线,并通过 filelist 兼容检查、Gateway 全量 Go 与 Web 测试、构建、OpenAPI、迁移、发布脚本、govulncheck 和 pnpm audit。审计结果为 0 个已知漏洞。
2026-07-28 04:04:55 +08:00
wangbo 31565af07a refactor(runner): 收敛上游失败决策与轮转策略
将同平台重试、跨平台动作、健康副作用和冷却排队统一到单一失败决策,避免旧降级策略与 failover 重复执行。\n\n增加事务级单源保护、候选实时刷新、冷却错误契约、管理接口严格校验及兼容策略只读展示。\n\n验证:真实 Gateway HTTP/PostgreSQL 接受测试 10 项通过;go test ./...、pnpm openapi、pnpm lint、pnpm test、pnpm build 均通过。
2026-07-27 23:50:17 +08:00
wangbo 046c16fc69 fix(audio): 兼容百炼 OpenAI 音频输入
为阿里云百炼 OpenAI-compatible Chat 统一保留标准 input_audio,并将旧版 audio_url 转换为 input_audio,同时根据 MIME 或扩展名补齐音频格式。

补充请求资产水合测试,确认 input_audio.data 会转换为裸 Base64 且保留 format。

验证:在 apps/api 执行 env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1 全部通过。
2026-07-27 23:50:17 +08:00
wangbo 95776c84f6 fix(media): 规范比例与 OpenAI 图像尺寸参数
统一过滤非 x:x 比例,避免默认值参与候选匹配或透传上游。

支持 xK 分辨率归一化,并按 Gemini、Volces 和 OpenAI 的参数规范生成上游请求;OpenAI 显式 WxH 始终保留用户原始尺寸。

验证:在独立临时 worktree 中执行 apps/api 全量 go test ./... -count=1 通过,真实 OpenAI 请求已到达上游但受本地无效凭据阻断。
2026-07-27 23:50:17 +08:00
wangbo 039220c835 feat(models): 注册 Kimi K3 和 Qwen3.8 2026-07-27 23:50:17 +08:00
chengcheng 31c32690b2 feat(models): 支持按使用场景筛选模型
Gateway 模型接口校验 usage_scene,并与 server-main 严格场景目录取交集;显式场景查询在上游不可用时关闭失败,避免返回不应暴露的模型。\n\n验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1
2026-07-27 13:15:23 +08:00
wangbo de10875439 feat(admin): 优化后台运维与任务管理
完善平台、基准模型、定价规则和实时负载的紧凑查询与滚动展示,并固定关键操作列。

新增管理员任务记录、批量计费结算和用户组限流、额度及模型权限维护能力,同时补充任务脱敏、查询索引与 OpenAPI 契约。

验证:pnpm openapi、Go 全量测试、pnpm lint、pnpm test、pnpm build、gofmt 与差异格式检查均通过。
2026-07-25 01:32:49 +08:00
wangbo b29d539d49 feat(database): 优化开发环境数据库连接配置逻辑
- 添加 database_url_targets_local_port 函数验证数据库URL端口匹配
- 新增 database_url_matches_container 变量标识数据库URL与容器匹配状态
- 更新 postgres_port 变量获取容器实际端口值
- 重构数据库URL配置条件判断逻辑
- 添加容器凭证刷新提示信息
- 优化Docker数据库创建条件判断流程
2026-07-24 23:54:31 +08:00
wangbo 152885bbb6 fix(routing): 对齐缓存亲和力硬规则与审计 2026-07-24 23:51:27 +08:00
easyai d7a0ec56f5 feat(web): 支持登录密码显示与隐藏
原因:登录时缺少密码可见性切换,长密码输入难以核对。

影响:新增可复用 PasswordInput,在登录表单提供眼睛按钮,并补充禁用态、焦点样式和无障碍标签。

风险:密码仅在用户主动点击时于当前输入框显示,不改变提交、存储或日志行为。

验证:前端 112 项测试通过;pnpm lint 通过;@easyai-ai-gateway/web 生产构建通过。
2026-07-24 23:27:28 +08:00
easyai 1fc80ffe23 test(http): 防止集成测试误用非测试数据库
原因:共享迁移辅助函数会直接对 AI_GATEWAY_TEST_DATABASE_URL 指向的数据库执行迁移,需要在任何写入前阻止误连非测试库。

影响:HTTP 集成测试仅允许名称明确包含 test 边界的数据库;不影响未配置数据库环境变量时的普通单元测试。

风险:自定义测试库名称若不符合规则会被拒绝,需要改用 test、test_*、*_test 或包含 _test_ 的名称。

验证:gofmt 无差异;env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1 通过。
2026-07-24 23:27:17 +08:00
wangbo a2ed32447a fix(storage): 补齐嵌套结果二进制转存
生产真实 Gemini 图片任务可能在标准 data 之外保留嵌套 Base64,导致最终持久化关口拒绝任务成功。\n\n在所有生成结果返回前统一执行递归转存,并对仍未被常规媒体转换识别的二进制回退到本地结果物化,确保 PostgreSQL 不接收原始二进制。\n\n验证:go test ./... -count=1;go vet ./...;使用生产 Gemini Flash Image 平台配置完成本地真实图片编辑请求,文件哈希一致且持久化结果无内联二进制。
2026-07-24 21:57:52 +08:00
wangbo 4f163ea6d7 perf(storage): 拦截任务二进制并本地暂存结果
原因:任务标准结果中的 Base64、Data URI 和 Buffer 会进入 PostgreSQL JSON,导致 TOAST 与备份体积快速增长。

影响:新增统一 JSON 持久化关口;upload_none 将二进制原子写入本地结果目录,数据库仅保存带 SHA-256 的有界占位符;任务详情、同步响应、异步查询和兼容协议按需校验恢复。补充 24 小时清理、容量上限、历史小批量治理命令及管理端说明。

风险:本地结果超过 TTL、丢失或损坏时分别返回明确的 410/500;空间不足时返回 503 且不重试上游。未自动执行历史治理。

验证:三种真实图片模型同步/异步与幂等重放通过;Go vet/全量测试、前端 111 测试、lint/typecheck/build、OpenAPI、迁移安全、govulncheck、依赖审计、手工发布测试及 Linux amd64 构建通过。
2026-07-24 21:13:09 +08:00
wangbo 2457de6a56 chore(release): 重新生成生产发布基线
首次发布时当前会话未加载生产 SSH 身份,导致 release manifest 无法记录现网基线。创建空提交生成新的不可变镜像 Tag,并在读取真实生产 release 后重新执行完整发布门禁;不包含源码变化。
2026-07-24 18:35:33 +08:00
wangbo 810dcfeee6 perf(storage): 极简化任务历史并增加保留治理
停止持久化 provider 原始响应、兼容响应快照、attempt/event/outbox 重复 JSON,并由标准任务结果动态生成 Kling/Keling/Volces 兼容响应。

增加事件去重与预算、极简 callback 投递、7/30 天分批清理、安全删除条件、并发迁移索引及可实际恢复的任务域排除备份。历史清理默认关闭,待兼容协议和异步恢复在线验证后单独启用。

验证:Go 全量测试与 go vet、PostgreSQL 18 集成与实际备份恢复、迁移安全测试、bash -n、ShellCheck、Compose 配置和人工发布脚本测试均通过。
2026-07-24 18:23:22 +08:00
wangbo 09375bfae7 fix(queue): 移除未配置的 GPT 图像并发上限 2026-07-24 17:03:28 +08:00
wangbo 20945dbbcb fix(runner): 延长媒体上游请求超时
图像生成上游同步响应可能超过 120 秒,原共享 HTTP 客户端会把仍在处理的提交标记为结果不确定。将客户端超时调整为 10 分钟,并增加回归测试;任务执行仍由现有租约、任务超时和取消机制约束。\n\n验证:\n- env -u AI_GATEWAY_TEST_DATABASE_URL go test ./internal/runner -run 'TestProviderHTTPClientTimeoutAllowsLongRunningMediaRequests|TestPlatformProxyMode' -count=1\n- env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1
2026-07-24 13:42:47 +08:00
wangbo 6c5daf29ca perf(queue): 按策略动态扩缩异步 Worker
将 River 执行容量按平台模型与用户组的有效并发策略动态调整,并为长任务续租、并发租约原子抢占和限流退避补充保护。\n\n统一平台模型限流继承语义,兼容历史 platformLimits/modelLimits,并为三个迁移图像模型建立独立并发租约。补充管理端显式继承/覆盖配置、指标、单元测试及隔离 PostgreSQL 验收。\n\n验证:go test ./... -count=1;pnpm lint;pnpm test;pnpm build;隔离 PostgreSQL 并发原子性/续租测试;128 任务与三分钟长任务动态 Worker 验收。
2026-07-24 12:26:56 +08:00
wangbo 290b8c1854 fix(queue): 扩展生产媒体任务吞吐 2026-07-24 10:37:20 +08:00
wangbo d80b2e5ebf fix(queue): 提升媒体任务执行并发 2026-07-24 10:31:18 +08:00
wangbo 36d1b18e10 fix(images): 规范 OpenAI 图像尺寸参数 2026-07-24 09:28:18 +08:00
wangbo 8df87875de fix(images): 优先转换编辑图片为 multipart 二进制
OpenAI 编辑即使声明支持 URL,也必须在 multipart 提交前转存为二进制文件。
2026-07-24 09:04:23 +08:00
wangbo db640c1211 fix(images): 兼容 OpenAI 多图编辑 multipart
- 将 OpenAI images/edits 请求构造成 multipart/form-data
- 在运行时安全转存 URL 图片为 data URL 后提交二进制文件
- 覆盖多图字段、真实上游模型名和内部元数据过滤测试
2026-07-24 08:59:42 +08:00
wangbo 1da712854f fix(migrations): 停用旧模型冲突别名
以可回滚的状态更新替代删除,满足生产迁移安全门禁。
2026-07-24 08:42:59 +08:00
wangbo 8b8e09cc22 fix(media): 支持大尺寸图像响应并补齐迁移契约
- 为图像协议放宽 JSON 响应上限并显式处理读取与超限错误
- 合并基础模型与平台能力,避免运行时丢失分辨率和比例约束
- 固化稳定 Gemini 图像模型的能力、计价和 preview 兼容别名
2026-07-24 08:40:30 +08:00
wangbo fb7e08fe5c fix(runtime): honor requested platform for routing 2026-07-24 07:56:48 +08:00
wangbo 22ee5bb624 fix(release): 修复首次发布基线解析 2026-07-23 23:54:31 +08:00
wangbo 76a7702925 fix(media): 对齐 EasyAI 媒体响应与转存策略
统一媒体异步提交和轮询响应,扩展 /ai/result 到当前用户的 Gateway 媒体任务,并保持既有 Gateway 字段兼容。\n\n启用转存但无可用渠道时回退到 24 小时本地静态资源;关闭转存时保留图片、音频等上游 Base64 字段。同步更新文件上传兼容结构、OpenAPI、管理端说明和回归测试。\n\n验证:cd apps/api && env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;gofmt -l 无输出;pnpm nx run web:typecheck。
2026-07-23 22:35:41 +08:00
wangbo 8b7d3e9c9a feat(seedance): 同步输入图片约束与自动转换 2026-07-23 21:55:42 +08:00
wangbo 46cdb1a288 fix(api): 缺失输出能力配置时透传请求 2026-07-23 11:07:31 +08:00
wangbo 8d86c4b3b3 feat(catalog): 统一模型调用身份并梳理生命周期
新增官方调用名、供应商真实名、显示名和兼容别名的独立契约,调整路由、目录聚合、管理端与 OpenAPI。

增加 Gemini、Qwen、DeepSeek、Claude 和 MiniMax 生命周期迁移、别名观测及引用保护,并补充数据库与聚合测试。
2026-07-22 20:23:36 +08:00
easyai e3f6c0fa3e Merge pull request #27: 分类展示公开 API 文档并修复 DEV 启动
合并 API 文档分类、Compose 数据库 DEV 启动修复与人工发布协作规则。
2026-07-22 17:23:39 +08:00
easyai fac6ec95da docs(agent): 明确取消自动 CI/CD 后的合并规则
补充本地验证证据要求,明确经用户授权可合并 PR 或直接更新 main,并禁止等待或伪造已经取消的 CI 状态。\n\n同时继续隔离 Git 变更授权与 publish、deploy 两阶段生产授权。\n\n验证:tests/release/manual-release-test.sh
2026-07-22 17:22:50 +08:00
easyai c305603589 fix(dev): 兼容 Compose 数据库容器启动
保留显式数据库连接配置,自动识别 Compose PostgreSQL 容器及宿主机端口,并在容器不存在时跳过不适用的 Docker 建库。\n\n验证:bash -n;ShellCheck;scripts/create-database.sh;pnpm dev 真实启动 API 与 Web
2026-07-22 17:12:48 +08:00
easyai 3de26d5157 feat(web): 分类展示公开 API 文档
将通用开放接口与兼容接口分层展示,补充完整接口目录、分类路由、搜索能力和接入约定。\n\n验证:pnpm nx run web:test;pnpm nx run web:build
2026-07-22 17:12:48 +08:00
easyai fbe8d1d3ec fix(openapi): 补回图片异步响应契约
通过 PR #26 修复 Images 异步 202 OpenAPI 契约回归。
2026-07-22 15:44:48 +08:00
easyai 2812cadd4f fix(openapi): 补回图片异步响应契约
为 Images 生成与编辑接口重新声明 X-Async 参数和 202 TaskAcceptedResponse,保持 OpenAPI 与真实运行时行为一致。

已重新生成 Swagger,并通过 HTTP API 测试和桌面端生成客户端类型检查验证。
2026-07-22 15:43:32 +08:00
easyai 0b7c46cca9 feat(api): 合并官方兼容接口响应协议 (#24)
合并兼容协议响应统一、提交状态修正、兼容任务元数据迁移及相应契约测试。\n\n已在最新 main 上完成完整本地门禁验证。
2026-07-22 15:38:46 +08:00
easyai e00851d7e6 fix(db): 调整兼容元数据迁移为可空列
生产迁移安全策略禁止在已有任务表上直接新增非空列。兼容提交响应读取已使用 COALESCE,因此将 JSONB 元数据列改为可空,不改变接口行为。

验证:ci-validate-migrations.mjs;migrations-test.sh;go test ./internal/store ./internal/httpapi。
2026-07-22 15:35:00 +08:00
easyai e07a997aa9 feat(api): 统一官方兼容接口响应协议
兼容接口现在以入口协议作为最终响应协议,同协议保留官方 Wire 响应,跨协议统一转换成功、任务状态与错误结构。

同时修正异步提交状态边界,持久化兼容公开任务标识和官方提交响应,并新增迁移、流式响应及协议契约测试。

验证:go vet ./...;go test ./...;govulncheck ./...;pnpm lint;pnpm test;pnpm build;pnpm audit --audit-level high;pnpm openapi;全部 CI 脚本。
2026-07-22 15:34:59 +08:00
easyai 42e8b517fd feat(gateway): 补齐桌面端高级媒体直连接口
通过 PR #25 合并桌面端 AI Gateway 直连接口、计费、安全隔离及真实 DEV 验收实现。
2026-07-22 15:29:54 +08:00
easyai 762d61c9cf docs(release): 补充仓库 Actions 停用验收
记录仓库级 has_actions=false 和 main 无保护规则的只读验收命令。
2026-07-22 15:17:11 +08:00
easyai dae5d16a58 refactor(release): 改为 Agent 双阶段人工发布
删除 Gitea Actions、Tag/Main 自动流水线和旧 Runner 配置,取消 Git 操作与发布授权的绑定。\n\n新增本地镜像发布、固定生产部署助手、digest manifest、迁移安全检查、simulation 冒烟及显式回滚流程。\n\n验证:pnpm lint、pnpm test、pnpm build、Go 全量测试、ShellCheck、Compose 配置、人工发布测试和 linux/amd64 完整栈冒烟。
2026-07-22 15:13:40 +08:00
easyai 3056cf8fca feat(gateway): 补齐桌面端高级媒体直连接口
ci / verify (pull_request) Successful in 15m34s
新增图片矢量化、视频超分、每日用量、计价与任务隔离能力,并通过环境变量解析平台凭据。

已通过 Go 全量门禁、迁移检查、镜像构建以及 Vectorizer 五格式和 Topaz 3 秒视频真实 DEV 验收。
2026-07-22 14:02:53 +08:00
easyai cbebfd7baa chore(ci): 推进生产迁移基线至 v0.4.4 (#23)
ci / verify (push) Successful in 11m18s
生产发布、健康检查与隔离恢复演练均已通过;仅推进生产迁移基线。
2026-07-22 11:19:59 +08:00
576 changed files with 103342 additions and 7080 deletions
+1
View File
@@ -8,6 +8,7 @@
.env
*.log
.local-secrets
node_modules
**/node_modules
+17
View File
@@ -29,6 +29,11 @@ IDENTITY_MODE=hybrid
# - hold: reject new production generation before any upstream request; existing settlements continue.
BILLING_ENGINE_MODE=observe
# River 执行 worker 会按平台模型和活跃用户组中更严格的 concurrent 总量每 5 秒重算。
# hard limit 是单进程安全边界;平台模型/用户组 concurrency lease 才是业务并发真值。
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=2048
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS=5
# Unified identity business settings are managed in System Settings > Unified
# Identity. Deployment only supplies the SecretStore and infrastructure timing.
AI_GATEWAY_PUBLIC_BASE_URL=http://localhost:8088
@@ -68,6 +73,18 @@ TASK_PROGRESS_CALLBACK_ENABLED=true
TASK_PROGRESS_CALLBACK_URL=http://localhost:3000/internal/platform/task-progress-callbacks
TASK_PROGRESS_CALLBACK_TIMEOUT_MS=5000
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS=10
# First deploy with cleanup disabled; enable only after compatibility and
# asynchronous-resume verification has passed.
AI_GATEWAY_TASK_CLEANUP_ENABLED=false
AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS=24
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES=10737418240
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES=268435456
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES=536870912
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED=true
CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178
VITE_GATEWAY_API_BASE_URL=http://localhost:8088
-129
View File
@@ -1,129 +0,0 @@
name: ci
on:
push:
branches: [main]
pull_request:
branches: [main]
jobs:
verify:
runs-on: easyai-gateway-ci-unprivileged-v2
services:
postgres:
image: docker.io/library/postgres:16-alpine@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777
env:
POSTGRES_USER: easyai_test
POSTGRES_HOST_AUTH_METHOD: trust
POSTGRES_DB: easyai_gateway_test
options: >-
--health-cmd "pg_isready -U easyai_test -d easyai_gateway_test"
--health-interval 2s
--health-timeout 5s
--health-retries 30
env:
TRIVY_DB_REPOSITORY: ghcr.m.daocloud.io/aquasecurity/trivy-db:2
AI_GATEWAY_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
AI_GATEWAY_TEST_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
steps:
- name: Checkout without external Actions
env:
CI_REPOSITORY: ${{ github.repository }}
CI_SERVER_URL: ${{ github.server_url }}
CI_SHA: ${{ github.sha }}
CI_JOB_TOKEN: ${{ github.token }}
CI_EVENT_BEFORE: ${{ github.event.before }}
CI_PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
set -eu
test -n "$CI_JOB_TOKEN"
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
git init .
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$CI_SHA"
for comparison_sha in "$CI_EVENT_BEFORE" "$CI_PR_BASE_SHA"; do
case "$comparison_sha" in
[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]*)
if test "${#comparison_sha}" -eq 40 && \
test "$comparison_sha" != 0000000000000000000000000000000000000000; then
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$comparison_sha"
fi
;;
esac
done
unset authorization CI_JOB_TOKEN
test ! -f .git/shallow
git checkout --detach "$CI_SHA"
test "$(git rev-parse HEAD)" = "$CI_SHA"
- name: Verify pinned host toolchains
run: |
go version
node --version
pnpm --version
docker-compose version
shellcheck --version
trivy --version
govulncheck -version
- name: Verify production migration safety
env:
CI_EVENT_NAME: ${{ github.event_name }}
CI_EVENT_BEFORE: ${{ github.event.before }}
CI_PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
production_base=$(cat deploy/ci/production-migration-base)
immutable_base=$CI_EVENT_BEFORE
if test "$CI_EVENT_NAME" = pull_request; then
immutable_base=$CI_PR_BASE_SHA
fi
case "$immutable_base" in
[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]*)
test "${#immutable_base}" -eq 40
test "$immutable_base" != 0000000000000000000000000000000000000000
;;
*) exit 1 ;;
esac
git merge-base --is-ancestor "$immutable_base" HEAD
node ./scripts/ci-validate-migrations.mjs \
"$production_base" "$immutable_base"
- name: Verify Go formatting
run: |
unformatted=$(gofmt -l apps/api)
test -z "$unformatted" || {
printf 'Go files require gofmt:\n%s\n' "$unformatted" >&2
exit 1
}
- name: Migrate PostgreSQL 16 integration database
working-directory: apps/api
run: go run ./cmd/migrate
- name: Verify Go code
working-directory: apps/api
env:
GOFLAGS: "-p=1"
GOMAXPROCS: "1"
run: |
go vet ./...
go test ./...
govulncheck ./...
- run: pnpm install --frozen-lockfile
- run: pnpm lint
- run: pnpm test
- run: pnpm build
- name: Audit JavaScript dependencies
run: pnpm audit --audit-level high
- name: Validate deployment configuration
run: |
docker-compose -f docker-compose.yml config --quiet
shellcheck scripts/ci-build-images.sh scripts/ci-validate-semver.sh \
scripts/provision-ci-runner.sh tests/ci/ci-build-images-test.sh \
tests/ci/migrations-test.sh tests/ci/pipeline-test.sh \
tests/ci/semver-test.sh
./tests/ci/ci-build-images-test.sh
./tests/ci/migrations-test.sh
./tests/ci/pipeline-test.sh
./tests/ci/semver-test.sh
- name: Scan repository
run: |
trivy fs --scanners vuln,secret,misconfig --severity HIGH,CRITICAL \
--ignore-unfixed --exit-code 1 --timeout 15m --skip-dirs .git \
--skip-dirs node_modules .
-114
View File
@@ -1,114 +0,0 @@
name: release-ci
on:
push:
tags: ['v*']
jobs:
verify-tag:
runs-on: easyai-gateway-ci-unprivileged-v2
services:
postgres:
image: docker.io/library/postgres:16-alpine@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777
env:
POSTGRES_USER: easyai_test
POSTGRES_HOST_AUTH_METHOD: trust
POSTGRES_DB: easyai_gateway_test
options: >-
--health-cmd "pg_isready -U easyai_test -d easyai_gateway_test"
--health-interval 2s
--health-timeout 5s
--health-retries 30
env:
TRIVY_DB_REPOSITORY: ghcr.m.daocloud.io/aquasecurity/trivy-db:2
AI_GATEWAY_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
AI_GATEWAY_TEST_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
steps:
- name: Checkout without external Actions
env:
CI_REPOSITORY: ${{ github.repository }}
CI_SERVER_URL: ${{ github.server_url }}
CI_SHA: ${{ github.sha }}
CI_JOB_TOKEN: ${{ github.token }}
run: |
set -eu
test -n "$CI_JOB_TOKEN"
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
git init .
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$CI_SHA"
unset authorization CI_JOB_TOKEN
test ! -f .git/shallow
git checkout --detach FETCH_HEAD
- name: Verify pinned host toolchains
run: |
go version
node --version
pnpm --version
docker-compose version
shellcheck --version
trivy --version
govulncheck -version
- name: Verify release tag ancestry
env:
CI_REPOSITORY: ${{ github.repository }}
CI_SERVER_URL: ${{ github.server_url }}
CI_SHA: ${{ github.sha }}
CI_JOB_TOKEN: ${{ github.token }}
run: |
set -eu
tag_name=${GITHUB_REF#refs/tags/}
./scripts/ci-validate-semver.sh "$tag_name"
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" \
+refs/heads/main:refs/remotes/origin/main
unset authorization CI_JOB_TOKEN
test ! -f .git/shallow
test "$(git rev-parse HEAD)" = "$CI_SHA"
git merge-base --is-ancestor "$CI_SHA" refs/remotes/origin/main
- name: Verify production migration safety
run: |
production_base=$(cat deploy/ci/production-migration-base)
node ./scripts/ci-validate-migrations.mjs "$production_base"
- name: Verify Go formatting
run: |
unformatted=$(gofmt -l apps/api)
test -z "$unformatted" || {
printf 'Go files require gofmt:\n%s\n' "$unformatted" >&2
exit 1
}
- name: Migrate PostgreSQL 16 integration database
working-directory: apps/api
run: go run ./cmd/migrate
- name: Verify Go code
working-directory: apps/api
env:
GOFLAGS: "-p=1"
GOMAXPROCS: "1"
run: |
go vet ./...
go test ./...
govulncheck ./...
- run: pnpm install --frozen-lockfile
- run: pnpm lint
- run: pnpm test
- run: pnpm build
- name: Audit JavaScript dependencies
run: pnpm audit --audit-level high
- name: Validate deployment configuration
run: |
docker-compose -f docker-compose.yml config --quiet
shellcheck scripts/ci-build-images.sh scripts/ci-validate-semver.sh \
scripts/provision-ci-runner.sh tests/ci/ci-build-images-test.sh \
tests/ci/migrations-test.sh tests/ci/pipeline-test.sh \
tests/ci/semver-test.sh
./tests/ci/ci-build-images-test.sh
./tests/ci/migrations-test.sh
./tests/ci/pipeline-test.sh
./tests/ci/semver-test.sh
- name: Scan repository
run: |
trivy fs --scanners vuln,secret,misconfig --severity HIGH,CRITICAL \
--ignore-unfixed --exit-code 1 --timeout 15m --skip-dirs .git \
--skip-dirs node_modules .
+35 -14
View File
@@ -9,19 +9,28 @@
3. `type` 使用小写英文,可选值为 `feat``fix``docs``test``refactor``perf``build``ci``chore``revert`
4. 提交摘要和正文必须使用中文;专有名词、协议名称、命令、路径、代码标识符和第三方原始错误可保留原文。
5. 摘要应简洁明确,末尾不加句号,不得使用“更新代码”“修复问题”等无法说明意图的模糊描述。
6. 非简单变更应在提交正文或 PR 描述中用中文说明原因、影响、风险和验证结果。
6. 非简单变更应在提交正文中用中文说明原因、影响、风险和验证结果。
## 仓库边界
1. 后端位于 `apps/api`,前端位于 `apps/web`,共享 TypeScript 契约位于 `packages/contracts`
2. 修改 HTTP 接口、请求或响应类型后,必须执行 `pnpm openapi` 并提交匹配的 OpenAPI 产物。
3. 数据库迁移只能新增,禁止修改已经进入生产基线的历史迁移;迁移必须通过生产迁移安全检查。
4. 不得把 `.env`、密码、Secret、Token、授权码、私钥或生产凭据提交到 Git、日志、测试输出或验收证据中。
3. 数据库迁移只能新增,禁止修改或删除已经进入 Git 历史迁移;发布前必须相对当前线上 SHA 执行迁移安全检查。
4. 不得把 `.env`、密码、Secret、Token、授权码、私钥或生产凭据提交到 Git、日志、测试输出、release manifest 或验收证据中。
5. 当前工作区存在用户改动时必须保留,不得覆盖、清理、重置或混入当前任务提交;需要隔离时使用独立分支、克隆或 worktree。
## 验证要求
根据改动范围执行最小充分验证;准备合并或发布时执行完整门禁:
根据改动范围执行最小充分验证。修改 Shell 脚本后必须执行 `bash -n` 和 ShellCheck;修改 Go 文件后必须确认 `gofmt -l` 没有输出。
本地发布的快速强制门禁由 `scripts/publish-release-images.sh` 固定执行:
```bash
cd apps/api && env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1
node scripts/ci-validate-migrations.mjs <当前线上完整 Git SHA>
```
完整测试、审计和镜像扫描不再由 Gitea 自动运行。需要合并大范围重构、准备安全审计或人工要求完整验证时执行:
```bash
cd apps/api && go vet ./... && go test ./... && govulncheck ./...
@@ -31,18 +40,30 @@ pnpm test
pnpm build
pnpm audit --audit-level high
docker compose -f docker-compose.yml config --quiet
./tests/ci/ci-build-images-test.sh
./tests/ci/migrations-test.sh
./tests/ci/pipeline-test.sh
./tests/ci/semver-test.sh
./tests/release/manual-release-test.sh
```
修改 Shell 脚本后还必须执行 `bash -n` 和 ShellCheck。修改 Go 文件后必须确认 `gofmt -l` 没有输出
仓库没有自动 CI 兜底,合并或推送前必须以本地命令的实际结果为验收依据;报告中应明确区分已运行、命中缓存、未运行和无法运行的项目,不得把未执行的检查表述为通过
## Git 与 CI/CD
## Git 与人工发布
1. 一个提交只包含一个逻辑变更,提交前检查暂存差异并确认不含敏感信息
2. `main` 只能通过短生命周期分支和 PR 合并,禁止直接推送或强制推送
3. PR 必须通过精确的 `ci / verify (pull_request)` 状态后才能合并;合并后还要确认同一 SHA 的 `ci / verify (push)` 成功
4. 生产版本仅使用稳定 SemVer `vMAJOR.MINOR.PATCH` 标签,并同时要求 `release-ci / verify-tag (push)` 成功
5. 未验证 protected tag、发布账本、数据库备份、健康检查和回滚路径时,不得声称 CI/CD 或生产发布已经完成。
1. 仓库不使用 Gitea Actions,不存在 Push、PR、Tag、Webhook、轮询或定时触发的自动构建和自动部署
2. `main` 不设置受保护分支或 required statusAgent 在获得用户明确授权后,可以合并已验证的 PR,也可以直接提交并推送 `main`。操作前必须获取最新 `origin/main`、确认工作区边界、检查提交差异和敏感信息,并执行与风险匹配的本地验证
3. 不得等待、伪造或手工补写已经取消的 CI status,也不得为了满足旧流程擅自恢复 Gitea Actions。Git 提交、合并、Push 和 Tag 授权只覆盖源码历史变更,永远不等于 publish 或 deploy 授权
4. 发布只能使用工作区干净、已经提交并属于 `origin/main` 历史的完整 SHA。镜像 Tag 使用完整 SHA,线上只能使用 Registry digest,禁止使用 `latest`
5. 本地镜像发布必须由用户明确要求后执行:
```bash
./scripts/publish-release-images.sh --components auto
```
该命令只能构建、冒烟、推送镜像和生成 `dist/releases/<SHA>.json`,不得修改生产。用户只要求“构建镜像”“推送镜像”或明确限定为 publish 时,执行到此为止。
6. 用户明确要求“发布”“上线”或“部署到线上”时,该次授权同时覆盖 publish 和 deploy。publish 成功并核验 manifest、组件、digest 和验证结果后,应直接继续执行下列部署命令,无需再次请求确认:
```bash
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
```
7. deploy 前后必须验证线上基线、镜像架构/revision、数据库备份条件、内部及公网健康检查和自动应用回滚。未实际完成这些验证时不得声称生产发布成功。
8. 回滚也必须由用户明确要求,只能选择服务器已有的历史 manifest;应用回滚不自动恢复数据库。
+28 -6
View File
@@ -6,8 +6,8 @@ ARG NODE_BUILD_IMAGE=node:${NODE_VERSION}-alpine
ARG WEB_RUNTIME_IMAGE=nginx:1.27-alpine
FROM --platform=$BUILDPLATFORM ${GO_BUILD_IMAGE} AS api-builder
ARG TARGETOS=linux
ARG TARGETARCH=amd64
ARG TARGETOS
ARG TARGETARCH
ARG GOPROXY=https://goproxy.cn,direct
ENV GOPROXY=$GOPROXY
@@ -32,16 +32,28 @@ RUN --mount=type=cache,target=/go/pkg/mod \
--mount=type=cache,target=/root/.cache/go-build \
cd apps/api && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway ./cmd/gateway && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-backfill-binary-results ./cmd/backfill-binary-results && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-emulator ./cmd/acceptance-emulator && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-callback-collector ./cmd/acceptance-callback-collector && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-load ./cmd/acceptance-load && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-snapshot ./cmd/acceptance-snapshot && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-bootstrap ./cmd/acceptance-bootstrap
FROM ${API_RUNTIME_IMAGE} AS api
RUN apk add --no-cache ca-certificates tzdata wget && \
RUN apk add --no-cache ca-certificates tzdata wget ffmpeg && \
adduser -D -H -u 10001 appuser
WORKDIR /app
COPY --from=api-builder /out/easyai-ai-gateway /app/easyai-ai-gateway
COPY --from=api-builder /out/easyai-ai-gateway-migrate /app/easyai-ai-gateway-migrate
COPY --from=api-builder /out/easyai-ai-gateway-backfill-binary-results /app/easyai-ai-gateway-backfill-binary-results
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-emulator /app/easyai-ai-gateway-acceptance-emulator
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-callback-collector /app/easyai-ai-gateway-acceptance-callback-collector
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-load /app/easyai-ai-gateway-acceptance-load
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-snapshot /app/easyai-ai-gateway-acceptance-snapshot
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-bootstrap /app/easyai-ai-gateway-acceptance-bootstrap
COPY apps/api/migrations /app/migrations
RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \
@@ -57,15 +69,16 @@ ENV APP_ENV=production \
CMD ["/app/easyai-ai-gateway"]
FROM --platform=$BUILDPLATFORM ${NODE_BUILD_IMAGE} AS web-builder
ARG NPM_CONFIG_REGISTRY=https://registry.npmjs.org
WORKDIR /src
RUN npm install -g pnpm@10.18.1
RUN npm install --registry "$NPM_CONFIG_REGISTRY" -g pnpm@10.18.1
COPY package.json pnpm-lock.yaml pnpm-workspace.yaml nx.json ./
COPY apps/web/package.json apps/web/package.json
COPY packages/contracts/package.json packages/contracts/package.json
RUN --mount=type=cache,id=pnpm-store,target=/root/.local/share/pnpm/store \
pnpm install --frozen-lockfile
pnpm install --frozen-lockfile --registry "$NPM_CONFIG_REGISTRY"
COPY packages packages
COPY apps/web apps/web
@@ -82,3 +95,12 @@ COPY docker/nginx.conf /etc/nginx/conf.d/default.conf
COPY --from=web-builder /src/apps/web/dist /usr/share/nginx/html
EXPOSE 80
FROM alpine:3.22 AS acceptance-netem
RUN apk add --no-cache ca-certificates curl iproute2 socat
COPY scripts/acceptance/netem-entrypoint.sh /usr/local/bin/netem-entrypoint
RUN chmod 0555 /usr/local/bin/netem-entrypoint
USER 65534:65534
ENTRYPOINT ["/usr/local/bin/netem-entrypoint"]
+43 -17
View File
@@ -59,9 +59,9 @@ OIDC 用户通过签名、Issuer、Audience、`tid`、Scope 和应用角色校
后端热更新可通过 `GO_WATCH_SHUTDOWN_GRACE_MS``GO_WATCH_RESTART_DELAY_MS` 调整旧进程退出等待时间与重启间隔。
## Docker Compose 一键部署
## Docker Compose 本地运行
仓库内提供了面向 `linux/amd64` 的 Docker Compose 构建和部署脚本,会自动构建 API/Web 镜像、启动 PostgreSQL、执行数据库迁移,并验证 API 与 Web 是否可访问:
仓库内的 Compose 脚本只用于本地构建和运行:启动 PostgreSQL、执行数据库迁移,并验证 API 与 Web 是否可访问。它不再包含 Registry 推送或生产部署能力
```bash
scripts/deploy-compose.sh
@@ -78,28 +78,21 @@ scripts/deploy-compose.sh
常用覆盖项:
```bash
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 scripts/deploy-compose.sh
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 AI_GATEWAY_PUSH=1 scripts/deploy-compose.sh
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 scripts/deploy-compose.sh push
AI_GATEWAY_IMAGE_TAG=local-test scripts/deploy-compose.sh
AI_GATEWAY_WEB_PORT=8080 AI_GATEWAY_API_PORT=18088 scripts/deploy-compose.sh
AI_GATEWAY_GO_PROXY='https://proxy.golang.org,direct' scripts/deploy-compose.sh
AI_GATEWAY_NPM_REGISTRY='https://registry.npmmirror.com' scripts/deploy-compose.sh
AI_GATEWAY_SKIP_BUILD=1 scripts/deploy-compose.sh
scripts/deploy-compose.sh down
scripts/deploy-compose.sh clean
```
默认镜像地址为:
默认本地镜像地址为:
- API: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway:latest`
- Web: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway-web:latest`
- API: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway:local`
- Web: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway-web:local`
执行 `scripts/deploy-compose.sh push` 或设置 `AI_GATEWAY_PUSH=1` 时,会同时推送当前版本 tag 和 `latest`。当前版本 tag 优先使用 `AI_GATEWAY_IMAGE_TAG`;如果没有设置,则使用根 `package.json` 里的 `version`
推送前需要先登录阿里云镜像仓库:
```bash
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
```
生产镜像禁止使用 `latest`,只能通过后文的人工 publish 命令推送完整 Git SHA Tag
Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.conf](docker/nginx.conf),可直接修改该文件调整静态资源、规范 `/api/v1` 公开入口和旧 `/gateway-api` 兼容反向代理。修改后执行以下命令使配置生效:
@@ -107,9 +100,40 @@ Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.
docker compose -f docker-compose.yml restart web
```
## 生产 CI/CD
## 生产人工发布
Gitea Actions 会在隔离的 rootless DinD Runner 中对 Pull Request、`main` Push 和版本 Tag 执行完整质量门禁;Tag 使用独立的 `release-ci / verify-tag (push)` context,不能复用旧的 `main` 成功状态。源码 Job 没有宿主 Docker、`sudo` 或生产部署权限。部署仓的 root-owned dispatcher 只在相同 SHA 的 `main` 与 Tag context 都成功后,用固定命令构建并扫描镜像,再以 Registry digest 发布 `ai.51easyai.com`。安装 Runner、Fork PR 审批、发布验证和回滚步骤见 [生产 CI/CD 运行手册](docs/runbooks/production-ci-cd.md),信任边界见 [ADR-001](docs/decisions/001-production-cicd.md)
本仓库没有 Gitea Actions、Webhook、Tag、`main` Push、轮询或定时发布。`main` 不使用受保护分支限制;commit、push 和 Tag 都不会构建镜像或更新生产
生产发布由 Agent 在用户明确指令下按两个阶段连续执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
```bash
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
./scripts/publish-release-images.sh --components auto
```
用户要求“发布”“上线”或“部署到线上”时,Agent 核验并报告 `dist/releases/<SHA>.json`、组件和 digest 后直接执行第二步,无需再次确认;用户仅要求构建或推送镜像时才在第一步后停止:
```bash
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
```
查看生产版本和显式回滚:
```bash
./scripts/deploy-production-release.sh --status
./scripts/deploy-production-release.sh --rollback <历史完整 Git SHA>
```
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md)
三节点 K3s、CloudNativePG、双 NGINX、备份恢复和跨节点文件验收见
[K3s 高可用运行手册](docs/operations/k3s-ha-runbook.md);正式流量接入前的 Gemini
带图、多参考图视频、Worker 强杀和容量阶梯验收见
[生产同构验收手册](docs/operations/production-acceptance.md)。决策背景见
[ADR-003](docs/decisions/003-manual-agent-release.md)。
生产验收前先运行[本地三节点同构验收](docs/operations/local-isomorphic-acceptance.md)
本地原生高并发、精确 amd64 制品冒烟、线上模拟、真实金丝雀和人工开闸使用同一
`acceptance-report/v1` 身份链路。
Compose 默认使用独立容器数据库 `postgres:18-alpine`,数据卷会保留在 `postgres_data``api_data`。为避免本地开发 `.env` 中的 `localhost` 数据库地址污染容器部署,compose 使用 `AI_GATEWAY_COMPOSE_*` 变量作为容器部署专用覆盖,例如:
@@ -144,6 +168,8 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
异步队列 worker 不使用固定业务并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的集群容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`,限制策略推导出的集群目标;`AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT` 默认 `32`,按单 Worker 的内存安全容量限制实例分配,即使其他实例失活也不会突破。Worker 的 `AI_GATEWAY_DATABASE_MAX_CONNS` 必须高于实例执行容量,为心跳、选主、租约续期和健康检查保留连接;生产环境按每实例执行容量 `24` 配置连接池上限 `32``AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE` 默认 `8`,把短任务准入、容量租约和唯一 River job 合并为有界原子微批次,摊薄同步提交,同时禁止重新形成整窗大事务。异步任务首次排队时会持久化已经鉴权的候选和 admission scope 快照;启用 `AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED=true` 的进程批量读取快照并仅刷新动态执行容量,只有显式重新选路才重新水化媒体和计算候选。生产 K3s 在两地 API 启用 dispatcher、在 Worker 禁用,使当前 PostgreSQL 主库同站点 API 负责低延迟准入,远端 Worker 只执行 River job;数据库任务锁和 scope 锁保证主库切换时另一地 API 可安全接管。未显式配置该变量时保持兼容行为,由执行 Worker 同时运行 dispatcher。`AI_GATEWAY_DATABASE_MIN_IDLE_CONNS` 控制启动预热连接数,默认 `0`,生产 K3s 配置为 `4``AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS` 生产配置为 `300` 秒。平台模型和用户组的 PostgreSQL concurrency lease 仍是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
## 迁移原则
1. 新服务先并行运行,不直接删除 `easyai-server-main` 内现有模块。
+628
View File
@@ -0,0 +1,628 @@
package main
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"flag"
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
var (
fullSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
)
type options struct {
clusterID string
releaseSHA string
apiImageDigest string
workerDigest string
emulatorBaseURL string
callbackURL string
output string
identityShards int
}
type runtimeFile struct {
SchemaVersion string `json:"schemaVersion"`
LocalClusterID string `json:"localClusterId"`
RunID string `json:"runId"`
RunToken string `json:"runToken"`
APIKeys []string `json:"apiKeys"`
Participants []runtimeParticipant `json:"participants"`
GeminiModel string `json:"geminiModel"`
VideoModel string `json:"videoModel"`
EmulatorBaseURL string `json:"emulatorBaseUrl"`
CallbackURL string `json:"callbackUrl"`
ReleaseSHA string `json:"releaseSha"`
APIImageDigest string `json:"apiImageDigest"`
WorkerImageDigest string `json:"workerImageDigest"`
SnapshotConfigHash string `json:"snapshotConfigHash"`
SnapshotSHA256 string `json:"snapshotSha256"`
CreatedAt time.Time `json:"createdAt"`
}
type runtimeParticipant struct {
APIKeyID string `json:"apiKeyId"`
UserID string `json:"userId"`
}
func main() {
opts, err := parseOptions()
if err != nil {
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
os.Exit(64)
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
if err := run(ctx, opts); err != nil {
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
os.Exit(1)
}
}
func parseOptions() (options, error) {
var opts options
flag.StringVar(&opts.clusterID, "local-cluster-id", "", "required local cluster marker")
flag.StringVar(&opts.releaseSHA, "release-sha", "", "full source Git SHA")
flag.StringVar(&opts.apiImageDigest, "api-image-digest", "", "immutable API image digest")
flag.StringVar(&opts.workerDigest, "worker-image-digest", "", "immutable Worker image digest")
flag.StringVar(&opts.emulatorBaseURL, "emulator-base-url", "", "in-cluster protocol emulator URL")
flag.StringVar(&opts.callbackURL, "callback-url", "", "in-cluster callback collector URL")
flag.StringVar(&opts.output, "output", "", "private runtime output file")
flag.IntVar(&opts.identityShards, "identity-shards", 32, "isolated acceptance identities")
flag.Parse()
opts.clusterID = strings.TrimSpace(opts.clusterID)
opts.releaseSHA = strings.ToLower(strings.TrimSpace(opts.releaseSHA))
opts.apiImageDigest = strings.ToLower(strings.TrimSpace(opts.apiImageDigest))
opts.workerDigest = strings.ToLower(strings.TrimSpace(opts.workerDigest))
opts.emulatorBaseURL = strings.TrimRight(strings.TrimSpace(opts.emulatorBaseURL), "/")
opts.callbackURL = strings.TrimSpace(opts.callbackURL)
if opts.clusterID == "" || !fullSHAPattern.MatchString(opts.releaseSHA) ||
!digestPattern.MatchString(opts.apiImageDigest) || !digestPattern.MatchString(opts.workerDigest) {
return options{}, errors.New("local cluster ID, full release SHA, and immutable image digests are required")
}
if opts.emulatorBaseURL == "" || opts.callbackURL == "" {
return options{}, errors.New("emulator and callback URLs are required")
}
if opts.identityShards < 1 || opts.identityShards > 128 {
return options{}, errors.New("identity shards must be between 1 and 128")
}
if strings.TrimSpace(opts.output) == "" {
return options{}, errors.New("private runtime output path is required")
}
if flag.NArg() != 0 {
return options{}, errors.New("unexpected positional arguments")
}
return opts, nil
}
func run(ctx context.Context, opts options) error {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
database, err := store.ConnectWithMaxConns(ctx, databaseURL, 4)
if err != nil {
return err
}
defer database.Close()
if err := verifyLocalClusterMarker(ctx, database, opts.clusterID); err != nil {
return err
}
if err := resetPreviousLocalRun(ctx, database, opts); err != nil {
return err
}
if err := verifyLocalDatabase(ctx, database, opts.clusterID); err != nil {
return err
}
snapshotConfigHash, snapshotSHA, err := importedSnapshot(ctx, database)
if err != nil {
return err
}
groupID, tenantID, err := ensureAcceptanceIdentityDomain(ctx, database)
if err != nil {
return err
}
participants := make([]runtimeParticipant, 0, opts.identityShards)
apiKeys := make([]string, 0, opts.identityShards)
for ordinal := 0; ordinal < opts.identityShards; ordinal++ {
userID, err := ensureAcceptanceUser(ctx, database, groupID, tenantID, ordinal)
if err != nil {
return err
}
keyID, secret, err := ensureAcceptanceAPIKey(ctx, database, userID, groupID, tenantID, ordinal)
if err != nil {
return err
}
if _, err := database.SetUserWalletBalance(ctx, store.WalletBalanceAdjustmentInput{
GatewayUserID: userID,
Currency: "resource",
BalanceText: "1000000000",
Reason: "local acceptance isolated wallet",
}); err != nil && !errors.Is(err, store.ErrWalletBalanceUnchanged) {
return err
}
participants = append(participants, runtimeParticipant{APIKeyID: keyID, UserID: userID})
apiKeys = append(apiKeys, secret)
}
if err := ensureAcceptanceAccessRules(ctx, database, groupID); err != nil {
return err
}
if err := ensureAcceptanceObjectStorage(ctx, database, opts.emulatorBaseURL); err != nil {
return err
}
geminiModel, videoModel, err := selectedModels(ctx, database)
if err != nil {
return err
}
runToken, err := randomToken()
if err != nil {
return err
}
run, err := database.CreateAcceptanceRun(ctx, store.CreateAcceptanceRunInput{
ReleaseSHA: opts.releaseSHA,
APIImageDigest: opts.apiImageDigest,
WorkerImageDigest: opts.workerDigest,
APIKeyID: participants[0].APIKeyID,
UserID: participants[0].UserID,
Token: runToken,
EmulatorBaseURL: opts.emulatorBaseURL,
CallbackURL: opts.callbackURL,
CapacityProfile: "P24",
Config: map[string]any{
"workloads": []any{"gemini_image_edit", "multi_reference_video"},
"participants": participants,
"localClusterId": opts.clusterID,
"snapshotConfigHash": snapshotConfigHash,
"snapshotSha256": snapshotSHA,
},
})
if err != nil {
return err
}
if _, err := database.ActivateAcceptanceRun(ctx, run.ID); err != nil {
return err
}
outputPath, err := privateOutputPath(opts.output)
if err != nil {
return err
}
payload, err := json.MarshalIndent(runtimeFile{
SchemaVersion: "acceptance-runtime/v1",
LocalClusterID: opts.clusterID,
RunID: run.ID,
RunToken: runToken,
APIKeys: apiKeys,
Participants: participants,
GeminiModel: geminiModel,
VideoModel: videoModel,
EmulatorBaseURL: opts.emulatorBaseURL,
CallbackURL: opts.callbackURL,
ReleaseSHA: opts.releaseSHA,
APIImageDigest: opts.apiImageDigest,
WorkerImageDigest: opts.workerDigest,
SnapshotConfigHash: snapshotConfigHash,
SnapshotSHA256: snapshotSHA,
CreatedAt: time.Now().UTC(),
}, "", " ")
if err != nil {
return err
}
outputFile, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
if err != nil {
return err
}
if _, err := outputFile.Write(append(payload, '\n')); err != nil {
_ = outputFile.Close()
return err
}
if err := outputFile.Close(); err != nil {
return err
}
fmt.Printf(
"acceptance_bootstrap=PASS cluster_id=%s run_id=%s identities=%d config_hash=%s runtime_file=%s\n",
opts.clusterID,
run.ID,
len(participants),
snapshotConfigHash,
outputPath,
)
return nil
}
func ensureAcceptanceObjectStorage(ctx context.Context, database *store.Store, emulatorBaseURL string) error {
baseURL := strings.TrimRight(strings.TrimSpace(emulatorBaseURL), "/")
if baseURL == "" {
return errors.New("acceptance object storage requires the emulator URL")
}
channels := []struct {
key string
name string
provider string
endpoint string
priority int
pathStyle bool
}{
{
key: "local-acceptance-oss", name: "Local Acceptance Aliyun OSS", provider: "aliyun_oss",
endpoint: baseURL + "/storage/oss/bucket", priority: 10, pathStyle: true,
},
{
key: "local-acceptance-s3", name: "Local Acceptance S3", provider: "s3",
endpoint: baseURL + "/storage/s3", priority: 20, pathStyle: true,
},
}
for _, channel := range channels {
_, err := database.Pool().Exec(ctx, `
INSERT INTO file_storage_channels (
channel_key, name, provider, credentials, config, retry_policy, priority, status
)
VALUES (
$1, $2, $3,
'{"accessKeyId":"local-acceptance","accessKeySecret":"local-acceptance-secret"}'::jsonb,
jsonb_build_object(
'endpoint', $4::text,
'region', 'local-acceptance-1',
'bucket', 'bucket',
'objectPrefix', 'acceptance/media',
'accessScope', 'private',
'forcePathStyle', $5::boolean,
'scenes', jsonb_build_array('upload', 'image_result', 'request_asset'),
'acceptanceEmulatorOnly', true
),
'{"enabled":true,"maxRetries":2,"backoffSeconds":[0.25,1],"strategy":"exponential"}'::jsonb,
$6, 'enabled'
)
ON CONFLICT (channel_key) DO UPDATE
SET name = EXCLUDED.name,
provider = EXCLUDED.provider,
credentials = EXCLUDED.credentials,
config = EXCLUDED.config,
retry_policy = EXCLUDED.retry_policy,
priority = EXCLUDED.priority,
status = 'enabled',
deleted_at = NULL,
last_error = NULL,
updated_at = now()`, channel.key, channel.name, channel.provider, channel.endpoint, channel.pathStyle, channel.priority)
if err != nil {
return fmt.Errorf("configure %s object storage channel: %w", channel.provider, err)
}
}
_, err := database.Pool().Exec(ctx, `
INSERT INTO system_settings (setting_key, value)
VALUES ('file_storage', '{"resultUploadPolicy":"default"}'::jsonb)
ON CONFLICT (setting_key) DO UPDATE
SET value = jsonb_set(COALESCE(system_settings.value, '{}'::jsonb), '{resultUploadPolicy}', '"default"'::jsonb, true),
updated_at = now()`)
return err
}
func resetPreviousLocalRun(ctx context.Context, database *store.Store, opts options) error {
mode, err := database.GetGatewayTrafficMode(ctx)
if err != nil {
return err
}
if mode.Mode == "live" {
return nil
}
if mode.Mode != "validation" || strings.TrimSpace(mode.RunID) == "" {
return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode)
}
run, err := database.GetAcceptanceRun(ctx, mode.RunID)
if err != nil {
return err
}
if run.Config["localClusterId"] != opts.clusterID {
return errors.New("refusing to replace an acceptance run owned by another cluster")
}
var outstanding int
if err := database.Pool().QueryRow(ctx, `
SELECT count(*)
FROM gateway_tasks
WHERE acceptance_run_id = $1::uuid
AND status IN ('queued','running')`, run.ID).Scan(&outstanding); err != nil {
return err
}
if outstanding != 0 {
return fmt.Errorf("refusing to replace local acceptance run with %d outstanding tasks", outstanding)
}
_, err = database.AbortAcceptanceRun(ctx, store.PromoteAcceptanceRunInput{
RunID: mode.RunID,
Revision: mode.Revision,
ReleaseSHA: mode.ReleaseSHA,
APIImageDigest: mode.APIImageDigest,
WorkerImageDigest: mode.WorkerImageDigest,
})
return err
}
func verifyLocalDatabase(ctx context.Context, database *store.Store, clusterID string) error {
if err := verifyLocalClusterMarker(ctx, database, clusterID); err != nil {
return err
}
mode, err := database.GetGatewayTrafficMode(ctx)
if err != nil {
return err
}
if mode.Mode != "live" {
return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode)
}
return nil
}
func verifyLocalClusterMarker(ctx context.Context, database *store.Store, clusterID string) error {
var marker string
err := database.Pool().QueryRow(ctx, `
SELECT COALESCE(value->>'clusterId', '')
FROM system_settings
WHERE setting_key = $1`, acceptancesnapshot.LocalClusterSettingKey).Scan(&marker)
if err != nil {
return fmt.Errorf("read local acceptance cluster marker: %w", err)
}
if marker != clusterID {
return errors.New("refusing bootstrap: local cluster marker mismatch")
}
return nil
}
func importedSnapshot(ctx context.Context, database *store.Store) (string, string, error) {
var configHash, snapshotSHA string
err := database.Pool().QueryRow(ctx, `
SELECT COALESCE(value->>'configHash', ''),
COALESCE(value->>'snapshotSha256', '')
FROM system_settings
WHERE setting_key = 'acceptance_snapshot'`).Scan(&configHash, &snapshotSHA)
if err != nil {
return "", "", fmt.Errorf("read imported acceptance snapshot: %w", err)
}
if len(configHash) != 64 || len(snapshotSHA) != 64 {
return "", "", errors.New("imported acceptance snapshot hashes are invalid")
}
return configHash, snapshotSHA, nil
}
func ensureAcceptanceIdentityDomain(ctx context.Context, database *store.Store) (string, string, error) {
var groupID string
err := database.Pool().QueryRow(ctx, `
INSERT INTO gateway_user_groups (
group_key, name, description, source, priority,
recharge_discount_policy, billing_discount_policy,
rate_limit_policy, quota_policy, metadata, status
)
VALUES (
'local-acceptance', 'Local Acceptance', 'Isolated local homologous acceptance identities',
'gateway', 1, '{"discountFactor":1}'::jsonb, '{"discountFactor":1}'::jsonb,
'{"rules":[]}'::jsonb, '{}'::jsonb,
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
)
ON CONFLICT (group_key) DO UPDATE
SET name = EXCLUDED.name,
description = EXCLUDED.description,
priority = EXCLUDED.priority,
rate_limit_policy = EXCLUDED.rate_limit_policy,
quota_policy = EXCLUDED.quota_policy,
metadata = EXCLUDED.metadata,
status = 'active',
updated_at = now()
RETURNING id::text`).Scan(&groupID)
if err != nil {
return "", "", err
}
var tenantID string
err = database.Pool().QueryRow(ctx, `
INSERT INTO gateway_tenants (
tenant_key, source, external_tenant_id, name, default_user_group_id,
billing_profile, rate_limit_policy, auth_policy, metadata, status
)
VALUES (
'local-acceptance', 'gateway', 'local-acceptance', 'Local Acceptance',
$1::uuid, '{}'::jsonb, '{"rules":[]}'::jsonb, '{}'::jsonb,
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
)
ON CONFLICT (tenant_key) DO UPDATE
SET default_user_group_id = EXCLUDED.default_user_group_id,
metadata = EXCLUDED.metadata,
status = 'active',
updated_at = now()
RETURNING id::text`, groupID).Scan(&tenantID)
return groupID, tenantID, err
}
func ensureAcceptanceUser(
ctx context.Context,
database *store.Store,
groupID string,
tenantID string,
ordinal int,
) (string, error) {
userKey := fmt.Sprintf("local-acceptance-%03d", ordinal)
roles := `["user"]`
if ordinal == 0 {
roles = `["manager"]`
}
var userID string
err := database.Pool().QueryRow(ctx, `
INSERT INTO gateway_users (
user_key, source, external_user_id, username, display_name,
gateway_tenant_id, tenant_id, tenant_key, default_user_group_id,
roles, auth_profile, metadata, status
)
VALUES (
$1, 'gateway', $1, $1, $2,
$3::uuid, 'local-acceptance', 'local-acceptance', $4::uuid,
$5::jsonb, '{}'::jsonb,
jsonb_build_object('purpose','local_acceptance','ordinal',$6::integer,'isolated',true),
'active'
)
ON CONFLICT (user_key) DO UPDATE
SET gateway_tenant_id = EXCLUDED.gateway_tenant_id,
tenant_id = EXCLUDED.tenant_id,
tenant_key = EXCLUDED.tenant_key,
default_user_group_id = EXCLUDED.default_user_group_id,
roles = EXCLUDED.roles,
metadata = EXCLUDED.metadata,
status = 'active',
deleted_at = NULL,
updated_at = now()
RETURNING id::text`,
userKey,
fmt.Sprintf("Local Acceptance %03d", ordinal),
tenantID,
groupID,
roles,
ordinal,
).Scan(&userID)
return userID, err
}
func ensureAcceptanceAPIKey(
ctx context.Context,
database *store.Store,
userID string,
groupID string,
tenantID string,
ordinal int,
) (string, string, error) {
name := fmt.Sprintf("Local Acceptance %03d", ordinal)
var keyID, secret string
err := database.Pool().QueryRow(ctx, `
SELECT id::text, COALESCE(key_secret, '')
FROM gateway_api_keys
WHERE gateway_user_id = $1::uuid
AND name = $2
AND status = 'active'
AND deleted_at IS NULL
AND COALESCE(key_secret, '') <> ''
ORDER BY created_at
LIMIT 1`, userID, name).Scan(&keyID, &secret)
if err == nil {
return keyID, secret, nil
}
if !store.IsNotFound(err) {
return "", "", err
}
created, err := database.CreateAPIKey(ctx, store.CreateAPIKeyInput{
Name: name,
Scopes: []string{"image", "video"},
}, &auth.User{
ID: userID,
Username: fmt.Sprintf("local-acceptance-%03d", ordinal),
GatewayUserID: userID,
GatewayTenantID: tenantID,
TenantID: "local-acceptance",
TenantKey: "local-acceptance",
UserGroupID: groupID,
Source: "gateway",
})
if err != nil {
return "", "", err
}
if _, err := database.Pool().Exec(ctx, `
UPDATE gateway_api_keys
SET user_group_id = $2::uuid,
rate_limit_policy = '{"rules":[]}'::jsonb,
quota_policy = '{}'::jsonb,
updated_at = now()
WHERE id = $1::uuid`, created.APIKey.ID, groupID); err != nil {
return "", "", err
}
return created.APIKey.ID, created.Secret, nil
}
func ensureAcceptanceAccessRules(ctx context.Context, database *store.Store, groupID string) error {
_, err := database.Pool().Exec(ctx, `
WITH selected AS (
SELECT 'platform'::text AS resource_type, platform.id AS resource_id
FROM integration_platforms platform
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
UNION
SELECT 'platform_model', model.id
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
UNION
SELECT 'base_model', model.base_model_id
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.base_model_id IS NOT NULL
)
INSERT INTO gateway_access_rules (
subject_type, subject_id, resource_type, resource_id, effect,
priority, min_permission_level, conditions, metadata, status
)
SELECT 'user_group', $1::uuid, selected.resource_type, selected.resource_id,
'allow', 1, 0, '{}'::jsonb,
'{"purpose":"local_acceptance"}'::jsonb, 'active'
FROM selected
ON CONFLICT (subject_type, subject_id, resource_type, resource_id, effect)
DO UPDATE SET status = 'active', priority = 1, updated_at = now()`, groupID)
return err
}
func selectedModels(ctx context.Context, database *store.Store) (string, string, error) {
var gemini, video string
err := database.Pool().QueryRow(ctx, `
SELECT base_model.invocation_name
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
WHERE platform.status = 'enabled'
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.enabled = true
AND model.model_type @> '["image_edit"]'::jsonb
ORDER BY platform.priority, model.created_at
LIMIT 1`).Scan(&gemini)
if err != nil {
return "", "", err
}
err = database.Pool().QueryRow(ctx, `
SELECT base_model.invocation_name
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
WHERE platform.status = 'enabled'
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
AND model.enabled = true
AND model.model_type @> '["omni_video"]'::jsonb
ORDER BY platform.priority, model.created_at
LIMIT 1`).Scan(&video)
return gemini, video, err
}
func randomToken() (string, error) {
payload := make([]byte, 32)
if _, err := rand.Read(payload); err != nil {
return "", err
}
return hex.EncodeToString(payload), nil
}
func privateOutputPath(path string) (string, error) {
absolute, err := filepath.Abs(strings.TrimSpace(path))
if err != nil {
return "", err
}
if _, err := os.Lstat(absolute); err == nil {
return "", errors.New("runtime output already exists")
} else if !os.IsNotExist(err) {
return "", err
}
if err := os.MkdirAll(filepath.Dir(absolute), 0o700); err != nil {
return "", err
}
return absolute, nil
}
@@ -0,0 +1,120 @@
package main
import (
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"os"
"strconv"
"strings"
"sync"
"time"
)
const maxCallbackBodyBytes = 1 << 20
type collector struct {
mu sync.Mutex
deliveries int64
duplicates int64
invalid int64
seen map[string]int
}
func main() {
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
if address == "" {
address = ":8091"
}
state := &collector{seen: map[string]int{}}
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
})
mux.HandleFunc("GET /report", state.report)
mux.HandleFunc("POST /callbacks", state.callback)
server := &http.Server{
Addr: address,
Handler: mux,
ReadHeaderTimeout: 10 * time.Second,
}
slog.Info("acceptance callback collector started", "address", address)
if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
slog.Error("acceptance callback collector stopped", "error", err)
os.Exit(1)
}
}
func (c *collector) callback(w http.ResponseWriter, request *http.Request) {
body, err := io.ReadAll(io.LimitReader(request.Body, maxCallbackBodyBytes+1))
if err != nil || len(body) > maxCallbackBodyBytes {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback body"})
return
}
decoder := json.NewDecoder(strings.NewReader(string(body)))
decoder.UseNumber()
var payload map[string]any
if err := decoder.Decode(&payload); err != nil {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback JSON"})
return
}
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
seq, _ := strconv.ParseInt(stringValue(payload["seq"]), 10, 64)
idempotency := strings.TrimSpace(request.Header.Get("Idempotency-Key"))
expected := taskID + ":" + strconv.FormatInt(seq, 10)
if taskID == "" || seq <= 0 || idempotency != expected {
c.recordInvalid()
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "callback identity is invalid"})
return
}
c.mu.Lock()
c.deliveries++
c.seen[expected]++
if c.seen[expected] > 1 {
c.duplicates++
}
c.mu.Unlock()
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (c *collector) report(w http.ResponseWriter, _ *http.Request) {
c.mu.Lock()
report := map[string]any{
"schemaVersion": "acceptance-callback-report/v1",
"deliveries": c.deliveries,
"duplicates": c.duplicates,
"invalid": c.invalid,
"uniqueEvents": len(c.seen),
}
c.mu.Unlock()
writeJSON(w, http.StatusOK, report)
}
func (c *collector) recordInvalid() {
c.mu.Lock()
c.invalid++
c.mu.Unlock()
}
func stringValue(value any) string {
switch typed := value.(type) {
case string:
return typed
case json.Number:
return typed.String()
case float64:
return strconv.FormatFloat(typed, 'f', -1, 64)
default:
return ""
}
}
func writeJSON(w http.ResponseWriter, status int, value any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(value)
}
+28
View File
@@ -0,0 +1,28 @@
package main
import (
"log/slog"
"net/http"
"os"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceemulator"
)
func main() {
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
if address == "" {
address = ":8090"
}
server := &http.Server{
Addr: address,
Handler: acceptanceemulator.New(acceptanceemulator.Config{}).Handler(),
ReadHeaderTimeout: 10 * time.Second,
}
slog.Info("acceptance protocol emulator started", "address", address)
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
slog.Error("acceptance protocol emulator stopped", "error", err)
os.Exit(1)
}
}
File diff suppressed because it is too large Load Diff
+430
View File
@@ -0,0 +1,430 @@
package main
import (
"bytes"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
return f(request)
}
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
payload := paddedPNG(4 << 20)
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
base64.StdEncoding.EncodeToString(payload))
size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
if err != nil {
t.Fatalf("stream Gemini output: %v", err)
}
expected := sha256.Sum256(payload)
if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
t.Fatalf("size=%d digest=%s", size, digest)
}
}
func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
first := paddedPNGVariant(256<<10, 1)
second := paddedPNGVariant(256<<10, 2)
if len(first) != 256<<10 || len(second) != 256<<10 {
t.Fatalf("variant sizes=%d/%d", len(first), len(second))
}
firstHash := sha256.Sum256(first)
secondHash := sha256.Sum256(second)
if firstHash == secondHash {
t.Fatal("distinct variants have the same SHA-256")
}
}
func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
var body struct {
Contents []struct {
Parts []struct {
InlineData *struct {
MIMEType string `json:"mimeType"`
Data string `json:"data"`
} `json:"inlineData"`
} `json:"parts"`
} `json:"contents"`
GenerationConfig struct {
ResponseModalities []string `json:"responseModalities"`
} `json:"generationConfig"`
}
stream := streamGeminiRequestBody(inputs)
defer stream.Close()
if err := json.NewDecoder(stream).Decode(&body); err != nil {
t.Fatalf("decode streamed Gemini request: %v", err)
}
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
t.Fatalf("unexpected Gemini body structure: %+v", body)
}
for index, input := range inputs {
inlineData := body.Contents[0].Parts[index+1].InlineData
if inlineData == nil || inlineData.MIMEType != "image/png" {
t.Fatalf("image %d inline data=%+v", index, inlineData)
}
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
if err != nil {
t.Fatalf("decode streamed input %d: %v", index, err)
}
if !bytes.Equal(decoded, input) {
t.Fatalf("streamed input %d differs from source", index)
}
}
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
}
}
func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
inputs := geminiInputs(3, 768<<10, "replay-test", 91)
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
if err != nil {
t.Fatalf("new Gemini request: %v", err)
}
if request.GetBody == nil {
t.Fatal("streaming Gemini request does not provide GetBody")
}
first, err := io.ReadAll(request.Body)
if err != nil {
t.Fatalf("read first request body: %v", err)
}
_ = request.Body.Close()
replayedBody, err := request.GetBody()
if err != nil {
t.Fatalf("replay Gemini request body: %v", err)
}
replayed, err := io.ReadAll(replayedBody)
if err != nil {
t.Fatalf("read replayed request body: %v", err)
}
_ = replayedBody.Close()
if !bytes.Equal(first, replayed) {
t.Fatal("replayed Gemini request body differs from the original")
}
if request.ContentLength != int64(len(first)) {
t.Fatalf("content length=%d, want %d", request.ContentLength, len(first))
}
}
func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
"https://fixtures.example/one.png",
"https://fixtures.example/two.png",
"https://fixtures.example/three.png",
})
if err != nil {
t.Fatalf("new multi-image request: %v", err)
}
var body struct {
Contents []struct {
Parts []struct {
FileData *struct {
FileURI string `json:"fileUri"`
} `json:"fileData"`
} `json:"parts"`
} `json:"contents"`
}
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
t.Fatalf("decode multi-image request: %v", err)
}
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
t.Fatalf("unexpected multi-image body: %+v", body)
}
for index := 1; index < 4; index++ {
if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
t.Fatalf("missing fileData at part %d", index)
}
}
}
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
images := make([]string, 16)
for index := range images {
images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
}
combinations := videoCombinations(images, 128)
if len(combinations) != 128 {
t.Fatalf("combinations=%d", len(combinations))
}
seen := map[string]struct{}{}
for _, combination := range combinations {
seen[fmt.Sprint(combination)] = struct{}{}
if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
t.Fatalf("invalid combination size=%d", len(combination))
}
}
if len(seen) != 128 {
t.Fatalf("unique combinations=%d", len(seen))
}
if got := requestedVideoCombinationCount(combinations, "video-capacity", 144, 0, 1); got != 131 {
t.Fatalf("requested combinations=%d, want 131", got)
}
}
func TestVideoProviderQuotaCombinationsUseDistinctNormalSizedImages(t *testing.T) {
images := []string{
"https://fixtures.example/image-00.png",
"https://fixtures.example/image-01.png",
"https://fixtures.example/image-02.png",
"https://fixtures.example/image-03.png",
}
combinations := videoProviderQuotaCombinations(images, 24)
if got := requestedVideoCombinationCount(combinations, "video-provider-quota", 24, 0, 1); got != 24 {
t.Fatalf("quota combinations=%d, want 24", got)
}
for _, combination := range combinations {
if len(combination) != 3 {
t.Fatalf("invalid quota combination: %v", combination)
}
for _, imageURL := range combination {
if !strings.Contains(imageURL, "acceptance_quota_variant=") {
t.Fatalf("quota image lacks unique variant: %s", imageURL)
}
}
}
}
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
output := paddedPNG(256 << 10)
encoded := base64.StdEncoding.EncodeToString(output)
var firstCalls atomic.Int64
var secondCalls atomic.Int64
var firstKeyCalls atomic.Int64
var secondKeyCalls atomic.Int64
newGateway := func(calls *atomic.Int64) *httptest.Server {
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls.Add(1)
switch r.Header.Get("Authorization") {
case "Bearer key-1":
firstKeyCalls.Add(1)
case "Bearer key-2":
secondKeyCalls.Add(1)
default:
t.Errorf("unexpected authorization header")
}
if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
t.Errorf("missing acceptance headers")
}
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Errorf("decode Gemini body: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"candidates": []any{map[string]any{"content": map[string]any{
"parts": []any{map[string]any{"inlineData": map[string]any{
"mimeType": "image/png", "data": encoded,
}}},
}}},
})
}))
}
first := newGateway(&firstCalls)
defer first.Close()
second := newGateway(&secondCalls)
defer second.Close()
opts := options{
gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
runToken: "token-1", geminiModel: "gemini-image-test",
}
result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1)
if result.err != nil || result.report.Completed != 8 {
t.Fatalf("result=%+v err=%v", result.report, result.err)
}
if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
}
if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
}
}
func TestDistributedShardIndexesCoverWorkloadWithoutOverlap(t *testing.T) {
first := options{shardIndex: 0, shardCount: 2}
second := options{shardIndex: 1, shardCount: 2}
if first.shardRequestCount(1001) != 501 || second.shardRequestCount(1001) != 500 {
t.Fatalf("shard counts=%d/%d", first.shardRequestCount(1001), second.shardRequestCount(1001))
}
seen := map[int]bool{}
for local := 0; local < first.shardRequestCount(1001); local++ {
seen[first.logicalRequestIndex(local)] = true
}
for local := 0; local < second.shardRequestCount(1001); local++ {
index := second.logicalRequestIndex(local)
if seen[index] {
t.Fatalf("duplicate logical index %d", index)
}
seen[index] = true
}
if len(seen) != 1001 {
t.Fatalf("covered indexes=%d", len(seen))
}
}
func TestGeminiInputVariantIsRunScoped(t *testing.T) {
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-b", 7) {
t.Fatal("different Run IDs produced the same input variant")
}
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-a", 8) {
t.Fatal("different logical indexes produced the same input variant")
}
}
func TestAcceptanceIdempotencyKeyIncludesExecutionAndLogicalIndex(t *testing.T) {
request := httptest.NewRequest(http.MethodPost, "https://gateway.example/v1", nil)
opts := options{
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1", executionID: "p24-2-gemini-baseline",
}
opts.setHeaders(request, 17, false)
if got := request.Header.Get("Idempotency-Key"); got != "acceptance-run-1-p24-2-gemini-baseline-17" {
t.Fatalf("idempotency key=%q", got)
}
poll := httptest.NewRequest(http.MethodGet, "https://gateway.example/result", nil)
opts.setHeaders(poll, 17, false)
if got := poll.Header.Get("Idempotency-Key"); got != "" {
t.Fatalf("GET request unexpectedly has idempotency key %q", got)
}
}
func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Header.Get("Authorization") != "" {
t.Fatal("acceptance credentials leaked to external media host")
}
w.Header().Set("Content-Type", "video/mp4")
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
}))
defer server.Close()
if err := validateVideoAsset(t.Context(), server.Client(), options{}, server.URL+"/result.mp4", 0); err != nil {
t.Fatalf("validate video asset: %v", err)
}
if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
t.Fatalf("media URL=%q", got)
}
}
func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "video/mp4")
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
}))
defer server.Close()
poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("gateway-only transport must not be used for external media")
})}
opts := options{gatewayTLSName: "gateway.easyai.local"}
if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
t.Fatalf("validate external video with independent transport: %v", err)
}
}
func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {
t.Fatalf("host=%q authorization=%q", request.Host, request.Header.Get("Authorization"))
}
if request.URL.Path != "/static/generated/result.mp4" {
t.Fatalf("path=%q", request.URL.Path)
}
w.Header().Set("Content-Type", "video/mp4")
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
}))
defer server.Close()
opts := options{
gateways: []string{server.URL}, apiKeys: []string{"key-1"},
runID: "run-1", runToken: "token-1", gatewayTLSName: "gateway.easyai.local",
}
if err := validateVideoAsset(t.Context(), server.Client(), opts, "/static/generated/result.mp4", 0); err != nil {
t.Fatalf("validate materialized video: %v", err)
}
got := findMediaURL(map[string]any{
"raw": map[string]any{"video_url": "http://internal.invalid/video.mp4"},
"data": []any{map[string]any{"video_url": "/static/generated/result.mp4"}},
})
if got != "/static/generated/result.mp4" {
t.Fatalf("preferred media URL=%q", got)
}
}
func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
got := redactError(
`token-1 failed at https://example.invalid/video.mp4?token=signed`,
options{runToken: "token-1"},
)
if got != `[REDACTED] failed at [REDACTED_URL]` {
t.Fatalf("redacted error=%q", got)
}
}
func TestAcceptanceReportCannotOverwriteExistingRunArtifact(t *testing.T) {
path := filepath.Join(t.TempDir(), "load-report.json")
if err := writeReportExclusive(path, []byte(`{"runId":"first"}`)); err != nil {
t.Fatalf("write first report: %v", err)
}
if err := writeReportExclusive(path, []byte(`{"runId":"second"}`)); err == nil {
t.Fatal("second report unexpectedly overwrote the first report")
}
payload, err := os.ReadFile(path)
if err != nil {
t.Fatalf("read report: %v", err)
}
if string(payload) != "{\"runId\":\"first\"}\n" {
t.Fatalf("report changed after rejected overwrite: %s", payload)
}
}
func TestAcceptanceFailurePreservesOperationAndHTTPStatus(t *testing.T) {
err := withOperation("video_poll", &httpStatusError{Status: http.StatusUnauthorized, Body: "unauthorized"})
var operationErr *operationError
if !errors.As(err, &operationErr) || operationErr.Operation != "video_poll" {
t.Fatalf("operation error=%#v", operationErr)
}
var statusErr *httpStatusError
if !errors.As(err, &statusErr) || statusErr.Status != http.StatusUnauthorized {
t.Fatalf("HTTP status error=%#v", statusErr)
}
}
func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
opts := options{
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
gatewayTLSName: "ai.example.com",
}
opts.setHeaders(request, 0, false)
if request.Host != "ai.example.com" {
t.Fatalf("Host=%q", request.Host)
}
}
func TestAcceptanceRootCAsRejectsSymlink(t *testing.T) {
root := t.TempDir()
target := filepath.Join(root, "ca.pem")
if err := os.WriteFile(target, []byte("not a certificate"), 0o600); err != nil {
t.Fatal(err)
}
link := filepath.Join(root, "ca-link.pem")
if err := os.Symlink(target, link); err != nil {
t.Fatal(err)
}
if _, err := acceptanceRootCAs(link); err == nil {
t.Fatal("symlink CA file was accepted")
}
}
+195
View File
@@ -0,0 +1,195 @@
package main
import (
"context"
"errors"
"flag"
"fmt"
"os"
"path/filepath"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
"github.com/jackc/pgx/v5/pgxpool"
)
func main() {
if len(os.Args) < 2 {
usage()
os.Exit(64)
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
var err error
switch os.Args[1] {
case "export":
err = runExport(ctx, os.Args[2:])
case "validate":
err = runValidate(os.Args[2:])
case "import":
err = runImport(ctx, os.Args[2:])
default:
usage()
os.Exit(64)
}
if err != nil {
fmt.Fprintln(os.Stderr, "acceptance snapshot:", err)
os.Exit(1)
}
}
func usage() {
fmt.Fprintln(os.Stderr, `Usage:
easyai-ai-gateway-acceptance-snapshot export --release-sha <SHA> --output <file>
easyai-ai-gateway-acceptance-snapshot validate --input <file>
easyai-ai-gateway-acceptance-snapshot import --input <file> --local-cluster-id <id>
export is read-only. import refuses to write unless system_settings contains a
matching acceptance_local_cluster_id marker.`)
}
func runExport(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("export", flag.ContinueOnError)
releaseSHA := flags.String("release-sha", "", "full source release SHA")
output := flags.String("output", "", "secret-safe snapshot output path")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected export arguments")
}
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
outputPath, err := safeOutputPath(*output)
if err != nil {
return err
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
return err
}
defer pool.Close()
snapshot, err := acceptancesnapshot.Export(ctx, pool, acceptancesnapshot.ExportOptions{
ReleaseSHA: *releaseSHA,
})
if err != nil {
return err
}
payload, err := acceptancesnapshot.Encode(snapshot)
if err != nil {
return err
}
if err := os.WriteFile(outputPath, payload, 0o600); err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_export=PASS schema=%s config_hash=%s snapshot_sha256=%s\n",
snapshot.SchemaVersion,
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func runValidate(args []string) error {
flags := flag.NewFlagSet("validate", flag.ContinueOnError)
input := flags.String("input", "", "snapshot input path")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected validate arguments")
}
snapshot, err := readSnapshot(*input)
if err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_validate=PASS schema=%s release=%s config_hash=%s snapshot_sha256=%s\n",
snapshot.SchemaVersion,
snapshot.Source.ReleaseSHA,
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func runImport(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("import", flag.ContinueOnError)
input := flags.String("input", "", "snapshot input path")
clusterID := flags.String("local-cluster-id", "", "required local cluster marker")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected import arguments")
}
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
if databaseURL == "" {
return errors.New("AI_GATEWAY_DATABASE_URL is required")
}
snapshot, err := readSnapshot(*input)
if err != nil {
return err
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
return err
}
defer pool.Close()
if err := acceptancesnapshot.Import(ctx, pool, snapshot, *clusterID); err != nil {
return err
}
fmt.Printf(
"acceptance_snapshot_import=PASS cluster_id=%s config_hash=%s snapshot_sha256=%s\n",
strings.TrimSpace(*clusterID),
snapshot.Source.ConfigHash,
snapshot.SnapshotSHA256,
)
return nil
}
func readSnapshot(path string) (acceptancesnapshot.Snapshot, error) {
path = strings.TrimSpace(path)
if path == "" {
return acceptancesnapshot.Snapshot{}, errors.New("snapshot path is required")
}
info, err := os.Lstat(path)
if err != nil {
return acceptancesnapshot.Snapshot{}, err
}
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return acceptancesnapshot.Snapshot{}, errors.New("snapshot must be a regular non-symlink file")
}
payload, err := os.ReadFile(path)
if err != nil {
return acceptancesnapshot.Snapshot{}, err
}
return acceptancesnapshot.Decode(payload)
}
func safeOutputPath(path string) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", errors.New("output path is required")
}
absolute, err := filepath.Abs(path)
if err != nil {
return "", err
}
parent := filepath.Dir(absolute)
if err := os.MkdirAll(parent, 0o700); err != nil {
return "", err
}
if info, err := os.Lstat(absolute); err == nil {
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return "", errors.New("output must be a regular non-symlink file")
}
} else if !os.IsNotExist(err) {
return "", err
}
return absolute, nil
}
+174
View File
@@ -0,0 +1,174 @@
package main
import (
"context"
"errors"
"flag"
"fmt"
"os"
"path/filepath"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/accessruleaudit"
"github.com/jackc/pgx/v5/pgxpool"
)
func main() {
if len(os.Args) < 2 {
usage()
os.Exit(64)
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
var err error
switch os.Args[1] {
case "export":
err = runExport(ctx, os.Args[2:])
case "verify":
err = runVerify(ctx, os.Args[2:])
default:
usage()
os.Exit(64)
}
if err != nil {
fmt.Fprintln(os.Stderr, "access-rule audit:", err)
os.Exit(1)
}
}
func usage() {
fmt.Fprintln(os.Stderr, `Usage:
easyai-ai-gateway-access-rule-audit export --output <before.json>
easyai-ai-gateway-access-rule-audit verify --before <before.json> --output <after.json>
Both commands are database read-only and only emit grouped counts and SHA-256
digests. Use a SELECT-only AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL role.`)
}
func runExport(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("export", flag.ContinueOnError)
output := flags.String("output", "", "secret-safe access-rule audit output")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected export arguments")
}
snapshot, outputPath, err := exportSnapshot(ctx, *output)
if err != nil {
return err
}
if err := writeSnapshot(outputPath, snapshot); err != nil {
return err
}
fmt.Printf("access_rule_audit_export=PASS total=%d allow=%d deny=%d live_sha256=%s\n", snapshot.Live.Total, snapshot.Live.AllowCount, snapshot.Live.DenyCount, snapshot.Live.SHA256)
return nil
}
func runVerify(ctx context.Context, args []string) error {
flags := flag.NewFlagSet("verify", flag.ContinueOnError)
beforePathValue := flags.String("before", "", "pre-migration audit snapshot")
output := flags.String("output", "", "post-migration audit output")
if err := flags.Parse(args); err != nil {
return err
}
if flags.NArg() != 0 {
return errors.New("unexpected verify arguments")
}
beforePath, err := regularInputPath(*beforePathValue)
if err != nil {
return err
}
after, outputPath, err := exportSnapshot(ctx, *output)
if err != nil {
return err
}
if beforePath == outputPath {
return errors.New("verification output must not overwrite the pre-migration snapshot")
}
payload, err := os.ReadFile(beforePath)
if err != nil {
return err
}
before, err := accessruleaudit.Decode(payload)
if err != nil {
return err
}
if err := accessruleaudit.VerifyMigration(before, after); err != nil {
return err
}
if err := writeSnapshot(outputPath, after); err != nil {
return err
}
fmt.Printf("access_rule_audit_verify=PASS archived_allow=%d live_allow=%d live_deny=%d deny_sha256=%s\n", after.Archive.ArchivedAllowCount, after.Live.AllowCount, after.Live.DenyCount, after.Live.DenySHA256)
return nil
}
func exportSnapshot(ctx context.Context, output string) (accessruleaudit.Snapshot, string, error) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL"))
if databaseURL == "" {
return accessruleaudit.Snapshot{}, "", errors.New("AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL is required")
}
outputPath, err := safeOutputPath(output)
if err != nil {
return accessruleaudit.Snapshot{}, "", err
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
return accessruleaudit.Snapshot{}, "", err
}
defer pool.Close()
snapshot, err := accessruleaudit.Export(ctx, pool, time.Now())
return snapshot, outputPath, err
}
func writeSnapshot(path string, snapshot accessruleaudit.Snapshot) error {
payload, err := accessruleaudit.Encode(snapshot)
if err != nil {
return err
}
return os.WriteFile(path, payload, 0o600)
}
func regularInputPath(path string) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", errors.New("input path is required")
}
absolute, err := filepath.Abs(path)
if err != nil {
return "", err
}
info, err := os.Lstat(absolute)
if err != nil {
return "", err
}
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return "", errors.New("input must be a regular non-symlink file")
}
return absolute, nil
}
func safeOutputPath(path string) (string, error) {
path = strings.TrimSpace(path)
if path == "" {
return "", errors.New("output path is required")
}
absolute, err := filepath.Abs(path)
if err != nil {
return "", err
}
parent := filepath.Dir(absolute)
if err := os.MkdirAll(parent, 0o700); err != nil {
return "", err
}
if info, err := os.Lstat(absolute); err == nil {
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
return "", errors.New("output must be a regular non-symlink file")
}
} else if !os.IsNotExist(err) {
return "", err
}
return absolute, nil
}
@@ -0,0 +1,133 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func main() {
apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
requireClean := flag.Bool("require-clean", false, "exit non-zero unless a complete dry-run finds no results requiring URL migration")
batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
afterID := flag.String("after-id", "", "resume after this task UUID")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg := config.Load()
if err := cfg.Validate(); err != nil {
logger.Error("invalid gateway configuration", "error", err)
os.Exit(1)
}
if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
os.Exit(1)
}
if *apply && *requireClean {
logger.Error("--apply and --require-clean cannot be combined")
os.Exit(1)
}
ctx := context.Background()
db, err := store.Connect(ctx, cfg.DatabaseURL)
if err != nil {
logger.Error("connect postgres failed", "error", err)
os.Exit(1)
}
defer db.Close()
service := runner.New(cfg, db, logger)
cursor := *afterID
scanned := 0
matched := 0
blockingMatched := 0
updated := 0
expired := 0
expiredLocalPlaceholders := 0
complete := false
for batch := 0; batch < *maxBatches; batch++ {
items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
if err != nil {
logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
os.Exit(1)
}
if len(items) == 0 {
complete = true
break
}
for _, item := range items {
cursor = item.ID
scanned++
if !runner.TaskResultNeedsURLMigration(item.Result) {
continue
}
matched++
isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
hasLocalPlaceholder := runner.TaskResultHasLocalPlaceholder(item.Result)
if isExpired && hasLocalPlaceholder {
expiredLocalPlaceholders++
if *apply {
logger.Warn("skip expired local result placeholder without overwriting stored result", "taskId", item.ID)
}
continue
}
blockingMatched++
if !*apply {
continue
}
persistent, changed, err := service.MigrateTaskResultToURLs(ctx, item.ID, item.Result)
if err != nil {
logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
os.Exit(1)
}
if !changed {
continue
}
ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
if err != nil {
logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
os.Exit(1)
}
if !ok {
continue
}
updated++
if isExpired {
expired++
}
}
if len(items) < *batchSize {
complete = true
break
}
}
logger.Info("binary result backfill completed",
"apply", *apply,
"scanned", scanned,
"matched", matched,
"blockingMatched", blockingMatched,
"updated", updated,
"expired", expired,
"expiredLocalPlaceholders", expiredLocalPlaceholders,
"complete", complete,
"resumeAfterId", cursor,
)
if *requireClean && (!complete || blockingMatched > 0) {
logger.Error("binary result URL migration gate failed", "complete", complete, "blockingMatched", blockingMatched, "expiredLocalPlaceholders", expiredLocalPlaceholders, "resumeAfterId", cursor)
os.Exit(1)
}
}
func localResultTTLHours(cfg config.Config) int {
if cfg.LocalResultTTLHours <= 0 {
return 24
}
return cfg.LocalResultTTLHours
}
+79 -14
View File
@@ -10,6 +10,7 @@ import (
"syscall"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/capacitycontroller"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/httpapi"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -36,30 +37,94 @@ func main() {
}
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
db, err := store.Connect(ctx, cfg.DatabaseURL)
executionMaxConns := cfg.DatabaseMaxConns
if cfg.DatabaseCriticalMaxConns > 0 {
executionMaxConns -= cfg.DatabaseCriticalMaxConns
}
if cfg.DatabaseRiverMaxConns > 0 {
executionMaxConns -= cfg.DatabaseRiverMaxConns
}
db, err := store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: executionMaxConns,
MinIdleConns: cfg.DatabaseMinIdleConns,
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect postgres failed", "error", err)
os.Exit(1)
}
defer db.Close()
if recovery, err := db.RecoverInterruptedRuntimeState(ctx); err != nil {
logger.Error("recover interrupted runtime state failed", "error", err)
os.Exit(1)
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 {
logger.Warn("interrupted runtime state recovered",
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
"releasedRateReservations", recovery.ReleasedRateReservations,
"failedAttempts", recovery.FailedAttempts,
"failedTasks", recovery.FailedTasks,
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
)
coordinationDB := db
if cfg.DatabaseCriticalMaxConns > 0 {
coordinationDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseCriticalMaxConns,
MinIdleConns: min(cfg.DatabaseCriticalMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect critical postgres pool failed", "error", err)
os.Exit(1)
}
defer coordinationDB.Close()
}
riverDB := db
if cfg.DatabaseRiverMaxConns > 0 {
riverDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
MaxConns: cfg.DatabaseRiverMaxConns,
MinIdleConns: min(cfg.DatabaseRiverMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
})
if err != nil {
stop()
logger.Error("connect River postgres pool failed", "error", err)
os.Exit(1)
}
defer riverDB.Close()
}
// Cancel background LISTEN and worker goroutines before closing the pool.
// This also prevents a startup panic from waiting forever in pgxpool.Close.
defer stop()
if cfg.RunsBackgroundWorkers() {
if recovery, recoveryErr := coordinationDB.RecoverInterruptedRuntimeState(ctx); recoveryErr != nil {
logger.Error("recover interrupted runtime state failed", "error", recoveryErr)
os.Exit(1)
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 {
logger.Warn("interrupted runtime state recovered",
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
"releasedRateReservations", recovery.ReleasedRateReservations,
"failedAttempts", recovery.FailedAttempts,
"failedTasks", recovery.FailedTasks,
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
"cleanedTaskAdmissions", recovery.CleanedTaskAdmissions,
)
}
}
var handler http.Handler
if cfg.RunsCapacityController() {
orchestrator, orchestratorErr := capacitycontroller.NewConfiguredAdapter(cfg)
if orchestratorErr != nil {
logger.Error("initialize capacity orchestrator adapter failed", "error", orchestratorErr)
os.Exit(1)
}
controller := capacitycontroller.New(cfg, coordinationDB, orchestrator, logger)
go controller.Run(ctx)
handler = controller.Handler()
} else {
handler = httpapi.NewServerWithStores(ctx, cfg, db, coordinationDB, riverDB, logger)
}
server := &http.Server{
Addr: cfg.HTTPAddr,
Handler: httpapi.NewServerWithContext(ctx, cfg, db, logger),
Handler: handler,
ReadHeaderTimeout: 10 * time.Second,
}
+89 -10
View File
@@ -6,6 +6,7 @@ import (
"log/slog"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
@@ -13,6 +14,16 @@ import (
"github.com/jackc/pgx/v5"
)
const (
noTransactionMigrationMarker = "-- easyai:migration:no-transaction"
migrationStatementSeparator = "-- easyai:migration:statement"
acceptanceImportReplayMarker = "-- easyai:migration:reapply-after-acceptance-import"
acceptanceImportReplayEnvironment = "AI_GATEWAY_MIGRATION_REAPPLY_ACCEPTANCE_IMPORT"
acceptanceLocalClusterSettingKey = "acceptance_local_cluster_id"
)
var acceptanceLocalClusterMarkerPattern = regexp.MustCompile(`^easyai-local-[0-9a-f]{24}$`)
func main() {
cfg := config.Load()
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
@@ -37,6 +48,21 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
logger.Error("ensure schema_migrations failed", "error", err)
os.Exit(1)
}
reapplyAcceptanceImport := strings.EqualFold(strings.TrimSpace(os.Getenv(acceptanceImportReplayEnvironment)), "true")
if reapplyAcceptanceImport {
var localClusterID string
if err := conn.QueryRow(ctx, `
SELECT COALESCE(value->>'clusterId', '')
FROM system_settings
WHERE setting_key = $1`, acceptanceLocalClusterSettingKey).Scan(&localClusterID); err != nil {
logger.Error("verify local acceptance database failed", "error", err)
os.Exit(1)
}
if !acceptanceLocalClusterMarkerPattern.MatchString(localClusterID) {
logger.Error("refusing acceptance migration replay outside a marked local cluster")
os.Exit(1)
}
}
files, err := filepath.Glob("migrations/*.sql")
if err != nil {
@@ -52,16 +78,38 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
logger.Error("check migration failed", "version", version, "error", err)
os.Exit(1)
}
if exists {
logger.Info("migration skipped", "version", version)
continue
}
sqlBytes, err := os.ReadFile(file)
if err != nil {
logger.Error("read migration file failed", "file", file, "error", err)
os.Exit(1)
}
replaying := exists && reapplyAcceptanceImport && hasMigrationMarker(string(sqlBytes), acceptanceImportReplayMarker)
if exists && !replaying {
logger.Info("migration skipped", "version", version)
continue
}
noTransaction, statements := migrationStatements(string(sqlBytes))
if noTransaction {
for _, statement := range statements {
if _, err := conn.Exec(ctx, statement); err != nil {
logger.Error("execute non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
}
if !exists {
if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
logger.Error("record non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
}
if replaying {
logger.Info("acceptance migration replayed", "version", version, "transactional", false)
} else {
logger.Info("migration applied", "version", version, "transactional", false)
}
continue
}
tx, err := conn.Begin(ctx)
if err != nil {
@@ -81,17 +129,48 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
logger.Error("execute migration failed", "version", version, "error", err)
os.Exit(1)
}
if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
_ = tx.Rollback(ctx)
logger.Error("record migration failed", "version", version, "error", err)
os.Exit(1)
if !exists {
if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
_ = tx.Rollback(ctx)
logger.Error("record migration failed", "version", version, "error", err)
os.Exit(1)
}
}
if err := tx.Commit(ctx); err != nil {
logger.Error("commit migration failed", "version", version, "error", err)
os.Exit(1)
}
logger.Info("migration applied", "version", version)
if replaying {
logger.Info("acceptance migration replayed", "version", version)
} else {
logger.Info("migration applied", "version", version)
}
}
fmt.Println("migrations complete")
}
func hasMigrationMarker(sql string, marker string) bool {
for _, line := range strings.Split(sql, "\n") {
if strings.TrimSpace(line) == marker {
return true
}
}
return false
}
func migrationStatements(sql string) (bool, []string) {
trimmed := strings.TrimSpace(sql)
if !strings.HasPrefix(trimmed, noTransactionMigrationMarker) {
return false, []string{sql}
}
trimmed = strings.TrimSpace(strings.TrimPrefix(trimmed, noTransactionMigrationMarker))
parts := strings.Split(trimmed, migrationStatementSeparator)
statements := make([]string, 0, len(parts))
for _, part := range parts {
if statement := strings.TrimSpace(part); statement != "" {
statements = append(statements, statement)
}
}
return true, statements
}
+124
View File
@@ -14,6 +14,102 @@ import (
"github.com/jackc/pgx/v5/pgxpool"
)
func TestMigrationStatementsSupportsConcurrentIndexes(t *testing.T) {
payload := `-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS first_index ON first_table(id);
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS second_index ON second_table(id);
`
noTransaction, statements := migrationStatements(payload)
if !noTransaction {
t.Fatal("expected a non-transaction migration")
}
if len(statements) != 2 {
t.Fatalf("statements=%d, want 2", len(statements))
}
if !strings.Contains(statements[0], "first_index") || !strings.Contains(statements[1], "second_index") {
t.Fatalf("unexpected statements: %#v", statements)
}
}
func TestAcceptanceImportReplayMarkerRequiresAnExactCommentLine(t *testing.T) {
if !hasMigrationMarker(
"-- preface\n"+acceptanceImportReplayMarker+"\nSELECT 1;\n",
acceptanceImportReplayMarker,
) {
t.Fatal("expected exact acceptance import replay marker")
}
if hasMigrationMarker(
"-- mentions "+acceptanceImportReplayMarker+" in prose\nSELECT 1;\n",
acceptanceImportReplayMarker,
) {
t.Fatal("prose mention must not enable acceptance import replay")
}
}
func TestSeedanceInputImageConstraintMigrationKeepsCatalogSnapshotsInSync(t *testing.T) {
payload, err := os.ReadFile("../../migrations/0078_seedance_input_image_constraints.sql")
if err != nil {
t.Fatal(err)
}
content := string(payload)
for _, required := range []string{
"easyai:豆包Seedance-2.0",
"easyai:豆包Seedance-2.0-fast",
"input_image_resolution_range",
"input_image_aspect_ratio_range",
"'{metadata,rawModel,capabilities}'",
"UPDATE platform_models",
} {
if !strings.Contains(content, required) {
t.Fatalf("Seedance input image migration is missing %q", required)
}
}
}
func TestVolcesSeedanceInputImageConstraintMigrationUsesDocumentedBounds(t *testing.T) {
payload, err := os.ReadFile("../../migrations/0098_volces_seedance20_input_image_constraints.sql")
if err != nil {
t.Fatal(err)
}
content := string(payload)
for _, required := range []string{
"volces:doubao-seedance-2-0-260128",
"volces:doubao-seedance-2-0-fast-260128",
"volces:doubao-seedance-2-0-mini-260615",
`"long_edge":300`,
`"long_edge":6000`,
`'[0.4,2.5]'::jsonb`,
"'{metadata,rawModel,capabilities}'",
"UPDATE platform_models",
} {
if !strings.Contains(content, required) {
t.Fatalf("Volces Seedance input image migration is missing %q", required)
}
}
}
func TestVolcesSeedanceAcceptanceReconciliationIsReplayable(t *testing.T) {
payload, err := os.ReadFile("../../migrations/0099_reconcile_volces_seedance20_acceptance_constraints.sql")
if err != nil {
t.Fatal(err)
}
content := string(payload)
for _, required := range []string{
acceptanceImportReplayMarker,
"volces:doubao-seedance-2-0-260128",
"volces:doubao-seedance-2-0-fast-260128",
"volces:doubao-seedance-2-0-mini-260615",
`"long_edge":300`,
`"long_edge":6000`,
`'[0.4,2.5]'::jsonb`,
} {
if !strings.Contains(content, required) {
t.Fatalf("Volces acceptance reconciliation is missing %q", required)
}
}
}
func TestSecurityEventSchemaMigrationsDefineCurrentLifecycle(t *testing.T) {
streamPayload, err := os.ReadFile("../../migrations/0063_oidc_security_events.sql")
if err != nil {
@@ -279,6 +375,34 @@ SET last_error_category='still-invalid' WHERE id=$1::uuid`, pairingID)
}
}
func TestOIDCSessionClientIdentityRepairMigrationUpgradesApplied0090Schema(t *testing.T) {
pool := newIdentityMigrationPostgresTestSchema(t)
ctx := context.Background()
if _, err := pool.Exec(ctx, `
CREATE TABLE gateway_oidc_sessions (
id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
oidc_user_binding_id uuid
)`); err != nil {
t.Fatalf("create pre-repair OIDC session schema: %v", err)
}
applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0091_oidc_session_client_identity.sql")
applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0091_oidc_session_client_identity.sql")
var dataType string
if err := pool.QueryRow(ctx, `
SELECT data_type
FROM information_schema.columns
WHERE table_schema=current_schema()
AND table_name='gateway_oidc_sessions'
AND column_name='oidc_client_id'`).Scan(&dataType); err != nil {
t.Fatalf("read repaired OIDC session client identity column: %v", err)
}
if dataType != "text" {
t.Fatalf("oidc_client_id data type=%q, want text", dataType)
}
}
func applyIdentityMigrationTestFile(t *testing.T, ctx context.Context, pool *pgxpool.Pool, path string) {
t.Helper()
payload, err := os.ReadFile(path)
+3637 -205
View File
File diff suppressed because it is too large Load Diff
+2476 -175
View File
File diff suppressed because it is too large Load Diff
+2
View File
@@ -15,6 +15,8 @@ require (
github.com/riverqueue/river/riverdriver/riverpgxv5 v0.24.0
github.com/riverqueue/river/rivertype v0.24.0
golang.org/x/crypto v0.52.0
golang.org/x/image v0.43.0
golang.org/x/net v0.54.0
golang.org/x/oauth2 v0.36.0
)
+4
View File
@@ -71,6 +71,10 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
golang.org/x/image v0.43.0 h1:FLxcP4ec2350nTfOC8ysKtqYSIFbk/QGjw1ZHNP4tsY=
golang.org/x/image v0.43.0/go.mod h1:rrpelvGFt+kLPAjPM4HeWPgrl0FtafueU//e5N0qk/Q=
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
@@ -0,0 +1,913 @@
package acceptanceemulator
import (
"bytes"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/binary"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"hash/crc32"
"image"
"image/color"
"image/jpeg"
"image/png"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
_ "golang.org/x/image/webp"
)
const (
maxProtocolBodyBytes = 128 << 20
maxImageReferenceBytes = 32 << 20
)
type Config struct {
Now func() time.Time
HTTPClient *http.Client
Wait func(context.Context, time.Duration) error
}
type Server struct {
now func() time.Time
httpClient *http.Client
wait func(context.Context, time.Duration) error
mu sync.RWMutex
tasks map[string]videoTask
fixtures map[string]fixture
nextID atomic.Uint64
report Report
callbacks map[string]map[int64]int
videoByIdempotency map[string]string
geminiIdempotency map[string]struct{}
storageObjects map[string]fixture
storageAttempts map[string]int
pngSmall string
pngLarge string
pngPeak string
}
type Report struct {
GeminiRequests int64 `json:"geminiRequests"`
GeminiInvalid int64 `json:"geminiInvalid"`
GeminiInputImages int64 `json:"geminiInputImages"`
GeminiInputBytes int64 `json:"geminiInputBytes"`
GeminiOutputBytes int64 `json:"geminiOutputBytes"`
VideoSubmissions int64 `json:"videoSubmissions"`
VideoPolls int64 `json:"videoPolls"`
VideoInvalid int64 `json:"videoInvalid"`
VideoReferenceCounts map[string]int64 `json:"videoReferenceCounts"`
VideoRoleCounts map[string]int64 `json:"videoRoleCounts"`
UniqueImageHashes int `json:"uniqueImageHashes"`
UniqueVideoTasks int `json:"uniqueVideoTasks"`
ForcedConversions int64 `json:"forcedConversionRequests"`
VerifiedConversions int64 `json:"verifiedConversions"`
CallbackEvents int64 `json:"callbackEvents"`
DuplicateCallbacks int64 `json:"duplicateCallbacks"`
MissingIdempotency int64 `json:"missingIdempotencyKeys"`
DuplicateSubmissions int64 `json:"duplicateSubmissionAttempts"`
StoragePuts int64 `json:"storagePuts"`
StorageGets int64 `json:"storageGets"`
StorageHeads int64 `json:"storageHeads"`
StorageDeletes int64 `json:"storageDeletes"`
StorageFailures int64 `json:"storageFailures"`
}
type videoTask struct {
ID string
Model string
CreatedAt time.Time
ReadyAt time.Time
UsageTokens int
ImageRefs []imageReference
}
type imageReference struct {
Role string
SHA256 string
ContentType string
Width int
Height int
}
type fixture struct {
ContentType string
Payload []byte
}
func New(config Config) *Server {
now := config.Now
if now == nil {
now = time.Now
}
client := config.HTTPClient
if client == nil {
client = &http.Client{Timeout: 30 * time.Second}
}
waiter := config.Wait
if waiter == nil {
waiter = wait
}
return &Server{
now: now, httpClient: client, wait: waiter, tasks: map[string]videoTask{}, fixtures: buildFixtures(),
callbacks: map[string]map[int64]int{},
videoByIdempotency: map[string]string{},
geminiIdempotency: map[string]struct{}{},
storageObjects: map[string]fixture{},
storageAttempts: map[string]int{},
report: Report{
VideoReferenceCounts: map[string]int64{},
VideoRoleCounts: map[string]int64{},
},
pngSmall: base64.StdEncoding.EncodeToString(paddedPNG(256 << 10)),
pngLarge: base64.StdEncoding.EncodeToString(paddedPNG(4 << 20)),
pngPeak: base64.StdEncoding.EncodeToString(paddedPNG(8 << 20)),
}
}
func (s *Server) Handler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", s.health)
mux.HandleFunc("GET /report", s.getReport)
mux.HandleFunc("POST /v1beta/models/", s.geminiGenerateContent)
mux.HandleFunc("POST /v1/models/", s.geminiGenerateContent)
mux.HandleFunc("POST /contents/generations/tasks", s.submitVideo)
mux.HandleFunc("GET /contents/generations/tasks/{taskID}", s.getVideo)
mux.HandleFunc("DELETE /contents/generations/tasks/{taskID}", s.deleteVideo)
mux.HandleFunc("GET /media/{asset}", s.getMedia)
mux.HandleFunc("GET /fixtures/{asset}", s.getFixture)
mux.HandleFunc("POST /callbacks", s.collectCallback)
mux.HandleFunc("PUT /storage/{profile}/{object...}", s.objectStorage)
mux.HandleFunc("GET /storage/{profile}/{object...}", s.objectStorage)
mux.HandleFunc("HEAD /storage/{profile}/{object...}", s.objectStorage)
mux.HandleFunc("DELETE /storage/{profile}/{object...}", s.objectStorage)
return mux
}
func (s *Server) objectStorage(w http.ResponseWriter, r *http.Request) {
profile := strings.ToLower(strings.TrimSpace(r.PathValue("profile")))
object := strings.TrimLeft(strings.TrimSpace(r.PathValue("object")), "/")
if object == "" {
http.NotFound(w, r)
return
}
baseProfile := strings.TrimSuffix(strings.TrimSuffix(strings.TrimSuffix(profile, "-transient"), "-auth"), "-fail")
if baseProfile != "oss" && baseProfile != "s3" {
http.NotFound(w, r)
return
}
key := profile + "/" + object
s.mu.Lock()
s.storageAttempts[key]++
attempt := s.storageAttempts[key]
if strings.HasSuffix(profile, "-auth") {
s.report.StorageFailures++
s.mu.Unlock()
http.Error(w, "forbidden", http.StatusForbidden)
return
}
if strings.HasSuffix(profile, "-fail") || (strings.HasSuffix(profile, "-transient") && attempt == 1) {
s.report.StorageFailures++
s.mu.Unlock()
http.Error(w, "temporary storage failure", http.StatusServiceUnavailable)
return
}
switch r.Method {
case http.MethodPut:
s.mu.Unlock()
payload, err := io.ReadAll(io.LimitReader(r.Body, maxProtocolBodyBytes+1))
if err != nil || len(payload) > maxProtocolBodyBytes {
http.Error(w, "invalid object", http.StatusBadRequest)
return
}
s.mu.Lock()
s.storageObjects[key] = fixture{ContentType: firstNonEmpty(r.Header.Get("Content-Type"), "application/octet-stream"), Payload: payload}
s.report.StoragePuts++
s.mu.Unlock()
w.WriteHeader(http.StatusOK)
case http.MethodGet:
item, ok := s.storageObjects[key]
if ok {
s.report.StorageGets++
}
s.mu.Unlock()
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", item.ContentType)
w.WriteHeader(http.StatusOK)
_, _ = w.Write(item.Payload)
case http.MethodHead:
item, ok := s.storageObjects[key]
if ok {
s.report.StorageHeads++
}
s.mu.Unlock()
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", item.ContentType)
w.Header().Set("Content-Length", strconv.Itoa(len(item.Payload)))
w.WriteHeader(http.StatusOK)
case http.MethodDelete:
_, ok := s.storageObjects[key]
delete(s.storageObjects, key)
if ok {
s.report.StorageDeletes++
}
s.mu.Unlock()
if !ok {
http.NotFound(w, r)
return
}
w.WriteHeader(http.StatusNoContent)
default:
s.mu.Unlock()
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
}
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if value = strings.TrimSpace(value); value != "" {
return value
}
}
return ""
}
func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
if !strings.HasSuffix(r.URL.Path, ":generateContent") {
http.NotFound(w, r)
return
}
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
if idempotencyKey == "" {
s.mu.Lock()
s.report.MissingIdempotency++
s.mu.Unlock()
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
return
}
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordGeminiInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
inputBytes, inputImages, err := validateGeminiImageRequest(body)
if err != nil {
s.recordGeminiInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
outputBytes, outputBase64, delay := s.geminiProfile(inputBytes)
if err := s.wait(r.Context(), delay); err != nil {
return
}
s.mu.Lock()
if _, exists := s.geminiIdempotency[idempotencyKey]; exists {
s.report.DuplicateSubmissions++
} else {
s.geminiIdempotency[idempotencyKey] = struct{}{}
}
s.report.GeminiRequests++
s.report.GeminiInputImages += int64(inputImages)
s.report.GeminiInputBytes += int64(inputBytes)
s.report.GeminiOutputBytes += int64(outputBytes)
s.mu.Unlock()
w.Header().Set("X-Request-ID", "acceptance-gemini-"+strconv.FormatUint(s.nextID.Add(1), 10))
writeJSON(w, http.StatusOK, map[string]any{
"candidates": []any{map[string]any{
"index": 0,
"content": map[string]any{
"role": "model",
"parts": []any{map[string]any{"inlineData": map[string]any{
"mimeType": "image/png",
"data": outputBase64,
}}},
},
"finishReason": "STOP",
}},
"usageMetadata": map[string]any{
"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2,
},
})
}
func (s *Server) geminiProfile(inputBytes int) (int, string, time.Duration) {
switch {
case inputBytes > acceptanceworkload.GeminiLarge.InputBytes:
profile := acceptanceworkload.GeminiPeak
return profile.OutputBytes, s.pngPeak, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
case inputBytes > acceptanceworkload.GeminiBaseline.InputBytes:
profile := acceptanceworkload.GeminiLarge
return profile.OutputBytes, s.pngLarge, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
default:
profile := acceptanceworkload.GeminiBaseline
return profile.OutputBytes, s.pngSmall, profile.DelayMin
}
}
func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
if idempotencyKey == "" {
s.mu.Lock()
s.report.MissingIdempotency++
s.mu.Unlock()
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
return
}
var body map[string]any
if err := decodeJSON(r, &body); err != nil {
s.recordVideoInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
refs, longRun, forceConversion, err := s.validateVideoRequest(r.Context(), body)
if err != nil {
s.recordVideoInvalid()
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
delay := videoDelay(body, longRun)
id := "acceptance-video-" + strconv.FormatUint(s.nextID.Add(1), 10)
task := videoTask{
ID: id, Model: strings.TrimSpace(stringValue(body["model"])),
CreatedAt: s.now(), ReadyAt: s.now().Add(delay), UsageTokens: videoUsageTokens(body), ImageRefs: refs,
}
s.mu.Lock()
if existingID := s.videoByIdempotency[idempotencyKey]; existingID != "" {
task = s.tasks[existingID]
s.report.DuplicateSubmissions++
s.mu.Unlock()
w.Header().Set("X-Request-ID", task.ID)
writeJSON(w, http.StatusOK, map[string]any{
"id": task.ID, "model": task.Model, "status": "queued",
"created_at": task.CreatedAt.Unix(),
})
return
}
s.tasks[id] = task
s.videoByIdempotency[idempotencyKey] = id
s.report.VideoSubmissions++
s.report.VideoReferenceCounts[strconv.Itoa(len(refs))]++
if forceConversion {
s.report.ForcedConversions++
s.report.VerifiedConversions++
}
for _, ref := range refs {
s.report.VideoRoleCounts[ref.Role]++
}
s.mu.Unlock()
w.Header().Set("X-Request-ID", id)
writeJSON(w, http.StatusOK, map[string]any{
"id": id, "model": task.Model, "status": "queued",
"created_at": task.CreatedAt.Unix(),
})
}
func (s *Server) getVideo(w http.ResponseWriter, r *http.Request) {
id := strings.TrimSpace(r.PathValue("taskID"))
s.mu.Lock()
task, ok := s.tasks[id]
if ok {
s.report.VideoPolls++
}
s.mu.Unlock()
if !ok {
writeProtocolError(w, http.StatusNotFound, "video task not found")
return
}
status := "running"
content := map[string]any{}
if !s.now().Before(task.ReadyAt) {
status = "succeeded"
content["video_url"] = requestOrigin(r) + "/media/" + url.PathEscape(id) + ".mp4"
}
w.Header().Set("X-Request-ID", id)
writeJSON(w, http.StatusOK, map[string]any{
"id": id, "model": task.Model, "status": status, "content": content,
"created_at": task.CreatedAt.Unix(),
"usage": map[string]any{"completion_tokens": task.UsageTokens, "total_tokens": task.UsageTokens},
})
}
func (s *Server) deleteVideo(w http.ResponseWriter, r *http.Request) {
id := strings.TrimSpace(r.PathValue("taskID"))
s.mu.Lock()
_, ok := s.tasks[id]
delete(s.tasks, id)
s.mu.Unlock()
if !ok {
writeProtocolError(w, http.StatusNotFound, "video task not found")
return
}
writeJSON(w, http.StatusOK, map[string]any{"id": id, "status": "cancelled"})
}
func (s *Server) getMedia(w http.ResponseWriter, r *http.Request) {
asset := strings.TrimSuffix(strings.TrimSpace(r.PathValue("asset")), ".mp4")
s.mu.RLock()
_, ok := s.tasks[asset]
s.mu.RUnlock()
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "video/mp4")
w.WriteHeader(http.StatusOK)
_, _ = w.Write(minimalMP4)
}
func (s *Server) getFixture(w http.ResponseWriter, r *http.Request) {
fixture, ok := s.fixtures[strings.TrimSpace(r.PathValue("asset"))]
if !ok {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", fixture.ContentType)
w.Header().Set("Cache-Control", "public, max-age=3600")
w.Header().Set("X-Content-SHA256", payloadSHA256(fixture.Payload))
w.WriteHeader(http.StatusOK)
_, _ = w.Write(fixture.Payload)
}
func (s *Server) getReport(w http.ResponseWriter, _ *http.Request) {
s.mu.RLock()
report := s.report
tasks := make([]videoTask, 0, len(s.tasks))
for _, task := range s.tasks {
tasks = append(tasks, task)
}
s.mu.RUnlock()
hashes := map[string]struct{}{}
for _, task := range tasks {
for _, ref := range task.ImageRefs {
hashes[ref.SHA256] = struct{}{}
}
}
report.UniqueImageHashes = len(hashes)
report.UniqueVideoTasks = len(tasks)
writeJSON(w, http.StatusOK, report)
}
func (s *Server) collectCallback(w http.ResponseWriter, r *http.Request) {
var payload map[string]any
if err := decodeJSON(r, &payload); err != nil {
writeProtocolError(w, http.StatusBadRequest, err.Error())
return
}
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
seq := int64(0)
switch value := payload["seq"].(type) {
case json.Number:
seq, _ = value.Int64()
case float64:
seq = int64(value)
}
if taskID == "" || seq <= 0 {
writeProtocolError(w, http.StatusBadRequest, "callback taskId and seq are required")
return
}
s.mu.Lock()
seen := s.callbacks[taskID]
if seen == nil {
seen = map[int64]int{}
s.callbacks[taskID] = seen
}
seen[seq]++
s.report.CallbackEvents++
if seen[seq] > 1 {
s.report.DuplicateCallbacks++
}
s.mu.Unlock()
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}
func (s *Server) validateVideoRequest(ctx context.Context, body map[string]any) ([]imageReference, bool, bool, error) {
content, _ := body["content"].([]any)
refs := make([]imageReference, 0, 9)
firstFrames := 0
lastFrames := 0
longRun := false
forceConversion := false
for _, raw := range content {
item, _ := raw.(map[string]any)
switch strings.TrimSpace(stringValue(item["type"])) {
case "text":
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-long-recovery") {
longRun = true
}
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-force-conversion") {
forceConversion = true
}
case "image_url":
role := strings.TrimSpace(stringValue(item["role"]))
switch role {
case "reference_image":
case "first_frame":
firstFrames++
case "last_frame":
lastFrames++
default:
return nil, false, false, fmt.Errorf("unsupported image role %q", role)
}
nested, _ := item["image_url"].(map[string]any)
rawURL := strings.TrimSpace(stringValue(nested["url"]))
ref, err := s.fetchImageReference(ctx, role, rawURL)
if err != nil {
return nil, false, false, err
}
refs = append(refs, ref)
}
}
if len(refs) != 3 && len(refs) != 6 && len(refs) != 9 {
return nil, false, false, fmt.Errorf("expected 3, 6, or 9 image references, got %d", len(refs))
}
if firstFrames > 1 || lastFrames > 1 {
return nil, false, false, errors.New("first_frame and last_frame must be unique")
}
if forceConversion {
for _, ref := range refs {
if !seedanceImageWithinOfficialInputRange(ref.Width, ref.Height) {
return nil, false, false, errors.New("forced oversized image did not pass through automatic normalization")
}
}
}
return refs, longRun, forceConversion, nil
}
func seedanceImageWithinOfficialInputRange(width int, height int) bool {
if width < 300 || width > 6000 || height < 300 || height > 6000 {
return false
}
ratio := float64(width) / float64(height)
return ratio >= 0.4 && ratio <= 2.5
}
func (s *Server) fetchImageReference(ctx context.Context, role string, rawURL string) (imageReference, error) {
if rawURL == "" {
return imageReference{}, errors.New("image reference URL is required")
}
payload, err := s.readImageReference(ctx, rawURL)
if err != nil {
return imageReference{}, err
}
config, format, err := image.DecodeConfig(bytes.NewReader(payload))
if err != nil {
return imageReference{}, fmt.Errorf("decode image reference: %w", err)
}
sum := sha256.Sum256(payload)
return imageReference{
Role: role, SHA256: hex.EncodeToString(sum[:]), ContentType: "image/" + format,
Width: config.Width, Height: config.Height,
}, nil
}
func (s *Server) readImageReference(ctx context.Context, rawURL string) ([]byte, error) {
if strings.HasPrefix(strings.ToLower(rawURL), "data:") {
return decodeImageDataURL(rawURL)
}
request, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return nil, err
}
response, err := s.httpClient.Do(request)
if err != nil {
return nil, fmt.Errorf("fetch image reference: %w", err)
}
defer response.Body.Close()
if response.StatusCode < 200 || response.StatusCode >= 300 {
return nil, fmt.Errorf("fetch image reference: status %d", response.StatusCode)
}
payload, err := io.ReadAll(io.LimitReader(response.Body, maxImageReferenceBytes+1))
if err != nil {
return nil, err
}
if len(payload) == 0 {
return nil, errors.New("image reference is empty")
}
if len(payload) > maxImageReferenceBytes {
return nil, errors.New("image reference exceeds 32 MiB")
}
return payload, nil
}
func decodeImageDataURL(raw string) ([]byte, error) {
header, encoded, ok := strings.Cut(raw, ",")
if !ok || len(header) <= len("data:") || encoded == "" {
return nil, errors.New("image data URL is malformed")
}
metadata := strings.Split(header[len("data:"):], ";")
if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(metadata[0])), "image/") {
return nil, errors.New("image data URL must declare an image media type")
}
base64Encoded := false
for _, item := range metadata[1:] {
if strings.EqualFold(strings.TrimSpace(item), "base64") {
base64Encoded = true
break
}
}
if !base64Encoded {
return nil, errors.New("image data URL must use base64 encoding")
}
if base64.StdEncoding.DecodedLen(len(encoded)) > maxImageReferenceBytes {
return nil, errors.New("image reference exceeds 32 MiB")
}
payload, err := base64.StdEncoding.DecodeString(encoded)
if err != nil {
return nil, errors.New("image data URL contains invalid base64")
}
if len(payload) == 0 {
return nil, errors.New("image reference is empty")
}
return payload, nil
}
func validateGeminiImageRequest(body map[string]any) (int, int, error) {
generationConfig, _ := body["generationConfig"].(map[string]any)
modalities, _ := generationConfig["responseModalities"].([]any)
hasImageModality := false
for _, modality := range modalities {
if strings.EqualFold(strings.TrimSpace(stringValue(modality)), "IMAGE") {
hasImageModality = true
}
}
if !hasImageModality {
return 0, 0, errors.New("generationConfig.responseModalities must contain IMAGE")
}
total := 0
images := 0
contents, _ := body["contents"].([]any)
for _, rawContent := range contents {
content, _ := rawContent.(map[string]any)
parts, _ := content["parts"].([]any)
for _, rawPart := range parts {
part, _ := rawPart.(map[string]any)
inline, _ := part["inlineData"].(map[string]any)
if inline == nil {
inline, _ = part["inline_data"].(map[string]any)
}
encoded := strings.TrimSpace(stringValue(inline["data"]))
if encoded == "" {
fileData, _ := part["fileData"].(map[string]any)
if fileData == nil {
fileData, _ = part["file_data"].(map[string]any)
}
fileURI := strings.TrimSpace(stringValue(fileData["fileUri"]))
if fileURI == "" {
fileURI = strings.TrimSpace(stringValue(fileData["file_uri"]))
}
if fileURI == "" {
fileURI = strings.TrimSpace(stringValue(fileData["uri"]))
}
if fileURI == "" {
continue
}
if !strings.HasPrefix(fileURI, "http://") && !strings.HasPrefix(fileURI, "https://") {
return 0, 0, errors.New("Gemini fileData URI must use HTTP or HTTPS")
}
images++
continue
}
decoded, err := base64.StdEncoding.DecodeString(encoded)
if err != nil {
return 0, 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err)
}
if len(decoded) == 0 {
return 0, 0, errors.New("Gemini inlineData is empty")
}
total += len(decoded)
images++
}
}
if images == 0 {
return 0, 0, errors.New("Gemini request has no inlineData or fileData image")
}
return total, images, nil
}
func (s *Server) recordGeminiInvalid() {
s.mu.Lock()
s.report.GeminiInvalid++
s.mu.Unlock()
}
func (s *Server) recordVideoInvalid() {
s.mu.Lock()
s.report.VideoInvalid++
s.mu.Unlock()
}
func decodeJSON(r *http.Request, target any) error {
reader := io.LimitReader(r.Body, maxProtocolBodyBytes+1)
payload, err := io.ReadAll(reader)
if err != nil {
return err
}
if len(payload) > maxProtocolBodyBytes {
return errors.New("protocol request body is too large")
}
decoder := json.NewDecoder(bytes.NewReader(payload))
decoder.UseNumber()
return decoder.Decode(target)
}
func writeJSON(w http.ResponseWriter, status int, value any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(value)
}
func writeProtocolError(w http.ResponseWriter, status int, message string) {
writeJSON(w, status, map[string]any{"error": map[string]any{
"code": "acceptance_protocol_invalid", "message": message,
}})
}
func wait(ctx context.Context, delay time.Duration) error {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
func scaledDelay(minimum time.Duration, maximum time.Duration, seed int) time.Duration {
if maximum <= minimum {
return minimum
}
windowSeconds := int64((maximum - minimum) / time.Second)
if windowSeconds <= 0 {
return minimum
}
return minimum + time.Duration(int64(seed)%(windowSeconds+1))*time.Second
}
func videoDelay(body map[string]any, longRun bool) time.Duration {
seed := int64(0)
switch value := body["seed"].(type) {
case json.Number:
seed, _ = value.Int64()
case float64:
seed = int64(value)
case int:
seed = int64(value)
}
if seed < 0 {
seed = -seed
}
if longRun {
return 2*time.Minute + time.Duration(seed%61)*time.Second
}
if videoPromptContains(body, "acceptance-provider-quota") {
return 5*time.Second + time.Duration(seed%6)*time.Second
}
if videoPromptContains(body, "acceptance-capacity-ladder") {
return 30*time.Second + time.Duration(seed%16)*time.Second
}
return 5*time.Second + time.Duration(seed%11)*time.Second
}
func videoUsageTokens(body map[string]any) int {
if videoPromptContains(body, "acceptance-provider-quota") {
return 7
}
return 1
}
func videoPromptContains(body map[string]any, marker string) bool {
content, _ := body["content"].([]any)
for _, raw := range content {
item, _ := raw.(map[string]any)
if strings.TrimSpace(stringValue(item["type"])) == "text" &&
strings.Contains(strings.ToLower(stringValue(item["text"])), marker) {
return true
}
}
return false
}
func requestOrigin(r *http.Request) string {
scheme := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
if scheme == "" {
scheme = "http"
}
host := strings.TrimSpace(r.Host)
return scheme + "://" + host
}
func stringValue(value any) string {
switch typed := value.(type) {
case nil:
return ""
case string:
return typed
case json.Number:
return typed.String()
default:
return fmt.Sprint(value)
}
}
func paddedPNG(size int) []byte {
base, _ := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")
if size <= len(base) {
return base
}
iendOffset := len(base) - 12
paddingLength := size - len(base) - 12
if paddingLength < 0 {
paddingLength = 0
}
chunk := make([]byte, 12+paddingLength)
binary.BigEndian.PutUint32(chunk[:4], uint32(paddingLength))
copy(chunk[4:8], []byte("teST"))
binary.BigEndian.PutUint32(chunk[8+paddingLength:], crc32.ChecksumIEEE(chunk[4:8+paddingLength]))
out := make([]byte, 0, len(base)+len(chunk))
out = append(out, base[:iendOffset]...)
out = append(out, chunk...)
out = append(out, base[iendOffset:]...)
return out
}
func buildFixtures() map[string]fixture {
fixtures := map[string]fixture{}
for index := 0; index < 4; index++ {
imageValue := patternedImage(768, 512, index+1)
var encoded bytes.Buffer
_ = png.Encode(&encoded, imageValue)
fixtures[fmt.Sprintf("image-%02d.png", index)] = fixture{ContentType: "image/png", Payload: encoded.Bytes()}
}
for index := 0; index < 4; index++ {
imageValue := patternedImage(1024, 768, index+11)
var encoded bytes.Buffer
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 88})
fixtures[fmt.Sprintf("image-%02d.jpg", index+4)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
}
webpFixtures := []string{
"524946465a00000057454250565038204e000000f003009d012a400040003e9148a04c25a42322220800b012096900d3ca8000103b93c116da67710000feeea63fff80dd7c5b4cbfff7381ff7381ff7381fc6d4326178e9af219737c3e2990000000",
"5249464650000000574542505650382044000000d003009d012a400040003e9148a04c25a42322220800b012096900760000206ea6a00af10b720000feeed3dfffc5b9fb75d21ffff8b43ad0eb43447bcc63798400000000",
"52494646500000005745425056503820440000009003009d012a400040003e9148a04c25a42322220800b0120969000010685716c42e300a2000fef558bffffb9c0fffd9c0fffd9c0fe3afff5ea96b93a3f9696c1c000000",
"524946464e0000005745425056503820420000009003009d012a400040003e9148a04c25a42322220800b0120969000010375350057885b90000fef4351ffffba7afffda7afffda7afe60fff239de951800000000000",
}
for index, encoded := range webpFixtures {
payload, _ := hex.DecodeString(encoded)
fixtures[fmt.Sprintf("image-%02d.webp", index+8)] = fixture{ContentType: "image/webp", Payload: payload}
}
for index := 0; index < 4; index++ {
imageValue := patternedImage(6144, 2160, index+21)
var encoded bytes.Buffer
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 90})
fixtures[fmt.Sprintf("image-%02d-oversized.jpg", index+12)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
}
return fixtures
}
func patternedImage(width int, height int, seed int) image.Image {
out := image.NewNRGBA(image.Rect(0, 0, width, height))
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
out.SetNRGBA(x, y, color.NRGBA{
R: uint8((x + seed*31) % 256),
G: uint8((y + seed*47) % 256),
B: uint8((x/8 + y/8 + seed*59) % 256),
A: 255,
})
}
}
return out
}
func payloadSHA256(payload []byte) string {
sum := sha256.Sum256(payload)
return hex.EncodeToString(sum[:])
}
var minimalMP4 = []byte{
0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm',
0, 0, 0, 0, 'i', 's', 'o', 'm', 'm', 'p', '4', '2',
0, 0, 0, 8, 'm', 'd', 'a', 't',
}
@@ -0,0 +1,406 @@
package acceptanceemulator
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"image"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
)
func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
var mu sync.Mutex
now := time.Unix(1_800_000_000, 0)
server := New(Config{
Now: func() time.Time {
mu.Lock()
defer mu.Unlock()
return now
},
Wait: func(context.Context, time.Duration) error { return nil },
})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
input := paddedPNG(256 << 10)
geminiBody, _ := json.Marshal(map[string]any{
"contents": []any{map[string]any{"parts": []any{
map[string]any{"text": "edit image"},
map[string]any{"inlineData": map[string]any{
"mimeType": "image/png", "data": base64.StdEncoding.EncodeToString(input),
}},
}}},
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
})
response, err := postIdempotent(
httpServer.URL+"/v1beta/models/gemini-test:generateContent",
geminiBody,
"gemini-task",
)
if err != nil {
t.Fatalf("Gemini request: %v", err)
}
if response.StatusCode != http.StatusOK {
payload, _ := io.ReadAll(response.Body)
t.Fatalf("Gemini status=%d body=%s", response.StatusCode, payload)
}
var geminiResult map[string]any
if err := json.NewDecoder(response.Body).Decode(&geminiResult); err != nil {
t.Fatalf("decode Gemini response: %v", err)
}
_ = response.Body.Close()
fixtureResponse, err := http.Get(httpServer.URL + "/fixtures/image-12-oversized.jpg")
if err != nil {
t.Fatalf("get oversized fixture: %v", err)
}
fixturePayload, readErr := io.ReadAll(fixtureResponse.Body)
_ = fixtureResponse.Body.Close()
if readErr != nil {
t.Fatalf("read oversized fixture: %v", readErr)
}
config, format, err := image.DecodeConfig(bytes.NewReader(fixturePayload))
if err != nil || format != "jpeg" || config.Width != 6144 || config.Height != 2160 {
t.Fatalf("oversized fixture format=%s config=%+v err=%v", format, config, err)
}
if _, format, err = image.Decode(bytes.NewReader(fixturePayload)); err != nil || format != "jpeg" {
t.Fatalf("fully decode oversized fixture format=%s err=%v", format, err)
}
content := []any{map[string]any{"type": "text", "text": "video"}}
for index, name := range []string{"image-00.png", "image-04.jpg", "image-08.webp"} {
role := "reference_image"
if index == 0 {
role = "first_frame"
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
})
}
videoBody, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": 1})
response, err = postIdempotent(
httpServer.URL+"/contents/generations/tasks",
videoBody,
"video-task",
)
if err != nil {
t.Fatalf("submit video: %v", err)
}
if response.StatusCode != http.StatusOK {
t.Fatalf("submit video status=%d", response.StatusCode)
}
var submitted map[string]any
_ = json.NewDecoder(response.Body).Decode(&submitted)
_ = response.Body.Close()
taskID := strings.TrimSpace(stringValue(submitted["id"]))
if taskID == "" {
t.Fatal("video emulator returned no task ID")
}
mu.Lock()
now = now.Add(20 * time.Second)
mu.Unlock()
response, err = http.Get(httpServer.URL + "/contents/generations/tasks/" + taskID)
if err != nil {
t.Fatalf("poll video: %v", err)
}
var polled map[string]any
_ = json.NewDecoder(response.Body).Decode(&polled)
_ = response.Body.Close()
if polled["status"] != "succeeded" {
t.Fatalf("video status=%v", polled["status"])
}
response, err = http.Get(httpServer.URL + "/report")
if err != nil {
t.Fatalf("get report: %v", err)
}
var report Report
_ = json.NewDecoder(response.Body).Decode(&report)
_ = response.Body.Close()
if report.GeminiRequests != 1 || report.VideoSubmissions != 1 || report.VideoInvalid != 0 ||
report.VideoReferenceCounts["3"] != 1 || report.UniqueImageHashes != 3 {
t.Fatalf("unexpected emulator report: %+v", report)
}
}
func TestForcedConversionRejectsImagesOutsideOfficialSeedanceRange(t *testing.T) {
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
content := []any{map[string]any{"type": "text", "text": "acceptance-force-conversion"}}
for _, name := range []string{"image-12-oversized.jpg", "image-00.png", "image-04.jpg"} {
content = append(content, map[string]any{
"type": "image_url", "role": "reference_image",
"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
})
}
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content})
response, err := postIdempotent(httpServer.URL+"/contents/generations/tasks", body, "oversized")
if err != nil {
t.Fatal(err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusBadRequest {
t.Fatalf("status=%d, want %d", response.StatusCode, http.StatusBadRequest)
}
}
func TestGeminiProtocolAcceptsMultipleFileDataImages(t *testing.T) {
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
parts := []any{map[string]any{"text": "combine references"}}
for index := 0; index < 3; index++ {
parts = append(parts, map[string]any{"fileData": map[string]any{
"mimeType": "image/png",
"fileUri": fmt.Sprintf("https://fixtures.example/%d.png", index),
}})
}
body, _ := json.Marshal(map[string]any{
"contents": []any{map[string]any{"parts": parts}},
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
})
response, err := postIdempotent(httpServer.URL+"/v1beta/models/gemini-test:generateContent", body, "gemini-multi")
if err != nil {
t.Fatalf("Gemini multi-image request: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
payload, _ := io.ReadAll(response.Body)
t.Fatalf("Gemini multi-image status=%d body=%s", response.StatusCode, payload)
}
reportResponse, err := http.Get(httpServer.URL + "/report")
if err != nil {
t.Fatalf("get report: %v", err)
}
defer reportResponse.Body.Close()
var report Report
if err := json.NewDecoder(reportResponse.Body).Decode(&report); err != nil {
t.Fatalf("decode report: %v", err)
}
if report.GeminiRequests != 1 || report.GeminiInputImages != 3 || report.GeminiInvalid != 0 {
t.Fatalf("unexpected multi-image report: %+v", report)
}
}
func TestImageFixturesFullyDecode(t *testing.T) {
for name, fixture := range buildFixtures() {
if !strings.HasPrefix(fixture.ContentType, "image/") {
continue
}
decoded, format, err := image.Decode(bytes.NewReader(fixture.Payload))
if err != nil {
t.Fatalf("fixture %s does not fully decode: %v", name, err)
}
bounds := decoded.Bounds()
if format == "" || bounds.Dx() <= 0 || bounds.Dy() <= 0 {
t.Fatalf("fixture %s format=%q bounds=%v", name, format, bounds)
}
}
}
func TestPaddedPNGIsValidAndExact(t *testing.T) {
for _, size := range []int{256 << 10, 4 << 20, 8 << 20} {
payload := paddedPNG(size)
if len(payload) != size {
t.Fatalf("PNG bytes=%d, want=%d", len(payload), size)
}
if _, _, err := image.DecodeConfig(bytes.NewReader(payload)); err != nil {
t.Fatalf("decode padded PNG: %v", err)
}
}
}
func TestScaledDelayStaysOnWholeSecondsWithinProfile(t *testing.T) {
delay := scaledDelay(8*time.Second, 15*time.Second, 2<<20)
if delay < 8*time.Second || delay > 15*time.Second || delay%time.Second != 0 {
t.Fatalf("scaled delay=%s", delay)
}
}
func TestVideoCapacityDelay(t *testing.T) {
body := map[string]any{
"seed": json.Number("7"),
"content": []any{map[string]any{
"type": "text",
"text": "capacity acceptance-capacity-ladder",
}},
}
delay := videoDelay(body, false)
if delay < 30*time.Second || delay > 45*time.Second {
t.Fatalf("video capacity delay=%s, want 30s..45s", delay)
}
}
func TestVideoProviderQuotaDelayAndUsage(t *testing.T) {
body := map[string]any{
"seed": json.Number("5"),
"content": []any{map[string]any{
"type": "text",
"text": "acceptance-provider-quota",
}},
}
delay := videoDelay(body, false)
if delay < 5*time.Second || delay > 10*time.Second {
t.Fatalf("video provider quota delay=%s, want 5s..10s", delay)
}
if tokens := videoUsageTokens(body); tokens != 7 {
t.Fatalf("video provider quota usage tokens=%d, want 7", tokens)
}
}
func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
httpServer := httptest.NewServer(server.Handler())
defer httpServer.Close()
dataImage := "data:image/png;base64," + base64.StdEncoding.EncodeToString(paddedPNG(0))
for _, imageCount := range []int{3, 6, 9} {
content := []any{map[string]any{"type": "text", "text": "multi-reference video"}}
for index := 0; index < imageCount; index++ {
role := "reference_image"
if index == 0 {
role = "first_frame"
}
if imageCount > 3 && index == imageCount-1 {
role = "last_frame"
}
imageURL := fmt.Sprintf("%s/fixtures/image-%02d.png", httpServer.URL, index%4)
if index%2 == 0 {
imageURL = dataImage
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{"url": imageURL},
})
}
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": imageCount})
response, err := postIdempotent(
httpServer.URL+"/contents/generations/tasks",
body,
fmt.Sprintf("video-%d", imageCount),
)
if err != nil {
t.Fatalf("%d-image submit: %v", imageCount, err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("%d-image status=%d", imageCount, response.StatusCode)
}
}
response, err := http.Get(httpServer.URL + "/report")
if err != nil {
t.Fatalf("get report: %v", err)
}
defer response.Body.Close()
var report Report
if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
t.Fatalf("decode report: %v", err)
}
for _, imageCount := range []string{"3", "6", "9"} {
if report.VideoReferenceCounts[imageCount] != 1 {
t.Fatalf("reference counts=%v", report.VideoReferenceCounts)
}
}
}
func TestObjectStorageEmulationSupportsRetryLifecycleAndFaults(t *testing.T) {
server := httptest.NewServer(New(Config{}).Handler())
defer server.Close()
objectURL := server.URL + "/storage/s3-transient/bucket/media/result.png"
request, _ := http.NewRequest(http.MethodPut, objectURL, strings.NewReader("image-bytes"))
request.Header.Set("Content-Type", "image/png")
response, err := http.DefaultClient.Do(request)
if err != nil {
t.Fatal(err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusServiceUnavailable {
t.Fatalf("first transient PUT status=%d", response.StatusCode)
}
request, _ = http.NewRequest(http.MethodPut, objectURL, strings.NewReader("image-bytes"))
request.Header.Set("Content-Type", "image/png")
response, err = http.DefaultClient.Do(request)
if err != nil {
t.Fatal(err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("retried PUT status=%d", response.StatusCode)
}
response, err = http.Get(objectURL)
if err != nil {
t.Fatal(err)
}
payload, _ := io.ReadAll(response.Body)
_ = response.Body.Close()
if response.StatusCode != http.StatusOK || string(payload) != "image-bytes" {
t.Fatalf("GET status=%d payload=%q", response.StatusCode, payload)
}
request, _ = http.NewRequest(http.MethodHead, objectURL, nil)
response, err = http.DefaultClient.Do(request)
if err != nil {
t.Fatal(err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("HEAD status=%d", response.StatusCode)
}
request, _ = http.NewRequest(http.MethodDelete, objectURL, nil)
response, err = http.DefaultClient.Do(request)
if err != nil {
t.Fatal(err)
}
_ = response.Body.Close()
if response.StatusCode != http.StatusNoContent {
t.Fatalf("DELETE status=%d", response.StatusCode)
}
for profile, wantStatus := range map[string]int{"oss-auth": http.StatusForbidden, "s3-fail": http.StatusServiceUnavailable} {
request, _ = http.NewRequest(http.MethodPut, server.URL+"/storage/"+profile+"/probe.bin", strings.NewReader("probe"))
response, err = http.DefaultClient.Do(request)
if err != nil {
t.Fatal(err)
}
_ = response.Body.Close()
if response.StatusCode != wantStatus {
t.Fatalf("%s status=%d, want %d", profile, response.StatusCode, wantStatus)
}
}
response, err = http.Get(server.URL + "/report")
if err != nil {
t.Fatal(err)
}
defer response.Body.Close()
var report Report
if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
t.Fatal(err)
}
if report.StoragePuts != 1 || report.StorageGets != 1 || report.StorageHeads != 1 || report.StorageDeletes != 1 || report.StorageFailures != 3 {
t.Fatalf("unexpected storage report: %+v", report)
}
}
func postIdempotent(url string, body []byte, key string) (*http.Response, error) {
request, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Idempotency-Key", key)
return http.DefaultClient.Do(request)
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,180 @@
package acceptancesnapshot
import (
"encoding/json"
"testing"
"time"
)
func TestSnapshotValidationRejectsSecretsAndDetectsTampering(t *testing.T) {
snapshot := validSnapshot(t)
if err := Validate(snapshot); err != nil {
t.Fatalf("validate fixture: %v", err)
}
tampered := snapshot
tampered.Candidates = append([]Candidate(nil), snapshot.Candidates...)
tampered.Candidates[0].Platform.Name = "tampered"
if err := Validate(tampered); err == nil {
t.Fatal("tampered snapshot passed validation")
}
unsafe := snapshot
unsafe.Candidates = append([]Candidate(nil), snapshot.Candidates...)
unsafe.Candidates[0].Platform.Config = map[string]any{"apiToken": "must-not-leak"}
refreshHashes(t, &unsafe)
if err := Validate(unsafe); err == nil {
t.Fatal("secret-like snapshot field passed validation")
}
}
func TestSanitizeMapRemovesCredentialAndProxyFields(t *testing.T) {
input := map[string]any{
"specType": "gemini",
"credentialEnv": "PRODUCTION_TOKEN",
"networkProxy": map[string]any{
"proxyUrl": "http://user:password@example.invalid",
},
"requestAssetImageURLFormat": "base64",
}
got := sanitizeMap(input)
if got["specType"] != "gemini" || got["requestAssetImageURLFormat"] != "base64" {
t.Fatalf("safe protocol fields were removed: %#v", got)
}
if _, ok := got["credentialEnv"]; ok {
t.Fatal("credentialEnv was not removed")
}
nested, _ := got["networkProxy"].(map[string]any)
if _, ok := nested["proxyUrl"]; ok {
t.Fatal("nested proxy URL was not removed")
}
}
func TestDecodeRejectsUnknownSnapshotFields(t *testing.T) {
snapshot := validSnapshot(t)
payload, err := json.Marshal(snapshot)
if err != nil {
t.Fatal(err)
}
var object map[string]any
if err := json.Unmarshal(payload, &object); err != nil {
t.Fatal(err)
}
object["unexpected"] = true
payload, _ = json.Marshal(object)
if _, err := Decode(payload); err == nil {
t.Fatal("snapshot with unknown field decoded successfully")
}
}
func validSnapshot(t *testing.T) Snapshot {
t.Helper()
candidates := []Candidate{
{
Workload: "gemini_image_edit",
Provider: Provider{
ProviderKey: "gemini", ProviderCode: "google-gemini",
DisplayName: "Gemini", ProviderType: "gemini",
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
},
BaseModel: BaseModel{
CanonicalModelKey: "gemini:image", InvocationName: "gemini-image",
ProviderModelName: "gemini-image", ModelType: []string{"image_edit"},
DisplayName: "Gemini Image", Capabilities: map[string]any{},
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
},
Platform: Platform{
Provider: "gemini", PlatformKey: "gemini-production", Name: "Gemini Production",
BaseURLPath: "/", AuthType: "api_key", Config: map[string]any{"specType": "gemini"},
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
},
PlatformModel: PlatformModel{
ModelName: "gemini-image", ProviderModelName: "gemini-image",
ModelType: []string{"image_edit"}, DisplayName: "Gemini Image",
CapabilityOverride: map[string]any{}, Capabilities: map[string]any{},
PricingMode: "inherit_discount", BillingConfigOverride: map[string]any{},
BillingConfig: map[string]any{}, PermissionConfig: map[string]any{},
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{},
},
RuntimePolicy: RuntimePolicy{
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
},
PricingRules: []PricingRule{},
Metadata: CandidateSource{
PlatformID: "00000000-0000-0000-0000-000000000001",
PlatformModelID: "00000000-0000-0000-0000-000000000002",
BaseModelID: "00000000-0000-0000-0000-000000000003",
},
},
{
Workload: "multi_reference_video",
Provider: Provider{
ProviderKey: "volces", ProviderCode: "volces",
DisplayName: "Volces", ProviderType: "volces",
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
},
BaseModel: BaseModel{
CanonicalModelKey: "volces:seedance", InvocationName: "seedance-2-fast",
ProviderModelName: "seedance-2-fast", ModelType: []string{"omni_video"},
DisplayName: "Seedance", Capabilities: map[string]any{"omni_video": map[string]any{"max_images": 9}},
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
},
Platform: Platform{
Provider: "volces", PlatformKey: "volces-production", Name: "Volces Production",
BaseURLPath: "/", AuthType: "bearer", Config: map[string]any{"specType": "volces"},
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
},
PlatformModel: PlatformModel{
ModelName: "seedance-2-fast", ProviderModelName: "seedance-2-fast",
ModelType: []string{"omni_video"}, DisplayName: "Seedance",
CapabilityOverride: map[string]any{"omni_video": map[string]any{"max_images": 9}},
Capabilities: map[string]any{}, PricingMode: "inherit_discount",
BillingConfigOverride: map[string]any{}, BillingConfig: map[string]any{},
PermissionConfig: map[string]any{}, RetryPolicy: map[string]any{},
RateLimitPolicy: map[string]any{}, RuntimePolicyOverride: map[string]any{},
},
RuntimePolicy: RuntimePolicy{
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
},
PricingRules: []PricingRule{},
Metadata: CandidateSource{
PlatformID: "00000000-0000-0000-0000-000000000004",
PlatformModelID: "00000000-0000-0000-0000-000000000005",
BaseModelID: "00000000-0000-0000-0000-000000000006",
},
},
}
snapshot := Snapshot{
SchemaVersion: SchemaVersion,
Source: Source{
ReleaseSHA: "0123456789abcdef0123456789abcdef01234567",
CreatedAt: time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC),
},
Candidates: candidates,
SecretSafe: true,
}
refreshHashes(t, &snapshot)
return snapshot
}
func refreshHashes(t *testing.T, snapshot *Snapshot) {
t.Helper()
configHash, err := candidateConfigHash(snapshot.Candidates)
if err != nil {
t.Fatal(err)
}
snapshot.Source.ConfigHash = configHash
snapshot.SnapshotSHA256 = ""
hash, err := snapshotHash(*snapshot)
if err != nil {
t.Fatal(err)
}
snapshot.SnapshotSHA256 = hash
}
@@ -0,0 +1,59 @@
package acceptanceworkload
import "time"
type GeminiProfile struct {
Name string
Requests int
InputImages int
InputBytes int
OutputBytes int
DelayMin time.Duration
DelayMax time.Duration
}
var (
GeminiBaseline = GeminiProfile{
Name: "gemini-baseline", Requests: 1000, InputImages: 1, InputBytes: 256 << 10, OutputBytes: 256 << 10,
DelayMin: 4 * time.Second, DelayMax: 4 * time.Second,
}
GeminiMultiImage = GeminiProfile{
Name: "gemini-multi-image", Requests: 192, InputImages: 3, InputBytes: 768 << 10, OutputBytes: 0,
DelayMin: 8 * time.Second, DelayMax: 15 * time.Second,
}
GeminiLarge = GeminiProfile{
Name: "gemini-large", Requests: 128, InputImages: 1, InputBytes: 2 << 20, OutputBytes: 4 << 20,
DelayMin: 8 * time.Second, DelayMax: 15 * time.Second,
}
GeminiPeak = GeminiProfile{
Name: "gemini-peak", Requests: 32, InputImages: 1, InputBytes: 8 << 20, OutputBytes: 8 << 20,
DelayMin: 15 * time.Second, DelayMax: 30 * time.Second,
}
)
func GeminiProfileByName(name string) (GeminiProfile, bool) {
switch name {
case GeminiBaseline.Name:
return GeminiBaseline, true
case GeminiMultiImage.Name:
return GeminiMultiImage, true
case GeminiLarge.Name:
return GeminiLarge, true
case GeminiPeak.Name:
return GeminiPeak, true
default:
return GeminiProfile{}, false
}
}
// VideoImageCount provides a deterministic 60/30/10 split for 3/6/9-image tasks.
func VideoImageCount(index int) int {
switch index % 10 {
case 0, 1, 2, 3, 4, 5:
return 3
case 6, 7, 8:
return 6
default:
return 9
}
}
@@ -0,0 +1,26 @@
package acceptanceworkload
import "testing"
func TestGeminiProfilesAndVideoDistribution(t *testing.T) {
for _, profile := range []GeminiProfile{GeminiBaseline, GeminiMultiImage, GeminiLarge, GeminiPeak} {
resolved, ok := GeminiProfileByName(profile.Name)
if !ok || resolved != profile {
t.Fatalf("profile %q resolved to %+v, ok=%v", profile.Name, resolved, ok)
}
if profile.Requests <= 0 || profile.InputImages <= 0 || profile.InputBytes < profile.InputImages || profile.OutputBytes < 0 ||
profile.DelayMin <= 0 || profile.DelayMax < profile.DelayMin {
t.Fatalf("invalid Gemini profile: %+v", profile)
}
if profile.Name != GeminiMultiImage.Name && profile.OutputBytes == 0 {
t.Fatalf("fixed-output Gemini profile has no output size: %+v", profile)
}
}
counts := map[int]int{}
for index := 0; index < 100; index++ {
counts[VideoImageCount(index)]++
}
if counts[3] != 60 || counts[6] != 30 || counts[9] != 10 {
t.Fatalf("video image distribution=%v", counts)
}
}
+323
View File
@@ -0,0 +1,323 @@
package accessruleaudit
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
const (
SchemaVersion = "access-rule-audit/v1"
MigrationBatch = "0102_access_rule_allow_whitelist_semantics"
)
type Snapshot struct {
SchemaVersion string `json:"schemaVersion"`
GeneratedAt time.Time `json:"generatedAt"`
SecretSafe bool `json:"secretSafe"`
Live RuleSet `json:"live"`
Archive *ArchiveManifest `json:"archive,omitempty"`
}
type RuleSet struct {
Total int64 `json:"total"`
SHA256 string `json:"sha256"`
AllowCount int64 `json:"allowCount"`
AllowSHA256 string `json:"allowSha256"`
DenyCount int64 `json:"denyCount"`
DenySHA256 string `json:"denySha256"`
Counts []RuleCount `json:"counts"`
}
type RuleCount struct {
SubjectType string `json:"subjectType"`
Effect string `json:"effect"`
ResourceType string `json:"resourceType"`
Status string `json:"status"`
Count int64 `json:"count"`
}
type ArchiveManifest struct {
MigrationBatch string `json:"migrationBatch"`
ManifestAllowCount int64 `json:"manifestAllowCount"`
ManifestAllowSHA string `json:"manifestAllowSha256"`
ArchivedAllowCount int64 `json:"archivedAllowCount"`
ArchivedAllowSHA string `json:"archivedAllowSha256"`
ManifestDenyCount int64 `json:"manifestDenyCount"`
ManifestDenySHA string `json:"manifestDenySha256"`
Consistent bool `json:"consistent"`
}
type ruleDigestRow struct {
id string
subjectType string
subjectID string
resourceType string
resourceID string
effect string
priority string
minPermissionLevel string
conditions string
metadata string
status string
createdEpoch string
updatedEpoch string
}
func Export(ctx context.Context, pool *pgxpool.Pool, now time.Time) (Snapshot, error) {
live, err := loadLiveRules(ctx, pool)
if err != nil {
return Snapshot{}, err
}
archive, err := loadArchiveManifest(ctx, pool)
if err != nil {
return Snapshot{}, err
}
if now.IsZero() {
now = time.Now()
}
snapshot := Snapshot{
SchemaVersion: SchemaVersion,
GeneratedAt: now.UTC(),
SecretSafe: true,
Live: live,
Archive: archive,
}
if err := Validate(snapshot); err != nil {
return Snapshot{}, err
}
return snapshot, nil
}
func Validate(snapshot Snapshot) error {
if snapshot.SchemaVersion != SchemaVersion {
return fmt.Errorf("unsupported access-rule audit schema %q", snapshot.SchemaVersion)
}
if !snapshot.SecretSafe {
return errors.New("access-rule audit snapshot is not marked secret-safe")
}
for name, value := range map[string]string{
"live": snapshot.Live.SHA256, "allow": snapshot.Live.AllowSHA256, "deny": snapshot.Live.DenySHA256,
} {
if !validSHA256(value) {
return fmt.Errorf("%s SHA-256 is invalid", name)
}
}
if snapshot.Live.Total != snapshot.Live.AllowCount+snapshot.Live.DenyCount {
return errors.New("live access-rule counts are inconsistent")
}
var groupedTotal int64
for _, count := range snapshot.Live.Counts {
if count.SubjectType == "" || count.Effect == "" || count.ResourceType == "" || count.Status == "" || count.Count < 1 {
return errors.New("access-rule grouped count is invalid")
}
groupedTotal += count.Count
}
if groupedTotal != snapshot.Live.Total {
return errors.New("access-rule grouped counts do not match total")
}
if snapshot.Archive != nil {
archive := snapshot.Archive
if archive.MigrationBatch != MigrationBatch || !validSHA256(archive.ManifestAllowSHA) || !validSHA256(archive.ArchivedAllowSHA) || !validSHA256(archive.ManifestDenySHA) {
return errors.New("access-rule archive manifest is invalid")
}
consistent := archive.ManifestAllowCount == archive.ArchivedAllowCount && archive.ManifestAllowSHA == archive.ArchivedAllowSHA
if archive.Consistent != consistent {
return errors.New("access-rule archive consistency marker is incorrect")
}
}
return nil
}
func VerifyMigration(before Snapshot, after Snapshot) error {
if err := Validate(before); err != nil {
return fmt.Errorf("before snapshot: %w", err)
}
if err := Validate(after); err != nil {
return fmt.Errorf("after snapshot: %w", err)
}
if after.Archive == nil || !after.Archive.Consistent {
return errors.New("verified allow archive is unavailable")
}
if after.Live.AllowCount != 0 {
return fmt.Errorf("live legacy allow count=%d, want 0", after.Live.AllowCount)
}
if before.Live.AllowCount != after.Archive.ManifestAllowCount || before.Live.AllowSHA256 != after.Archive.ManifestAllowSHA {
return errors.New("pre-migration allow snapshot does not match archived manifest")
}
if before.Live.DenyCount != after.Live.DenyCount || before.Live.DenySHA256 != after.Live.DenySHA256 {
return errors.New("deny rules changed during allow migration")
}
return nil
}
func Encode(snapshot Snapshot) ([]byte, error) {
if err := Validate(snapshot); err != nil {
return nil, err
}
payload, err := json.MarshalIndent(snapshot, "", " ")
if err != nil {
return nil, err
}
return append(payload, '\n'), nil
}
func Decode(payload []byte) (Snapshot, error) {
decoder := json.NewDecoder(strings.NewReader(string(payload)))
decoder.DisallowUnknownFields()
var snapshot Snapshot
if err := decoder.Decode(&snapshot); err != nil {
return Snapshot{}, err
}
if err := Validate(snapshot); err != nil {
return Snapshot{}, err
}
return snapshot, nil
}
func loadLiveRules(ctx context.Context, pool *pgxpool.Pool) (RuleSet, error) {
rows, err := pool.Query(ctx, `
SELECT id::text, subject_type, subject_id::text, resource_type,
resource_id::text, effect, priority::text,
min_permission_level::text, conditions::text, metadata::text,
status, extract(epoch FROM created_at)::text,
extract(epoch FROM updated_at)::text
FROM gateway_access_rules
ORDER BY id`)
if err != nil {
return RuleSet{}, err
}
defer rows.Close()
allDigests := make([]string, 0)
allowDigests := make([]string, 0)
denyDigests := make([]string, 0)
counts := map[string]int64{}
for rows.Next() {
var row ruleDigestRow
if err := rows.Scan(
&row.id, &row.subjectType, &row.subjectID, &row.resourceType,
&row.resourceID, &row.effect, &row.priority,
&row.minPermissionLevel, &row.conditions, &row.metadata,
&row.status, &row.createdEpoch, &row.updatedEpoch,
); err != nil {
return RuleSet{}, err
}
digest := rowSHA256(row)
allDigests = append(allDigests, digest)
if row.effect == "allow" {
allowDigests = append(allowDigests, digest)
}
if row.effect == "deny" {
denyDigests = append(denyDigests, digest)
}
counts[strings.Join([]string{row.subjectType, row.effect, row.resourceType, row.status}, "\x00")]++
}
if err := rows.Err(); err != nil {
return RuleSet{}, err
}
grouped := make([]RuleCount, 0, len(counts))
for key, count := range counts {
parts := strings.Split(key, "\x00")
grouped = append(grouped, RuleCount{SubjectType: parts[0], Effect: parts[1], ResourceType: parts[2], Status: parts[3], Count: count})
}
sort.Slice(grouped, func(i, j int) bool {
left := grouped[i]
right := grouped[j]
return strings.Join([]string{left.SubjectType, left.Effect, left.ResourceType, left.Status}, "\x00") <
strings.Join([]string{right.SubjectType, right.Effect, right.ResourceType, right.Status}, "\x00")
})
return RuleSet{
Total: int64(len(allDigests)),
SHA256: aggregateSHA256(allDigests),
AllowCount: int64(len(allowDigests)),
AllowSHA256: aggregateSHA256(allowDigests),
DenyCount: int64(len(denyDigests)),
DenySHA256: aggregateSHA256(denyDigests),
Counts: grouped,
}, nil
}
func loadArchiveManifest(ctx context.Context, pool *pgxpool.Pool) (*ArchiveManifest, error) {
var available bool
if err := pool.QueryRow(ctx, `
SELECT to_regclass('gateway_access_rule_migration_batches') IS NOT NULL
AND to_regclass('gateway_access_rule_allow_archive') IS NOT NULL`).Scan(&available); err != nil {
return nil, err
}
if !available {
return nil, nil
}
manifest := &ArchiveManifest{MigrationBatch: MigrationBatch}
if err := pool.QueryRow(ctx, `
SELECT allow_count, allow_sha256, deny_count, deny_sha256
FROM gateway_access_rule_migration_batches
WHERE migration_batch = $1`, MigrationBatch).Scan(
&manifest.ManifestAllowCount,
&manifest.ManifestAllowSHA,
&manifest.ManifestDenyCount,
&manifest.ManifestDenySHA,
); err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, nil
}
return nil, err
}
rows, err := pool.Query(ctx, `
SELECT row_sha256
FROM gateway_access_rule_allow_archive
WHERE migration_batch = $1
ORDER BY rule_id`, MigrationBatch)
if err != nil {
return nil, err
}
defer rows.Close()
digests := make([]string, 0)
for rows.Next() {
var digest string
if err := rows.Scan(&digest); err != nil {
return nil, err
}
digests = append(digests, digest)
}
if err := rows.Err(); err != nil {
return nil, err
}
manifest.ArchivedAllowCount = int64(len(digests))
manifest.ArchivedAllowSHA = aggregateSHA256(digests)
manifest.Consistent = manifest.ManifestAllowCount == manifest.ArchivedAllowCount && manifest.ManifestAllowSHA == manifest.ArchivedAllowSHA
return manifest, nil
}
func rowSHA256(row ruleDigestRow) string {
return sha256Hex(strings.Join([]string{
row.id, row.subjectType, row.subjectID, row.resourceType,
row.resourceID, row.effect, row.priority, row.minPermissionLevel,
row.conditions, row.metadata, row.status, row.createdEpoch, row.updatedEpoch,
}, "\x1f"))
}
func aggregateSHA256(rowDigests []string) string {
return sha256Hex(strings.Join(rowDigests, "\n"))
}
func sha256Hex(value string) string {
digest := sha256.Sum256([]byte(value))
return hex.EncodeToString(digest[:])
}
func validSHA256(value string) bool {
if len(value) != 64 {
return false
}
_, err := hex.DecodeString(value)
return err == nil && value == strings.ToLower(value)
}
@@ -0,0 +1,69 @@
package accessruleaudit
import (
"testing"
"time"
)
func TestVerifyMigrationAcceptsMatchingArchiveAndUnchangedDeny(t *testing.T) {
before := auditSnapshot(3, "allow-before", 2, "deny-before", nil)
after := auditSnapshot(0, "", 2, "deny-before", &ArchiveManifest{
MigrationBatch: MigrationBatch,
ManifestAllowCount: 3,
ManifestAllowSHA: sha256Hex("allow-before"),
ArchivedAllowCount: 3,
ArchivedAllowSHA: sha256Hex("allow-before"),
ManifestDenyCount: 2,
ManifestDenySHA: sha256Hex("deny-before"),
Consistent: true,
})
if err := VerifyMigration(before, after); err != nil {
t.Fatalf("verify matching migration: %v", err)
}
}
func TestVerifyMigrationRejectsChangedDenyOrMissingArchive(t *testing.T) {
before := auditSnapshot(1, "allow", 1, "deny", nil)
withoutArchive := auditSnapshot(0, "", 1, "deny", nil)
if err := VerifyMigration(before, withoutArchive); err == nil {
t.Fatal("missing archive must fail verification")
}
changedDeny := auditSnapshot(0, "", 1, "changed-deny", &ArchiveManifest{
MigrationBatch: MigrationBatch,
ManifestAllowCount: 1,
ManifestAllowSHA: sha256Hex("allow"),
ArchivedAllowCount: 1,
ArchivedAllowSHA: sha256Hex("allow"),
ManifestDenyCount: 1,
ManifestDenySHA: sha256Hex("deny"),
Consistent: true,
})
if err := VerifyMigration(before, changedDeny); err == nil {
t.Fatal("changed deny hash must fail verification")
}
}
func auditSnapshot(allowCount int64, allowSeed string, denyCount int64, denySeed string, archive *ArchiveManifest) Snapshot {
counts := make([]RuleCount, 0, 2)
if allowCount > 0 {
counts = append(counts, RuleCount{SubjectType: "api_key", Effect: "allow", ResourceType: "platform_model", Status: "active", Count: allowCount})
}
if denyCount > 0 {
counts = append(counts, RuleCount{SubjectType: "api_key", Effect: "deny", ResourceType: "platform_model", Status: "active", Count: denyCount})
}
return Snapshot{
SchemaVersion: SchemaVersion,
GeneratedAt: time.Date(2026, 8, 3, 0, 0, 0, 0, time.UTC),
SecretSafe: true,
Live: RuleSet{
Total: allowCount + denyCount,
SHA256: sha256Hex("all-" + allowSeed + "-" + denySeed),
AllowCount: allowCount,
AllowSHA256: sha256Hex(allowSeed),
DenyCount: denyCount,
DenySHA256: sha256Hex(denySeed),
Counts: counts,
},
Archive: archive,
}
}
+49 -25
View File
@@ -34,27 +34,36 @@ const (
)
type User struct {
ID string `json:"sub"`
Username string `json:"username"`
Roles []string `json:"role,omitempty"`
TenantID string `json:"tenantId,omitempty"`
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
TenantKey string `json:"tenantKey,omitempty"`
SSOID string `json:"sso_id,omitempty"`
Source string `json:"source,omitempty"`
GatewayUserID string `json:"gatewayUserId,omitempty"`
UserGroupID string `json:"userGroupId,omitempty"`
UserGroupKey string `json:"userGroupKey,omitempty"`
UserGroupKeys []string `json:"userGroupKeys,omitempty"`
APIKeyID string `json:"apiKeyId,omitempty"`
APIKeySecret string `json:"apiKeySecret,omitempty"`
APIKeyName string `json:"apiKeyName,omitempty"`
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
APIKeyScopes []string `json:"apiKeyScopes,omitempty"`
TokenExpiresAt time.Time `json:"-"`
TokenIssuedAt time.Time `json:"-"`
Issuer string `json:"-"`
TokenPurpose string `json:"-"`
ID string `json:"sub"`
Username string `json:"username"`
DisplayName string `json:"-"`
Email string `json:"-"`
Phone string `json:"-"`
AvatarURL string `json:"-"`
Roles []string `json:"role,omitempty"`
ContextType string `json:"contextType,omitempty"`
TenantID string `json:"tenantId,omitempty"`
TenantName string `json:"tenantName,omitempty"`
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
TenantKey string `json:"tenantKey,omitempty"`
SSOID string `json:"sso_id,omitempty"`
Source string `json:"source,omitempty"`
GatewayUserID string `json:"gatewayUserId,omitempty"`
UserGroupID string `json:"userGroupId,omitempty"`
UserGroupKey string `json:"userGroupKey,omitempty"`
UserGroupKeys []string `json:"userGroupKeys,omitempty"`
APIKeyID string `json:"apiKeyId,omitempty"`
APIKeySecret string `json:"apiKeySecret,omitempty"`
APIKeyName string `json:"apiKeyName,omitempty"`
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
APIKeyScopes []string `json:"apiKeyScopes,omitempty"`
TokenExpiresAt time.Time `json:"-"`
TokenIssuedAt time.Time `json:"-"`
Issuer string `json:"-"`
ApplicationID string `json:"-"`
OIDCClientID string `json:"-"`
OIDCUserBindingID string `json:"-"`
TokenPurpose string `json:"-"`
}
type contextKey string
@@ -64,9 +73,10 @@ const userContextKey contextKey = "easyai-auth-user"
var ErrUnauthorized = errors.New("unauthorized")
type RequestAuthError struct {
Status int
Code string
Message string
Status int
Code string
Message string
RetryAfterSeconds int
}
func (e *RequestAuthError) Error() string { return e.Code }
@@ -75,6 +85,12 @@ func NewRequestAuthError(status int, code, message string) error {
return &RequestAuthError{Status: status, Code: code, Message: message}
}
func NewRetryableRequestAuthError(status int, code, message string, retryAfterSeconds int) error {
return &RequestAuthError{
Status: status, Code: code, Message: message, RetryAfterSeconds: retryAfterSeconds,
}
}
type Authenticator struct {
JWTSecret string
ServerMainBaseURL string
@@ -122,6 +138,9 @@ func (a *Authenticator) Require(permission Permission, next http.Handler) http.H
}
var requestError *RequestAuthError
if errors.As(err, &requestError) {
if requestError.RetryAfterSeconds > 0 {
w.Header().Set("Retry-After", strconv.Itoa(requestError.RetryAfterSeconds))
}
w.Header().Set("Content-Type", "application/json; charset=utf-8")
w.Header().Set("Cache-Control", "no-store")
w.WriteHeader(requestError.Status)
@@ -313,7 +332,12 @@ func (a *Authenticator) verifyAPIKey(ctx context.Context, apiKey string) (*User,
return user, nil
}
if !errors.Is(err, ErrUnauthorized) {
return nil, err
return nil, NewRetryableRequestAuthError(
http.StatusServiceUnavailable,
"AUTH_STORE_UNAVAILABLE",
"authentication service temporarily unavailable",
2,
)
}
if strings.HasPrefix(apiKey, "sk-gw-") {
return nil, ErrUnauthorized
@@ -0,0 +1,31 @@
package auth
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
)
func TestAPIKeyStoreFailureReturnsRetryableServiceUnavailable(t *testing.T) {
authenticator := New("test-secret", "", "")
authenticator.LocalAPIKeyVerifier = func(context.Context, string) (*User, error) {
return nil, errors.New("database unavailable")
}
handler := authenticator.Require(PermissionBasic, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("unavailable authentication store must not call the protected handler")
}))
request := httptest.NewRequest(http.MethodGet, "/protected", nil)
request.Header.Set("Authorization", "Bearer sk-gw-local")
recorder := httptest.NewRecorder()
handler.ServeHTTP(recorder, request)
if recorder.Code != http.StatusServiceUnavailable {
t.Fatalf("status=%d, want 503; body=%s", recorder.Code, recorder.Body.String())
}
if recorder.Header().Get("Retry-After") != "2" {
t.Fatalf("Retry-After=%q, want 2", recorder.Header().Get("Retry-After"))
}
}
@@ -0,0 +1,63 @@
package auth
import (
"os"
"path/filepath"
"runtime"
"slices"
"strings"
"testing"
)
func TestGatewaySourceDoesNotDependOnAuthenticationCoreVocabulary(t *testing.T) {
_, currentFile, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("resolve current test file")
}
repositoryRoot := filepath.Clean(filepath.Join(filepath.Dir(currentFile), "..", "..", "..", ".."))
appsRoot := filepath.Join(repositoryRoot, "apps")
forbiddenMarkers := []string{
strings.ToLower("key" + "cloak"),
"/rea" + "lms/",
}
sourceExtensions := []string{".go", ".js", ".json", ".ts", ".tsx", ".yaml", ".yml"}
skippedDirectories := []string{"coverage", "dist", "node_modules", "test-results"}
var violations []string
err := filepath.WalkDir(appsRoot, func(filePath string, entry os.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
if entry.IsDir() {
if filePath != appsRoot && slices.Contains(skippedDirectories, entry.Name()) {
return filepath.SkipDir
}
return nil
}
if !slices.Contains(sourceExtensions, strings.ToLower(filepath.Ext(entry.Name()))) {
return nil
}
content, readErr := os.ReadFile(filePath)
if readErr != nil {
return readErr
}
lowerContent := strings.ToLower(string(content))
for _, marker := range forbiddenMarkers {
if strings.Contains(lowerContent, marker) {
relativePath, relativeErr := filepath.Rel(repositoryRoot, filePath)
if relativeErr != nil {
return relativeErr
}
violations = append(violations, filepath.ToSlash(relativePath))
break
}
}
return nil
})
if err != nil {
t.Fatalf("scan Gateway sources: %v", err)
}
if len(violations) > 0 {
t.Fatalf("Gateway source contains authentication-core-specific vocabulary: %s", strings.Join(violations, ", "))
}
}
+164 -29
View File
@@ -20,6 +20,7 @@ import (
"time"
"github.com/golang-jwt/jwt/v5"
"github.com/google/uuid"
)
const maxOIDCResponseBytes = 1 << 20
@@ -30,6 +31,9 @@ type OIDCConfig struct {
Issuer string
Audience string
TenantID string
TenantMode string
ApplicationID string
ClientID string
RolePrefix string
RequiredScopes []string
JWKSCacheTTL time.Duration
@@ -44,10 +48,12 @@ type OIDCConfig struct {
}
type OIDCSecurityEventIdentity struct {
Issuer string
TenantID string
Subject string
IssuedAt time.Time
Issuer string
ApplicationID string
ContextType string
TenantID string
Subject string
IssuedAt time.Time
}
type OIDCSecurityEventEvaluation struct {
@@ -91,8 +97,16 @@ func NewOIDCVerifier(config OIDCConfig) (*OIDCVerifier, error) {
config.Issuer = strings.TrimRight(strings.TrimSpace(config.Issuer), "/")
config.Audience = strings.TrimSpace(config.Audience)
config.TenantID = strings.TrimSpace(config.TenantID)
config.TenantMode = strings.TrimSpace(config.TenantMode)
config.ApplicationID = strings.TrimSpace(config.ApplicationID)
config.ClientID = strings.TrimSpace(config.ClientID)
config.RolePrefix = strings.TrimSpace(config.RolePrefix)
if err := validatePublicURL(config.Issuer, config.AppEnv); err != nil || config.Audience == "" || config.TenantID == "" || config.RolePrefix == "" {
if config.TenantMode == "" {
config.TenantMode = "single_tenant"
}
validTenantMode := config.TenantMode == "single_tenant" && config.TenantID != "" ||
config.TenantMode == "multi_tenant" && config.TenantID == "" && config.ApplicationID != ""
if err := validatePublicURL(config.Issuer, config.AppEnv); err != nil || config.Audience == "" || !validTenantMode || config.RolePrefix == "" {
return nil, errors.New("issuer, audience, tenant and role prefix are required")
}
if config.IntrospectionEnabled && config.IntrospectionCredentialProvider == nil &&
@@ -132,15 +146,28 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
parser := jwt.NewParser(jwt.WithValidMethods([]string{"RS256", "ES256"}))
unverified, _, err := parser.ParseUnverified(raw, jwt.MapClaims{})
if err != nil || unverified == nil {
return nil, oidcUnauthorized("token envelope is invalid", err)
return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", err)
}
unverifiedClaims, ok := unverified.Claims.(jwt.MapClaims)
if !ok {
return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", nil)
}
if stringClaim(unverifiedClaims, "iss") == "" {
return nil, oidcUnauthorized("ISSUER_MISSING", "issuer is missing", nil)
}
if len(stringSliceClaim(unverifiedClaims, "aud")) == 0 {
return nil, oidcUnauthorized("AUDIENCE_MISSING", "audience is missing", nil)
}
if _, exists := unverifiedClaims["exp"]; !exists {
return nil, oidcUnauthorized("EXP_MISSING", "exp is missing", nil)
}
kid, _ := unverified.Header["kid"].(string)
if kid == "" {
return nil, oidcUnauthorized("kid is missing", nil)
return nil, oidcUnauthorized("KID_MISSING", "kid is missing", nil)
}
key, err := v.key(ctx, kid)
if err != nil {
return nil, oidcUnauthorized("signing key lookup failed", err)
return nil, oidcUnauthorized("SIGNING_KEY_LOOKUP_FAILED", "signing key lookup failed", err)
}
token, err := jwt.Parse(raw, func(token *jwt.Token) (any, error) {
if token.Header["kid"] != kid {
@@ -150,43 +177,60 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
}, jwt.WithValidMethods([]string{"RS256", "ES256"}), jwt.WithIssuer(v.config.Issuer),
jwt.WithAudience(v.config.Audience), jwt.WithExpirationRequired(), jwt.WithLeeway(30*time.Second))
if err != nil || !token.Valid {
return nil, oidcUnauthorized("signature or registered claims are invalid", err)
return nil, oidcUnauthorized(registeredClaimsValidationCategory(err), "signature or registered claims are invalid", err)
}
claims, ok := token.Claims.(jwt.MapClaims)
if !ok || stringClaim(claims, "sub") == "" || stringClaim(claims, "tid") != v.config.TenantID {
return nil, oidcUnauthorized("stable identity claims are invalid", nil)
contextType := stringClaim(claims, "context_type")
tenantID := stringClaim(claims, "tid")
validContext := contextType == "tenant" && tenantID == v.config.TenantID
if v.config.TenantMode == "multi_tenant" {
switch contextType {
case "platform":
validContext = tenantID == ""
case "tenant":
validContext = uuid.Validate(tenantID) == nil
default:
validContext = false
}
}
clientID := stringClaim(claims, "client_id")
if !ok || stringClaim(claims, "sub") == "" || !validContext ||
v.config.ClientID != "" && clientID != v.config.ClientID {
return nil, oidcUnauthorized("STABLE_IDENTITY_CLAIMS_INVALID", "stable identity claims are invalid", nil)
}
expiresAt, ok := numericDateClaim(claims["exp"])
if !ok {
return nil, oidcUnauthorized("exp is invalid", nil)
return nil, oidcUnauthorized("EXP_INVALID", "exp is invalid", nil)
}
if _, ok := numericDateClaim(claims["nbf"]); !ok {
return nil, oidcUnauthorized("nbf is missing", nil)
return nil, oidcUnauthorized("NBF_MISSING", "nbf is missing", nil)
}
issuedAt, hasIssuedAt := numericDateClaim(claims["iat"])
scopes := scopeClaims(claims)
if !containsAll(scopes, v.config.RequiredScopes) {
return nil, oidcUnauthorized("required scope is missing", nil)
return nil, oidcUnauthorized("REQUIRED_SCOPE_MISSING", "required scope is missing", nil)
}
roles := gatewayRoles(stringSliceClaim(claims, "roles"), v.config.RolePrefix)
if len(roles) == 0 {
roles = gatewayRoles(stringSliceClaim(claims, "role"), v.config.RolePrefix)
}
if len(roles) == 0 {
return nil, oidcUnauthorized("mapped role is missing", nil)
return nil, oidcUnauthorized("MAPPED_ROLE_MISSING", "mapped role is missing", nil)
}
if v.config.SecurityEventEvaluator != nil {
evaluation, evaluateErr := v.config.SecurityEventEvaluator(ctx, OIDCSecurityEventIdentity{
Issuer: v.config.Issuer, TenantID: v.config.TenantID, Subject: stringClaim(claims, "sub"), IssuedAt: issuedAt,
Issuer: v.config.Issuer, ApplicationID: v.config.ApplicationID,
ContextType: contextType, TenantID: tenantID,
Subject: stringClaim(claims, "sub"), IssuedAt: issuedAt,
})
if evaluateErr != nil {
return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_SECURITY_EVENT_STATE_UNAVAILABLE", "认证撤销状态暂时不可用")
}
if evaluation.Enabled && !hasIssuedAt {
return nil, oidcUnauthorized("iat is required when security events are enabled", nil)
return nil, oidcUnauthorized("SECURITY_EVENT_IAT_MISSING", "iat is required when security events are enabled", nil)
}
if evaluation.Revoked {
return nil, oidcUnauthorized("token was issued before the revocation watermark", nil)
return nil, oidcUnauthorized("TOKEN_REVOKED", "token was issued before the revocation watermark", nil)
}
if evaluation.RequireIntrospection {
active, introspectionErr := v.introspect(ctx, raw)
@@ -194,35 +238,126 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_INTROSPECTION_UNAVAILABLE", "认证中心内省暂时不可用")
}
if !active {
return nil, oidcUnauthorized("token is inactive", nil)
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", nil)
}
} else if !evaluation.Enabled && v.config.IntrospectionEnabled {
active, introspectionErr := v.introspect(ctx, raw)
if introspectionErr != nil || !active {
return nil, oidcUnauthorized("token is inactive", introspectionErr)
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", introspectionErr)
}
}
} else if v.config.IntrospectionEnabled {
active, err := v.introspect(ctx, raw)
if err != nil || !active {
return nil, oidcUnauthorized("token is inactive", err)
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", err)
}
}
username := stringClaim(claims, "preferred_username")
username := oidcProfileText(claims, "preferred_username", 320)
if username == "" {
username = stringClaim(claims, "username")
username = oidcProfileText(claims, "username", 320)
}
return &User{
ID: stringClaim(claims, "sub"), Username: username, Roles: roles,
TenantID: v.config.TenantID, Source: "oidc", TokenExpiresAt: expiresAt, TokenIssuedAt: issuedAt, Issuer: v.config.Issuer,
ID: stringClaim(claims, "sub"), Username: username,
DisplayName: oidcProfileText(claims, "name", 200),
Email: oidcVerifiedProfileText(claims, "email", "email_verified", 320),
Phone: oidcVerifiedProfileText(claims, "phone_number", "phone_number_verified", 64),
AvatarURL: safeOIDCProfileURL(oidcProfileText(claims, "picture", 2048)),
Roles: roles,
ContextType: contextType, TenantID: tenantID, Source: "oidc",
TokenExpiresAt: expiresAt, TokenIssuedAt: issuedAt, Issuer: v.config.Issuer,
ApplicationID: v.config.ApplicationID, OIDCClientID: clientID,
}, nil
}
func oidcUnauthorized(reason string, cause error) error {
if cause != nil {
return fmt.Errorf("%w: %s: %v", ErrUnauthorized, reason, cause)
func oidcProfileText(claims jwt.MapClaims, key string, limit int) string {
value := strings.TrimSpace(stringClaim(claims, key))
if value == "" {
return ""
}
runes := []rune(value)
if len(runes) > limit {
return string(runes[:limit])
}
return value
}
func oidcVerifiedProfileText(claims jwt.MapClaims, key, verifiedKey string, limit int) string {
verified, ok := claims[verifiedKey].(bool)
if !ok || !verified {
return ""
}
return oidcProfileText(claims, key, limit)
}
func safeOIDCProfileURL(value string) string {
parsed, err := url.Parse(value)
if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil {
return ""
}
return parsed.String()
}
type oidcValidationError struct {
category string
reason string
cause error
}
func (err *oidcValidationError) Error() string {
if err.cause != nil {
return fmt.Sprintf("%v: %s: %v", ErrUnauthorized, err.reason, err.cause)
}
return fmt.Sprintf("%v: %s", ErrUnauthorized, err.reason)
}
func (err *oidcValidationError) Unwrap() []error {
if err.cause == nil {
return []error{ErrUnauthorized}
}
return []error{ErrUnauthorized, err.cause}
}
func OIDCValidationCategory(err error) string {
var validationError *oidcValidationError
if errors.As(err, &validationError) {
return validationError.category
}
return "UNCLASSIFIED"
}
func registeredClaimsValidationCategory(err error) string {
switch {
case errors.Is(err, jwt.ErrTokenRequiredClaimMissing):
return "REQUIRED_CLAIM_MISSING"
case errors.Is(err, jwt.ErrTokenInvalidAudience):
return "AUDIENCE_INVALID"
case errors.Is(err, jwt.ErrTokenInvalidIssuer):
return "ISSUER_INVALID"
case errors.Is(err, jwt.ErrTokenExpired):
return "TOKEN_EXPIRED"
case errors.Is(err, jwt.ErrTokenNotValidYet):
return "TOKEN_NOT_VALID_YET"
case errors.Is(err, jwt.ErrTokenUsedBeforeIssued):
return "TOKEN_USED_BEFORE_ISSUED"
case errors.Is(err, jwt.ErrTokenSignatureInvalid):
return "SIGNATURE_INVALID"
case errors.Is(err, jwt.ErrTokenMalformed):
return "TOKEN_MALFORMED"
case errors.Is(err, jwt.ErrTokenUnverifiable):
return "TOKEN_UNVERIFIABLE"
case errors.Is(err, jwt.ErrTokenInvalidClaims):
return "CLAIMS_INVALID"
default:
return "REGISTERED_CLAIMS_INVALID"
}
}
func oidcUnauthorized(category, reason string, cause error) error {
return &oidcValidationError{
category: category,
reason: reason,
cause: cause,
}
return fmt.Errorf("%w: %s", ErrUnauthorized, reason)
}
func (v *OIDCVerifier) key(ctx context.Context, kid string) (any, error) {
+33 -9
View File
@@ -11,6 +11,7 @@ import (
"sync"
"github.com/coreos/go-oidc/v3/oidc"
"github.com/google/uuid"
"golang.org/x/oauth2"
)
@@ -77,15 +78,38 @@ func (c *OIDCPublicClient) ValidateConfiguration(ctx context.Context) error {
return err
}
func (c *OIDCPublicClient) AuthorizationURL(ctx context.Context, state, nonce, pkceVerifier string) (string, error) {
func (c *OIDCPublicClient) AuthorizationURL(
ctx context.Context,
state, nonce, pkceVerifier, contextType, tenantHint string,
) (string, error) {
if strings.TrimSpace(state) == "" || strings.TrimSpace(nonce) == "" || !validPKCEVerifier(pkceVerifier) {
return "", errors.New("state, nonce and PKCE verifier are required")
}
contextType = strings.TrimSpace(contextType)
if contextType != "" && contextType != "platform" && contextType != "tenant" {
return "", errors.New("context type must be empty, platform or tenant")
}
if strings.TrimSpace(tenantHint) != "" && uuid.Validate(strings.TrimSpace(tenantHint)) != nil {
return "", errors.New("tenant hint must be a UUID")
}
if contextType != "tenant" && strings.TrimSpace(tenantHint) != "" {
return "", errors.New("only tenant context can bind a tenant hint")
}
config, _, err := c.configuration(ctx)
if err != nil {
return "", err
}
return config.AuthCodeURL(state, oidc.Nonce(nonce), oauth2.S256ChallengeOption(pkceVerifier)), nil
options := []oauth2.AuthCodeOption{
oidc.Nonce(nonce),
oauth2.S256ChallengeOption(pkceVerifier),
}
if contextType != "" {
options = append(options, oauth2.SetAuthURLParam("context_type", contextType))
}
if strings.TrimSpace(tenantHint) != "" {
options = append(options, oauth2.SetAuthURLParam("tenant_hint", strings.TrimSpace(tenantHint)))
}
return config.AuthCodeURL(state, options...), nil
}
func (c *OIDCPublicClient) ExchangeCode(ctx context.Context, code, verifier string) (OIDCTokenResponse, error) {
@@ -121,30 +145,30 @@ func (c *OIDCPublicClient) Refresh(ctx context.Context, refreshToken string) (OI
func (c *OIDCPublicClient) VerifyIDToken(ctx context.Context, raw, expectedNonce string) (string, error) {
expectedNonce = strings.TrimSpace(expectedNonce)
if strings.TrimSpace(raw) == "" || expectedNonce == "" {
return "", oidcUnauthorized("ID token validation context is invalid", nil)
return "", oidcUnauthorized("ID_TOKEN_CONTEXT_INVALID", "ID token validation context is invalid", nil)
}
if _, _, err := c.configuration(ctx); err != nil {
return "", oidcUnauthorized("ID token provider discovery failed", err)
return "", oidcUnauthorized("ID_TOKEN_DISCOVERY_FAILED", "ID token provider discovery failed", err)
}
c.mu.Lock()
verifier := c.idTokenVerifier
c.mu.Unlock()
if verifier == nil {
return "", oidcUnauthorized("ID token verifier is unavailable", nil)
return "", oidcUnauthorized("ID_TOKEN_VERIFIER_UNAVAILABLE", "ID token verifier is unavailable", nil)
}
token, err := verifier.Verify(c.requestContext(ctx), raw)
if err != nil {
return "", oidcUnauthorized("ID token signature or registered claims are invalid", nil)
return "", oidcUnauthorized("ID_TOKEN_REGISTERED_CLAIMS_INVALID", "ID token signature or registered claims are invalid", nil)
}
if token.Subject == "" || token.Nonce != expectedNonce {
return "", oidcUnauthorized("ID token subject or nonce is invalid", nil)
return "", oidcUnauthorized("ID_TOKEN_SUBJECT_OR_NONCE_INVALID", "ID token subject or nonce is invalid", nil)
}
var claims map[string]any
if err := token.Claims(&claims); err != nil {
return "", oidcUnauthorized("ID token claims are invalid", nil)
return "", oidcUnauthorized("ID_TOKEN_CLAIMS_INVALID", "ID token claims are invalid", nil)
}
if _, ok := numericDateClaim(claims["nbf"]); !ok {
return "", oidcUnauthorized("ID token nbf is missing", nil)
return "", oidcUnauthorized("ID_TOKEN_NBF_MISSING", "ID token nbf is missing", nil)
}
return token.Subject, nil
}
+47 -1
View File
@@ -67,7 +67,10 @@ func TestOIDCPublicClientUsesAuthorizationCodePKCES256WithoutSecret(t *testing.T
if err := client.ValidateConfiguration(context.Background()); err != nil {
t.Fatalf("ValidateConfiguration() error = %v", err)
}
authorizationURL, err := client.AuthorizationURL(context.Background(), "state", "nonce", pkceVerifier)
tenantHint := "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"
authorizationURL, err := client.AuthorizationURL(
context.Background(), "state", "nonce", pkceVerifier, "tenant", tenantHint,
)
if err != nil {
t.Fatal(err)
}
@@ -78,11 +81,54 @@ func TestOIDCPublicClientUsesAuthorizationCodePKCES256WithoutSecret(t *testing.T
if query.Get("response_type") != "code" || query.Get("code_challenge_method") != "S256" || query.Get("code_challenge") != expectedChallenge {
t.Fatalf("authorization request is not PKCE S256: %v", query)
}
if query.Get("tenant_hint") != tenantHint {
t.Fatalf("tenant_hint = %q, want %q", query.Get("tenant_hint"), tenantHint)
}
if query.Get("context_type") != "tenant" {
t.Fatalf("context_type = %q, want tenant", query.Get("context_type"))
}
if _, err := client.ExchangeCode(context.Background(), "authorization-code", pkceVerifier); err != nil {
t.Fatal(err)
}
}
func TestOIDCPublicClientOmitsOptionalContextForUnifiedLogin(t *testing.T) {
var issuer string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/.well-known/openid-configuration" {
http.NotFound(w, r)
return
}
_ = json.NewEncoder(w).Encode(map[string]string{
"issuer": issuer, "authorization_endpoint": issuer + "/authorize",
"token_endpoint": issuer + "/token", "jwks_uri": issuer + "/jwks",
})
}))
defer server.Close()
issuer = server.URL
client, err := NewOIDCPublicClient(OIDCPublicClientConfig{
AppEnv: "test", Issuer: issuer, ClientID: "gateway-public",
RedirectURI: "https://gateway.example.com/callback",
PostLogoutRedirectURI: "https://gateway.example.com/",
HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
authorizationURL, err := client.AuthorizationURL(
context.Background(), "state", "nonce",
"test-pkce-verifier-with-at-least-43-characters-1234", "", "",
)
if err != nil {
t.Fatal(err)
}
parsed, _ := url.Parse(authorizationURL)
if parsed.Query().Has("context_type") || parsed.Query().Has("tenant_hint") {
t.Fatalf("unified login leaked optional context binding: %v", parsed.Query())
}
}
func TestOIDCPublicClientRejectsOfflineAccess(t *testing.T) {
_, err := NewOIDCPublicClient(OIDCPublicClientConfig{
AppEnv: "production",
+124 -8
View File
@@ -66,11 +66,27 @@ func TestOIDCVerifierAcceptsRS256AndES256StableClaims(t *testing.T) {
t.Fatal(err)
}
if user.ID != "platform-subject" || user.TenantID != "tenant-1" || user.Source != "oidc" ||
user.Username != "acceptance" || user.DisplayName != "王小明" ||
user.Email != "real.user@example.test" || user.Phone != "+8613800000000" ||
user.AvatarURL != "https://static.example.test/avatar.png" ||
len(user.Roles) != 1 || user.Roles[0] != "admin" {
t.Fatalf("unexpected OIDC user: %#v", user)
}
})
}
unverified := signedOIDCToken(t, issuer, "rsa-key", jwt.SigningMethodRS256, rsaKey, func(claims jwt.MapClaims) {
claims["email_verified"] = false
claims["phone_number_verified"] = false
claims["picture"] = "javascript:alert(1)"
})
user, err := verifier.Verify(context.Background(), unverified)
if err != nil {
t.Fatal(err)
}
if user.Email != "" || user.Phone != "" || user.AvatarURL != "" {
t.Fatalf("untrusted profile claims were accepted: %#v", user)
}
}
func TestOIDCVerifierRejectsMissingOrMismatchedSecurityClaims(t *testing.T) {
@@ -91,20 +107,116 @@ func TestOIDCVerifierRejectsMissingOrMismatchedSecurityClaims(t *testing.T) {
RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
})
tests := []struct {
name string
mutate func(jwt.MapClaims)
name, wantCategory string
mutate func(jwt.MapClaims)
}{
{"missing nbf", func(claims jwt.MapClaims) { delete(claims, "nbf") }},
{"wrong audience", func(claims jwt.MapClaims) { claims["aud"] = "other-api" }},
{"wrong tenant", func(claims jwt.MapClaims) { claims["tid"] = "tenant-2" }},
{"missing scope", func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
{"unmapped role", func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
{"missing nbf", "NBF_MISSING", func(claims jwt.MapClaims) { delete(claims, "nbf") }},
{"missing exp", "EXP_MISSING", func(claims jwt.MapClaims) { delete(claims, "exp") }},
{"missing issuer", "ISSUER_MISSING", func(claims jwt.MapClaims) { delete(claims, "iss") }},
{"missing audience", "AUDIENCE_MISSING", func(claims jwt.MapClaims) { delete(claims, "aud") }},
{"wrong audience", "AUDIENCE_INVALID", func(claims jwt.MapClaims) { claims["aud"] = "other-api" }},
{"missing context type", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { delete(claims, "context_type") }},
{"unknown context type", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { claims["context_type"] = "account" }},
{"wrong tenant", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { claims["tid"] = "tenant-2" }},
{"missing scope", "REQUIRED_SCOPE_MISSING", func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
{"unmapped role", "MAPPED_ROLE_MISSING", func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
raw := signedOIDCToken(t, issuer, "ec-key", jwt.SigningMethodES256, key, test.mutate)
if _, err := verifier.Verify(context.Background(), raw); err == nil {
t.Fatal("expected token to be rejected")
} else {
if !errors.Is(err, ErrUnauthorized) {
t.Fatalf("validation error no longer wraps ErrUnauthorized: %v", err)
}
if category := OIDCValidationCategory(err); category != test.wantCategory {
t.Fatalf("validation category=%q, want %q", category, test.wantCategory)
}
}
})
}
}
func TestOIDCVerifierAcceptsExplicitPlatformAndTenantContextsForMultiTenantApplication(t *testing.T) {
key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
var issuer string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.URL.Path == "/.well-known/openid-configuration" {
_ = json.NewEncoder(w).Encode(map[string]any{
"issuer": issuer,
"jwks_uri": issuer + "/jwks",
})
return
}
_ = json.NewEncoder(w).Encode(map[string]any{
"keys": []any{ecJWK("ec-key", &key.PublicKey)},
})
}))
defer server.Close()
issuer = server.URL
verifier, err := NewOIDCVerifier(OIDCConfig{
AppEnv: "test",
Issuer: issuer, Audience: "gateway-api",
TenantMode: "multi_tenant",
ApplicationID: "11111111-1111-4111-8111-111111111111",
RolePrefix: "gateway.", HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
platformToken := signedOIDCToken(
t,
issuer,
"ec-key",
jwt.SigningMethodES256,
key,
func(claims jwt.MapClaims) {
claims["context_type"] = "platform"
delete(claims, "tid")
},
)
platformUser, err := verifier.Verify(context.Background(), platformToken)
if err != nil {
t.Fatalf("platform context rejected: %v", err)
}
if platformUser.ContextType != "platform" || platformUser.TenantID != "" {
t.Fatalf("platform user=%#v", platformUser)
}
tenantID := "22222222-2222-4222-8222-222222222222"
tenantToken := signedOIDCToken(
t,
issuer,
"ec-key",
jwt.SigningMethodES256,
key,
func(claims jwt.MapClaims) {
claims["context_type"] = "tenant"
claims["tid"] = tenantID
},
)
tenantUser, err := verifier.Verify(context.Background(), tenantToken)
if err != nil {
t.Fatalf("tenant context rejected: %v", err)
}
if tenantUser.ContextType != "tenant" || tenantUser.TenantID != tenantID {
t.Fatalf("tenant user=%#v", tenantUser)
}
for name, mutate := range map[string]func(jwt.MapClaims){
"platform with tenant": func(claims jwt.MapClaims) {
claims["context_type"] = "platform"
},
"tenant without tenant": func(claims jwt.MapClaims) {
claims["context_type"] = "tenant"
delete(claims, "tid")
},
} {
t.Run(name, func(t *testing.T) {
raw := signedOIDCToken(
t, issuer, "ec-key", jwt.SigningMethodES256, key, mutate,
)
if _, err := verifier.Verify(context.Background(), raw); err == nil {
t.Fatal("context/tenant mismatch was accepted")
}
})
}
@@ -302,8 +414,12 @@ func signedOIDCToken(t *testing.T, issuer, kid string, method jwt.SigningMethod,
now := time.Now()
claims := jwt.MapClaims{
"iss": issuer, "aud": "gateway-api", "sub": "platform-subject", "tid": "tenant-1",
"context_type": "tenant",
"preferred_username": "acceptance", "roles": []string{"gateway.admin"},
"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
"name": "王小明", "email": "real.user@example.test", "email_verified": true,
"phone_number": "+8613800000000", "phone_number_verified": true,
"picture": "https://static.example.test/avatar.png",
"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
"exp": now.Add(time.Hour).Unix(),
}
if mutate != nil {
@@ -0,0 +1,40 @@
package capacitycontroller
import (
"errors"
"os"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
)
func NewConfiguredAdapter(cfg config.Config) (OrchestratorAdapter, error) {
switch strings.ToLower(strings.TrimSpace(cfg.CapacityOrchestratorAdapter)) {
case "static":
pools, err := cfg.CapacityPools()
if err != nil {
return nil, err
}
initial := make(map[string]int, len(pools))
for _, pool := range pools {
initial[pool.ID] = pool.BootstrapReplicas
}
return NewStaticAdapter(initial), nil
case "", "kubernetes":
return NewKubernetesClient(KubernetesConfig{
Namespace: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_NAMESPACE", adapterEnv("POD_NAMESPACE", "easyai")),
APIServer: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_API_SERVER", "https://kubernetes.default.svc"),
TokenFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_TOKEN_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/token"),
CAFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_CA_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"),
})
default:
return nil, errors.New("unsupported capacity orchestrator adapter")
}
}
func adapterEnv(name, fallback string) string {
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
return value
}
return fallback
}
@@ -0,0 +1,454 @@
package capacitycontroller
import (
"context"
"encoding/json"
"errors"
"log/slog"
"net/http"
"os"
"strings"
"sync"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/jackc/pgx/v5"
)
const (
controllerReconcileInterval = 10 * time.Second
controllerLeadershipRetry = 5 * time.Second
controllerDeletionCost = -1000
)
type capacityStore interface {
TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error)
WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error)
ListPoolQueueRuntime(context.Context) ([]store.PoolQueueRuntime, error)
ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
MarkWorkerDraining(context.Context, string) error
ReactivateWorkerInstance(context.Context, string) error
PublishDesiredCapacity(context.Context, executionpool.DesiredCapacity) error
}
type OrchestratorAdapter interface {
PoolState(context.Context, string, string) (PoolInfrastructureState, error)
ScalePool(context.Context, string, string, int) error
SetInstanceTerminationPriority(context.Context, string, int) error
}
type Status struct {
Leader bool `json:"leader"`
LastRunAt time.Time `json:"lastRunAt,omitempty"`
LastError string `json:"lastError,omitempty"`
Queue store.WorkerQueueRuntime `json:"queue"`
Plan Plan `json:"plan"`
ScaleActions uint64 `json:"scaleActions"`
}
type Controller struct {
cfg config.Config
store capacityStore
orchestrator OrchestratorAdapter
pools []config.ExecutionPoolCapacityConfig
logger *slog.Logger
expectedRevision string
now func() time.Time
highSince time.Time
lowSince time.Time
statusMu sync.RWMutex
status Status
}
func New(
cfg config.Config,
db capacityStore,
orchestrator OrchestratorAdapter,
logger *slog.Logger,
) *Controller {
return &Controller{
cfg: cfg, store: db, orchestrator: orchestrator, logger: logger,
expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
now: time.Now,
pools: mustCapacityPools(cfg),
}
}
func mustCapacityPools(cfg config.Config) []config.ExecutionPoolCapacityConfig {
pools, _ := cfg.CapacityPools()
return pools
}
func (controller *Controller) Run(ctx context.Context) {
for ctx.Err() == nil {
leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx)
if err != nil {
controller.setError(err)
if !waitContext(ctx, controllerLeadershipRetry) {
return
}
continue
}
if !acquired {
controller.setLeader(false)
controller.clearError()
if !waitContext(ctx, controllerLeadershipRetry) {
return
}
continue
}
controller.setLeader(true)
controller.clearError()
controller.runLeader(ctx, leadership)
leadership.Release()
controller.setLeader(false)
}
}
func (controller *Controller) runLeader(ctx context.Context, leadership store.Leadership) {
ticker := time.NewTicker(controllerReconcileInterval)
defer ticker.Stop()
if err := controller.reconcile(ctx); err != nil {
controller.logError("initial capacity reconciliation failed", err)
}
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
keepAliveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
err := leadership.KeepAlive(keepAliveCtx)
cancel()
if err != nil {
controller.logError("capacity controller leadership lost", err)
return
}
if err := controller.reconcile(ctx); err != nil {
controller.logError("capacity reconciliation failed", err)
}
}
}
}
func (controller *Controller) reconcile(ctx context.Context) error {
queue, err := controller.store.WorkerQueueRuntime(ctx)
if err != nil {
return err
}
poolQueues, err := controller.store.ListPoolQueueRuntime(ctx)
if err != nil {
return err
}
demandByPool := make(map[string]int, len(poolQueues))
for _, poolQueue := range poolQueues {
demandByPool[poolQueue.PoolID] = poolQueue.Queued + poolQueue.Running
}
instances, err := controller.store.ListWorkerInstanceRuntime(ctx)
if err != nil {
return err
}
database, err := controller.store.CapacityDatabaseHealth(ctx)
if err != nil {
return err
}
now := controller.now()
poolResources := make([]PoolResources, 0, len(controller.pools))
infrastructureStates := make(map[string]PoolInfrastructureState, len(controller.pools))
for _, pool := range controller.pools {
if pool.MaxReplicas == 0 {
continue
}
state, stateErr := controller.orchestrator.PoolState(ctx, pool.ID, pool.AdapterRef)
if stateErr != nil {
return stateErr
}
infrastructureStates[pool.ID] = state
poolResources = append(poolResources, PoolResources{
PoolID: pool.ID, CurrentReplicas: state.CurrentReplicas,
MinReplicas: pool.MinReplicas, MaxReplicas: pool.MaxReplicas,
Demand: demandByPool[pool.ID],
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
UsedMemoryBytes: state.UsedMemoryBytes,
WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes,
AllocatableMilliCPU: state.AllocatableMilliCPU,
UsedMilliCPU: state.UsedMilliCPU,
WorkerRequestMilliCPU: state.WorkerRequestMilliCPU,
MemoryPressure: state.MemoryPressure,
Nodes: nodeResources(state.Nodes),
})
}
currentTotal := 0
for _, pool := range poolResources {
currentTotal += pool.CurrentReplicas
}
target := controller.cfg.WorkerTargetOutstandingPerReplica
if target < 1 {
target = 2 * controller.cfg.AsyncWorkerInstanceHardLimit
}
rawDesired := (max(queue.Queued+queue.Running, 0) + target - 1) / target
if rawDesired > currentTotal {
if controller.highSince.IsZero() {
controller.highSince = now
}
controller.lowSince = time.Time{}
} else if rawDesired < currentTotal && queue.Queued == 0 {
if controller.lowSince.IsZero() {
controller.lowSince = now
}
controller.highSince = time.Time{}
} else {
controller.highSince = time.Time{}
controller.lowSince = time.Time{}
}
revisionHealthy := controller.revisionsMatch(instances, infrastructureStates)
scaleUpEligible := !controller.highSince.IsZero() &&
now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second &&
revisionHealthy
scaleDownEligible := !controller.lowSince.IsZero() &&
now.Sub(controller.lowSince) >= time.Duration(controller.cfg.WorkerScaleDownStabilizationSeconds)*time.Second &&
revisionHealthy
plan := CalculatePlan(PlanInput{
Queued: queue.Queued, Running: queue.Running,
InstanceSlots: controller.cfg.AsyncWorkerInstanceHardLimit,
TargetOutstandingPerReplica: target,
MemoryTargetPercent: controller.cfg.NodeMemoryTargetPercent,
MemoryHardPercent: controller.cfg.NodeMemoryHardPercent,
CPUTargetPercent: controller.cfg.NodeCPUTargetPercent,
DatabaseConnections: database.Connections,
DatabaseConnectionBudget: controller.cfg.PostgresConnectionBudget,
NonWorkerConnectionBudget: controller.cfg.PostgresNonWorkerConnectionBudget,
WorkerDatabasePoolMax: controller.cfg.WorkerDatabaseMaxConns,
SynchronousDatabasePeers: database.SynchronousPeers,
ScaleUpEligible: scaleUpEligible,
ScaleDownEligible: scaleDownEligible,
Pools: poolResources,
})
if !revisionHealthy && plan.FrozenReason == "" {
plan.FrozenReason = "release_revision_mismatch"
}
for _, poolPlan := range plan.Pools {
reason := "capacity_plan"
if plan.FrozenReason != "" {
reason = "frozen:" + plan.FrozenReason
}
if err := controller.store.PublishDesiredCapacity(ctx, executionpool.DesiredCapacity{
PoolID: poolPlan.PoolID, Desired: poolPlan.DesiredReplicas,
Reason: reason, ValidUntil: now.Add(2 * controllerReconcileInterval),
}); err != nil {
return err
}
}
if controller.cfg.WorkerAutoscalingEnabled {
for _, poolPlan := range plan.Pools {
switch {
case poolPlan.DesiredReplicas > poolPlan.CurrentReplicas:
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.DesiredReplicas); err != nil {
return err
}
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.DesiredReplicas, "scale_up")
case poolPlan.DesiredReplicas < poolPlan.CurrentReplicas:
if err := controller.reconcileScaleDown(ctx, now, poolPlan, instances); err != nil {
return err
}
}
}
}
controller.statusMu.Lock()
controller.status.LastRunAt = now
controller.status.LastError = ""
controller.status.Queue = queue
controller.status.Plan = plan
controller.statusMu.Unlock()
return nil
}
func nodeResources(states []KubernetesNodeState) []NodeResources {
nodes := make([]NodeResources, 0, len(states))
for _, state := range states {
nodes = append(nodes, NodeResources{
NodeName: state.NodeName,
CurrentReplicas: state.CurrentReplicas,
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
UsedMemoryBytes: state.UsedMemoryBytes,
WorkerUsedMemoryBytes: state.WorkerUsedMemoryBytes,
AllocatableMilliCPU: state.AllocatableMilliCPU,
UsedMilliCPU: state.UsedMilliCPU,
WorkerUsedMilliCPU: state.WorkerUsedMilliCPU,
MemoryPressure: state.MemoryPressure,
})
}
return nodes
}
func (controller *Controller) reconcileScaleDown(
ctx context.Context,
now time.Time,
poolPlan PoolPlan,
instances []store.WorkerInstanceRuntime,
) error {
for _, instance := range instances {
if workerPoolID(instance) != poolPlan.PoolID || instance.Status != "draining" {
continue
}
if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
if strings.TrimSpace(instance.OrchestratorInstanceRef) == "" {
return errors.New("worker orchestrator instance reference is required for scale down")
}
if err := controller.orchestrator.SetInstanceTerminationPriority(ctx, instance.OrchestratorInstanceRef, controllerDeletionCost); err != nil {
return err
}
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.CurrentReplicas-1); err != nil {
return err
}
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas-1, "drained_scale_down")
return nil
}
if instance.DrainingAt != nil &&
now.Sub(*instance.DrainingAt) >= time.Duration(controller.cfg.WorkerDrainTimeoutSeconds)*time.Second {
if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
return err
}
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_timeout")
}
return nil
}
var candidate *store.WorkerInstanceRuntime
for index := range instances {
instance := &instances[index]
if workerPoolID(*instance) != poolPlan.PoolID || instance.Status != "active" {
continue
}
if candidate == nil ||
instance.RunningTasks+instance.ActiveLeases < candidate.RunningTasks+candidate.ActiveLeases {
candidate = instance
}
}
if candidate == nil {
return nil
}
if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
return err
}
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_started")
return nil
}
func workerPoolID(instance store.WorkerInstanceRuntime) string {
if strings.TrimSpace(instance.PoolID) != "" {
return instance.PoolID
}
return instance.Site
}
func (controller *Controller) adapterRef(poolID string) string {
for _, pool := range controller.pools {
if pool.ID == poolID {
return pool.AdapterRef
}
}
return poolID
}
func (controller *Controller) revisionsMatch(
instances []store.WorkerInstanceRuntime,
states map[string]PoolInfrastructureState,
) bool {
if controller.expectedRevision == "" {
return true
}
for _, state := range states {
if state.Revision != "" && state.Revision != controller.expectedRevision {
return false
}
}
for _, instance := range instances {
if instance.Revision != "" && instance.Revision != controller.expectedRevision {
return false
}
}
return true
}
func (controller *Controller) Status() Status {
controller.statusMu.RLock()
defer controller.statusMu.RUnlock()
return controller.status
}
func (controller *Controller) Handler() http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"ok":true,"service":"easyai-capacity-controller"}`))
})
mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, _ *http.Request) {
status := controller.Status()
// Followers are intentionally idle while the database advisory lock is
// held by the elected leader. They are still ready to take over and must
// not make a two-replica Deployment permanently fail its rollout.
if status.LastError != "" {
http.Error(w, `{"ok":false}`, http.StatusServiceUnavailable)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"ok":true}`))
})
mux.HandleFunc("GET /status", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(controller.Status())
})
return mux
}
func (controller *Controller) setLeader(leader bool) {
controller.statusMu.Lock()
controller.status.Leader = leader
controller.statusMu.Unlock()
}
func (controller *Controller) setError(err error) {
controller.statusMu.Lock()
controller.status.LastError = err.Error()
controller.statusMu.Unlock()
}
func (controller *Controller) clearError() {
controller.statusMu.Lock()
controller.status.LastError = ""
controller.statusMu.Unlock()
}
func (controller *Controller) logError(message string, err error) {
controller.setError(err)
if controller.logger != nil {
controller.logger.Error(message, "error", err)
}
}
func (controller *Controller) observeScale(site string, from int, to int, reason string) {
controller.statusMu.Lock()
controller.status.ScaleActions++
controller.statusMu.Unlock()
if controller.logger != nil {
controller.logger.Info("worker capacity action",
"site", site,
"fromReplicas", from,
"toReplicas", to,
"reason", reason,
)
}
}
func waitContext(ctx context.Context, duration time.Duration) bool {
timer := time.NewTimer(duration)
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}
@@ -0,0 +1,482 @@
package capacitycontroller
import (
"bytes"
"context"
"crypto/tls"
"crypto/x509"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strconv"
"strings"
"time"
)
type KubernetesConfig struct {
Namespace string
APIServer string
TokenFile string
CAFile string
HTTPClient *http.Client
}
type KubernetesClient struct {
namespace string
baseURL string
tokenFile string
client *http.Client
}
type PoolInfrastructureState struct {
PoolID string
NodeName string
CurrentReplicas int
Revision string
Image string
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerRequestMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerRequestMilliCPU int64
MemoryPressure bool
Nodes []PoolNodeState
}
type PoolNodeState struct {
NodeName string
CurrentReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerUsedMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerUsedMilliCPU int64
MemoryPressure bool
}
type KubernetesSiteState = PoolInfrastructureState
type KubernetesNodeState = PoolNodeState
var _ OrchestratorAdapter = (*KubernetesClient)(nil)
func (client *KubernetesClient) PoolState(ctx context.Context, poolID, adapterRef string) (PoolInfrastructureState, error) {
state, err := client.SiteState(ctx, adapterRef)
state.PoolID = poolID
return state, err
}
func (client *KubernetesClient) ScalePool(ctx context.Context, _, adapterRef string, replicas int) error {
return client.ScaleWorkerDeployment(ctx, adapterRef, replicas)
}
func (client *KubernetesClient) SetInstanceTerminationPriority(ctx context.Context, instanceRef string, priority int) error {
return client.SetPodDeletionCost(ctx, instanceRef, priority)
}
func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) {
if strings.TrimSpace(config.Namespace) == "" {
return nil, errors.New("capacity controller Kubernetes namespace is required")
}
if config.APIServer == "" {
config.APIServer = "https://kubernetes.default.svc"
}
if config.TokenFile == "" {
config.TokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
}
parsed, err := url.Parse(strings.TrimRight(config.APIServer, "/"))
if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
(parsed.Scheme != "https" && config.HTTPClient == nil) {
return nil, errors.New("capacity controller Kubernetes API server URL is invalid")
}
client := config.HTTPClient
if client == nil {
if config.CAFile == "" {
config.CAFile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
}
certificate, readErr := os.ReadFile(config.CAFile)
if readErr != nil {
return nil, fmt.Errorf("read Kubernetes service account CA: %w", readErr)
}
roots := x509.NewCertPool()
if !roots.AppendCertsFromPEM(certificate) {
return nil, errors.New("Kubernetes service account CA is invalid")
}
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DisableCompression = true
transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
client = &http.Client{
Timeout: 10 * time.Second,
Transport: transport,
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
return &KubernetesClient{
namespace: config.Namespace,
baseURL: strings.TrimRight(config.APIServer, "/"),
tokenFile: config.TokenFile,
client: client,
}, nil
}
func (client *KubernetesClient) SiteState(ctx context.Context, site string) (KubernetesSiteState, error) {
site = strings.TrimSpace(site)
if site == "" {
return KubernetesSiteState{}, errors.New("site is required")
}
var deployment struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Spec struct {
Replicas int `json:"replicas"`
Template struct {
Spec struct {
Containers []struct {
Name string `json:"name"`
Image string `json:"image"`
Env []struct {
Name string `json:"name"`
Value string `json:"value"`
} `json:"env"`
Resources struct {
Requests map[string]string `json:"requests"`
} `json:"resources"`
} `json:"containers"`
} `json:"spec"`
} `json:"template"`
} `json:"spec"`
}
deploymentPath := fmt.Sprintf(
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s",
url.PathEscape(client.namespace),
url.PathEscape(site),
)
if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil {
return KubernetesSiteState{}, err
}
state := KubernetesSiteState{PoolID: site, CurrentReplicas: deployment.Spec.Replicas}
for _, container := range deployment.Spec.Template.Spec.Containers {
if container.Name != "worker" {
continue
}
state.Image = container.Image
state.WorkerRequestMemoryBytes, _ = parseBinaryQuantity(container.Resources.Requests["memory"])
state.WorkerRequestMilliCPU, _ = parseCPUQuantity(container.Resources.Requests["cpu"])
for _, environment := range container.Env {
if environment.Name == "AI_GATEWAY_REVISION" {
state.Revision = environment.Value
}
}
}
var nodes struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Status struct {
Allocatable map[string]string `json:"allocatable"`
Conditions []struct {
Type string `json:"type"`
Status string `json:"status"`
} `json:"conditions"`
} `json:"status"`
} `json:"items"`
}
nodePath := "/api/v1/nodes?labelSelector=" + url.QueryEscape(
"easyai.io/site="+site+",easyai.io/worker=true",
)
if err := client.getJSON(ctx, nodePath, &nodes); err != nil {
return KubernetesSiteState{}, err
}
readyNodes := nodes.Items[:0]
for _, node := range nodes.Items {
ready := false
for _, condition := range node.Status.Conditions {
if condition.Type == "Ready" && condition.Status == "True" {
ready = true
break
}
}
if ready {
readyNodes = append(readyNodes, node)
}
}
nodes.Items = readyNodes
if len(nodes.Items) == 0 {
// A zero-replica Deployment with no eligible nodes represents an
// explicitly disabled site. It must remain visible to the planner, but it
// has no node or metrics budget to collect. A non-zero Deployment still
// fails closed so replicas cannot disappear from capacity accounting.
if state.CurrentReplicas == 0 {
return state, nil
}
return KubernetesSiteState{}, fmt.Errorf("site %s has no ready matching nodes", site)
}
var pods struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
DeletionTimestamp *time.Time `json:"deletionTimestamp"`
} `json:"metadata"`
Spec struct {
NodeName string `json:"nodeName"`
} `json:"spec"`
} `json:"items"`
}
podPath := fmt.Sprintf(
"/api/v1/namespaces/%s/pods?labelSelector=%s",
url.PathEscape(client.namespace),
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
)
if err := client.getJSON(ctx, podPath, &pods); err != nil {
return KubernetesSiteState{}, err
}
replicasByNode := make(map[string]int, len(nodes.Items))
nodeByPod := make(map[string]string, len(pods.Items))
for _, pod := range pods.Items {
if pod.Metadata.DeletionTimestamp == nil && pod.Spec.NodeName != "" {
replicasByNode[pod.Spec.NodeName]++
nodeByPod[pod.Metadata.Name] = pod.Spec.NodeName
}
}
var podMetrics struct {
Items []struct {
Metadata struct {
Name string `json:"name"`
} `json:"metadata"`
Containers []struct {
Usage map[string]string `json:"usage"`
} `json:"containers"`
} `json:"items"`
}
podMetricsPath := fmt.Sprintf(
"/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods?labelSelector=%s",
url.PathEscape(client.namespace),
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
)
if err := client.getJSON(ctx, podMetricsPath, &podMetrics); err != nil {
return KubernetesSiteState{}, err
}
workerMemoryByNode := make(map[string]int64, len(nodes.Items))
workerCPUByNode := make(map[string]int64, len(nodes.Items))
measuredPods := make(map[string]bool, len(podMetrics.Items))
for _, pod := range podMetrics.Items {
nodeName := nodeByPod[pod.Metadata.Name]
if nodeName == "" {
continue
}
measuredPods[pod.Metadata.Name] = true
for _, container := range pod.Containers {
memory, memoryErr := parseBinaryQuantity(container.Usage["memory"])
cpu, cpuErr := parseCPUQuantity(container.Usage["cpu"])
if memoryErr != nil || cpuErr != nil {
return KubernetesSiteState{}, fmt.Errorf(
"site %s pod %s returned incomplete resource metrics",
site,
pod.Metadata.Name,
)
}
workerMemoryByNode[nodeName] += memory
workerCPUByNode[nodeName] += cpu
}
}
for podName := range nodeByPod {
if !measuredPods[podName] {
return KubernetesSiteState{}, fmt.Errorf(
"site %s pod %s has no resource metrics yet",
site,
podName,
)
}
}
for _, node := range nodes.Items {
nodeState := KubernetesNodeState{
NodeName: node.Metadata.Name,
CurrentReplicas: replicasByNode[node.Metadata.Name],
WorkerUsedMemoryBytes: workerMemoryByNode[node.Metadata.Name],
WorkerUsedMilliCPU: workerCPUByNode[node.Metadata.Name],
}
nodeState.AllocatableMemoryBytes, _ = parseBinaryQuantity(node.Status.Allocatable["memory"])
nodeState.AllocatableMilliCPU, _ = parseCPUQuantity(node.Status.Allocatable["cpu"])
for _, condition := range node.Status.Conditions {
if condition.Type == "MemoryPressure" && condition.Status == "True" {
nodeState.MemoryPressure = true
state.MemoryPressure = true
}
}
var metrics struct {
Usage map[string]string `json:"usage"`
}
metricsPath := "/apis/metrics.k8s.io/v1beta1/nodes/" + url.PathEscape(nodeState.NodeName)
if err := client.getJSON(ctx, metricsPath, &metrics); err != nil {
return KubernetesSiteState{}, err
}
nodeState.UsedMemoryBytes, _ = parseBinaryQuantity(metrics.Usage["memory"])
nodeState.UsedMilliCPU, _ = parseCPUQuantity(metrics.Usage["cpu"])
if nodeState.AllocatableMemoryBytes <= 0 || nodeState.AllocatableMilliCPU <= 0 {
return KubernetesSiteState{}, fmt.Errorf(
"site %s node %s returned incomplete resource quantities",
site,
nodeState.NodeName,
)
}
state.Nodes = append(state.Nodes, nodeState)
state.AllocatableMemoryBytes += nodeState.AllocatableMemoryBytes
state.AllocatableMilliCPU += nodeState.AllocatableMilliCPU
state.UsedMemoryBytes += nodeState.UsedMemoryBytes
state.UsedMilliCPU += nodeState.UsedMilliCPU
}
if len(state.Nodes) == 1 {
state.NodeName = state.Nodes[0].NodeName
}
if state.WorkerRequestMemoryBytes <= 0 || state.WorkerRequestMilliCPU <= 0 ||
state.AllocatableMemoryBytes <= 0 || state.AllocatableMilliCPU <= 0 {
return KubernetesSiteState{}, fmt.Errorf("site %s returned incomplete resource quantities", site)
}
return state, nil
}
func (client *KubernetesClient) ScaleWorkerDeployment(ctx context.Context, site string, replicas int) error {
if replicas < 0 || replicas > 64 {
return errors.New("worker replica target must be between 0 and 64")
}
path := fmt.Sprintf(
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s/scale",
url.PathEscape(client.namespace),
url.PathEscape(strings.TrimSpace(site)),
)
return client.patchJSON(ctx, path, map[string]any{"spec": map[string]any{"replicas": replicas}})
}
func (client *KubernetesClient) SetPodDeletionCost(ctx context.Context, podName string, cost int) error {
podName = strings.TrimSpace(podName)
if podName == "" {
return errors.New("worker pod name is required")
}
path := fmt.Sprintf(
"/api/v1/namespaces/%s/pods/%s",
url.PathEscape(client.namespace),
url.PathEscape(podName),
)
return client.patchJSON(ctx, path, map[string]any{
"metadata": map[string]any{
"annotations": map[string]any{
"controller.kubernetes.io/pod-deletion-cost": strconv.Itoa(cost),
},
},
})
}
func (client *KubernetesClient) getJSON(ctx context.Context, path string, target any) error {
request, err := client.request(ctx, http.MethodGet, path, nil)
if err != nil {
return err
}
response, err := client.client.Do(request)
if err != nil {
return fmt.Errorf("request Kubernetes API: %w", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
}
decoder := json.NewDecoder(io.LimitReader(response.Body, 4*1024*1024))
if err := decoder.Decode(target); err != nil {
return fmt.Errorf("decode Kubernetes API %s: %w", path, err)
}
return nil
}
func (client *KubernetesClient) patchJSON(ctx context.Context, path string, payload any) error {
encoded, err := json.Marshal(payload)
if err != nil {
return err
}
request, err := client.request(ctx, http.MethodPatch, path, bytes.NewReader(encoded))
if err != nil {
return err
}
request.Header.Set("Content-Type", "application/merge-patch+json")
response, err := client.client.Do(request)
if err != nil {
return fmt.Errorf("patch Kubernetes API: %w", err)
}
defer response.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
if response.StatusCode < 200 || response.StatusCode >= 300 {
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
}
return nil
}
func (client *KubernetesClient) request(ctx context.Context, method string, path string, body io.Reader) (*http.Request, error) {
token, err := os.ReadFile(client.tokenFile)
if err != nil || strings.TrimSpace(string(token)) == "" {
clear(token)
return nil, errors.New("Kubernetes service account token is unavailable")
}
request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, body)
if err != nil {
clear(token)
return nil, err
}
request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token)))
request.Header.Set("Accept", "application/json")
clear(token)
return request, nil
}
func parseBinaryQuantity(value string) (int64, error) {
value = strings.TrimSpace(value)
if value == "" {
return 0, errors.New("quantity is empty")
}
units := []struct {
suffix string
multiplier int64
}{
{"Ei", 1 << 60},
{"Pi", 1 << 50},
{"Ti", 1 << 40},
{"Gi", 1 << 30},
{"Mi", 1 << 20},
{"Ki", 1 << 10},
{"G", 1_000_000_000},
{"M", 1_000_000},
{"K", 1_000},
}
for _, unit := range units {
if strings.HasSuffix(value, unit.suffix) {
number, err := strconv.ParseFloat(strings.TrimSuffix(value, unit.suffix), 64)
return int64(number * float64(unit.multiplier)), err
}
}
return strconv.ParseInt(value, 10, 64)
}
func parseCPUQuantity(value string) (int64, error) {
value = strings.TrimSpace(value)
if strings.HasSuffix(value, "n") {
nanocores, err := strconv.ParseInt(strings.TrimSuffix(value, "n"), 10, 64)
return nanocores / 1_000_000, err
}
if strings.HasSuffix(value, "u") {
microcores, err := strconv.ParseInt(strings.TrimSuffix(value, "u"), 10, 64)
return microcores / 1_000, err
}
if strings.HasSuffix(value, "m") {
return strconv.ParseInt(strings.TrimSuffix(value, "m"), 10, 64)
}
cores, err := strconv.ParseFloat(value, 64)
return int64(cores * 1000), err
}
@@ -0,0 +1,172 @@
package capacitycontroller
import (
"context"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
)
func TestKubernetesClientReadsSiteAndMutatesOnlyWorkerScale(t *testing.T) {
var patches []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Header.Get("Authorization") != "Bearer test-token" {
http.Error(w, "unauthorized", http.StatusUnauthorized)
return
}
switch {
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-hongkong"):
_, _ = w.Write([]byte(`{
"spec":{"replicas":1,"template":{"spec":{"containers":[{
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
"resources":{"requests":{"memory":"512Mi","cpu":"250m"}}
}]}}}
}`))
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
if request.URL.Query().Get("labelSelector") != "easyai.io/site=hongkong,easyai.io/worker=true" {
http.Error(w, "unexpected Worker node selector", http.StatusBadRequest)
return
}
_, _ = w.Write([]byte(`{"items":[
{"metadata":{"name":"easyai-hongkong"},"status":{
"allocatable":{"memory":"8Gi","cpu":"4"},
"conditions":[{"type":"Ready","status":"False"}]
}},
{"metadata":{"name":"easyai-hongkong-worker-2"},"status":{
"allocatable":{"memory":"8Gi","cpu":"4"},
"conditions":[{"type":"Ready","status":"True"},{"type":"MemoryPressure","status":"False"}]
}}]}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/api/v1/namespaces/easyai/pods"):
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"worker-hongkong-1"},"spec":{"nodeName":"easyai-hongkong-worker-2"}}]}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong-worker-2"):
_, _ = w.Write([]byte(`{"usage":{"memory":"3Gi","cpu":"500m"}}`))
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong"):
t.Error("NotReady node metrics must not be queried")
http.Error(w, "NotReady node", http.StatusInternalServerError)
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
strings.HasSuffix(request.URL.Path, "/pods"):
_, _ = w.Write([]byte(`{"items":[{
"metadata":{"name":"worker-hongkong-1"},
"containers":[{"usage":{"memory":"384Mi","cpu":"125m"}}]
}]}`))
case request.Method == http.MethodPatch:
patches = append(patches, request.URL.Path)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{}`))
default:
http.NotFound(w, request)
}
}))
defer server.Close()
tokenFile := filepath.Join(t.TempDir(), "token")
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
t.Fatal(err)
}
client, err := NewKubernetesClient(KubernetesConfig{
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
state, err := client.SiteState(context.Background(), "hongkong")
if err != nil {
t.Fatal(err)
}
if state.CurrentReplicas != 1 || state.NodeName != "easyai-hongkong-worker-2" ||
state.WorkerRequestMemoryBytes != 512<<20 || state.WorkerRequestMilliCPU != 250 ||
state.AllocatableMemoryBytes != 8<<30 || state.UsedMemoryBytes != 3<<30 ||
state.Nodes[0].WorkerUsedMemoryBytes != 384<<20 || state.Nodes[0].WorkerUsedMilliCPU != 125 {
t.Fatalf("site state=%+v", state)
}
if err := client.SetPodDeletionCost(context.Background(), "worker-pod", -1000); err != nil {
t.Fatal(err)
}
if err := client.ScaleWorkerDeployment(context.Background(), "hongkong", 2); err != nil {
t.Fatal(err)
}
if len(patches) != 2 ||
!strings.Contains(patches[0], "/pods/worker-pod") ||
!strings.Contains(patches[1], "/deployments/easyai-worker-hongkong/scale") {
t.Fatalf("patches=%v", patches)
}
}
func TestKubernetesClientAllowsDisabledSiteWithoutWorkerNodes(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
switch {
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
_, _ = w.Write([]byte(`{
"spec":{"replicas":0,"template":{"spec":{"containers":[{
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
"resources":{"requests":{"memory":"1536Mi","cpu":"500m"}}
}]}}}
}`))
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
_, _ = w.Write([]byte(`{"items":[]}`))
default:
http.NotFound(w, request)
}
}))
defer server.Close()
tokenFile := filepath.Join(t.TempDir(), "token")
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
t.Fatal(err)
}
client, err := NewKubernetesClient(KubernetesConfig{
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
state, err := client.SiteState(context.Background(), "ningbo")
if err != nil {
t.Fatal(err)
}
if state.PoolID != "ningbo" || state.CurrentReplicas != 0 || len(state.Nodes) != 0 {
t.Fatalf("site state=%+v", state)
}
}
func TestKubernetesClientRejectsActiveSiteWithoutWorkerNodes(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
switch {
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
_, _ = w.Write([]byte(`{"spec":{"replicas":1,"template":{"spec":{"containers":[]}}}}`))
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
_, _ = w.Write([]byte(`{"items":[]}`))
default:
http.NotFound(w, request)
}
}))
defer server.Close()
tokenFile := filepath.Join(t.TempDir(), "token")
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
t.Fatal(err)
}
client, err := NewKubernetesClient(KubernetesConfig{
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
})
if err != nil {
t.Fatal(err)
}
if _, err := client.SiteState(context.Background(), "ningbo"); err == nil ||
!strings.Contains(err.Error(), "site ningbo has no ready matching nodes") {
t.Fatalf("expected missing Worker node error, got %v", err)
}
}
func TestQuantityParsing(t *testing.T) {
if value, err := parseBinaryQuantity("1536Mi"); err != nil || value != 1536<<20 {
t.Fatalf("memory=%d err=%v", value, err)
}
if value, err := parseCPUQuantity("250000000n"); err != nil || value != 250 {
t.Fatalf("cpu=%d err=%v", value, err)
}
}
@@ -0,0 +1,289 @@
package capacitycontroller
import (
"math"
"sort"
)
type PoolResources struct {
PoolID string
Demand int
CurrentReplicas int
MinReplicas int
MaxReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerRequestMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerRequestMilliCPU int64
MemoryPressure bool
Nodes []NodeResources
}
type NodeResources struct {
NodeName string
CurrentReplicas int
AllocatableMemoryBytes int64
UsedMemoryBytes int64
WorkerUsedMemoryBytes int64
AllocatableMilliCPU int64
UsedMilliCPU int64
WorkerUsedMilliCPU int64
MemoryPressure bool
}
type PlanInput struct {
Queued int
Running int
InstanceSlots int
TargetOutstandingPerReplica int
MemoryTargetPercent int
MemoryHardPercent int
CPUTargetPercent int
DatabaseConnections int
DatabaseConnectionBudget int
NonWorkerConnectionBudget int
WorkerDatabasePoolMax int
SynchronousDatabasePeers int
ScaleUpEligible bool
ScaleDownEligible bool
Pools []PoolResources
}
type PoolPlan struct {
PoolID string `json:"poolId"`
Demand int `json:"demand"`
CurrentReplicas int `json:"currentReplicas"`
DesiredReplicas int `json:"desiredReplicas"`
ResourceMax int `json:"resourceMax"`
MemoryPercent float64 `json:"memoryPercent"`
CPUPercent float64 `json:"cpuPercent"`
}
type Plan struct {
RawDesired int `json:"rawDesired"`
DesiredTotal int `json:"desiredTotal"`
CurrentTotal int `json:"currentTotal"`
FrozenReason string `json:"frozenReason,omitempty"`
Pools []PoolPlan `json:"pools"`
}
func CalculatePlan(input PlanInput) Plan {
if input.InstanceSlots < 1 {
input.InstanceSlots = 1
}
target := input.TargetOutstandingPerReplica
if target < 1 {
target = 2 * input.InstanceSlots
}
if input.MemoryTargetPercent < 1 {
input.MemoryTargetPercent = 75
}
if input.MemoryHardPercent < 1 {
input.MemoryHardPercent = 85
}
if input.CPUTargetPercent < 1 {
input.CPUTargetPercent = 70
}
rawDesired := int(math.Ceil(float64(max(input.Queued+input.Running, 0)) / float64(target)))
plan := Plan{RawDesired: rawDesired}
minTotal := 0
resourceMaxTotal := 0
for _, pool := range input.Pools {
for _, node := range pool.Nodes {
if node.MemoryPressure {
pool.MemoryPressure = true
}
}
current := max(pool.CurrentReplicas, 0)
minReplicas := max(pool.MinReplicas, 0)
configMax := max(pool.MaxReplicas, minReplicas)
resourceMax, memoryPercent, cpuPercent := poolResourceMaximum(
pool,
input.MemoryTargetPercent,
input.CPUTargetPercent,
)
resourceMax = min(resourceMax, configMax)
if resourceMax < minReplicas && plan.FrozenReason == "" {
plan.FrozenReason = "pool_resource_budget"
}
resourceMax = max(resourceMax, minReplicas)
plan.Pools = append(plan.Pools, PoolPlan{
PoolID: pool.PoolID, CurrentReplicas: current, DesiredReplicas: minReplicas,
Demand: pool.Demand,
ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
})
plan.CurrentTotal += current
minTotal += minReplicas
resourceMaxTotal += resourceMax
if pool.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) {
plan.FrozenReason = "node_memory_pressure"
}
}
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
workerReplicaBudget := max(
(input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax,
0,
)
capPoolResourceMaxima(plan.Pools, workerReplicaBudget)
resourceMaxTotal = 0
for _, pool := range plan.Pools {
resourceMaxTotal += pool.ResourceMax
}
if workerReplicaBudget < minTotal && plan.FrozenReason == "" {
plan.FrozenReason = "database_connection_budget"
}
}
desired := max(rawDesired, minTotal)
desired = min(desired, resourceMaxTotal)
if input.SynchronousDatabasePeers < 1 && plan.FrozenReason == "" {
plan.FrozenReason = "database_not_synchronous"
}
if desired > plan.CurrentTotal {
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
additional := max(
(input.DatabaseConnectionBudget-input.DatabaseConnections)/input.WorkerDatabasePoolMax,
0,
)
desired = min(desired, plan.CurrentTotal+additional)
}
if !input.ScaleUpEligible || plan.FrozenReason != "" ||
(input.DatabaseConnectionBudget > 0 && input.DatabaseConnections >= input.DatabaseConnectionBudget) {
desired = plan.CurrentTotal
if plan.FrozenReason == "" {
if !input.ScaleUpEligible {
plan.FrozenReason = "scale_up_window"
} else {
plan.FrozenReason = "database_connection_budget"
}
}
} else {
desired = min(desired, max(plan.CurrentTotal*2, plan.CurrentTotal+2))
}
} else if desired < plan.CurrentTotal && !input.ScaleDownEligible {
desired = plan.CurrentTotal
if plan.FrozenReason == "" {
plan.FrozenReason = "scale_down_stabilization"
}
}
desired = max(desired, minTotal)
plan.DesiredTotal = desired
distributeDesiredReplicas(plan.Pools, desired)
return plan
}
func capPoolResourceMaxima(pools []PoolPlan, total int) {
minimum := 0
original := make(map[string]int, len(pools))
for index := range pools {
original[pools[index].PoolID] = pools[index].ResourceMax
pools[index].ResourceMax = pools[index].DesiredReplicas
minimum += pools[index].ResourceMax
}
target := max(total, minimum)
assigned := minimum
for assigned < target {
sort.SliceStable(pools, func(left, right int) bool {
leftRoom := original[pools[left].PoolID] - pools[left].ResourceMax
rightRoom := original[pools[right].PoolID] - pools[right].ResourceMax
if leftRoom != rightRoom {
return leftRoom > rightRoom
}
return pools[left].PoolID < pools[right].PoolID
})
if original[pools[0].PoolID] <= pools[0].ResourceMax {
break
}
pools[0].ResourceMax++
assigned++
}
}
func poolResourceMaximum(pool PoolResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
if len(pool.Nodes) > 0 {
if pool.WorkerRequestMemoryBytes <= 0 || pool.WorkerRequestMilliCPU <= 0 {
return 0, 0, 0
}
memoryMax := 0
cpuMax := 0
memoryPercent := float64(0)
cpuPercent := float64(0)
for _, node := range pool.Nodes {
nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes)
nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU)
memoryPercent = max(memoryPercent, nodeMemoryPercent)
cpuPercent = max(cpuPercent, nodeCPUPercent)
nonWorkerMemory := max(
node.UsedMemoryBytes-node.WorkerUsedMemoryBytes,
0,
)
memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory
memoryMax += int(max(memoryBudget, 0) / pool.WorkerRequestMemoryBytes)
nonWorkerCPU := max(
node.UsedMilliCPU-node.WorkerUsedMilliCPU,
0,
)
cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
cpuMax += int(max(cpuBudget, 0) / pool.WorkerRequestMilliCPU)
}
return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}
memoryPercent := usagePercent(pool.UsedMemoryBytes, pool.AllocatableMemoryBytes)
cpuPercent := usagePercent(pool.UsedMilliCPU, pool.AllocatableMilliCPU)
memoryMax := pool.MaxReplicas
if pool.AllocatableMemoryBytes > 0 && pool.WorkerRequestMemoryBytes > 0 {
nonWorker := max(pool.UsedMemoryBytes-int64(pool.CurrentReplicas)*pool.WorkerRequestMemoryBytes, 0)
budget := pool.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
memoryMax = int(max(budget, 0) / pool.WorkerRequestMemoryBytes)
}
cpuMax := pool.MaxReplicas
if pool.AllocatableMilliCPU > 0 && pool.WorkerRequestMilliCPU > 0 {
nonWorker := max(pool.UsedMilliCPU-int64(pool.CurrentReplicas)*pool.WorkerRequestMilliCPU, 0)
budget := pool.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
cpuMax = int(max(budget, 0) / pool.WorkerRequestMilliCPU)
}
return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}
func distributeDesiredReplicas(pools []PoolPlan, desired int) {
if len(pools) == 0 {
return
}
assigned := 0
for index := range pools {
assigned += pools[index].DesiredReplicas
}
for assigned < desired {
sort.SliceStable(pools, func(left, right int) bool {
leftDemandPerReplica := float64(max(pools[left].Demand, 0)) / float64(max(pools[left].DesiredReplicas, 1))
rightDemandPerReplica := float64(max(pools[right].Demand, 0)) / float64(max(pools[right].DesiredReplicas, 1))
if leftDemandPerReplica != rightDemandPerReplica {
return leftDemandPerReplica > rightDemandPerReplica
}
leftRoom := pools[left].ResourceMax - pools[left].DesiredReplicas
rightRoom := pools[right].ResourceMax - pools[right].DesiredReplicas
if leftRoom != rightRoom {
return leftRoom > rightRoom
}
if pools[left].DesiredReplicas != pools[right].DesiredReplicas {
return pools[left].DesiredReplicas < pools[right].DesiredReplicas
}
return pools[left].PoolID < pools[right].PoolID
})
if pools[0].DesiredReplicas >= pools[0].ResourceMax {
break
}
pools[0].DesiredReplicas++
assigned++
}
sort.Slice(pools, func(left, right int) bool { return pools[left].PoolID < pools[right].PoolID })
}
func usagePercent(used int64, allocatable int64) float64 {
if allocatable <= 0 {
return 0
}
return float64(used) * 100 / float64(allocatable)
}
@@ -0,0 +1,164 @@
package capacitycontroller
import (
"encoding/json"
"strings"
"testing"
)
func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) {
input := PlanInput{
Queued: 1000, Running: 48, InstanceSlots: 24,
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Pools: []PoolResources{
{
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30,
WorkerRequestMemoryBytes: 1 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500,
},
{
PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
WorkerRequestMemoryBytes: 1 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500,
},
},
}
plan := CalculatePlan(input)
if plan.RawDesired != 22 {
t.Fatalf("raw desired=%d, want 22", plan.RawDesired)
}
if plan.DesiredTotal != 4 {
t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan)
}
if plan.Pools[0].PoolID != "hongkong" || plan.Pools[0].DesiredReplicas != 3 {
t.Fatalf("pool allocation=%+v, want hongkong=3 ningbo=1", plan.Pools)
}
}
func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 200, InstanceSlots: 24,
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Pools: []PoolResources{{
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30,
WorkerRequestMemoryBytes: 1 << 30,
}},
})
if plan.DesiredTotal != 1 || plan.FrozenReason != "node_memory_pressure" {
t.Fatalf("plan=%+v, want frozen at one replica", plan)
}
}
func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 200, InstanceSlots: 24,
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Pools: []PoolResources{{
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29,
WorkerRequestMemoryBytes: 2 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500,
}},
})
if plan.DesiredTotal != 1 || plan.FrozenReason != "pool_resource_budget" {
t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan)
}
}
func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) {
input := PlanInput{
InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150,
WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1,
ScaleDownEligible: false,
Pools: []PoolResources{
{PoolID: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
{PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
},
}
plan := CalculatePlan(input)
if plan.DesiredTotal != 4 || plan.FrozenReason != "scale_down_stabilization" {
t.Fatalf("plan=%+v, want current replicas during stabilization", plan)
}
input.ScaleDownEligible = true
plan = CalculatePlan(input)
if plan.DesiredTotal != 2 {
t.Fatalf("desired total=%d, want min total 2", plan.DesiredTotal)
}
}
func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 500, InstanceSlots: 24,
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Pools: []PoolResources{{
PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8,
WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500,
Nodes: []NodeResources{
{
NodeName: "hk-a", CurrentReplicas: 1,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 1000,
},
{
NodeName: "hk-b", CurrentReplicas: 1,
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 2 << 30,
AllocatableMilliCPU: 4000, UsedMilliCPU: 500,
},
},
}},
})
if plan.Pools[0].ResourceMax != 7 {
t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Pools[0].ResourceMax, plan)
}
if plan.DesiredTotal != 4 {
t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal)
}
}
func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) {
plan := CalculatePlan(PlanInput{
Queued: 500, InstanceSlots: 24,
DatabaseConnections: 20, DatabaseConnectionBudget: 150,
NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32,
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
Pools: []PoolResources{
{PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
{PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
},
})
if plan.DesiredTotal != 2 {
t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan)
}
resourceMax := 0
for _, pool := range plan.Pools {
resourceMax += pool.ResourceMax
}
if resourceMax != 2 {
t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan)
}
}
func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) {
payload, err := json.Marshal(Plan{
RawDesired: 3,
Pools: []PoolPlan{{PoolID: "hongkong", ResourceMax: 2}},
})
if err != nil {
t.Fatal(err)
}
text := string(payload)
for _, field := range []string{`"rawDesired":3`, `"pools"`, `"poolId":"hongkong"`, `"resourceMax":2`} {
if !strings.Contains(text, field) {
t.Fatalf("plan JSON %s does not contain %s", text, field)
}
}
}
@@ -0,0 +1,42 @@
package capacitycontroller
import (
"context"
"sync"
)
// StaticAdapter proves that capacity coordination does not require Kubernetes.
// It reports configured fixed replicas and intentionally ignores scale writes.
type StaticAdapter struct {
mu sync.RWMutex
replicas map[string]int
}
var _ OrchestratorAdapter = (*StaticAdapter)(nil)
func NewStaticAdapter(initial map[string]int) *StaticAdapter {
copyOfInitial := make(map[string]int, len(initial))
for poolID, replicas := range initial {
copyOfInitial[poolID] = replicas
}
return &StaticAdapter{replicas: copyOfInitial}
}
func (adapter *StaticAdapter) PoolState(_ context.Context, poolID, _ string) (PoolInfrastructureState, error) {
adapter.mu.RLock()
replicas := adapter.replicas[poolID]
adapter.mu.RUnlock()
return PoolInfrastructureState{
PoolID: poolID, CurrentReplicas: replicas,
AllocatableMemoryBytes: 1 << 60, WorkerRequestMemoryBytes: 1,
AllocatableMilliCPU: 1 << 50, WorkerRequestMilliCPU: 1,
}, nil
}
func (adapter *StaticAdapter) ScalePool(context.Context, string, string, int) error {
return nil
}
func (adapter *StaticAdapter) SetInstanceTerminationPriority(context.Context, string, int) error {
return nil
}
@@ -0,0 +1,24 @@
package capacitycontroller
import (
"context"
"testing"
)
func TestStaticAdapterDoesNotRequireKubernetes(t *testing.T) {
adapter := NewStaticAdapter(map[string]int{"pool-a": 2})
state, err := adapter.PoolState(context.Background(), "pool-a", "ignored")
if err != nil {
t.Fatal(err)
}
if state.PoolID != "pool-a" || state.CurrentReplicas != 2 {
t.Fatalf("state=%+v", state)
}
if err := adapter.ScalePool(context.Background(), "pool-a", "ignored", 4); err != nil {
t.Fatal(err)
}
state, _ = adapter.PoolState(context.Background(), "pool-a", "ignored")
if state.CurrentReplicas != 2 {
t.Fatalf("static adapter scaled to %d", state.CurrentReplicas)
}
}
@@ -0,0 +1,282 @@
package clients
import (
"context"
"encoding/json"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestVectorizerClientMultipartAndPrivateResumeState(t *testing.T) {
var receivedPath string
var receivedFields map[string]string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedPath = r.URL.Path
if username, password, ok := r.BasicAuth(); !ok || username != "id" || password != "secret" {
t.Fatalf("unexpected vectorizer auth")
}
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart: %v", err)
}
receivedFields = map[string]string{}
for key, values := range r.MultipartForm.Value {
receivedFields[key] = values[0]
}
w.Header().Set("Content-Type", "image/svg+xml")
w.Header().Set("X-Image-Token", "private-image-token")
w.Header().Set("X-Receipt", "private-receipt")
_, _ = io.WriteString(w, `<svg xmlns="http://www.w3.org/2000/svg"/>`)
}))
defer server.Close()
privateState := map[string]any{}
response, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP("93.184.216.34")}}, nil
}}).Run(context.Background(), Request{
Kind: "images.vectorize", Model: "easy-image-vectorizer-1",
Body: map[string]any{"source": map[string]any{"url": "https://example.com/input.png"}, "format": "svg", "maxColors": 16, "cleanupLevel": "strong"},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"accessKey": "id", "secretKey": "secret"}),
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error { privateState = payload; return nil },
})
if err != nil {
t.Fatalf("run vectorizer: %v", err)
}
if receivedPath != "/vectorize" || receivedFields["image.url"] == "" || receivedFields["output.file_format"] != "svg" || receivedFields["processing.max_colors"] != "16" {
t.Fatalf("unexpected multipart request: path=%s fields=%+v", receivedPath, receivedFields)
}
if privateState["imageToken"] != "private-image-token" || privateState["receipt"] != "private-receipt" {
t.Fatalf("private resume state missing: %+v", privateState)
}
serialized := strings.ToLower(toJSONForTest(response.Result))
if strings.Contains(serialized, "private-image-token") || strings.Contains(serialized, "private-receipt") {
t.Fatalf("private vectorizer state leaked into response: %s", serialized)
}
}
func TestVectorizerClientRejectsPrivateSourceBeforeProviderCall(t *testing.T) {
providerCalls := 0
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { providerCalls++ }))
defer server.Close()
_, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP("127.0.0.1")}}, nil
}}).Run(context.Background(), Request{
Body: map[string]any{"source": map[string]any{"url": "https://assets.example.test/input.png"}},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
})
if err == nil || !strings.Contains(err.Error(), "blocked network") {
t.Fatalf("expected blocked source error, got %v", err)
}
if providerCalls != 0 {
t.Fatalf("provider received %d calls for blocked source", providerCalls)
}
}
func TestVectorizerClientReusesImageTokenForExtraFormat(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/download" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatal(err)
}
if r.FormValue("image.token") != "token" || r.FormValue("receipt") != "receipt" || r.FormValue("output.file_format") != "pdf" {
t.Fatalf("unexpected reuse form: %+v", r.MultipartForm.Value)
}
w.Header().Set("Content-Type", "application/pdf")
_, _ = w.Write([]byte("%PDF-1.7\n"))
}))
defer server.Close()
response, err := (VectorizerClient{}).Run(context.Background(), Request{
Body: map[string]any{"_vectorizer_image_token": "token", "_vectorizer_receipt": "receipt", "format": "pdf"},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
})
if err != nil {
t.Fatal(err)
}
item := response.Result["data"].([]any)[0].(map[string]any)
if item["type"] != "file" || item["mime_type"] != "application/pdf" {
t.Fatalf("unexpected pdf result: %+v", item)
}
}
func TestTopazClientUploadsPartsPollsAndValidatesOutput(t *testing.T) {
var mu sync.Mutex
called := map[string]int{}
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mu.Lock()
called[r.Method+" "+r.URL.Path]++
mu.Unlock()
switch {
case r.Method == http.MethodGet && r.URL.Path == "/source.mp4":
_, _ = w.Write([]byte("small-video-source"))
case r.Method == http.MethodPost && r.URL.Path == "/video/":
if r.Header.Get("X-API-Key") != "topaz-key" {
t.Fatalf("missing Topaz API key")
}
_, _ = io.WriteString(w, `{"requestId":"job-1"}`)
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/accept":
_, _ = io.WriteString(w, `{"urls":["`+server.URL+`/upload/1","`+server.URL+`/upload/2"]}`)
case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/upload/"):
w.Header().Set("ETag", `"etag-`+strings.TrimPrefix(r.URL.Path, "/upload/")+`"`)
w.WriteHeader(http.StatusOK)
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/complete-upload/":
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `{}`)
case r.Method == http.MethodGet && r.URL.Path == "/video/job-1/status":
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/output.mp4"}}`)
default:
http.NotFound(w, r)
}
}))
defer server.Close()
phases := []string{}
client := TopazClient{Probe: func(_ context.Context, target string) (TopazVideoMetadata, error) {
if strings.HasPrefix(target, "http") {
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24, HasAudio: true}, nil
}
return TopazVideoMetadata{Width: 320, Height: 180, Duration: 3, FrameRate: 24, FrameCount: 72, HasAudio: true}, nil
}}
response, err := client.Run(context.Background(), Request{
Model: "easy-proteus-standard-4",
Body: map[string]any{"video_url": server.URL + "/source.mp4", "output_width": 640, "output_height": 360, "preserve_audio": true},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"allowPrivateSourceDownloads": true, "pollIntervalMs": 1, "pollTimeoutMs": 100}),
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error {
phases = append(phases, payload["phase"].(string))
return nil
},
OnRemoteTaskPolled: func(_ string, payload map[string]any) error {
phases = append(phases, payload["phase"].(string))
return nil
},
})
if err != nil {
t.Fatalf("run Topaz: %v", err)
}
if response.RequestID != "job-1" || !containsStringForTest(phases, "created") || !containsStringForTest(phases, "uploaded") {
t.Fatalf("unexpected Topaz response/phases: response=%+v phases=%+v", response, phases)
}
if called["PUT /upload/1"] != 1 || called["PUT /upload/2"] != 1 || called["PATCH /video/job-1/complete-upload/"] != 1 {
t.Fatalf("upload lifecycle incomplete: %+v", called)
}
}
func TestTopazClientResumesUploadedTaskWithoutDownloadingSource(t *testing.T) {
var sourceCalls int
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/source.mp4" {
sourceCalls++
}
if r.URL.Path == "/video/job-resume/status" {
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/result.mp4"}}`)
return
}
http.NotFound(w, r)
}))
defer server.Close()
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24}, nil
}}).Run(context.Background(), Request{
RemoteTaskID: "job-resume", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 100}),
})
if err != nil {
t.Fatal(err)
}
if sourceCalls != 0 {
t.Fatalf("resume downloaded source %d times", sourceCalls)
}
}
func TestTopazClientReportsProviderFailureAndPollingTimeout(t *testing.T) {
tests := []struct {
name string
statusBody string
timeoutMS int
want string
}{
{name: "provider failure", statusBody: `{"status":"failed","message":"upstream rejected"}`, timeoutMS: 100, want: "upstream rejected"},
{name: "poll timeout", statusBody: `{"status":"processing"}`, timeoutMS: 3, want: "polling timed out"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/video/job/status" {
http.NotFound(w, r)
return
}
_, _ = io.WriteString(w, test.statusBody)
}))
defer server.Close()
_, err := (TopazClient{}).Run(context.Background(), Request{
RemoteTaskID: "job", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
Body: map[string]any{"video_url": "https://assets.example.test/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
"pollIntervalMs": 1, "pollTimeoutMs": test.timeoutMS,
}),
})
if err == nil || !strings.Contains(err.Error(), test.want) {
t.Fatalf("expected %q error, got %v", test.want, err)
}
})
}
}
func TestTopazClientRejectsOversizedInputBeforeProbe(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/source.mp4" {
_, _ = io.WriteString(w, "too-large")
return
}
http.NotFound(w, r)
}))
defer server.Close()
probeCalls := 0
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
probeCalls++
return TopazVideoMetadata{}, nil
}}).Run(context.Background(), Request{
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
"allowPrivateSourceDownloads": true, "maxInputBytes": 3,
}),
})
if err == nil || !strings.Contains(err.Error(), "size limit") {
t.Fatalf("expected input size limit error, got %v", err)
}
if probeCalls != 0 {
t.Fatalf("oversized input was probed %d times", probeCalls)
}
}
func storeCandidate(baseURL, specType string, credentials map[string]any) store.RuntimeModelCandidate {
return storeCandidateWithConfig(baseURL, specType, credentials, nil)
}
func storeCandidateWithConfig(baseURL, specType string, credentials, config map[string]any) store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{BaseURL: baseURL, SpecType: specType, Provider: specType, Credentials: credentials, PlatformConfig: config}
}
func toJSONForTest(value any) string {
raw, _ := json.Marshal(value)
return string(raw)
}
func containsStringForTest(values []string, expected string) bool {
for _, value := range values {
if value == expected {
return true
}
}
return false
}
+328 -3
View File
@@ -12,7 +12,8 @@ import (
const OpenAIReasoningEffortValidationMessage = "reasoning_effort must be one of: none, minimal, low, medium, high, xhigh, max"
var (
openAIReasoningEfforts = map[string]struct{}{
chatReasoningEffortOrder = []string{"none", "minimal", "low", "medium", "high", "xhigh", "max"}
openAIReasoningEfforts = map[string]struct{}{
"none": {},
"minimal": {},
"low": {},
@@ -49,9 +50,48 @@ func ValidateOpenAIReasoningEffort(value any) error {
}
func applyOpenAIChatReasoningParams(body map[string]any, candidate store.RuntimeModelCandidate) {
_ = applyOpenAIChatReasoningParamsWithSource(body, candidate, nil)
}
func applyOpenAIChatReasoningParamsWithSource(body map[string]any, candidate store.RuntimeModelCandidate, original map[string]any) error {
effort := normalizedReasoningString(body["reasoning_effort"])
_, explicitEffort := original["reasoning_effort"]
_, explicitTemperature := original["temperature"]
model := chatReasoningModelName(body, candidate)
if isAliyunBailianOpenAI(candidate) && isAliyunQwen38MaxPreview(model) {
if explicitEffort && qwen38MaxPreviewReasoningEffort(effort) != effort {
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream only supports low, medium, or xhigh without changing the requested reasoning semantics")
}
if temperature, ok := finiteFloatFromAny(body["temperature"]); explicitTemperature && ok && temperature < 0.6 {
return explicitParameterAdaptationError("temperature", "the selected upstream requires temperature >= 0.6")
}
applyAliyunQwen38Reasoning(body, effort)
if temperature, ok := finiteFloatFromAny(body["temperature"]); ok && temperature < 0.6 {
body["temperature"] = 0.6
}
return nil
}
if effort == "" || !isOpenAIReasoningEffort(effort) {
return
return nil
}
resolved, state := resolveCandidateReasoningEffort(effort, candidate)
if state == reasoningCapabilityUnsupported {
if explicitEffort {
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream does not support reasoning_effort")
}
delete(body, "reasoning_effort")
return nil
}
if resolved != "" {
if explicitEffort && resolved != effort {
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream does not support the requested reasoning effort exactly")
}
effort = resolved
}
if explicitEffort {
if err := validateExplicitProviderReasoningEffort(effort, candidate, model); err != nil {
return err
}
}
body["reasoning_effort"] = effort
@@ -65,10 +105,239 @@ func applyOpenAIChatReasoningParams(body map[string]any, candidate store.Runtime
case isVolcesOpenAI(candidate):
applyVolcesReasoning(body, candidate, effort)
}
return nil
}
func validateExplicitProviderReasoningEffort(effort string, candidate store.RuntimeModelCandidate, model string) error {
if effort == "none" {
return nil
}
unsupported := func() error {
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream cannot preserve the requested reasoning effort exactly")
}
switch {
case isAliyunBailianOpenAI(candidate):
if !isAliyunHighMaxReasoningModel(model) || highMaxReasoningEffort(effort) != effort {
return unsupported()
}
case isDeepSeekOpenAI(candidate):
if highMaxReasoningEffort(effort) != effort {
return unsupported()
}
case isZhipuOpenAI(candidate):
if !isZhipuReasoningEffortModel(model) || zhipuReasoningEffort(effort) != effort {
return unsupported()
}
case isVolcesOpenAI(candidate):
if !isVolcesReasoningEffortModel(model) || volcesChatReasoningEffort(effort) != effort {
return unsupported()
}
}
return nil
}
func applyOpenAIResponsesReasoningParams(body map[string]any, candidate store.RuntimeModelCandidate) {
_ = applyOpenAIResponsesReasoningParamsWithSource(body, candidate, nil)
}
func applyOpenAIResponsesReasoningParamsWithSource(body map[string]any, candidate store.RuntimeModelCandidate, original map[string]any) error {
reasoning, _ := body["reasoning"].(map[string]any)
if reasoning != nil {
reasoning = cloneBody(reasoning)
}
effort := ""
if reasoning != nil {
effort = normalizedReasoningString(reasoning["effort"])
}
originalReasoning, _ := original["reasoning"].(map[string]any)
_, explicitEffort := originalReasoning["effort"]
_, explicitTemperature := original["temperature"]
model := chatReasoningModelName(body, candidate)
qwen38 := isAliyunBailianOpenAI(candidate) && isAliyunQwen38MaxPreview(model)
if qwen38 {
if effort == "" {
if enabled, ok := body["enable_thinking"].(bool); ok && !enabled {
effort = "low"
}
} else {
mapped := qwen38MaxPreviewReasoningEffort(effort)
if explicitEffort && mapped != effort {
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream only supports low, medium, or xhigh without changing the requested reasoning semantics")
}
effort = mapped
}
} else if effort != "" && isOpenAIReasoningEffort(effort) {
resolved, state := resolveCandidateReasoningEffort(effort, candidate)
if state == reasoningCapabilityUnsupported {
if explicitEffort {
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream does not support reasoning.effort")
}
effort = ""
} else if resolved != "" {
if explicitEffort && resolved != effort {
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream does not support the requested reasoning effort exactly")
}
effort = resolved
}
}
if effort != "" {
if reasoning == nil {
reasoning = map[string]any{}
}
reasoning["effort"] = effort
} else if reasoning != nil {
delete(reasoning, "effort")
}
if len(reasoning) > 0 {
body["reasoning"] = reasoning
} else {
delete(body, "reasoning")
}
delete(body, "reasoning_effort")
delete(body, "enable_thinking")
if qwen38 {
if temperature, ok := finiteFloatFromAny(body["temperature"]); ok && temperature < 0.6 {
if explicitTemperature {
return explicitParameterAdaptationError("temperature", "the selected upstream requires temperature >= 0.6")
}
body["temperature"] = 0.6
}
}
return nil
}
func explicitParameterAdaptationError(param string, message string) error {
return &ClientError{Code: "invalid_parameter", Message: message, Param: param, StatusCode: 400, Retryable: false}
}
func applyAliyunQwen38Reasoning(body map[string]any, effort string) {
defer delete(body, "thinking_budget_tokens")
requestedDisable := false
if enabled, ok := body["enable_thinking"].(bool); ok && !enabled {
requestedDisable = true
}
body["enable_thinking"] = true
budget, hasBudget := nonNegativeIntFromAny(body["thinking_budget"])
if !hasBudget {
budget, hasBudget = nonNegativeIntFromAny(body["thinking_budget_tokens"])
}
if hasBudget {
if budget > 262144 {
budget = 262144
}
body["thinking_budget"] = budget
delete(body, "reasoning_effort")
return
}
delete(body, "thinking_budget")
if effort == "" {
if requestedDisable {
body["reasoning_effort"] = "low"
}
return
}
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
}
func qwen38MaxPreviewReasoningEffort(effort string) string {
switch effort {
case "none", "minimal", "low":
return "low"
case "medium":
return "medium"
default:
return "xhigh"
}
}
type reasoningCapabilityState int
const (
reasoningCapabilityUnknown reasoningCapabilityState = iota
reasoningCapabilityUnsupported
reasoningCapabilitySupported
)
func resolveCandidateReasoningEffort(effort string, candidate store.RuntimeModelCandidate) (string, reasoningCapabilityState) {
supported := map[string]struct{}{}
declaredLevels := false
explicitlyUnsupported := false
explicitlySupported := false
for _, key := range []string{"text_generate", "tools_call", "image_analysis", "video_understanding", "audio_understanding", "omni"} {
capability, ok := candidate.Capabilities[key].(map[string]any)
if !ok {
continue
}
if value, ok := capability["supportThinking"].(bool); ok && !value {
explicitlyUnsupported = true
continue
}
if value, ok := capability["supportThinking"].(bool); ok && value {
explicitlySupported = true
}
rawLevels, ok := capability["thinkingEffortLevels"].([]any)
if !ok {
if stringsValue, stringsOK := capability["thinkingEffortLevels"].([]string); stringsOK {
rawLevels = make([]any, len(stringsValue))
for index, value := range stringsValue {
rawLevels[index] = value
}
ok = true
}
}
if !ok {
continue
}
declaredLevels = true
for _, raw := range rawLevels {
level := normalizedReasoningString(raw)
if isOpenAIReasoningEffort(level) {
supported[level] = struct{}{}
}
}
}
if len(supported) == 0 {
if declaredLevels || (explicitlyUnsupported && !explicitlySupported) {
return "", reasoningCapabilityUnsupported
}
return effort, reasoningCapabilityUnknown
}
requestedIndex := reasoningEffortIndex(effort)
best := ""
bestDistance := len(chatReasoningEffortOrder) + 1
for index, candidateEffort := range chatReasoningEffortOrder {
if _, ok := supported[candidateEffort]; !ok {
continue
}
distance := index - requestedIndex
if distance < 0 {
distance = -distance
}
if distance < bestDistance {
best = candidateEffort
bestDistance = distance
}
}
return best, reasoningCapabilitySupported
}
func reasoningEffortIndex(effort string) int {
for index, value := range chatReasoningEffortOrder {
if value == effort {
return index
}
}
return 0
}
func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandidate, effort string) {
defer delete(body, "thinking_budget_tokens")
budget, hasBudget := positiveIntFromAny(body["thinking_budget"])
if !hasBudget {
budget, hasBudget = positiveIntFromAny(body["thinking_budget_tokens"])
}
delete(body, "thinking_budget")
if effort == "none" {
body["enable_thinking"] = false
@@ -78,6 +347,15 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
body["enable_thinking"] = true
model := chatReasoningModelName(body, candidate)
if isAliyunQwen38MaxPreview(model) {
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
body["thinking_budget"] = budget
delete(body, "reasoning_effort")
return
}
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
return
}
if isAliyunHighMaxReasoningModel(model) {
body["reasoning_effort"] = highMaxReasoningEffort(effort)
return
@@ -85,7 +363,7 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
delete(body, "reasoning_effort")
if isAliyunThinkingBudgetModel(model) {
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
if hasBudget {
body["thinking_budget"] = budget
}
}
@@ -200,6 +478,10 @@ func isAliyunHighMaxReasoningModel(model string) bool {
return strings.Contains(model, "deepseek-v4") || strings.HasPrefix(model, "glm-")
}
func isAliyunQwen38MaxPreview(model string) bool {
return strings.Contains(model, "qwen3.8-max-preview")
}
func isAliyunThinkingBudgetModel(model string) bool {
return strings.Contains(model, "qwen") || strings.Contains(model, "qwq") || strings.Contains(model, "qvq") || strings.Contains(model, "kimi")
}
@@ -282,6 +564,49 @@ func positiveIntFromAny(value any) (int, bool) {
return int(math.Floor(number)), true
}
func nonNegativeIntFromAny(value any) (int, bool) {
number, ok := finiteFloatFromAny(value)
if !ok || number < 0 {
return 0, false
}
return int(math.Floor(number)), true
}
func finiteFloatFromAny(value any) (float64, bool) {
if value == nil {
return 0, false
}
var number float64
switch typed := value.(type) {
case int:
number = float64(typed)
case int64:
number = float64(typed)
case float64:
number = typed
case float32:
number = float64(typed)
case jsonNumber:
parsed, err := typed.Float64()
if err != nil {
return 0, false
}
number = parsed
case string:
parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64)
if err != nil {
return 0, false
}
number = parsed
default:
return 0, false
}
if math.IsNaN(number) || math.IsInf(number, 0) {
return 0, false
}
return number, true
}
type jsonNumber interface {
Float64() (float64, error)
}
@@ -0,0 +1,184 @@
package clients
import (
"encoding/json"
"os"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestQwen38MatchesCrossServiceBehaviorVectors(t *testing.T) {
raw, err := os.ReadFile("testdata/qwen38-reasoning-vectors.json")
if err != nil {
t.Fatalf("read behavior vectors: %v", err)
}
var vectors []struct {
Name string `json:"name"`
Input map[string]any `json:"input"`
Expected map[string]any `json:"expected"`
}
if err := json.Unmarshal(raw, &vectors); err != nil {
t.Fatalf("decode behavior vectors: %v", err)
}
candidate := store.RuntimeModelCandidate{Provider: "aliyun-bailian-openai", ProviderModelName: "qwen3.8-max-preview"}
for _, vector := range vectors {
t.Run(vector.Name, func(t *testing.T) {
body := cloneBody(vector.Input)
body["model"] = "qwen3.8-max-preview"
applyOpenAIChatReasoningParams(body, candidate)
for _, key := range []string{"enable_thinking", "reasoning_effort", "thinking_budget", "temperature"} {
expected, expectedPresent := vector.Expected[key]
actual, actualPresent := body[key]
if expectedPresent != actualPresent || (expectedPresent && normalizedJSONScalar(actual) != normalizedJSONScalar(expected)) {
t.Fatalf("%s mismatch: expected=%#v actual=%#v body=%+v", key, expected, actual, body)
}
}
})
}
}
func normalizedJSONScalar(value any) any {
if number, ok := value.(int); ok {
return float64(number)
}
return value
}
func TestQwen38MaxPreviewReasoningNormalization(t *testing.T) {
tests := []struct {
name string
body map[string]any
effort any
budget any
temp any
enabled any
noEffort bool
}{
{name: "none", body: map[string]any{"reasoning_effort": "none"}, effort: "low"},
{name: "minimal", body: map[string]any{"reasoning_effort": "minimal"}, effort: "low"},
{name: "medium", body: map[string]any{"reasoning_effort": "medium"}, effort: "medium"},
{name: "high", body: map[string]any{"reasoning_effort": "high"}, effort: "xhigh"},
{name: "max", body: map[string]any{"reasoning_effort": "max"}, effort: "xhigh"},
{name: "unspecified", body: map[string]any{}, noEffort: true},
{name: "forced on", body: map[string]any{"enable_thinking": false}, effort: "low"},
{name: "budget wins", body: map[string]any{"reasoning_effort": "high", "thinking_budget": 999999}, budget: 262144, noEffort: true},
{name: "zero budget", body: map[string]any{"reasoning_effort": "medium", "thinking_budget_tokens": 0}, budget: 0, noEffort: true},
{name: "temperature floor", body: map[string]any{"temperature": 0.1}, temp: 0.6, noEffort: true},
}
candidate := store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
ProviderModelName: "qwen3.8-max-preview",
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
body := cloneBody(test.body)
body["model"] = "qwen3.8-max-preview"
applyOpenAIChatReasoningParams(body, candidate)
if body["enable_thinking"] != true {
t.Fatalf("Qwen3.8 must always enable thinking: %+v", body)
}
if test.noEffort {
if _, ok := body["reasoning_effort"]; ok {
t.Fatalf("reasoning_effort must be omitted: %+v", body)
}
} else if body["reasoning_effort"] != test.effort {
t.Fatalf("expected effort %#v, got %+v", test.effort, body)
}
if test.budget != nil && body["thinking_budget"] != test.budget {
t.Fatalf("expected budget %#v, got %+v", test.budget, body)
}
if test.temp != nil && body["temperature"] != test.temp {
t.Fatalf("expected temperature %#v, got %+v", test.temp, body)
}
})
}
}
func TestGenericReasoningCapabilityUsesNearestTierAndIgnoresModeSwitch(t *testing.T) {
body := map[string]any{"reasoning_effort": "medium"}
applyOpenAIChatReasoningParams(body, store.RuntimeModelCandidate{
Provider: "openai",
Capabilities: map[string]any{
"text_generate": map[string]any{
"supportThinking": true,
"supportThinkingModeSwitch": false,
"thinkingEffortLevels": []any{"low", "high"},
},
},
})
if body["reasoning_effort"] != "low" {
t.Fatalf("equidistant effort must choose lower tier, got %+v", body)
}
}
func TestGenericReasoningKeepsUnknownAndRemovesExplicitUnsupported(t *testing.T) {
unknown := map[string]any{"reasoning_effort": "high"}
applyOpenAIChatReasoningParams(unknown, store.RuntimeModelCandidate{Provider: "openai"})
if unknown["reasoning_effort"] != "high" {
t.Fatalf("unknown capability should preserve effort: %+v", unknown)
}
unsupported := map[string]any{"reasoning_effort": "high"}
applyOpenAIChatReasoningParams(unsupported, store.RuntimeModelCandidate{
Provider: "openai",
Capabilities: map[string]any{
"text_generate": map[string]any{"supportThinking": false},
},
})
if _, ok := unsupported["reasoning_effort"]; ok {
t.Fatalf("explicit unsupported capability should remove effort: %+v", unsupported)
}
}
func TestQwen38ResponsesReasoningNormalization(t *testing.T) {
candidate := store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
ProviderModelName: "qwen3.8-max-preview",
}
tests := []struct {
name string
requested string
expected string
}{
{name: "none", requested: "none", expected: "low"},
{name: "minimal", requested: "minimal", expected: "low"},
{name: "medium", requested: "medium", expected: "medium"},
{name: "high", requested: "high", expected: "xhigh"},
{name: "max", requested: "max", expected: "xhigh"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
body := map[string]any{
"model": "qwen3.8-max-preview",
"reasoning": map[string]any{"effort": test.requested, "summary": "auto"},
"temperature": 0.1,
}
applyOpenAIResponsesReasoningParams(body, candidate)
reasoning := body["reasoning"].(map[string]any)
if reasoning["effort"] != test.expected || reasoning["summary"] != "auto" {
t.Fatalf("unexpected Responses reasoning: %+v", body)
}
if body["temperature"] != 0.6 {
t.Fatalf("Qwen3.8 Responses temperature was not floored: %+v", body)
}
})
}
}
func TestQwen38ResponsesForcedThinkingAndDefault(t *testing.T) {
candidate := store.RuntimeModelCandidate{Provider: "aliyun-bailian-openai", ProviderModelName: "qwen3.8-max-preview"}
forced := map[string]any{"enable_thinking": false}
applyOpenAIResponsesReasoningParams(forced, candidate)
if forced["reasoning"].(map[string]any)["effort"] != "low" {
t.Fatalf("disabled Qwen3.8 Responses request must become low: %+v", forced)
}
if _, ok := forced["enable_thinking"]; ok {
t.Fatalf("Chat-only alias leaked into Responses body: %+v", forced)
}
defaulted := map[string]any{"reasoning": map[string]any{"summary": "auto"}}
applyOpenAIResponsesReasoningParams(defaulted, candidate)
reasoning := defaulted["reasoning"].(map[string]any)
if _, ok := reasoning["effort"]; ok || reasoning["summary"] != "auto" {
t.Fatalf("unspecified effort must preserve upstream xhigh default: %+v", defaulted)
}
}
+600 -18
View File
@@ -5,6 +5,7 @@ import (
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
@@ -344,6 +345,453 @@ func TestOpenAIClientChatContract(t *testing.T) {
}
}
func TestOpenAIClientDelayedStream429ReturnsBeforeAnyDelta(t *testing.T) {
requestStarted := make(chan struct{})
releaseResponse := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
close(requestStarted)
<-releaseResponse
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusTooManyRequests)
_ = json.NewEncoder(w).Encode(map[string]any{
"error": map[string]any{"message": "controlled delayed 429"},
})
}))
defer server.Close()
release := func() {
select {
case <-releaseResponse:
default:
close(releaseResponse)
}
}
defer release()
type runResult struct {
response Response
err error
}
deltaCount := 0
done := make(chan runResult, 1)
go func() {
response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions",
Model: "delayed-429-model",
Stream: true,
Body: map[string]any{
"model": "delayed-429-model",
"stream": true,
"messages": []any{map[string]any{"role": "user", "content": "ping"}},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ProviderModelName: "delayed-429-model",
Credentials: map[string]any{"apiKey": "test-key"},
},
StreamDelta: func(StreamDeltaEvent) error {
deltaCount++
return nil
},
})
done <- runResult{response: response, err: err}
}()
select {
case <-requestStarted:
case <-time.After(time.Second):
t.Fatal("upstream request did not start")
}
select {
case result := <-done:
t.Fatalf("stream returned before upstream status arrived: response=%+v err=%v", result.response, result.err)
case <-time.After(50 * time.Millisecond):
}
release()
var result runResult
select {
case result = <-done:
case <-time.After(time.Second):
t.Fatal("stream did not return after delayed 429")
}
clientErr, ok := result.err.(*ClientError)
if !ok || clientErr.Code != "rate_limit" || clientErr.StatusCode != http.StatusTooManyRequests || !clientErr.Retryable {
t.Fatalf("unexpected delayed 429 error: %T %+v", result.err, result.err)
}
if deltaCount != 0 {
t.Fatalf("delayed 429 emitted %d stream deltas, want 0", deltaCount)
}
}
func TestOpenAIClientAliyunChatSendsStandardInputAudioAndNormalizesLegacyAudioURL(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatalf("decode request: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "chatcmpl-audio",
"object": "chat.completion",
"model": captured["model"],
"choices": []any{map[string]any{
"message": map[string]any{"role": "assistant", "content": "ok"},
}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions",
Model: "qwen-omni",
Body: map[string]any{
"model": "qwen-omni",
"messages": []any{
map[string]any{
"role": "user",
"content": []any{
map[string]any{
"type": "input_audio",
"input_audio": map[string]any{
"data": "https://cdn.example.com/audio/standard.wav",
"format": "wav",
},
},
map[string]any{
"type": "audio_url",
"audio_url": map[string]any{"url": "https://cdn.example.com/audio/legacy.m4a"},
},
},
},
},
},
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
BaseURL: server.URL,
ProviderModelName: "qwen-omni",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatalf("run aliyun openai-compatible client: %v", err)
}
messages, _ := captured["messages"].([]any)
message, _ := messages[0].(map[string]any)
content, _ := message["content"].([]any)
standard, _ := content[0].(map[string]any)
if standard["type"] != "input_audio" {
t.Fatalf("standard input_audio was not preserved: %+v", standard)
}
standardAudio, _ := standard["input_audio"].(map[string]any)
if standardAudio["data"] != "https://cdn.example.com/audio/standard.wav" || standardAudio["format"] != "wav" {
t.Fatalf("unexpected standard input_audio: %+v", standardAudio)
}
legacy, _ := content[1].(map[string]any)
if legacy["type"] != "input_audio" {
t.Fatalf("legacy audio_url was not normalized: %+v", legacy)
}
legacyAudio, _ := legacy["input_audio"].(map[string]any)
if legacyAudio["data"] != "https://cdn.example.com/audio/legacy.m4a" || legacyAudio["format"] != "m4a" {
t.Fatalf("unexpected normalized legacy audio: %+v", legacyAudio)
}
}
func TestOpenAIClientImageEditUsesMultipartWithMultipleImages(t *testing.T) {
var contentType string
var receivedFields map[string][]string
var receivedImageCount int
var receivedImageFields []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
contentType = r.Header.Get("Content-Type")
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart image edit: %v", err)
}
receivedFields = r.MultipartForm.Value
for field, files := range r.MultipartForm.File {
receivedImageFields = append(receivedImageFields, field)
receivedImageCount += len(files)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.edits",
ModelType: "image_edit",
Model: "gpt-image-2",
Body: map[string]any{
"model": "gpt-image-2",
"prompt": "combine the references",
"size": "1024x1536",
"quality": "medium",
"aspect_ratio": "2:3",
"resolution": "2K",
"width": 1024,
"height": 1536,
"images": []any{
"data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("image one")),
"data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("image two")),
},
"_metadata": map[string]any{"private": true},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2-vip",
ModelType: "image_edit",
Credentials: map[string]any{"apiKey": "openai-key"},
},
})
if err != nil {
t.Fatalf("run OpenAI image edit: %v", err)
}
if !strings.HasPrefix(contentType, "multipart/form-data; boundary=") {
t.Fatalf("image edit should use multipart/form-data, got %q", contentType)
}
if receivedImageCount != 2 || len(receivedImageFields) != 1 || receivedImageFields[0] != "image[]" {
t.Fatalf("unexpected multipart image files: fields=%v count=%d", receivedImageFields, receivedImageCount)
}
if receivedFields["model"][0] != "gpt-image-2-vip" ||
receivedFields["prompt"][0] != "combine the references" ||
receivedFields["size"][0] != "1024x1536" ||
receivedFields["quality"][0] != "medium" {
t.Fatalf("unexpected multipart fields: %+v", receivedFields)
}
if _, ok := receivedFields["_metadata"]; ok {
t.Fatalf("internal metadata must not be forwarded: %+v", receivedFields)
}
for _, field := range []string{"aspect_ratio", "resolution", "width", "height"} {
if _, ok := receivedFields[field]; ok {
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, receivedFields)
}
}
}
func TestOpenAIClientImageEditDownloadsURLWhenBuildingMultipart(t *testing.T) {
imagePayload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
fetchCount := 0
imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fetchCount++
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(imagePayload)
}))
defer imageServer.Close()
var uploaded []byte
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart image edit: %v", err)
}
files := r.MultipartForm.File["image"]
if len(files) != 1 {
t.Fatalf("expected one multipart image, got %+v", r.MultipartForm.File)
}
file, err := files[0].Open()
if err != nil {
t.Fatalf("open multipart image: %v", err)
}
defer file.Close()
uploaded, err = io.ReadAll(file)
if err != nil {
t.Fatalf("read multipart image: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer upstream.Close()
_, err := (OpenAIClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
Kind: "images.edits",
ModelType: "image_edit",
Body: map[string]any{
"image": imageServer.URL + "/source.png",
"prompt": "edit it",
},
Candidate: store.RuntimeModelCandidate{
BaseURL: upstream.URL,
Provider: "openai",
ProviderModelName: "gpt-image-compatible",
ModelType: "image_edit",
Credentials: map[string]any{"apiKey": "openai-key"},
},
})
if err != nil {
t.Fatalf("run OpenAI URL multipart image edit: %v", err)
}
if fetchCount != 1 || string(uploaded) != string(imagePayload) {
t.Fatalf("image URL should be fetched once during multipart construction: fetches=%d uploaded=%v", fetchCount, uploaded)
}
}
func TestKelingAndMinimaxMaterializeRemoteMediaInsideClient(t *testing.T) {
payload := []byte("remote media payload")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "audio/mpeg")
_, _ = w.Write(payload)
}))
defer server.Close()
encoded, err := kelingImageToBase64(context.Background(), Request{HTTPClient: server.Client()}, server.URL+"/image.png")
if err != nil {
t.Fatalf("materialize Keling image URL: %v", err)
}
if encoded != base64.StdEncoding.EncodeToString(payload) {
t.Fatalf("unexpected Keling image base64: %q", encoded)
}
audio, filename, contentType, err := minimaxVoiceCloneFilePayload(context.Background(), server.Client(), server.URL+"/audio.mp3", "voice_clone")
if err != nil {
t.Fatalf("materialize Minimax audio URL: %v", err)
}
if string(audio) != string(payload) || filename != "voice_clone.mp3" || contentType != "audio/mpeg" {
t.Fatalf("unexpected Minimax audio materialization: filename=%q contentType=%q payload=%q", filename, contentType, audio)
}
}
func TestOpenAIClientImageEditUsesJSONURLsOnlyWhenExplicitlyConfigured(t *testing.T) {
var contentType string
var received map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
contentType = r.Header.Get("Content-Type")
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
t.Fatalf("decode image edit JSON: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"url": "https://cdn.example.com/output.png"}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.edits",
ModelType: "image_edit",
Model: "compatible-image-edit",
Body: map[string]any{
"images": []any{
"https://cdn.example.com/first.png",
map[string]any{"url": "https://cdn.example.com/second.png"},
},
"mask_image": map[string]any{"image_url": "https://cdn.example.com/mask.png"},
"prompt": "combine the references",
"_metadata": map[string]any{"private": true},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "compatible-openai",
SpecType: "openai",
ProviderModelName: "compatible-image-edit-v2",
ModelType: "image_edit",
Credentials: map[string]any{"apiKey": "openai-key"},
PlatformConfig: map[string]any{"imageEditRequestFormat": "json_url"},
},
})
if err != nil {
t.Fatalf("run OpenAI-compatible JSON image edit: %v", err)
}
if !strings.HasPrefix(contentType, "application/json") {
t.Fatalf("explicit JSON URL mode should use application/json, got %q", contentType)
}
images, _ := received["image"].([]any)
if len(images) != 2 || images[0] != "https://cdn.example.com/first.png" || images[1] != "https://cdn.example.com/second.png" {
t.Fatalf("unexpected JSON image URLs: %+v", received)
}
if received["mask"] != "https://cdn.example.com/mask.png" || received["model"] != "compatible-image-edit-v2" {
t.Fatalf("unexpected JSON image edit body: %+v", received)
}
if _, ok := received["images"]; ok {
t.Fatalf("images alias must be normalized to image: %+v", received)
}
if _, ok := received["_metadata"]; ok {
t.Fatalf("internal metadata must not be forwarded: %+v", received)
}
}
func TestOpenAIClientImageEditJSONModeRejectsInlineImageData(t *testing.T) {
_, _, err := openAIRequestPayload(context.Background(), "images.edits", map[string]any{
"image": "data:image/png;base64,aW1hZ2U=",
}, store.RuntimeModelCandidate{PlatformConfig: map[string]any{"imageEditRequestFormat": "json_url"}})
clientErr, ok := err.(*ClientError)
if !ok || clientErr.Code != "invalid_parameter" || clientErr.Param != "image" || clientErr.StatusCode != http.StatusBadRequest {
t.Fatalf("expected non-retryable image URL validation error, got %#v", err)
}
}
func TestOpenAIImageEditFieldNameUsesArraySyntaxForMultipleImages(t *testing.T) {
tests := []struct {
name string
candidate store.RuntimeModelCandidate
imageCount int
want string
}{
{name: "compatible single image", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1"}, imageCount: 1, want: "image"},
{name: "compatible multiple images", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1"}, imageCount: 2, want: "image[]"},
{name: "official single image", candidate: store.RuntimeModelCandidate{BaseURL: "https://api.openai.com/v1"}, imageCount: 1, want: "image[]"},
{name: "explicit scalar override", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1", PlatformConfig: map[string]any{"imageEditMultipartFieldName": "image"}}, imageCount: 2, want: "image"},
{name: "explicit array override", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1", PlatformConfig: map[string]any{"imageEditMultipartFieldName": "image[]"}}, imageCount: 1, want: "image[]"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := openAIImageEditFieldName(test.candidate, test.imageCount); got != test.want {
t.Fatalf("field name = %q, want %q", got, test.want)
}
})
}
}
func TestOpenAIClientImageGenerationUsesSizeWithoutGenericGeometryFields(t *testing.T) {
var received map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
t.Fatalf("decode image generation body: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.generations",
ModelType: "image_generate",
Model: "gpt-image-2",
Body: map[string]any{
"model": "gpt-image-2",
"prompt": "draw",
"size": "2048x1152",
"quality": "high",
"aspect_ratio": "16:9",
"aspectRatio": "16:9",
"resolution": "2K",
"width": 2048,
"height": 1152,
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2-vip",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "openai-key"},
},
})
if err != nil {
t.Fatalf("run OpenAI image generation: %v", err)
}
if received["model"] != "gpt-image-2-vip" ||
received["prompt"] != "draw" ||
received["size"] != "2048x1152" ||
received["quality"] != "high" {
t.Fatalf("unexpected image generation fields: %+v", received)
}
for _, field := range []string{"aspect_ratio", "aspectRatio", "resolution", "width", "height"} {
if _, ok := received[field]; ok {
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, received)
}
}
}
func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
cases := []struct {
name string
@@ -405,6 +853,36 @@ func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
}
},
},
{
name: "aliyun qwen 3.8 preview preserves supported effort",
provider: "aliyun-bailian-openai",
model: "qwen3.8-max-preview",
body: map[string]any{
"reasoning_effort": "high",
},
assertion: func(t *testing.T, body map[string]any) {
if body["enable_thinking"] != true || body["reasoning_effort"] != "xhigh" {
t.Fatalf("aliyun qwen 3.8 should use xhigh reasoning effort, got %+v", body)
}
},
},
{
name: "aliyun qwen 3.8 preview uses explicit budget exclusively",
provider: "aliyun-bailian-openai",
model: "qwen3.8-max-preview",
body: map[string]any{
"reasoning_effort": "medium",
"thinking_budget_tokens": 16384,
},
assertion: func(t *testing.T, body map[string]any) {
if body["enable_thinking"] != true || body["thinking_budget"] != float64(16384) {
t.Fatalf("aliyun qwen 3.8 should use the explicit thinking budget, got %+v", body)
}
if _, ok := body["reasoning_effort"]; ok {
t.Fatalf("aliyun qwen 3.8 must not combine reasoning effort with thinking budget: %+v", body)
}
},
},
{
name: "same deepseek model on aliyun uses aliyun protocol",
provider: "aliyun-bailian-openai",
@@ -707,13 +1185,11 @@ func TestOpenAIClientChatRequestNormalizesToolContext(t *testing.T) {
}
assistant, _ := messages[0].(map[string]any)
if _, ok := assistant["functionCall"]; ok {
t.Fatalf("functionCall should be converted away: %+v", assistant)
t.Fatalf("functionCall alias should be converted away: %+v", assistant)
}
toolCalls, _ := assistant["tool_calls"].([]any)
toolCall, _ := toolCalls[0].(map[string]any)
function, _ := toolCall["function"].(map[string]any)
function, _ := assistant["function_call"].(map[string]any)
if function["name"] != "lookup" || function["arguments"] != `{"q":"weather"}` {
t.Fatalf("unexpected normalized tool call: %+v", assistant)
t.Fatalf("unexpected normalized legacy function call: %+v", assistant)
}
toolMessage, _ := messages[1].(map[string]any)
if toolMessage["tool_call_id"] != "call_0" || toolMessage["toolCallId"] != nil {
@@ -811,14 +1287,18 @@ func TestOpenAIClientChatResponseNormalizesToolCallFormats(t *testing.T) {
if message["content"] != "calling tools" {
t.Fatalf("tool_use block should be removed from content: %+v", message)
}
for _, key := range []string{"toolCalls", "function_call"} {
for _, key := range []string{"toolCalls"} {
if _, ok := message[key]; ok {
t.Fatalf("%s should be converted away: %+v", key, message)
}
}
legacyFunction, _ := message["function_call"].(map[string]any)
if legacyFunction["name"] != "legacy_lookup" || legacyFunction["arguments"] != "{\"city\":\"NYC\"}" {
t.Fatalf("canonical legacy function_call should be preserved: %+v", message)
}
toolCalls, _ := message["tool_calls"].([]any)
if len(toolCalls) != 3 {
t.Fatalf("expected 3 normalized tool calls, got %+v", message)
if len(toolCalls) != 2 {
t.Fatalf("expected 2 normalized tool calls plus legacy function_call, got %+v", message)
}
assertToolCall := func(index int, id string, name string, arguments string) {
t.Helper()
@@ -829,8 +1309,7 @@ func TestOpenAIClientChatResponseNormalizesToolCallFormats(t *testing.T) {
}
}
assertToolCall(0, "call_camel", "camel_lookup", "{\"city\":\"SF\"}")
assertToolCall(1, "call_1", "legacy_lookup", "{\"city\":\"NYC\"}")
assertToolCall(2, "toolu_1", "anthropic_lookup", "{\"city\":\"Boston\"}")
assertToolCall(1, "toolu_1", "anthropic_lookup", "{\"city\":\"Boston\"}")
}
func TestOpenAIClientChatStreamContract(t *testing.T) {
@@ -1088,19 +1567,21 @@ func TestOpenAIClientChatStreamNormalizesToolCallFormats(t *testing.T) {
if len(captured) != 3 {
t.Fatalf("unexpected captured events: %+v", captured)
}
for _, event := range captured {
for index, event := range captured {
choices, _ := event.Event["choices"].([]any)
choice, _ := choices[0].(map[string]any)
delta, _ := choice["delta"].(map[string]any)
if _, ok := delta["function_call"]; ok {
t.Fatalf("function_call should be converted away: %+v", event.Event)
}
if _, ok := delta["functionCall"]; ok {
t.Fatalf("functionCall should be converted away: %+v", event.Event)
}
if _, ok := delta["toolCall"]; ok {
t.Fatalf("toolCall should be converted away: %+v", event.Event)
}
if index < 2 {
if _, ok := delta["function_call"]; !ok {
t.Fatalf("canonical legacy function_call should be preserved: %+v", event.Event)
}
}
}
choices, _ := response.Result["choices"].([]any)
choice, _ := choices[0].(map[string]any)
@@ -1304,6 +1785,104 @@ func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
t.Fatalf("unexpected image response: %+v", response.Result)
}
if response.Wire == nil || len(response.Wire.RawJSON) != 0 {
t.Fatalf("Gemini image response retained duplicate raw JSON: %+v", response.Wire)
}
candidates, _ := response.Wire.Body["candidates"].([]any)
content, _ := candidates[0].(map[string]any)["content"].(map[string]any)
wireParts, _ := content["parts"].([]any)
inline, _ := wireParts[0].(map[string]any)["inlineData"].(map[string]any)
if inline["data"] != "aW1hZ2U=" {
t.Fatalf("decoded Gemini wire body lost Base64 output: %+v", response.Wire.Body)
}
}
func TestGeminiClientImageGenerateRejectsPromptBlockedWithoutImage(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("X-Request-Id", "gemini-blocked-request")
_ = json.NewEncoder(w).Encode(map[string]any{
"promptFeedback": map[string]any{
"blockReason": "SAFETY",
"blockReasonMessage": "The request was blocked by the image safety policy.",
},
"usageMetadata": map[string]any{"promptTokenCount": 12, "totalTokenCount": 12},
})
}))
defer server.Close()
_, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.generations",
ModelType: "image_generate",
Model: "gemini-image",
Body: map[string]any{"model": "gemini-image", "prompt": "blocked prompt"},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ProviderModelName: "gemini-image",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "gemini-key"},
},
})
if err == nil {
t.Fatal("missing Gemini image must fail instead of returning a placeholder")
}
if ErrorCode(err) != "gemini_image_output_missing" || IsRetryable(err) {
t.Fatalf("unexpected missing image error: code=%q retryable=%v err=%v", ErrorCode(err), IsRetryable(err), err)
}
meta := ErrorResponseMetadata(err)
if meta.RequestID != "gemini-blocked-request" || meta.ResponseDurationMS < 0 {
t.Fatalf("missing image error lost response metadata: %+v", meta)
}
details := ErrorDetails(err)
if details["reason"] != "prompt_blocked" || details["promptBlockReason"] != "SAFETY" || details["candidateCount"] != 0 {
t.Fatalf("unexpected prompt block details: %+v", details)
}
if !strings.Contains(err.Error(), "blockReason=SAFETY") || strings.Contains(err.Error(), "gemini-image-placeholder") {
t.Fatalf("missing image error should explain the upstream block without a placeholder: %v", err)
}
}
func TestGeminiMissingImageDetailsExplainEmptyImagePayload(t *testing.T) {
details := geminiMissingImageDetails(map[string]any{
"candidates": []any{map[string]any{
"finishReason": "STOP",
"content": map[string]any{"parts": []any{
map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": ""}},
map[string]any{"text": "The model completed without producing an image."},
}},
}},
})
message := geminiMissingImageMessage(details)
if details["reason"] != "empty_image_payload" || details["candidateCount"] != 1 || details["partCount"] != 2 {
t.Fatalf("unexpected empty image payload details: %+v", details)
}
if !strings.Contains(message, "missing inline data or file URI") ||
!strings.Contains(message, "finishReason=STOP") ||
!strings.Contains(message, "The model completed without producing an image.") {
t.Fatalf("missing image message should preserve the concrete extraction reason: %s", message)
}
}
func TestGeminiMissingImageDetailsPreserveUpstreamErrorEnvelope(t *testing.T) {
details := geminiMissingImageDetails(map[string]any{
"error": map[string]any{
"code": 500,
"status": "INTERNAL",
"message": "image backend returned no generated asset",
},
})
message := geminiMissingImageMessage(details)
if details["reason"] != "upstream_error" ||
details["upstreamErrorCode"] != "500" ||
details["upstreamErrorStatus"] != "INTERNAL" {
t.Fatalf("unexpected upstream error details: %+v", details)
}
if !strings.Contains(message, "upstreamErrorCode=500") ||
!strings.Contains(message, "upstreamErrorStatus=INTERNAL") ||
!strings.Contains(message, "image backend returned no generated asset") {
t.Fatalf("upstream error message should remain diagnostic: %s", message)
}
}
func TestGeminiGenerationConfigInitializesMissingImageConfig(t *testing.T) {
@@ -1977,7 +2556,7 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if item["url"] != "https://example.com/out.mp4" || response.Usage.TotalTokens != 9 {
if item["url"] != "https://example.com/out.mp4" || response.Result["content"] != nil || response.Result["usage"] != nil || response.Usage.TotalTokens != 9 {
t.Fatalf("unexpected response: %+v usage=%+v", response.Result, response.Usage)
}
}
@@ -2021,8 +2600,9 @@ func TestVolcesClientVideoRetriesTransientPollAndKeepsOfficialResult(t *testing.
if polls != 2 || len(persisted) != 1 || persisted[0] != "cgt-retry:succeeded" {
t.Fatalf("unexpected poll state polls=%d persisted=%+v", polls, persisted)
}
if response.Result["updated_at"] != float64(124) || response.Result["seed"] != float64(7) || response.Result["raw"] == nil {
t.Fatalf("official result fields lost: %+v", response.Result)
if response.Result["seed"] != float64(7) || response.Result["raw"] != nil || response.Result["updated_at"] != nil ||
response.Result["content"] != nil || response.Result["usage"] != nil || response.Result["upstream_task_id"] != nil {
t.Fatalf("provider response should be reduced to the canonical result: %+v", response.Result)
}
}
@@ -2109,6 +2689,8 @@ func TestVolcesVideoBodyAllowsOnlyTaskPayloadFields(t *testing.T) {
"tools": []any{map[string]any{"type": "web_search"}},
"task_id": "local-task-id",
"runMode": "simulation",
"modelType": "omni_video",
"model_type": "omni_video",
"fps": 24,
"content": []any{
map[string]any{"type": "text", "text": "Use <<<element_1>>> in a product reveal"},
@@ -2300,7 +2882,7 @@ func TestVolcesClientVideoResumePollsExistingTaskID(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "cgt-existing" || item["url"] != "https://example.com/resumed.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/resumed.mp4" {
t.Fatalf("unexpected resumed response: %+v", response.Result)
}
}
+614 -19
View File
@@ -1,8 +1,10 @@
package clients
import (
"bufio"
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"mime"
@@ -11,6 +13,8 @@ import (
"path"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type GeminiClient struct {
@@ -22,25 +26,73 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
if apiKey == "" {
return Response{}, &ClientError{Code: "missing_credentials", Message: "gemini api key is required", Retryable: false}
}
if geminiVeoRequest(request) {
return c.runVeo(ctx, request, apiKey)
}
body := geminiBody(request)
_, nativeContents := request.Body["contents"]
if !nativeContents && GeminiUsesOfficialAPI(request.Candidate) {
var err error
body, err = c.prepareOfficialGeminiImageFiles(ctx, request, apiKey, body)
if err != nil {
return Response{}, err
}
} else if !nativeContents {
var err error
body, err = prepareCompatibleGeminiInlineImages(ctx, body)
if err != nil {
return Response{}, err
}
}
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey), bytes.NewReader(raw))
endpoint := geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey)
if request.Stream {
endpoint = geminiStreamURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(raw))
if err != nil {
return Response{}, err
}
req.Header.Set("Content-Type", "application/json")
applyUpstreamIdempotency(req, request)
responseStartedAt := time.Now()
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return Response{}, transportClientError(err)
}
if err := notifyResponseReceived(request); err != nil {
return Response{}, err
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
var result map[string]any
var wire *WireResponse
if request.Stream && resp.StatusCode >= 200 && resp.StatusCode < 300 {
wire = &WireResponse{Protocol: ProtocolGeminiGenerateContent, StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
result, err = decodeGeminiStreamResponse(resp, request.StreamDelta, wire)
} else {
result, wire, err = decodeHTTPResponseForProtocol(resp, ProtocolGeminiGenerateContent)
}
responseFinishedAt := time.Now()
if err != nil {
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
}
output := geminiResult(request, result)
output, err := geminiResult(request, result)
if err != nil {
if requestID == "" {
requestID = firstNonEmptyString(result["responseId"], result["response_id"])
}
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
}
if wire != nil && geminiImageModelType(request.ModelType) {
// Gemini image responses can contain multi-megabyte Base64 fields. The
// decoded wire body is still used for native synchronous passthrough, so
// retaining the original JSON bytes would keep a second full copy alive
// until the client response finishes.
wire.RawJSON = nil
}
if requestID == "" {
requestID = requestIDFromResult(output)
}
@@ -52,10 +104,244 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
UpstreamProtocol: ProtocolGeminiGenerateContent,
Wire: wire,
}, nil
}
// GeminiUsesOfficialAPI distinguishes the native Google Gemini endpoint from
// Gemini-shaped compatibility services and Google's OpenAI compatibility path.
func GeminiUsesOfficialAPI(candidate store.RuntimeModelCandidate) bool {
base := strings.TrimSpace(candidate.BaseURL)
if base == "" {
return true
}
parsed, err := url.Parse(base)
if err != nil || !strings.EqualFold(parsed.Hostname(), "generativelanguage.googleapis.com") {
return false
}
pathValue := strings.ToLower(strings.TrimRight(parsed.Path, "/"))
return pathValue != "/openai" && !strings.HasSuffix(pathValue, "/openai")
}
func (c GeminiClient) prepareOfficialGeminiImageFiles(ctx context.Context, request Request, apiKey string, body map[string]any) (map[string]any, error) {
client := httpClient(request.HTTPClient, c.HTTPClient)
cache := map[string]map[string]any{}
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, request.Candidate, apiKey, body, cache)
if err != nil {
return nil, err
}
out, _ := prepared.(map[string]any)
if out == nil {
return body, nil
}
return out, nil
}
func prepareOfficialGeminiImageFileValue(ctx context.Context, client *http.Client, candidate store.RuntimeModelCandidate, apiKey string, value any, cache map[string]map[string]any) (any, error) {
switch typed := value.(type) {
case map[string]any:
next := make(map[string]any, len(typed))
for key, item := range typed {
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, candidate, apiKey, item, cache)
if err != nil {
return nil, err
}
next[key] = prepared
}
for _, key := range []string{"fileData", "file_data"} {
fileData, ok := next[key].(map[string]any)
if !ok {
continue
}
fileURI := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"])
mimeType := firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(fileURI))
if !requestAssetStringIsHTTPURL(fileURI) || !strings.HasPrefix(strings.ToLower(mimeType), "image/") || geminiOfficialFileURI(fileURI) {
continue
}
cacheKey := fileURI + "\x00" + mimeType
uploaded := cache[cacheKey]
if uploaded == nil {
var err error
uploaded, err = uploadOfficialGeminiImageFile(ctx, client, candidate, apiKey, fileURI, mimeType)
if err != nil {
fallback, fallbackErr := compatibleGeminiInlineImage(ctx, fileURI, mimeType)
if fallbackErr != nil {
return nil, err
}
delete(next, key)
next["inlineData"] = fallback
continue
}
cache[cacheKey] = uploaded
}
next[key] = cloneBody(uploaded)
}
return next, nil
case []any:
next := make([]any, 0, len(typed))
for _, item := range typed {
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, candidate, apiKey, item, cache)
if err != nil {
return nil, err
}
next = append(next, prepared)
}
return next, nil
default:
return value, nil
}
}
func prepareCompatibleGeminiInlineImages(ctx context.Context, body map[string]any) (map[string]any, error) {
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, body, map[string]map[string]any{})
if err != nil {
return nil, err
}
out, _ := prepared.(map[string]any)
if out == nil {
return body, nil
}
return out, nil
}
func prepareCompatibleGeminiInlineImageValue(ctx context.Context, value any, cache map[string]map[string]any) (any, error) {
switch typed := value.(type) {
case map[string]any:
next := make(map[string]any, len(typed))
for key, item := range typed {
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, item, cache)
if err != nil {
return nil, err
}
next[key] = prepared
}
for _, key := range []string{"fileData", "file_data"} {
fileData, ok := next[key].(map[string]any)
if !ok {
continue
}
fileURI := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"])
mimeType := firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(fileURI))
if !requestAssetStringIsHTTPURL(fileURI) || !strings.HasPrefix(strings.ToLower(mimeType), "image/") {
continue
}
cacheKey := fileURI + "\x00" + mimeType
inline := cache[cacheKey]
if inline == nil {
var err error
inline, err = compatibleGeminiInlineImage(ctx, fileURI, mimeType)
if err != nil {
return nil, err
}
cache[cacheKey] = inline
}
delete(next, key)
next["inlineData"] = cloneBody(inline)
}
return next, nil
case []any:
next := make([]any, 0, len(typed))
for _, item := range typed {
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, item, cache)
if err != nil {
return nil, err
}
next = append(next, prepared)
}
return next, nil
default:
return value, nil
}
}
func compatibleGeminiInlineImage(ctx context.Context, sourceURL string, declaredMimeType string) (map[string]any, error) {
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, sourceURL, 256<<20)
if err != nil {
return nil, err
}
return map[string]any{
"mimeType": geminiMediaMime(firstNonEmptyString(declaredMimeType, fetchedMimeType), "image"),
"data": base64.StdEncoding.EncodeToString(payload),
}, nil
}
func uploadOfficialGeminiImageFile(ctx context.Context, client *http.Client, candidate store.RuntimeModelCandidate, apiKey string, sourceURL string, declaredMimeType string) (map[string]any, error) {
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, sourceURL, 256<<20)
if err != nil {
return nil, err
}
mimeType := geminiMediaMime(firstNonEmptyString(declaredMimeType, fetchedMimeType), "image")
base := strings.TrimRight(strings.TrimSpace(candidate.BaseURL), "/")
if base == "" {
base = "https://generativelanguage.googleapis.com"
}
parsed, err := url.Parse(base)
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
return nil, &ClientError{Code: "invalid_parameter", Message: "invalid official Gemini base URL", Retryable: false}
}
uploadStartURL := parsed.Scheme + "://" + parsed.Host + "/upload/v1beta/files"
metadata, _ := json.Marshal(map[string]any{"file": map[string]any{"display_name": fmt.Sprintf("easyai-image-%d", time.Now().UnixMilli())}})
startRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadStartURL, bytes.NewReader(metadata))
if err != nil {
return nil, err
}
startRequest.Header.Set("x-goog-api-key", apiKey)
startRequest.Header.Set("X-Goog-Upload-Protocol", "resumable")
startRequest.Header.Set("X-Goog-Upload-Command", "start")
startRequest.Header.Set("X-Goog-Upload-Header-Content-Length", fmt.Sprintf("%d", len(payload)))
startRequest.Header.Set("X-Goog-Upload-Header-Content-Type", mimeType)
startRequest.Header.Set("Content-Type", "application/json")
startResponse, err := client.Do(startRequest)
if err != nil {
return nil, transportClientError(err)
}
if startResponse.StatusCode < http.StatusOK || startResponse.StatusCode >= http.StatusMultipleChoices {
_, responseErr := decodeHTTPResponse(startResponse)
return nil, responseErr
}
uploadURL := strings.TrimSpace(startResponse.Header.Get("X-Goog-Upload-URL"))
_ = startResponse.Body.Close()
if uploadURL == "" {
return nil, &ClientError{Code: "invalid_response", Message: "Gemini Files API did not return an upload URL", Retryable: false}
}
uploadRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadURL, bytes.NewReader(payload))
if err != nil {
return nil, err
}
uploadRequest.Header.Set("X-Goog-Upload-Offset", "0")
uploadRequest.Header.Set("X-Goog-Upload-Command", "upload, finalize")
uploadRequest.Header.Set("Content-Type", mimeType)
uploadResponse, err := client.Do(uploadRequest)
if err != nil {
return nil, transportClientError(err)
}
result, err := decodeHTTPResponse(uploadResponse)
if err != nil {
return nil, err
}
file := mapFromAny(result["file"])
fileURI := firstNonEmptyString(file["uri"], result["uri"])
if fileURI == "" {
return nil, &ClientError{Code: "invalid_response", Message: "Gemini Files API response did not include file.uri", Retryable: false}
}
return map[string]any{"fileUri": fileURI, "mimeType": firstNonEmptyString(file["mimeType"], file["mime_type"], mimeType)}, nil
}
func geminiOfficialFileURI(value string) bool {
parsed, err := url.Parse(strings.TrimSpace(value))
return err == nil && strings.EqualFold(parsed.Hostname(), "generativelanguage.googleapis.com") && strings.Contains(strings.ToLower(parsed.Path), "/files/")
}
func geminiURL(baseURL string, model string, apiKey string) string {
return geminiActionURL(baseURL, model, "generateContent", apiKey, false)
}
func geminiStreamURL(baseURL string, model string, apiKey string) string {
return geminiActionURL(baseURL, model, "streamGenerateContent", apiKey, true)
}
func geminiActionURL(baseURL string, model string, action string, apiKey string, stream bool) string {
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
if base == "" {
base = "https://generativelanguage.googleapis.com"
@@ -67,7 +353,53 @@ func geminiURL(baseURL string, model string, apiKey string) string {
base += "/v1beta"
}
escapedModel := url.PathEscape(model)
return fmt.Sprintf("%s/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
endpoint := fmt.Sprintf("%s/models/%s:%s?key=%s", base, escapedModel, action, url.QueryEscape(apiKey))
if stream {
endpoint += "&alt=sse"
}
return endpoint
}
func decodeGeminiStreamResponse(resp *http.Response, onDelta StreamDelta, wire *WireResponse) (map[string]any, error) {
defer resp.Body.Close()
scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), 16*1024*1024)
result := map[string]any{}
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if !strings.HasPrefix(line, "data:") {
continue
}
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "" || data == "[DONE]" {
continue
}
var event map[string]any
if err := json.Unmarshal([]byte(data), &event); err != nil {
return nil, &ClientError{Code: "invalid_response", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
}
result = event
if onDelta != nil {
if err := onDelta(StreamDeltaEvent{
Text: geminiText(event),
Event: event,
WireProtocol: ProtocolGeminiGenerateContent,
WireStatusCode: resp.StatusCode,
WireHeaders: compatibleResponseHeaders(resp.Header),
}); err != nil {
return nil, err
}
}
}
if err := scanner.Err(); err != nil {
clientErr := transportClientError(err)
clientErr.Wire = wire
return nil, clientErr
}
if len(result) == 0 {
return nil, &ClientError{Code: "invalid_response", Message: "gemini stream returned no events", StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
}
return result, nil
}
func geminiBody(request Request) map[string]any {
@@ -200,24 +532,39 @@ func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]
if out == nil {
out = map[string]any{}
}
if aspectRatio := firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"]); aspectRatio != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
if imageConfig != nil {
aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(imageConfig["aspectRatio"], imageConfig["aspect_ratio"]))
imageSize := normalizedProviderImageResolution(firstNonEmptyString(imageConfig["imageSize"], imageConfig["image_size"]))
delete(imageConfig, "aspectRatio")
delete(imageConfig, "aspect_ratio")
delete(imageConfig, "imageSize")
delete(imageConfig, "image_size")
if aspectRatio != "" {
imageConfig["aspectRatio"] = aspectRatio
}
if imageSize != "" {
imageConfig["imageSize"] = imageSize
}
}
if aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"])); aspectRatio != "" {
if imageConfig == nil {
imageConfig = map[string]any{}
}
imageConfig["aspectRatio"] = aspectRatio
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if imageSize := firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"]); imageSize != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
if imageSize := normalizedProviderImageResolution(firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"])); imageSize != "" {
if imageConfig == nil {
imageConfig = map[string]any{}
}
imageConfig["imageSize"] = imageSize
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if len(imageConfig) > 0 {
out["imageConfig"] = imageConfig
} else {
delete(out, "imageConfig")
}
delete(out, "image_config")
if count := intFromAny(firstPresent(body["n"], body["candidateCount"], body["candidate_count"])); count > 0 {
out["candidateCount"] = count
}
@@ -544,19 +891,24 @@ func geminiToolsFromOpenAITools(value any) []any {
return []any{map[string]any{"functionDeclarations": declarations}}
}
func geminiResult(request Request, raw map[string]any) map[string]any {
func geminiResult(request Request, raw map[string]any) (map[string]any, error) {
if geminiImageModelType(request.ModelType) {
data := geminiImageData(raw)
if len(data) == 0 {
data = []any{map[string]any{"url": "/static/provider/gemini-image-placeholder.png"}}
details := geminiMissingImageDetails(raw)
return nil, &ClientError{
Code: "gemini_image_output_missing",
Message: geminiMissingImageMessage(details),
Details: details,
Retryable: false,
}
}
return map[string]any{
"id": "gemini-image",
"created": nowUnix(),
"model": request.Model,
"data": data,
"raw": raw,
}
}, nil
}
message, finishReason := geminiChatMessage(raw)
return map[string]any{
@@ -570,8 +922,7 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
"message": message,
}},
"usage": geminiUsageMap(raw),
"raw": raw,
}
}, nil
}
func textFromMessages(body map[string]any) string {
@@ -686,6 +1037,250 @@ func geminiImageData(raw map[string]any) []any {
return out
}
func geminiMissingImageDetails(raw map[string]any) map[string]any {
details := map[string]any{
"provider": "gemini",
"resourceType": "image",
}
upstreamError, _ := raw["error"].(map[string]any)
upstreamErrorCode := stringFromPathValue(upstreamError["code"])
upstreamErrorStatus := strings.TrimSpace(stringFromAny(upstreamError["status"]))
upstreamErrorMessage := compactGeminiDiagnosticText(stringFromAny(upstreamError["message"]), 320)
if upstreamErrorCode != "" {
details["upstreamErrorCode"] = upstreamErrorCode
}
if upstreamErrorStatus != "" {
details["upstreamErrorStatus"] = upstreamErrorStatus
}
if upstreamErrorMessage != "" {
details["upstreamErrorMessage"] = upstreamErrorMessage
}
promptFeedback, _ := raw["promptFeedback"].(map[string]any)
if promptFeedback == nil {
promptFeedback, _ = raw["prompt_feedback"].(map[string]any)
}
promptBlockReason := firstNonEmptyString(promptFeedback["blockReason"], promptFeedback["block_reason"])
promptBlockMessage := compactGeminiDiagnosticText(firstNonEmptyString(
promptFeedback["blockReasonMessage"],
promptFeedback["block_reason_message"],
promptFeedback["message"],
), 320)
if promptBlockReason != "" {
details["promptBlockReason"] = promptBlockReason
}
if promptBlockMessage != "" {
details["promptBlockMessage"] = promptBlockMessage
}
candidates, _ := raw["candidates"].([]any)
details["candidateCount"] = len(candidates)
finishReasons := make([]string, 0)
safetyCategories := make([]string, 0)
partKinds := make([]string, 0)
providerMessages := make([]string, 0)
partCount := 0
emptyImagePayload := false
for _, rawCandidate := range candidates {
candidate, _ := rawCandidate.(map[string]any)
if finishReason := strings.TrimSpace(stringFromAny(firstPresent(candidate["finishReason"], candidate["finish_reason"]))); finishReason != "" {
finishReasons = appendGeminiDiagnosticValue(finishReasons, finishReason)
}
safetyRatings, _ := candidate["safetyRatings"].([]any)
if safetyRatings == nil {
safetyRatings, _ = candidate["safety_ratings"].([]any)
}
for _, rawRating := range safetyRatings {
rating, _ := rawRating.(map[string]any)
blocked, _ := firstPresent(rating["blocked"], rating["isBlocked"], rating["is_blocked"]).(bool)
if !blocked {
continue
}
if category := strings.TrimSpace(stringFromAny(rating["category"])); category != "" {
safetyCategories = appendGeminiDiagnosticValue(safetyCategories, category)
}
}
content, _ := candidate["content"].(map[string]any)
parts, _ := content["parts"].([]any)
partCount += len(parts)
for _, rawPart := range parts {
part, _ := rawPart.(map[string]any)
if text := compactGeminiDiagnosticText(stringFromAny(part["text"]), 320); text != "" {
partKinds = appendGeminiDiagnosticValue(partKinds, "text")
providerMessages = appendGeminiDiagnosticValue(providerMessages, text)
}
inline, inlinePresent := part["inlineData"].(map[string]any)
if !inlinePresent {
inline, inlinePresent = part["inline_data"].(map[string]any)
}
if inlinePresent {
partKinds = appendGeminiDiagnosticValue(partKinds, "inlineData")
if strings.TrimSpace(stringFromAny(inline["data"])) == "" {
emptyImagePayload = true
}
}
fileData, filePresent := part["fileData"].(map[string]any)
if !filePresent {
fileData, filePresent = part["file_data"].(map[string]any)
}
if filePresent {
partKinds = appendGeminiDiagnosticValue(partKinds, "fileData")
if firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]) == "" {
emptyImagePayload = true
}
}
for _, item := range []struct {
key string
kind string
}{
{key: "thought", kind: "thought"},
{key: "functionCall", kind: "functionCall"},
{key: "function_call", kind: "functionCall"},
{key: "executableCode", kind: "executableCode"},
{key: "executable_code", kind: "executableCode"},
{key: "codeExecutionResult", kind: "codeExecutionResult"},
{key: "code_execution_result", kind: "codeExecutionResult"},
} {
if _, ok := part[item.key]; ok {
partKinds = appendGeminiDiagnosticValue(partKinds, item.kind)
}
}
}
}
details["partCount"] = partCount
if len(finishReasons) > 0 {
details["candidateFinishReasons"] = finishReasons
}
if len(safetyCategories) > 0 {
details["blockedSafetyCategories"] = safetyCategories
}
if len(partKinds) > 0 {
details["observedPartKinds"] = partKinds
}
if len(providerMessages) > 0 {
details["providerMessages"] = providerMessages
}
reason := "no_supported_image_parts"
switch {
case upstreamErrorCode != "" || upstreamErrorStatus != "" || upstreamErrorMessage != "":
reason = "upstream_error"
case promptBlockReason != "":
reason = "prompt_blocked"
case len(safetyCategories) > 0 || geminiFinishReasonsIndicateFiltering(finishReasons):
reason = "candidate_filtered"
case len(candidates) == 0:
reason = "no_candidates"
case emptyImagePayload:
reason = "empty_image_payload"
case partCount == 0:
reason = "no_content_parts"
}
details["reason"] = reason
return details
}
func geminiMissingImageMessage(details map[string]any) string {
message := "Gemini image response contains no extractable image"
switch stringFromAny(details["reason"]) {
case "upstream_error":
message += ": upstream returned an error envelope"
case "prompt_blocked":
message += ": prompt blocked"
case "candidate_filtered":
message += ": candidate filtered"
case "no_candidates":
message += ": upstream response contains no candidates"
case "empty_image_payload":
message += ": image part is missing inline data or file URI"
case "no_content_parts":
message += ": candidate content contains no parts"
default:
message += ": candidate parts contain no supported image resource"
}
if reason := stringFromAny(details["promptBlockReason"]); reason != "" {
message += " (blockReason=" + reason + ")"
}
if code := stringFromAny(details["upstreamErrorCode"]); code != "" {
message += "; upstreamErrorCode=" + code
}
if status := stringFromAny(details["upstreamErrorStatus"]); status != "" {
message += "; upstreamErrorStatus=" + status
}
if values := geminiDiagnosticStrings(details["candidateFinishReasons"]); len(values) > 0 {
message += "; finishReason=" + strings.Join(values, ",")
}
if values := geminiDiagnosticStrings(details["blockedSafetyCategories"]); len(values) > 0 {
message += "; blockedSafetyCategories=" + strings.Join(values, ",")
}
if values := geminiDiagnosticStrings(details["providerMessages"]); len(values) > 0 {
message += "; providerMessage=" + values[0]
} else if upstreamMessage := stringFromAny(details["upstreamErrorMessage"]); upstreamMessage != "" {
message += "; providerMessage=" + upstreamMessage
} else if blockMessage := stringFromAny(details["promptBlockMessage"]); blockMessage != "" {
message += "; providerMessage=" + blockMessage
}
if values := geminiDiagnosticStrings(details["observedPartKinds"]); len(values) > 0 {
message += "; observedPartKinds=" + strings.Join(values, ",")
}
return compactGeminiDiagnosticText(message, 1800)
}
func geminiFinishReasonsIndicateFiltering(values []string) bool {
for _, value := range values {
normalized := strings.ToUpper(strings.TrimSpace(value))
switch normalized {
case "SAFETY", "IMAGE_SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII":
return true
}
}
return false
}
func geminiDiagnosticStrings(value any) []string {
switch typed := value.(type) {
case []string:
return typed
case []any:
values := make([]string, 0, len(typed))
for _, item := range typed {
if text := strings.TrimSpace(stringFromAny(item)); text != "" {
values = append(values, text)
}
}
return values
default:
return nil
}
}
func appendGeminiDiagnosticValue(values []string, value string) []string {
value = strings.TrimSpace(value)
if value == "" {
return values
}
for _, existing := range values {
if existing == value {
return values
}
}
if len(values) >= 8 {
return values
}
return append(values, value)
}
func compactGeminiDiagnosticText(value string, maxRunes int) string {
value = strings.Join(strings.Fields(value), " ")
if maxRunes <= 0 {
return ""
}
runes := []rune(value)
if len(runes) <= maxRunes {
return value
}
return string(runes[:maxRunes]) + "…"
}
func geminiUsage(raw map[string]any) Usage {
usageMap := geminiUsageMap(raw)
input := intFromAny(usageMap["prompt_tokens"])
+667
View File
@@ -0,0 +1,667 @@
package clients
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
)
const geminiVeoMaxVideoBytes int64 = 256 << 20
type geminiVeoImage struct {
URI string
MimeType string
}
func geminiVeoRequest(request Request) bool {
return request.Kind == "videos.generations" &&
strings.Contains(strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate))), "veo")
}
func (c GeminiClient) runVeo(ctx context.Context, request Request, apiKey string) (Response, error) {
startedAt := time.Now()
operationName := strings.TrimSpace(request.RemoteTaskID)
lastRequestID := operationName
var operation map[string]any
var wire *WireResponse
if operationName == "" {
body, err := geminiVeoBody(ctx, request)
if err != nil {
return Response{}, err
}
operation, lastRequestID, wire, err = c.geminiVeoPost(ctx, request, apiKey, body)
if err != nil {
return Response{}, annotateResponseError(err, lastRequestID, startedAt, time.Now())
}
operationName = strings.TrimSpace(stringFromAny(operation["name"]))
if err := validateGeminiVeoOperationName(operationName); err != nil {
return Response{}, err
}
if request.OnRemoteTaskSubmitted != nil {
if err := request.OnRemoteTaskSubmitted(operationName, geminiVeoOperationCheckpoint(operation)); err != nil {
return Response{}, err
}
}
} else if err := validateGeminiVeoOperationName(operationName); err != nil {
return Response{}, err
}
interval := durationFromConfig(request.Candidate.PlatformConfig, 10*time.Second, "pollIntervalMs", "poll_interval_ms")
timeout := durationFromConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
deadline := time.NewTimer(timeout)
defer deadline.Stop()
nextPoll := time.NewTimer(0)
defer nextPoll.Stop()
transientFailures := 0
for {
if operation != nil && boolFromAny(operation["done"]) {
return c.geminiVeoCompletedResponse(ctx, request, apiKey, operationName, operation, wire, lastRequestID, startedAt)
}
select {
case <-ctx.Done():
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: lastRequestID, Retryable: true}
case <-deadline.C:
return Response{}, &ClientError{
Code: "timeout",
Message: fmt.Sprintf("gemini Veo operation %s did not finish before timeout", operationName),
RequestID: lastRequestID,
Retryable: true,
}
case <-nextPoll.C:
pollStartedAt := time.Now()
result, requestID, pollWire, err := c.geminiVeoGetOperation(ctx, request, apiKey, operationName)
if requestID != "" {
lastRequestID = requestID
}
if err != nil {
err = annotateResponseError(err, lastRequestID, pollStartedAt, time.Now())
if !IsRetryable(err) {
return Response{}, err
}
transientFailures++
resetGeminiVeoPollTimer(nextPoll, geminiVeoRetryInterval(interval, transientFailures))
continue
}
transientFailures = 0
operation = result
wire = pollWire
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(operationName, geminiVeoOperationCheckpoint(operation)); err != nil {
return Response{}, err
}
}
if boolFromAny(operation["done"]) {
return c.geminiVeoCompletedResponse(ctx, request, apiKey, operationName, operation, wire, lastRequestID, startedAt)
}
resetGeminiVeoPollTimer(nextPoll, interval)
}
}
}
func geminiVeoBody(ctx context.Context, request Request) (map[string]any, error) {
prompt := firstNonEmptyPrompt(request.Body, "")
if prompt == "" {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo prompt is required", Param: "prompt", StatusCode: http.StatusBadRequest, Retryable: false}
}
firstFrame, lastFrame, references := geminiVeoImageInputs(request.Body)
instance := map[string]any{"prompt": prompt}
if firstFrame.URI != "" {
image, err := geminiVeoInlineImage(ctx, firstFrame)
if err != nil {
return nil, err
}
instance["image"] = image
}
if lastFrame.URI != "" {
image, err := geminiVeoInlineImage(ctx, lastFrame)
if err != nil {
return nil, err
}
instance["lastFrame"] = image
}
if len(references) > 3 {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo supports at most 3 reference images", Param: "reference_images", StatusCode: http.StatusBadRequest, Retryable: false}
}
if len(references) > 0 {
items := make([]any, 0, len(references))
for _, reference := range references {
image, err := geminiVeoInlineImage(ctx, reference)
if err != nil {
return nil, err
}
items = append(items, map[string]any{"image": image, "referenceType": "asset"})
}
instance["referenceImages"] = items
}
parameters, err := geminiVeoParameters(request.Body, lastFrame.URI != "", len(references) > 0)
if err != nil {
return nil, err
}
body := map[string]any{"instances": []any{instance}}
if len(parameters) > 0 {
body["parameters"] = parameters
}
return body, nil
}
func geminiVeoParameters(body map[string]any, hasLastFrame bool, hasReferences bool) (map[string]any, error) {
parameters := map[string]any{}
count := intFromAny(firstPresent(body["n"], body["numberOfVideos"], body["number_of_videos"]))
if count == 0 {
count = 1
}
if count != 1 {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo supports exactly 1 output video", Param: "n", StatusCode: http.StatusBadRequest, Retryable: false}
}
parameters["sampleCount"] = count
duration := intFromAny(firstPresent(body["duration"], body["duration_seconds"], body["durationSeconds"]))
if duration != 0 {
if duration != 4 && duration != 6 && duration != 8 {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo duration must be 4, 6, or 8 seconds", Param: "duration", StatusCode: http.StatusBadRequest, Retryable: false}
}
parameters["durationSeconds"] = duration
}
aspectRatio := strings.TrimSpace(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
if aspectRatio != "" {
if aspectRatio != "16:9" && aspectRatio != "9:16" {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo aspect ratio must be 16:9 or 9:16", Param: "aspect_ratio", StatusCode: http.StatusBadRequest, Retryable: false}
}
parameters["aspectRatio"] = aspectRatio
}
resolution := geminiVeoResolution(firstNonEmptyString(body["resolution"], body["size"]))
if resolution == "invalid" {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo resolution must be 720p, 1080p, or 4k", Param: "resolution", StatusCode: http.StatusBadRequest, Retryable: false}
}
if resolution != "" {
parameters["resolution"] = resolution
}
if (resolution == "1080p" || resolution == "4k" || hasLastFrame || hasReferences) && duration != 0 && duration != 8 {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo requires an 8-second duration for 1080p, 4k, last-frame, or reference-image generation", Param: "duration", StatusCode: http.StatusBadRequest, Retryable: false}
}
for _, item := range []struct {
to string
from []string
}{
{to: "personGeneration", from: []string{"personGeneration", "person_generation"}},
{to: "negativePrompt", from: []string{"negativePrompt", "negative_prompt"}},
} {
if value := strings.TrimSpace(firstNonEmptyStringValue(body, item.from...)); value != "" {
parameters[item.to] = value
}
}
for _, key := range []string{"enhancePrompt", "enhance_prompt"} {
if value, ok := body[key].(bool); ok {
parameters["enhancePrompt"] = value
break
}
}
return parameters, nil
}
func geminiVeoResolution(value string) string {
normalized := strings.ToLower(strings.TrimSpace(value))
switch normalized {
case "":
return ""
case "720p", "1080p", "4k":
return normalized
case "2160p", "3840x2160", "2160x3840":
return "4k"
default:
return "invalid"
}
}
func geminiVeoImageInputs(body map[string]any) (geminiVeoImage, geminiVeoImage, []geminiVeoImage) {
first := firstGeminiVeoImage(body["first_frame"], body["firstFrame"], body["first_frame_image"], body["firstFrameImage"])
last := firstGeminiVeoImage(body["last_frame"], body["lastFrame"], body["last_frame_image"], body["lastFrameImage"])
references := geminiVeoImagesFromValues(body["reference_images"], body["referenceImages"], body["reference_image"], body["referenceImage"])
images := geminiVeoImagesFromValues(body["image"], body["images"], body["image_url"], body["imageUrl"], body["image_urls"], body["imageUrls"])
if first.URI == "" && len(images) > 0 {
first = images[0]
images = images[1:]
}
references = append(references, images...)
for _, item := range contentItems(body["content"]) {
if !strings.Contains(strings.ToLower(strings.TrimSpace(stringFromAny(item["type"]))), "image") {
continue
}
image := firstGeminiVeoImage(item["image_url"], item["imageUrl"], item["url"], item["image"])
if image.URI == "" {
continue
}
if image.MimeType == "" {
image.MimeType = firstNonEmptyString(item["mime_type"], item["mimeType"])
}
switch strings.ToLower(strings.TrimSpace(stringFromAny(item["role"]))) {
case "first_frame", "firstframe":
if first.URI == "" {
first = image
}
case "last_frame", "lastframe":
if last.URI == "" {
last = image
}
default:
if first.URI == "" {
first = image
} else {
references = append(references, image)
}
}
}
return first, last, deduplicateGeminiVeoImages(references, first.URI, last.URI)
}
func firstGeminiVeoImage(values ...any) geminiVeoImage {
for _, value := range values {
if images := geminiVeoImagesFromAny(value); len(images) > 0 {
return images[0]
}
}
return geminiVeoImage{}
}
func geminiVeoImagesFromValues(values ...any) []geminiVeoImage {
out := make([]geminiVeoImage, 0)
for _, value := range values {
out = append(out, geminiVeoImagesFromAny(value)...)
}
return out
}
func geminiVeoImagesFromAny(value any) []geminiVeoImage {
switch typed := value.(type) {
case string:
if uri := strings.TrimSpace(typed); uri != "" {
return []geminiVeoImage{{URI: uri}}
}
case []any:
out := make([]geminiVeoImage, 0, len(typed))
for _, item := range typed {
out = append(out, geminiVeoImagesFromAny(item)...)
}
return out
case []string:
out := make([]geminiVeoImage, 0, len(typed))
for _, item := range typed {
out = append(out, geminiVeoImagesFromAny(item)...)
}
return out
case map[string]any:
if nested := firstPresent(typed["image_url"], typed["imageUrl"], typed["image"]); nested != nil {
if images := geminiVeoImagesFromAny(nested); len(images) > 0 {
if images[0].MimeType == "" {
images[0].MimeType = firstNonEmptyString(typed["mime_type"], typed["mimeType"])
}
return images
}
}
uri := firstNonEmptyString(typed["url"], typed["uri"], typed["data"], typed["bytesBase64Encoded"], typed["bytes_base64_encoded"])
if uri != "" {
return []geminiVeoImage{{URI: uri, MimeType: firstNonEmptyString(typed["mime_type"], typed["mimeType"])}}
}
}
return nil
}
func deduplicateGeminiVeoImages(images []geminiVeoImage, exclusions ...string) []geminiVeoImage {
seen := map[string]bool{}
for _, exclusion := range exclusions {
if exclusion = strings.TrimSpace(exclusion); exclusion != "" {
seen[exclusion] = true
}
}
out := make([]geminiVeoImage, 0, len(images))
for _, image := range images {
image.URI = strings.TrimSpace(image.URI)
if image.URI == "" || seen[image.URI] {
continue
}
seen[image.URI] = true
out = append(out, image)
}
return out
}
func geminiVeoInlineImage(ctx context.Context, image geminiVeoImage) (map[string]any, error) {
parsed := geminiDataURL(image.URI)
mimeType := strings.TrimSpace(image.MimeType)
data := ""
if parsed != nil {
data = parsed.data
mimeType = firstNonEmptyString(mimeType, parsed.mimeType)
} else if requestAssetStringIsHTTPURL(image.URI) {
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, image.URI, 256<<20)
if err != nil {
return nil, err
}
data = base64.StdEncoding.EncodeToString(payload)
mimeType = firstNonEmptyString(mimeType, fetchedMimeType)
} else if !requestLikeURL(image.URI) {
data = strings.TrimSpace(image.URI)
}
if data == "" {
return nil, &ClientError{
Code: "invalid_parameter",
Message: "gemini Veo image input must be hydrated as base64 data",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
if _, err := base64.StdEncoding.DecodeString(data); err != nil {
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo image input contains invalid base64 data", Param: "image", StatusCode: http.StatusBadRequest, Retryable: false}
}
if mimeType == "" {
mimeType = "image/png"
}
return map[string]any{"bytesBase64Encoded": data, "mimeType": geminiMediaMime(mimeType, "image")}, nil
}
func requestAssetStringIsHTTPURL(value string) bool {
parsed, err := url.Parse(strings.TrimSpace(value))
return err == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https"))
}
func requestLikeURL(value string) bool {
normalized := strings.ToLower(strings.TrimSpace(value))
return strings.HasPrefix(normalized, "http://") || strings.HasPrefix(normalized, "https://") || strings.HasPrefix(normalized, "/static/")
}
func (c GeminiClient) geminiVeoPost(ctx context.Context, request Request, apiKey string, body map[string]any) (map[string]any, string, *WireResponse, error) {
raw, _ := json.Marshal(body)
endpoint := geminiVeoActionURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate))
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(raw))
if err != nil {
return nil, "", nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-goog-api-key", apiKey)
applyUpstreamIdempotency(req, request)
if err := notifySubmissionStarted(request); err != nil {
return nil, "", nil, err
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", nil, transportClientError(err)
}
if err := notifyResponseReceived(request); err != nil {
resp.Body.Close()
return nil, "", nil, err
}
requestID := requestIDFromHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolGeminiVeo)
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
return result, requestID, wire, notifyErr
}
return result, requestID, wire, err
}
func (c GeminiClient) geminiVeoGetOperation(ctx context.Context, request Request, apiKey string, operationName string) (map[string]any, string, *WireResponse, error) {
endpoint, err := geminiVeoOperationURL(request.Candidate.BaseURL, operationName)
if err != nil {
return nil, "", nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
if err != nil {
return nil, "", nil, err
}
req.Header.Set("x-goog-api-key", apiKey)
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", nil, transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolGeminiVeo)
return result, requestID, wire, err
}
func (c GeminiClient) geminiVeoCompletedResponse(ctx context.Context, request Request, apiKey string, operationName string, operation map[string]any, wire *WireResponse, requestID string, startedAt time.Time) (Response, error) {
finishedAt := time.Now()
if operationError := mapFromAny(operation["error"]); len(operationError) > 0 {
code := firstNonEmptyString(operationError["status"], operationError["code"])
if code == "" {
code = "gemini_veo_failed"
}
message := strings.TrimSpace(stringFromAny(operationError["message"]))
if message == "" {
message = "gemini Veo operation failed"
}
return Response{}, &ClientError{
Code: strings.ToLower(code),
Message: message,
RequestID: requestID,
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
Retryable: false,
}
}
items, err := c.geminiVeoResultItems(ctx, request, apiKey, operation)
finishedAt = time.Now()
if err != nil {
return Response{}, annotateResponseError(err, requestID, startedAt, finishedAt)
}
result := map[string]any{
"id": operationName,
"object": "video.generation",
"created": nowUnix(),
"model": request.Model,
"upstream_task_id": operationName,
"data": items,
}
return Response{
Result: result,
RequestID: firstNonEmpty(requestID, operationName),
Progress: providerProgress(request),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
UpstreamProtocol: ProtocolGeminiVeo,
Wire: wire,
}, nil
}
func (c GeminiClient) geminiVeoResultItems(ctx context.Context, request Request, apiKey string, operation map[string]any) ([]any, error) {
response := mapFromAny(operation["response"])
generateResponse := mapFromAny(firstPresent(response["generateVideoResponse"], response["generate_video_response"]))
samples := mapListFromAny(firstPresent(generateResponse["generatedSamples"], generateResponse["generated_samples"], response["generatedVideos"], response["generated_videos"]))
if len(samples) == 0 {
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo operation returned no generated video", Retryable: false}
}
items := make([]any, 0, len(samples))
for _, sample := range samples {
video := mapFromAny(sample["video"])
if len(video) == 0 {
video = sample
}
mimeType := firstNonEmptyString(video["mimeType"], video["mime_type"], sample["mimeType"], sample["mime_type"])
if mimeType == "" {
mimeType = "video/mp4"
}
var payload []byte
if encoded := firstNonEmptyString(video["videoBytes"], video["video_bytes"], video["bytesBase64Encoded"], video["bytes_base64_encoded"]); encoded != "" {
decoded, err := base64.StdEncoding.DecodeString(encoded)
if err != nil {
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo returned invalid video base64 data", Retryable: false}
}
payload = decoded
} else if uri := strings.TrimSpace(firstNonEmptyString(video["uri"], video["url"])); uri != "" {
var err error
payload, mimeType, err = c.geminiVeoDownload(ctx, request, apiKey, uri, mimeType)
if err != nil {
return nil, err
}
}
if len(payload) == 0 {
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo generated video payload is missing", Retryable: false}
}
items = append(items, map[string]any{"type": "video", "video_bytes": payload, "mime_type": mimeType})
}
return items, nil
}
func (c GeminiClient) geminiVeoDownload(ctx context.Context, request Request, apiKey string, rawURI string, fallbackMimeType string) ([]byte, string, error) {
downloadURL, trustedHost, err := geminiVeoDownloadURL(request.Candidate.BaseURL, rawURI)
if err != nil {
return nil, "", err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
if err != nil {
return nil, "", err
}
if trustedHost {
req.Header.Set("x-goog-api-key", apiKey)
}
client := httpClient(request.HTTPClient, c.HTTPClient)
redirectClient := *client
originalRedirect := client.CheckRedirect
redirectClient.CheckRedirect = func(next *http.Request, via []*http.Request) error {
if originalRedirect != nil {
if err := originalRedirect(next, via); err != nil {
return err
}
}
if len(via) > 0 && !strings.EqualFold(next.URL.Hostname(), via[0].URL.Hostname()) {
next.Header.Del("x-goog-api-key")
}
return nil
}
resp, err := redirectClient.Do(req)
if err != nil {
return nil, "", transportClientError(err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
return nil, "", &ClientError{Code: statusCodeName(resp.StatusCode), Message: errorMessage(raw, resp.Status), StatusCode: resp.StatusCode, RequestID: requestIDFromHTTPResponse(resp), Retryable: HTTPRetryable(resp.StatusCode)}
}
payload, err := io.ReadAll(io.LimitReader(resp.Body, geminiVeoMaxVideoBytes+1))
if err != nil {
if timeoutErr := transportClientError(err); timeoutErr.Code == "timeout" {
timeoutErr.StatusCode = resp.StatusCode
timeoutErr.RequestID = requestIDFromHTTPResponse(resp)
return nil, "", timeoutErr
}
return nil, "", &ClientError{Code: "response_read_error", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: true}
}
if int64(len(payload)) > geminiVeoMaxVideoBytes {
return nil, "", &ClientError{Code: "response_too_large", Message: fmt.Sprintf("gemini Veo video exceeds %d bytes", geminiVeoMaxVideoBytes), StatusCode: resp.StatusCode, Retryable: false}
}
mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
if mimeType == "" || mimeType == "application/octet-stream" {
mimeType = fallbackMimeType
}
if mimeType == "" {
mimeType = "video/mp4"
}
return payload, mimeType, nil
}
func geminiVeoVersionedBaseURL(baseURL string) string {
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
if base == "" {
base = "https://generativelanguage.googleapis.com"
}
base = strings.TrimSuffix(base, "/openai")
if !strings.HasSuffix(base, "/v1") && !strings.HasSuffix(base, "/v1beta") && !strings.HasSuffix(base, "/v1alpha") {
base += "/v1beta"
}
return base
}
func geminiVeoActionURL(baseURL string, model string) string {
return fmt.Sprintf("%s/models/%s:predictLongRunning", geminiVeoVersionedBaseURL(baseURL), url.PathEscape(strings.TrimSpace(model)))
}
func geminiVeoOperationURL(baseURL string, operationName string) (string, error) {
if err := validateGeminiVeoOperationName(operationName); err != nil {
return "", err
}
segments := strings.Split(strings.Trim(operationName, "/"), "/")
for index, segment := range segments {
segments[index] = url.PathEscape(segment)
}
return geminiVeoVersionedBaseURL(baseURL) + "/" + strings.Join(segments, "/"), nil
}
func validateGeminiVeoOperationName(operationName string) error {
name := strings.TrimSpace(operationName)
if name == "" || strings.Contains(name, "..") || strings.ContainsAny(name, "?#\\") || strings.HasPrefix(name, "/") || strings.Contains(name, "://") {
return &ClientError{Code: "invalid_response", Message: "gemini Veo operation name is invalid", Retryable: false}
}
if !strings.Contains(name, "/operations/") && !strings.HasPrefix(name, "operations/") {
return &ClientError{Code: "invalid_response", Message: "gemini Veo operation name is invalid", Retryable: false}
}
return nil
}
func geminiVeoDownloadURL(baseURL string, rawURI string) (string, bool, error) {
base, err := url.Parse(geminiVeoVersionedBaseURL(baseURL))
if err != nil {
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo base URL is invalid", Retryable: false}
}
parsed, err := url.Parse(strings.TrimSpace(rawURI))
if err != nil {
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo video URI is invalid", Retryable: false}
}
if !parsed.IsAbs() {
parsed = base.ResolveReference(parsed)
}
if parsed.Scheme != "https" && parsed.Scheme != "http" {
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo video URI scheme is unsupported", Retryable: false}
}
hostname := strings.ToLower(strings.TrimSpace(parsed.Hostname()))
trusted := strings.EqualFold(hostname, base.Hostname()) || hostname == "generativelanguage.googleapis.com"
return parsed.String(), trusted, nil
}
func geminiVeoOperationCheckpoint(operation map[string]any) map[string]any {
checkpoint := map[string]any{"done": boolFromAny(operation["done"])}
if metadata := mapFromAny(operation["metadata"]); len(metadata) > 0 {
checkpoint["metadata"] = metadata
}
if operationError := mapFromAny(operation["error"]); len(operationError) > 0 {
checkpoint["error"] = operationError
}
return checkpoint
}
func resetGeminiVeoPollTimer(timer *time.Timer, duration time.Duration) {
if !timer.Stop() {
select {
case <-timer.C:
default:
}
}
timer.Reset(duration)
}
func geminiVeoRetryInterval(interval time.Duration, failures int) time.Duration {
if failures <= 0 {
return interval
}
multiplier := 1 << min(failures-1, 4)
result := interval * time.Duration(multiplier)
if result > 30*time.Second {
return 30 * time.Second
}
return result
}
@@ -0,0 +1,383 @@
package clients
import (
"context"
"encoding/base64"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"os"
"reflect"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestGeminiVeoRunUsesOfficialLongRunningProtocol(t *testing.T) {
const operationName = "models/veo-3.1-generate-preview/operations/test-operation"
videoPayload := []byte("test mp4 payload")
var server *httptest.Server
var pollCount atomic.Int32
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("x-goog-api-key") != "test-gemini-key" {
t.Errorf("missing Gemini API key header for %s", r.URL.Path)
}
switch {
case r.Method == http.MethodPost && r.URL.Path == "/v1beta/models/veo-3.1-generate-preview:predictLongRunning":
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode Veo body: %v", err)
}
instances := body["instances"].([]any)
if prompt := instances[0].(map[string]any)["prompt"]; prompt != "a paper boat on a lake" {
t.Errorf("unexpected prompt: %v", prompt)
}
parameters := body["parameters"].(map[string]any)
if parameters["durationSeconds"] != float64(4) || parameters["resolution"] != "720p" || parameters["aspectRatio"] != "16:9" {
t.Errorf("unexpected parameters: %+v", parameters)
}
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName})
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/"+operationName:
if pollCount.Add(1) == 1 {
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName, "done": false})
return
}
_ = json.NewEncoder(w).Encode(map[string]any{
"name": operationName,
"done": true,
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
map[string]any{"video": map[string]any{"uri": server.URL + "/v1beta/files/generated-video:download", "mimeType": "video/mp4"}},
}}},
})
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/files/generated-video:download":
w.Header().Set("Content-Type", "video/mp4")
_, _ = w.Write(videoPayload)
default:
http.NotFound(w, r)
}
}))
defer server.Close()
var submitted string
polled := 0
response, err := (GeminiClient{}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "veo-3.1",
Body: map[string]any{
"prompt": "a paper boat on a lake",
"duration": 4,
"resolution": "720p",
"aspect_ratio": "16:9",
},
Candidate: testGeminiVeoCandidate(server.URL),
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
submitted = remoteTaskID
return nil
},
OnRemoteTaskPolled: func(remoteTaskID string, _ map[string]any) error {
if remoteTaskID != operationName {
t.Errorf("unexpected polled operation: %s", remoteTaskID)
}
polled++
return nil
},
})
if err != nil {
t.Fatalf("run Gemini Veo: %v", err)
}
if submitted != operationName || polled != 2 || pollCount.Load() != 2 {
t.Fatalf("unexpected async callbacks: submitted=%q polled=%d requests=%d", submitted, polled, pollCount.Load())
}
if response.UpstreamProtocol != ProtocolGeminiVeo || response.RequestID != operationName {
t.Fatalf("unexpected response metadata: %+v", response)
}
data := response.Result["data"].([]any)
item := data[0].(map[string]any)
if got := item["video_bytes"].([]byte); !reflect.DeepEqual(got, videoPayload) {
t.Fatalf("unexpected video payload: %q", got)
}
if _, exposed := item["uri"]; exposed {
t.Fatalf("upstream authenticated URI must not be exposed: %+v", item)
}
}
func TestGeminiVeoRunResumesOperationWithoutSubmittingAgain(t *testing.T) {
const operationName = "models/veo-3.1-generate-preview/operations/resumed"
var postCount atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodPost {
postCount.Add(1)
http.Error(w, "unexpected submit", http.StatusInternalServerError)
return
}
_ = json.NewEncoder(w).Encode(map[string]any{
"name": operationName,
"done": true,
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
map[string]any{"video": map[string]any{"videoBytes": base64.StdEncoding.EncodeToString([]byte("resumed video")), "mimeType": "video/mp4"}},
}}},
})
}))
defer server.Close()
response, err := (GeminiClient{}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "veo-3.1",
Body: map[string]any{"prompt": "ignored on resume"},
Candidate: testGeminiVeoCandidate(server.URL),
RemoteTaskID: operationName,
})
if err != nil {
t.Fatalf("resume Gemini Veo: %v", err)
}
if postCount.Load() != 0 || response.Result["id"] != operationName {
t.Fatalf("resume submitted a new operation or returned wrong ID: posts=%d result=%+v", postCount.Load(), response.Result)
}
}
func TestGeminiVeoBodyMapsImageAndParameters(t *testing.T) {
first := "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("first"))
last := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("last"))
reference := "data:image/webp;base64," + base64.StdEncoding.EncodeToString([]byte("reference"))
body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
"prompt": "animate",
"first_frame": first,
"last_frame": last,
"reference_images": []any{reference},
"duration": 8,
"resolution": "2160p",
"aspect_ratio": "9:16",
"negative_prompt": "text",
"person_generation": "allow_adult",
"enhance_prompt": false,
}})
if err != nil {
t.Fatalf("build Gemini Veo body: %v", err)
}
instance := body["instances"].([]any)[0].(map[string]any)
image := instance["image"].(map[string]any)
lastFrame := instance["lastFrame"].(map[string]any)
references := instance["referenceImages"].([]any)
if image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString([]byte("first")) || image["mimeType"] != "image/png" {
t.Fatalf("unexpected first frame: %+v", image)
}
if lastFrame["mimeType"] != "image/jpeg" || len(references) != 1 {
t.Fatalf("unexpected last/reference images: last=%+v references=%+v", lastFrame, references)
}
parameters := body["parameters"].(map[string]any)
if parameters["resolution"] != "4k" || parameters["durationSeconds"] != 8 || parameters["enhancePrompt"] != false {
t.Fatalf("unexpected Gemini Veo parameters: %+v", parameters)
}
}
func TestGeminiVeoDownloadsImageURLWhenBuildingOfficialPayload(t *testing.T) {
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
fetchCount := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
fetchCount++
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(payload)
}))
defer server.Close()
body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
"prompt": "animate",
"first_frame": server.URL + "/input.png",
"duration": 8,
}})
if err != nil {
t.Fatalf("build Gemini Veo body from URL: %v", err)
}
instance := body["instances"].([]any)[0].(map[string]any)
image := instance["image"].(map[string]any)
if fetchCount != 1 || image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString(payload) || image["mimeType"] != "image/png" {
t.Fatalf("Veo URL should be fetched during official payload construction: fetches=%d image=%+v", fetchCount, image)
}
}
func TestGeminiOfficialAPIClassification(t *testing.T) {
tests := []struct {
baseURL string
want bool
}{
{baseURL: "", want: true},
{baseURL: "https://generativelanguage.googleapis.com", want: true},
{baseURL: "https://generativelanguage.googleapis.com/v1beta", want: true},
{baseURL: "https://generativelanguage.googleapis.com/v1beta/openai", want: false},
{baseURL: "https://gemini-compatible.example.com/v1beta", want: false},
}
for _, test := range tests {
if got := GeminiUsesOfficialAPI(store.RuntimeModelCandidate{BaseURL: test.baseURL}); got != test.want {
t.Fatalf("GeminiUsesOfficialAPI(%q) = %t, want %t", test.baseURL, got, test.want)
}
}
}
func TestGeminiCompatibleConvertsFileDataURLToInlineBase64(t *testing.T) {
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(payload)
}))
defer server.Close()
body := map[string]any{"contents": []any{map[string]any{"parts": []any{map[string]any{
"fileData": map[string]any{"fileUri": server.URL + "/source.png", "mimeType": "image/png"},
}}}}}
prepared, err := prepareCompatibleGeminiInlineImages(context.Background(), body)
if err != nil {
t.Fatalf("prepare compatible Gemini inline image: %v", err)
}
contents := prepared["contents"].([]any)
parts := contents[0].(map[string]any)["parts"].([]any)
part := parts[0].(map[string]any)
inline := part["inlineData"].(map[string]any)
if _, exists := part["fileData"]; exists || inline["mimeType"] != "image/png" || inline["data"] != base64.StdEncoding.EncodeToString(payload) {
t.Fatalf("compatible Gemini should use inlineData: %+v", part)
}
}
func TestGeminiOfficialImageUploadUsesFilesAPI(t *testing.T) {
imagePayload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(imagePayload)
}))
defer imageServer.Close()
var uploadServer *httptest.Server
startCount := 0
uploadCount := 0
uploadServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/upload/v1beta/files":
startCount++
if r.Header.Get("x-goog-api-key") != "gemini-key" || r.Header.Get("X-Goog-Upload-Protocol") != "resumable" || r.Header.Get("X-Goog-Upload-Header-Content-Type") != "image/png" {
t.Fatalf("unexpected Gemini Files start headers: %+v", r.Header)
}
w.Header().Set("X-Goog-Upload-URL", uploadServer.URL+"/upload/session")
w.WriteHeader(http.StatusOK)
case "/upload/session":
uploadCount++
raw, err := io.ReadAll(r.Body)
if err != nil {
t.Fatalf("read Gemini Files upload: %v", err)
}
if string(raw) != string(imagePayload) || r.Header.Get("X-Goog-Upload-Command") != "upload, finalize" {
t.Fatalf("unexpected Gemini Files upload body=%v headers=%+v", raw, r.Header)
}
_ = json.NewEncoder(w).Encode(map[string]any{"file": map[string]any{
"uri": "https://generativelanguage.googleapis.com/v1beta/files/easyai-image",
"mimeType": "image/png",
}})
default:
http.NotFound(w, r)
}
}))
defer uploadServer.Close()
fileData, err := uploadOfficialGeminiImageFile(context.Background(), uploadServer.Client(), store.RuntimeModelCandidate{BaseURL: uploadServer.URL + "/v1beta"}, "gemini-key", imageServer.URL+"/source.png", "image/png")
if err != nil {
t.Fatalf("upload official Gemini image: %v", err)
}
if startCount != 1 || uploadCount != 1 || fileData["fileUri"] != "https://generativelanguage.googleapis.com/v1beta/files/easyai-image" || fileData["mimeType"] != "image/png" {
t.Fatalf("unexpected Gemini Files result: starts=%d uploads=%d fileData=%+v", startCount, uploadCount, fileData)
}
}
func TestGeminiVeoRejectsInvalidDurationBeforeSubmission(t *testing.T) {
_, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{"prompt": "test", "duration": 5}})
if err == nil || ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "duration" || !strings.Contains(err.Error(), "4, 6, or 8") {
t.Fatalf("unexpected invalid duration error: %v", err)
}
}
func TestGeminiVeoDownloadURLOnlyTrustsProviderAndGoogleAPIs(t *testing.T) {
for _, test := range []struct {
uri string
trusted bool
}{
{uri: "https://gemini-proxy.example.com/v1beta/files/video", trusted: true},
{uri: "https://generativelanguage.googleapis.com/v1beta/files/video", trusted: true},
{uri: "https://storage.googleapis.com/signed/video", trusted: false},
{uri: "https://attacker.example.com/video", trusted: false},
} {
_, trusted, err := geminiVeoDownloadURL("https://gemini-proxy.example.com/v1beta", test.uri)
if err != nil {
t.Fatalf("parse %s: %v", test.uri, err)
}
if trusted != test.trusted {
t.Fatalf("unexpected trust for %s: got %t want %t", test.uri, trusted, test.trusted)
}
}
}
func TestGeminiVeoLive(t *testing.T) {
apiKey := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_API_KEY"))
if apiKey == "" {
t.Skip("GEMINI_VEO_LIVE_API_KEY is not configured")
}
baseURL := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_BASE_URL"))
model := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_MODEL"))
if model == "" {
model = "veo-3.1-generate-preview"
}
ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute)
defer cancel()
operationName := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_OPERATION"))
response, err := (GeminiClient{}).Run(ctx, Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: model,
Body: map[string]any{
"prompt": "A small red paper boat gently floating on calm water, static camera, natural daylight.",
"duration": 4,
"resolution": "720p",
"aspect_ratio": "16:9",
},
Candidate: store.RuntimeModelCandidate{
Provider: "gemini",
BaseURL: baseURL,
ProviderModelName: model,
Credentials: map[string]any{"apiKey": apiKey},
PlatformConfig: map[string]any{"pollIntervalMs": 10000, "pollTimeoutMs": 660000},
},
RemoteTaskID: operationName,
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
t.Logf("live Gemini Veo submitted operation %s", remoteTaskID)
return nil
},
})
if err != nil {
t.Fatalf("live Gemini Veo request failed: %v", err)
}
data, _ := response.Result["data"].([]any)
if len(data) != 1 {
t.Fatalf("live Gemini Veo returned %d videos", len(data))
}
item, _ := data[0].(map[string]any)
payload, _ := item["video_bytes"].([]byte)
if len(payload) < 1024 {
t.Fatalf("live Gemini Veo video is unexpectedly small: %d bytes", len(payload))
}
if mimeType := strings.TrimSpace(stringFromAny(item["mime_type"])); !strings.HasPrefix(mimeType, "video/") {
t.Fatalf("live Gemini Veo returned unexpected MIME type: %q", mimeType)
}
t.Logf("live Gemini Veo accepted operation %s and downloaded %d video bytes", response.RequestID, len(payload))
}
func testGeminiVeoCandidate(baseURL string) store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{
Provider: "gemini",
BaseURL: baseURL,
ProviderModelName: "veo-3.1-generate-preview",
Credentials: map[string]any{"apiKey": "test-gemini-key"},
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
}
}
+312 -41
View File
@@ -13,6 +13,11 @@ import (
"time"
)
const (
defaultMaxJSONResponseBytes int64 = 16 << 20
mediaMaxJSONResponseBytes int64 = 128 << 20
)
func credential(candidate map[string]any, keys ...string) string {
for _, key := range keys {
if value, ok := candidate[key].(string); ok && strings.TrimSpace(value) != "" {
@@ -44,25 +49,140 @@ func intValue(body map[string]any, key string, fallback int) int {
}
func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
result, _, err := decodeHTTPResponseForProtocol(resp, "")
return result, err
}
func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[string]any, *WireResponse, error) {
defer resp.Body.Close()
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
maxBytes := maxJSONResponseBytesForProtocol(protocol)
raw, readErr := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1))
tooLarge := int64(len(raw)) > maxBytes
if tooLarge {
raw = raw[:maxBytes]
}
wire := &WireResponse{
Protocol: strings.TrimSpace(protocol),
StatusCode: resp.StatusCode,
Headers: compatibleResponseHeaders(resp.Header),
// raw is already owned by this response. Reusing it avoids retaining a
// second copy of media responses that may be as large as 128 MiB.
RawJSON: raw,
}
if readErr != nil {
if len(raw) > 0 {
_ = json.Unmarshal(raw, &wire.Body)
}
if timeoutErr := transportClientError(readErr); timeoutErr.Code == "timeout" {
timeoutErr.StatusCode = resp.StatusCode
timeoutErr.RequestID = requestIDFromHTTPResponse(resp)
timeoutErr.Wire = wire
return nil, wire, timeoutErr
}
return nil, wire, &ClientError{
Code: "response_read_error",
Message: readErr.Error(),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: true,
Wire: wire,
}
}
if tooLarge {
if len(raw) > 0 {
_ = json.Unmarshal(raw, &wire.Body)
}
return nil, wire, &ClientError{
Code: "response_too_large",
Message: fmt.Sprintf("upstream JSON response exceeds %d bytes", maxBytes),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: false,
Wire: wire,
}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, &ClientError{
if len(raw) > 0 {
_ = json.Unmarshal(raw, &wire.Body)
}
return nil, wire, &ClientError{
Code: statusCodeName(resp.StatusCode),
Message: errorMessage(raw, resp.Status),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: HTTPRetryable(resp.StatusCode),
Wire: wire,
}
}
var out map[string]any
if len(raw) == 0 {
return map[string]any{}, nil
wire.Body = map[string]any{}
return map[string]any{}, wire, nil
}
if err := json.Unmarshal(raw, &out); err != nil {
return nil, &ClientError{Code: "invalid_response", Message: err.Error(), Retryable: false}
return nil, wire, &ClientError{Code: "invalid_response", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
}
return out, nil
// Result and wire metadata are read-only until the runner replaces the
// canonical result, so they can share the decoded object instead of keeping
// another full copy of a large Base64 media response.
wire.Body = out
return out, wire, nil
}
func maxJSONResponseBytesForProtocol(protocol string) int64 {
switch strings.TrimSpace(protocol) {
case ProtocolGeminiGenerateContent, ProtocolOpenAIImages:
return mediaMaxJSONResponseBytes
default:
return defaultMaxJSONResponseBytes
}
}
func compatibleResponseHeaders(source http.Header) map[string][]string {
if len(source) == 0 {
return nil
}
allowed := map[string]bool{
"content-type": true, "retry-after": true,
"x-request-id": true, "request-id": true,
"x-ratelimit-limit-requests": true, "x-ratelimit-remaining-requests": true,
"x-ratelimit-reset-requests": true, "x-ratelimit-limit-tokens": true,
"x-ratelimit-remaining-tokens": true, "x-ratelimit-reset-tokens": true,
"x-goog-request-id": true, "x-goog-upload-status": true,
"x-goog-upload-url": true, "x-goog-upload-size-received": true,
}
out := map[string][]string{}
for key, values := range source {
if !allowed[strings.ToLower(strings.TrimSpace(key))] {
continue
}
out[http.CanonicalHeaderKey(key)] = append([]string(nil), values...)
}
if len(out) == 0 {
return nil
}
return out
}
func notifySubmissionStarted(request Request) error {
if request.OnUpstreamSubmissionStarted == nil {
return nil
}
return request.OnUpstreamSubmissionStarted()
}
func notifyResponseReceived(request Request) error {
if request.OnUpstreamResponseReceived == nil {
return nil
}
return request.OnUpstreamResponseReceived()
}
func notifyWireResponse(request Request, wire *WireResponse) error {
if wire == nil || request.OnUpstreamWireResponse == nil {
return nil
}
return request.OnUpstreamWireResponse(wire)
}
func decodeOpenAIStreamResponse(resp *http.Response, onDelta StreamDelta) (map[string]any, error) {
@@ -88,6 +208,9 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
rawLines := make([]string, 0)
parts := make([]string, 0)
refusalParts := make([]string, 0)
streamLogprobs := make([]any, 0)
streamAnnotations := make([]any, 0)
reasoningParts := make([]string, 0)
var last map[string]any
var usage Usage
@@ -118,6 +241,11 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
if reasoningText != "" {
reasoningParts = append(reasoningParts, reasoningText)
}
if refusal := streamEventRefusal(event); refusal != "" {
refusalParts = append(refusalParts, refusal)
}
streamLogprobs = append(streamLogprobs, streamEventLogprobs(event)...)
streamAnnotations = append(streamAnnotations, streamEventAnnotations(event)...)
aggregateStreamToolCalls(event, toolCalls)
if reason := streamEventFinishReason(event); reason != "" {
finishReason = reason
@@ -145,7 +273,7 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
}
return out, true, nil
}
return buildOpenAIStreamResult(last, parts, reasoningParts, toolCalls, finishReason, usage), true, nil
return buildOpenAIStreamResult(last, parts, refusalParts, streamLogprobs, streamAnnotations, reasoningParts, toolCalls, finishReason, usage), true, nil
}
func decodeOpenAIStream(raw []byte) (map[string]any, bool) {
@@ -156,8 +284,8 @@ func decodeOpenAIStream(raw []byte) (map[string]any, bool) {
return result, ok && err == nil
}
func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts []string, toolCalls map[int]map[string]any, finishReason string, usage Usage) map[string]any {
if len(parts) == 0 && len(reasoningParts) == 0 && len(toolCalls) == 0 {
func buildOpenAIStreamResult(last map[string]any, parts []string, refusalParts []string, streamLogprobs []any, streamAnnotations []any, reasoningParts []string, toolCalls map[int]map[string]any, finishReason string, usage Usage) map[string]any {
if len(parts) == 0 && len(refusalParts) == 0 && len(streamLogprobs) == 0 && len(streamAnnotations) == 0 && len(reasoningParts) == 0 && len(toolCalls) == 0 {
return last
}
message := map[string]any{
@@ -167,6 +295,12 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
if len(reasoningParts) > 0 {
message["reasoning_content"] = strings.Join(reasoningParts, "")
}
if len(refusalParts) > 0 {
message["refusal"] = strings.Join(refusalParts, "")
}
if len(streamAnnotations) > 0 {
message["annotations"] = streamAnnotations
}
if len(toolCalls) > 0 {
message["tool_calls"] = sortedStreamToolCalls(toolCalls)
}
@@ -174,15 +308,17 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
finishReason = "stop"
}
var out map[string]any
choice := map[string]any{
"index": 0, "message": message, "finish_reason": finishReason,
}
if len(streamLogprobs) > 0 {
choice["logprobs"] = map[string]any{"content": streamLogprobs, "refusal": nil}
}
out = map[string]any{
"id": stringFromAny(firstPresent(last["id"], "chatcmpl-stream")),
"object": "chat.completion",
"model": stringFromAny(last["model"]),
"choices": []any{map[string]any{
"index": 0,
"message": message,
"finish_reason": finishReason,
}},
"id": stringFromAny(firstPresent(last["id"], "chatcmpl-stream")),
"object": "chat.completion",
"model": stringFromAny(last["model"]),
"choices": []any{choice},
}
if usage.TotalTokens > 0 {
usageMap := map[string]any{
@@ -200,6 +336,42 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
return out
}
func streamEventRefusal(event map[string]any) string {
choices, _ := event["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
delta, _ := choice["delta"].(map[string]any)
if refusal := stringFromAny(delta["refusal"]); refusal != "" {
return refusal
}
}
return ""
}
func streamEventLogprobs(event map[string]any) []any {
choices, _ := event["choices"].([]any)
out := make([]any, 0)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
logprobs, _ := choice["logprobs"].(map[string]any)
content, _ := logprobs["content"].([]any)
out = append(out, content...)
}
return out
}
func streamEventAnnotations(event map[string]any) []any {
choices, _ := event["choices"].([]any)
out := make([]any, 0)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
delta, _ := choice["delta"].(map[string]any)
annotations, _ := delta["annotations"].([]any)
out = append(out, annotations...)
}
return out
}
// NormalizeChatCompletionRequestBody 将后续请求里的工具调用上下文还原为
// OpenAI Chat Completions 标准格式,便于再次发送给 OpenAI-compatible 上游。
func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
@@ -219,7 +391,7 @@ func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
continue
}
copied := cloneMapAny(message)
normalizeToolCallsContainer(copied, false)
normalizeRequestToolCallsContainer(copied)
normalizeToolMessageFields(copied)
toolMessages, cleanContent, changed := toolResultMessagesFromContent(copied["content"])
if changed {
@@ -241,6 +413,67 @@ func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
return out
}
// normalizeRequestToolCallsContainer preserves canonical OpenAI Chat
// structures. Only known provider aliases are normalized. Legacy
// function_call and standard custom tool calls must remain intact.
func normalizeRequestToolCallsContainer(container map[string]any) {
if container == nil {
return
}
toolCalls := canonicalToolCalls(container["tool_calls"], false)
for _, key := range []string{"tool_call", "toolCall", "toolCalls"} {
if raw, ok := container[key]; ok {
for _, normalized := range normalizeRawToolCalls(raw, len(toolCalls), false) {
toolCalls = append(toolCalls, normalized)
}
delete(container, key)
}
}
if contentToolCalls, cleanContent, changed := toolCallsFromContent(container["content"], len(toolCalls), false); changed {
toolCalls = append(toolCalls, contentToolCalls...)
setNormalizedContent(container, cleanContent, false)
}
if partToolCalls := toolCallsFromParts(container["parts"], len(toolCalls), false); len(partToolCalls) > 0 {
toolCalls = append(toolCalls, partToolCalls...)
delete(container, "parts")
}
if len(toolCalls) > 0 {
container["tool_calls"] = toolCalls
}
if functionCall, ok := container["functionCall"]; ok {
if _, canonical := container["function_call"]; !canonical {
if normalized := normalizeToolCall(functionCall, 0, false); normalized != nil {
container["function_call"] = normalized["function"]
}
}
delete(container, "functionCall")
}
}
func canonicalToolCalls(value any, stream bool) []any {
items, ok := value.([]any)
if !ok {
if value == nil {
return nil
}
items = []any{value}
}
out := make([]any, 0, len(items))
for index, raw := range items {
toolCall, _ := raw.(map[string]any)
typeName := stringFromAny(toolCall["type"])
if (typeName == "custom" && len(mapFromAny(toolCall["custom"])) > 0) ||
((typeName == "" || typeName == "function") && len(mapFromAny(toolCall["function"])) > 0) {
out = append(out, cloneMapAny(toolCall))
continue
}
if normalized := normalizeToolCall(raw, index, stream); normalized != nil {
out = append(out, normalized)
}
}
return out
}
func cloneMapAny(source map[string]any) map[string]any {
if source == nil {
return nil
@@ -332,14 +565,27 @@ func normalizeToolCallsContainer(container map[string]any, stream bool) {
if container == nil {
return
}
toolCalls := make([]any, 0)
for _, rawToolCall := range rawToolCallValues(container) {
for _, normalized := range normalizeRawToolCalls(rawToolCall, len(toolCalls), stream) {
toolCalls = append(toolCalls, normalized)
toolCalls := canonicalToolCalls(container["tool_calls"], stream)
for _, key := range []string{"tool_call", "toolCall", "toolCalls"} {
if rawToolCall, ok := container[key]; ok {
for _, normalized := range normalizeRawToolCalls(rawToolCall, len(toolCalls), stream) {
toolCalls = append(toolCalls, normalized)
}
delete(container, key)
}
}
if functionCall, ok := container["functionCall"]; ok {
if _, canonical := container["function_call"]; !canonical {
if normalized := normalizeToolCall(functionCall, 0, stream); normalized != nil {
container["function_call"] = normalized["function"]
}
}
delete(container, "functionCall")
}
if contentToolCalls, cleanContent, changed := toolCallsFromContent(container["content"], len(toolCalls), stream); changed {
toolCalls = append(toolCalls, contentToolCalls...)
for _, normalized := range contentToolCalls {
toolCalls = append(toolCalls, normalized)
}
setNormalizedContent(container, cleanContent, stream)
}
if partToolCalls := toolCallsFromParts(container["parts"], len(toolCalls), stream); len(partToolCalls) > 0 {
@@ -349,9 +595,6 @@ func normalizeToolCallsContainer(container map[string]any, stream bool) {
if len(toolCalls) > 0 {
container["tool_calls"] = toolCalls
}
for _, key := range []string{"tool_call", "toolCall", "toolCalls", "function_call", "functionCall"} {
delete(container, key)
}
}
func normalizeToolMessageFields(message map[string]any) {
@@ -443,6 +686,20 @@ func normalizeToolCall(value any, index int, stream bool) map[string]any {
if len(source) == 0 {
return nil
}
if stringFromAny(source["type"]) == "custom" {
customSource := mapFromAny(source["custom"])
if len(customSource) == 0 {
return nil
}
toolCall := cloneMapAny(source)
toolCall["custom"] = cloneMapAny(customSource)
if stream {
if _, ok := toolCall["index"]; !ok {
toolCall["index"] = index
}
}
return toolCall
}
functionSource := mapFromAny(source["function"])
if len(functionSource) == 0 {
functionSource = mapFromAny(firstPresent(source["function_call"], source["functionCall"]))
@@ -862,7 +1119,7 @@ func aggregateStreamToolCalls(event map[string]any, toolCalls map[int]map[string
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
delta, _ := choice["delta"].(map[string]any)
rawToolCalls, _ := delta["tool_calls"].([]any)
rawToolCalls := streamToolCallsFromDelta(delta)
for _, rawToolCall := range rawToolCalls {
incoming, _ := rawToolCall.(map[string]any)
index := intFromAny(incoming["index"])
@@ -876,25 +1133,39 @@ func aggregateStreamToolCalls(event map[string]any, toolCalls map[int]map[string
current[key] = value
}
}
incomingFn, _ := incoming["function"].(map[string]any)
if len(incomingFn) == 0 {
continue
}
currentFn, _ := current["function"].(map[string]any)
if currentFn == nil {
currentFn = map[string]any{}
current["function"] = currentFn
}
if name, ok := incomingFn["name"].(string); ok && name != "" {
currentFn["name"] = stringFromAny(currentFn["name"]) + name
}
if arguments, ok := incomingFn["arguments"].(string); ok && arguments != "" {
currentFn["arguments"] = stringFromAny(currentFn["arguments"]) + arguments
}
aggregateStreamToolPayload(current, incoming, "function", "arguments")
aggregateStreamToolPayload(current, incoming, "custom", "input")
}
}
}
func streamToolCallsFromDelta(delta map[string]any) []any {
rawToolCalls, _ := delta["tool_calls"].([]any)
out := append([]any(nil), rawToolCalls...)
if functionCall, ok := delta["function_call"].(map[string]any); ok {
out = append(out, map[string]any{"index": 0, "type": "function", "function": functionCall})
}
return out
}
func aggregateStreamToolPayload(current map[string]any, incoming map[string]any, containerKey string, payloadKey string) {
incomingPayload, _ := incoming[containerKey].(map[string]any)
if len(incomingPayload) == 0 {
return
}
currentPayload, _ := current[containerKey].(map[string]any)
if currentPayload == nil {
currentPayload = map[string]any{}
current[containerKey] = currentPayload
}
if name, ok := incomingPayload["name"].(string); ok && name != "" {
currentPayload["name"] = stringFromAny(currentPayload["name"]) + name
}
if payload, ok := incomingPayload[payloadKey].(string); ok && payload != "" {
currentPayload[payloadKey] = stringFromAny(currentPayload[payloadKey]) + payload
}
}
func sortedStreamToolCalls(toolCalls map[int]map[string]any) []any {
indices := make([]int, 0, len(toolCalls))
for index := range toolCalls {
+51 -20
View File
@@ -80,7 +80,18 @@ func (c KelingClient) runVideo(ctx context.Context, request Request, token strin
}()
if upstreamTaskID == "" {
submitResult, requestID, err := c.postJSON(ctx, request, prepared.Endpoint, token, prepared.Payload)
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
submitResult, requestID, submitWire, err := c.postJSON(ctx, request, prepared.Endpoint, token, prepared.Payload)
if notifyErr := notifyWireResponse(request, submitWire); notifyErr != nil {
return Response{}, notifyErr
}
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
return Response{}, notifyErr
}
}
submitRequestID = requestID
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
@@ -122,6 +133,11 @@ func (c KelingClient) runVideo(ctx context.Context, request Request, token strin
return Response{}, annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
}
lastResult = pollResult
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(upstreamTaskID, pollResult); err != nil {
return Response{}, err
}
}
switch kelingTaskStatus(pollResult) {
case "succeed":
@@ -176,7 +192,18 @@ func (c KelingClient) runTaskAPIVideo(ctx context.Context, request Request, toke
if err != nil {
return Response{}, err
}
submitResult, requestID, err := c.postJSONAt(ctx, request, taskAPIBaseURL, endpoint, token, payload)
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
submitResult, requestID, submitWire, err := c.postJSONAt(ctx, request, taskAPIBaseURL, endpoint, token, payload, ProtocolKlingV2Omni)
if notifyErr := notifyWireResponse(request, submitWire); notifyErr != nil {
return Response{}, notifyErr
}
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
return Response{}, notifyErr
}
}
submitRequestID = requestID
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
@@ -224,6 +251,11 @@ func (c KelingClient) runTaskAPIVideo(ctx context.Context, request Request, toke
if err != nil {
return Response{}, annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
}
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(upstreamTaskID, pollResult); err != nil {
return Response{}, err
}
}
task := kelingTaskAPITask(pollResult, upstreamTaskID)
lastStatus = strings.ToLower(strings.TrimSpace(stringFromAny(task["status"])))
@@ -567,31 +599,31 @@ func (c KelingClient) kelingOmniElementList(ctx context.Context, request Request
return elements, createdIDs, nil
}
func (c KelingClient) postJSON(ctx context.Context, request Request, path string, token string, body map[string]any) (map[string]any, string, error) {
return c.postJSONAt(ctx, request, request.Candidate.BaseURL, path, token, body)
func (c KelingClient) postJSON(ctx context.Context, request Request, path string, token string, body map[string]any) (map[string]any, string, *WireResponse, error) {
return c.postJSONAt(ctx, request, request.Candidate.BaseURL, path, token, body, ProtocolKlingV1Omni)
}
func (c KelingClient) postJSONAt(ctx context.Context, request Request, baseURL string, path string, token string, body map[string]any) (map[string]any, string, error) {
func (c KelingClient) postJSONAt(ctx context.Context, request Request, baseURL string, path string, token string, body map[string]any, protocol string) (map[string]any, string, *WireResponse, error) {
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(baseURL, path), bytes.NewReader(raw))
if err != nil {
return nil, "", err
return nil, "", nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+token)
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", nil, transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, protocol)
if err != nil {
return result, requestID, err
return result, requestID, wire, err
}
if code := intFromAny(result["code"]); code != 0 {
return result, requestID, &ClientError{Code: kelingEnvelopeErrorCode(result), Message: kelingEnvelopeErrorMessage(result), RequestID: firstNonEmpty(requestID, stringFromAny(result["request_id"])), Retryable: false}
return result, requestID, wire, &ClientError{Code: kelingEnvelopeErrorCode(result), Message: kelingEnvelopeErrorMessage(result), RequestID: firstNonEmpty(requestID, stringFromAny(result["request_id"])), Retryable: false, Wire: wire}
}
return result, firstNonEmpty(requestID, stringFromAny(result["request_id"])), nil
return result, firstNonEmpty(requestID, stringFromAny(result["request_id"])), wire, nil
}
func (c KelingClient) getJSON(ctx context.Context, request Request, path string, token string) (map[string]any, string, error) {
@@ -606,7 +638,7 @@ func (c KelingClient) getJSONAt(ctx context.Context, request Request, baseURL st
req.Header.Set("Authorization", "Bearer "+token)
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
@@ -629,7 +661,7 @@ func (c KelingClient) createKelingElement(ctx context.Context, request Request,
req.Header.Set("Authorization", "Bearer "+token)
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return "", transportClientError(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
@@ -663,7 +695,7 @@ func (c KelingClient) cleanupKelingElements(ctx context.Context, request Request
if id == "" {
continue
}
_, _, _ = c.postJSON(ctx, request, "/general/delete-elements", token, map[string]any{"element_id": id})
_, _, _, _ = c.postJSON(ctx, request, "/general/delete-elements", token, map[string]any{"element_id": id})
}
return nil
}
@@ -713,7 +745,7 @@ func kelingImageToBase64(ctx context.Context, request Request, value string) (st
}
resp, err := httpClient(request.HTTPClient).Do(req)
if err != nil {
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return "", transportClientError(err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
@@ -722,7 +754,7 @@ func kelingImageToBase64(ctx context.Context, request Request, value string) (st
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
if err != nil {
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return "", transportClientError(err)
}
return base64.StdEncoding.EncodeToString(raw), nil
}
@@ -1274,7 +1306,6 @@ func kelingVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
@@ -1310,7 +1341,6 @@ func kelingTaskAPIVideoSuccessResult(request Request, upstreamTaskID string, tas
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
@@ -1343,9 +1373,10 @@ func kelingPollInterval(request Request) time.Duration {
}
func kelingPollTimeout(request Request) time.Duration {
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["kelingPollTimeoutSeconds"], request.Candidate.PlatformConfig["klingPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), 600)
fallbackSeconds := ProviderRequestTimeout(request.Kind).Seconds()
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["kelingPollTimeoutSeconds"], request.Candidate.PlatformConfig["klingPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), fallbackSeconds)
if seconds < 1 {
seconds = 600
seconds = fallbackSeconds
}
return time.Duration(seconds) * time.Second
}
+18 -18
View File
@@ -326,17 +326,17 @@ func (c MinimaxClient) runSpeech(ctx context.Context, request Request) (Response
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: "minimax speech audio hex is invalid: " + err.Error(), RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)), ResponseStartedAt: startedAt, ResponseFinishedAt: finishedAt, ResponseDurationMS: responseDurationMS(startedAt, finishedAt), Retryable: false}
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["raw_data"] = cloneMapAny(result)
normalized["data"] = []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}}
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}},
}
return Response{
Result: normalized,
RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)),
@@ -378,12 +378,12 @@ func (c MinimaxClient) runVoiceClone(ctx context.Context, request Request) (Resp
if isProviderTaskFailure(providerTaskSpec{Name: "minimax"}, result) {
return Response{}, providerTaskFailure(providerTaskSpec{Name: "minimax"}, result, firstNonEmptyString(requestID, uploadRequestID, requestIDFromResult(result)), startedAt)
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["voice_id"] = stringFromAny(payload["voice_id"])
normalized["raw_data"] = cloneMapAny(result)
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"voice_id": stringFromAny(payload["voice_id"]),
}
if demoAudio := firstNonEmptyString(valueAtPath(result, "demo_audio"), valueAtPath(result, "data.demo_audio")); demoAudio != "" {
normalized["demo_audio"] = demoAudio
normalized["data"] = []any{map[string]any{"type": "audio", "url": demoAudio}}
@@ -646,7 +646,7 @@ func providerPostMultipartFile(ctx context.Context, client *http.Client, url str
applyProviderAuth(req, credentials, auth)
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
@@ -0,0 +1,232 @@
package clients
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestGeminiGenerationConfigDropsDefaultRatioTokensAndNormalizesKSize(t *testing.T) {
config := geminiGenerationConfig(map[string]any{
"aspect_ratio": "auto",
"size": "2k",
"generationConfig": map[string]any{
"imageConfig": map[string]any{
"aspectRatio": "adaptive",
"imageSize": "2k",
},
},
}, true)
imageConfig := mapFromAny(config["imageConfig"])
if _, ok := imageConfig["aspectRatio"]; ok {
t.Fatalf("non-ratio tokens must not be sent to Gemini: %+v", config)
}
if imageConfig["imageSize"] != "2K" {
t.Fatalf("Gemini imageSize should use canonical K format: %+v", config)
}
}
func TestNormalizeOfficialOpenAIImageSizeUsesModelSpecificWireFormat(t *testing.T) {
legacy := map[string]any{
"model": "gpt-image-1",
"size": "2K",
"width": 2048,
"height": 1152,
"aspect_ratio": "16:9",
}
normalizeOfficialOpenAIImageSize(legacy, map[string]any{
"size": "2K",
"aspect_ratio": "16:9",
})
if legacy["size"] != "1536x1024" {
t.Fatalf("legacy GPT Image should use an official orientation size: %+v", legacy)
}
flexible := map[string]any{
"model": "gpt-image-2",
"size": "2K",
"resolution": "2k",
"aspect_ratio": "16:9",
}
normalizeOfficialOpenAIImageSize(flexible, map[string]any{
"size": "2K",
"aspect_ratio": "16:9",
})
if flexible["size"] != "2048x1152" {
t.Fatalf("GPT Image 2 should receive flexible pixel dimensions: %+v", flexible)
}
rounded := map[string]any{
"model": "gpt-image-2",
"size": "2K",
"resolution": "2K",
"aspect_ratio": "3:2",
}
normalizeOfficialOpenAIImageSize(rounded, map[string]any{
"size": "2K",
"aspect_ratio": "3:2",
})
if rounded["size"] != "2048x1360" {
t.Fatalf("GPT Image 2 dimensions should satisfy the official 16-pixel multiple: %+v", rounded)
}
explicit := map[string]any{
"model": "gpt-image-2",
"size": "1232x768",
"width": 1232,
"height": 768,
}
normalizeOfficialOpenAIImageSize(explicit, map[string]any{
"size": "1234x777",
})
if explicit["size"] != "1234x777" {
t.Fatalf("an explicit WxH size must be sent unchanged to OpenAI: %+v", explicit)
}
explicitLegacy := map[string]any{
"model": "gpt-image-1",
"size": "2048x1152",
}
normalizeOfficialOpenAIImageSize(explicitLegacy, map[string]any{
"size": "2048x1152",
})
if explicitLegacy["size"] != "2048x1152" {
t.Fatalf("explicit WxH must also win for legacy OpenAI image models: %+v", explicitLegacy)
}
}
func TestOpenAIClientUsesOriginalWxHAndConvertsKResolution(t *testing.T) {
received := make([]map[string]any, 0, 3)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode OpenAI image request: %v", err)
}
received = append(received, body)
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
candidate := store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "openai-key"},
}
requests := []Request{
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "explicit dimensions",
"size": "1232x768",
"width": 1232,
"height": 768,
},
OriginalBody: map[string]any{
"size": "1234x777",
},
Candidate: candidate,
},
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "resolution conversion",
"size": "2048x1365",
"resolution": "2K",
"aspect_ratio": "3:2",
"width": 2048,
"height": 1365,
},
OriginalBody: map[string]any{
"size": "2K",
"aspect_ratio": "3:2",
},
Candidate: candidate,
},
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "resolution field conversion",
"size": "1024x1024",
"resolution": "1K",
"aspect_ratio": "1:1",
"width": 1024,
"height": 1024,
},
OriginalBody: map[string]any{
"resolution": "1k",
"aspect_ratio": "1:1",
},
Candidate: candidate,
},
}
for _, request := range requests {
if _, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request); err != nil {
t.Fatalf("run OpenAI image request: %v", err)
}
}
if len(received) != 3 {
t.Fatalf("unexpected request count: %d", len(received))
}
if received[0]["size"] != "1234x777" {
t.Fatalf("explicit WxH should win over preprocessed dimensions: %+v", received[0])
}
if received[1]["size"] != "2048x1360" {
t.Fatalf("2K with 3:2 should use normalized OpenAI dimensions: %+v", received[1])
}
if received[2]["size"] != "1024x1024" {
t.Fatalf("resolution=1K should be converted to OpenAI dimensions: %+v", received[2])
}
for _, body := range received {
for _, key := range []string{"resolution", "aspect_ratio", "width", "height"} {
if _, ok := body[key]; ok {
t.Fatalf("generic geometry field %q must not reach OpenAI: %+v", key, body)
}
}
}
}
func TestVolcesImageBodySelectsConfiguredSizeFormat(t *testing.T) {
body := volcesImageBody(Request{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"model": "Seedream-5.0-Pro",
"resolution": "2k",
"size": "2048x1152",
"width": 2048,
"height": 1152,
"aspect_ratio": "16:9",
"platformId": "gateway-internal-platform",
},
Candidate: store.RuntimeModelCandidate{
ProviderModelName: "doubao-seedream-5-0-pro-260628",
ModelType: "image_generate",
Capabilities: map[string]any{
"image_generate": map[string]any{
"size_param_format": "resolution",
},
},
},
})
if body["size"] != "2K" {
t.Fatalf("resolution-only Volces model should receive canonical K size: %+v", body)
}
for _, key := range []string{"resolution", "width", "height", "aspect_ratio", "platformId"} {
if _, ok := body[key]; ok {
t.Fatalf("generic geometry field %q must not reach Volces: %+v", key, body)
}
}
}
@@ -0,0 +1,251 @@
package clients
import (
"math"
"strconv"
"strings"
)
const (
openAIGPTImage2MaxEdge = 3840
openAIGPTImage2MinPixels = 655360
openAIGPTImage2MaxPixels = 8294400
openAIGPTImage2MaxRatio = 3
openAIGPTImage2Multiple = 16
)
type providerImageDimensions struct {
width int
height int
}
func normalizedProviderAspectRatio(value any) string {
compact := strings.Map(func(r rune) rune {
switch r {
case ' ', '\t', '\r', '\n':
return -1
default:
return r
}
}, strings.TrimSpace(stringFromAny(value)))
parts := strings.Split(compact, ":")
if len(parts) != 2 {
return ""
}
width, widthErr := strconv.ParseFloat(parts[0], 64)
height, heightErr := strconv.ParseFloat(parts[1], 64)
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
return ""
}
return parts[0] + ":" + parts[1]
}
func normalizedProviderImageResolution(value any) string {
text := strings.ToLower(strings.TrimSpace(stringFromAny(value)))
if len(text) < 2 || text[len(text)-1] != 'k' {
return ""
}
magnitude, err := strconv.Atoi(text[:len(text)-1])
if err != nil || magnitude <= 0 {
return ""
}
return strconv.Itoa(magnitude) + "K"
}
func normalizeOfficialOpenAIImageSize(body map[string]any, originalBody map[string]any) {
if originalBody == nil {
originalBody = body
}
if width, height, ok := providerPixelDimensions(stringFromAny(originalBody["size"])); ok {
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
return
}
model := strings.ToLower(strings.TrimSpace(stringFromAny(body["model"])))
resolution := normalizedProviderImageResolution(firstNonEmptyString(
originalBody["resolution"],
originalBody["size"],
body["resolution"],
body["size"],
))
if resolution != "" {
width, height, ok := providerBodyPixelDimensions(body)
if !ok {
width, height, ok = providerResolutionDimensions(
resolution,
firstNonEmptyString(originalBody["aspect_ratio"], originalBody["aspectRatio"], originalBody["ratio"]),
)
}
if !ok {
delete(body, "size")
return
}
switch {
case strings.HasPrefix(model, "gpt-image-2"):
dimensions := normalizeOpenAIGPTImage2Dimensions(width, height)
body["size"] = strconv.Itoa(dimensions.width) + "x" + strconv.Itoa(dimensions.height)
case strings.HasPrefix(model, "gpt-image-"):
body["size"] = officialOpenAILegacyImageSize(width, height)
default:
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
return
}
if !strings.HasPrefix(model, "gpt-image-") {
return
}
size := strings.ToLower(strings.TrimSpace(stringFromAny(body["size"])))
if size == "auto" {
body["size"] = "auto"
return
}
if strings.HasPrefix(model, "gpt-image-2") {
if width, height, ok := providerPixelDimensions(size); ok {
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
return
}
switch size {
case "1024x1024", "1536x1024", "1024x1536":
body["size"] = size
return
}
width, height, ok := requestedProviderDimensions(body)
if !ok {
body["size"] = "auto"
return
}
body["size"] = officialOpenAILegacyImageSize(width, height)
}
func officialOpenAILegacyImageSize(width int, height int) string {
switch {
case width > height:
return "1536x1024"
case height > width:
return "1024x1536"
default:
return "1024x1024"
}
}
func providerBodyPixelDimensions(body map[string]any) (int, int, bool) {
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
return width, height, true
}
width := intFromAny(body["width"])
height := intFromAny(body["height"])
return width, height, width > 0 && height > 0
}
func providerResolutionDimensions(resolution string, aspectRatio any) (int, int, bool) {
resolution = normalizedProviderImageResolution(resolution)
if resolution == "" {
return 0, 0, false
}
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
if err != nil || magnitude <= 0 {
return 0, 0, false
}
baseSide := magnitude * 1024
ratio := providerAspectRatioNumber(aspectRatio)
width, height := baseSide, baseSide
if ratio > 1 {
height = max(1, int(math.Round(float64(baseSide)/ratio)))
} else if ratio > 0 && ratio < 1 {
width = max(1, int(math.Round(float64(baseSide)*ratio)))
}
return width, height, true
}
func requestedProviderDimensions(body map[string]any) (int, int, bool) {
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
return width, height, true
}
width := intFromAny(body["width"])
height := intFromAny(body["height"])
if width > 0 && height > 0 {
return width, height, true
}
ratio := providerAspectRatioNumber(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
if ratio > 1 {
return 2, 1, true
}
if ratio > 0 && ratio < 1 {
return 1, 2, true
}
if ratio == 1 {
return 1, 1, true
}
return 0, 0, false
}
func providerPixelDimensions(value string) (int, int, bool) {
parts := strings.Split(strings.ToLower(strings.TrimSpace(value)), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
return 0, 0, false
}
return width, height, true
}
func providerAspectRatioNumber(value any) float64 {
normalized := normalizedProviderAspectRatio(value)
if normalized == "" {
return 0
}
parts := strings.Split(normalized, ":")
width, _ := strconv.ParseFloat(parts[0], 64)
height, _ := strconv.ParseFloat(parts[1], 64)
return width / height
}
func normalizeOpenAIGPTImage2Dimensions(width int, height int) providerImageDimensions {
if width <= 0 || height <= 0 {
return providerImageDimensions{}
}
rawRatio := float64(width) / float64(height)
ratio := rawRatio
if ratio < 1 {
ratio = 1 / ratio
}
ratio = min(ratio, float64(openAIGPTImage2MaxRatio))
landscape := rawRatio >= 1
longEdge := min(max(width, height), openAIGPTImage2MaxEdge)
dimensions := buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
for dimensions.width*dimensions.height > openAIGPTImage2MaxPixels {
longEdge = max(openAIGPTImage2Multiple, max(dimensions.width, dimensions.height)-openAIGPTImage2Multiple)
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
}
for dimensions.width*dimensions.height < openAIGPTImage2MinPixels && max(dimensions.width, dimensions.height) < openAIGPTImage2MaxEdge {
longEdge = min(openAIGPTImage2MaxEdge, max(dimensions.width, dimensions.height)+openAIGPTImage2Multiple)
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, true)
}
return dimensions
}
func buildOpenAIGPTImage2Dimensions(longEdge int, ratio float64, landscape bool, roundUp bool) providerImageDimensions {
normalizedLongEdge := roundOpenAIGPTImage2Multiple(float64(longEdge), roundUp)
normalizedShortEdge := roundOpenAIGPTImage2Multiple(float64(normalizedLongEdge)/ratio, roundUp)
if float64(normalizedLongEdge)/float64(normalizedShortEdge) > openAIGPTImage2MaxRatio {
normalizedShortEdge += openAIGPTImage2Multiple
}
if landscape {
return providerImageDimensions{width: normalizedLongEdge, height: normalizedShortEdge}
}
return providerImageDimensions{width: normalizedShortEdge, height: normalizedLongEdge}
}
func roundOpenAIGPTImage2Multiple(value float64, roundUp bool) int {
rounded := math.Round(value / openAIGPTImage2Multiple)
if roundUp {
rounded = math.Ceil(value / openAIGPTImage2Multiple)
}
return max(openAIGPTImage2Multiple, int(rounded)*openAIGPTImage2Multiple)
}
@@ -0,0 +1,82 @@
package clients
import (
"context"
"io"
"net/http"
"testing"
"time"
)
type timeoutResponseBody struct {
read bool
}
func (body *timeoutResponseBody) Read(buffer []byte) (int, error) {
if body.read {
return 0, context.DeadlineExceeded
}
body.read = true
return copy(buffer, []byte(`{"partial":`)), nil
}
func (*timeoutResponseBody) Close() error { return nil }
func TestProviderRequestTimeoutUsesMediaDefaults(t *testing.T) {
for _, test := range []struct {
kind string
want time.Duration
}{
{kind: "images.generations", want: 20 * time.Minute},
{kind: "images.edits", want: 20 * time.Minute},
{kind: "videos.generations", want: 30 * time.Minute},
{kind: "chat.completions", want: 10 * time.Minute},
} {
if got := ProviderRequestTimeout(test.kind); got != test.want {
t.Fatalf("timeout for %s: got %s want %s", test.kind, got, test.want)
}
}
}
func TestTransportTimeoutUsesTerminalTimeoutCode(t *testing.T) {
err := transportClientError(context.DeadlineExceeded)
if err.Code != "timeout" || err.Retryable {
t.Fatalf("transport timeout classification = %+v, want terminal timeout", err)
}
}
func TestResponseBodyTimeoutUsesTerminalTimeoutCode(t *testing.T) {
response := &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(&timeoutResponseBody{}),
}
_, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIImages)
if err == nil {
t.Fatal("expected response body timeout")
}
clientErr, ok := err.(*ClientError)
if !ok || clientErr.Code != "timeout" || clientErr.Retryable {
t.Fatalf("response body timeout classification = %#v, want terminal timeout", err)
}
if clientErr.StatusCode != http.StatusOK || clientErr.Wire != wire {
t.Fatalf("response body timeout lost response metadata: %#v", clientErr)
}
}
func TestMediaPollTimeoutsUseTaskDefaultsWithoutPlatformOverride(t *testing.T) {
image := Request{Kind: "images.generations", Candidate: storeCandidateWithConfig("", "", nil, nil)}
video := Request{Kind: "videos.generations", Candidate: storeCandidateWithConfig("", "", nil, nil)}
if got := providerPollTimeout(image); got != 20*time.Minute {
t.Fatalf("image provider poll timeout: got %s want %s", got, 20*time.Minute)
}
if got := providerPollTimeout(video); got != 30*time.Minute {
t.Fatalf("video provider poll timeout: got %s want %s", got, 30*time.Minute)
}
if got := kelingPollTimeout(video); got != 30*time.Minute {
t.Fatalf("Keling video poll timeout: got %s want %s", got, 30*time.Minute)
}
if got := volcesPollTimeout(video); got != 30*time.Minute {
t.Fatalf("Volces video poll timeout: got %s want %s", got, 30*time.Minute)
}
}
+505 -18
View File
@@ -3,8 +3,14 @@ package clients
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/textproto"
"net/url"
"strings"
"time"
@@ -12,7 +18,8 @@ import (
)
type OpenAIClient struct {
HTTPClient *http.Client
HTTPClient *http.Client
Corrections *ParameterCorrectionCache
}
func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error) {
@@ -42,9 +49,19 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}
if endpointKind == "chat.completions" {
body = NormalizeChatCompletionRequestBody(body)
applyOpenAIChatReasoningParams(body, request.Candidate)
normalizedBody, normalizeErr := normalizeOpenAIChatAudioParts(body, request.Candidate)
if normalizeErr != nil {
return Response{}, normalizeErr
}
body = normalizedBody
if err := applyOpenAIChatReasoningParamsWithSource(body, request.Candidate, openAIAdaptationSource(request, endpointKind)); err != nil {
return Response{}, err
}
body = FilterOpenAIChatRequestBody(body)
} else if request.Kind == "responses" {
if err := applyOpenAIResponsesReasoningParamsWithSource(body, request.Candidate, request.OriginalBody); err != nil {
return Response{}, err
}
body = FilterOpenAIResponsesRequestBody(body)
if _, hasInput := body["input"]; !hasInput {
if messages, hasMessages := request.Body["messages"]; hasMessages {
@@ -59,34 +76,89 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}
}
body["model"] = upstreamModelName(request.Candidate)
normalizeOpenAIImageRequestBody(endpointKind, body, request.OriginalBody)
stream := openAIEndpointSupportsStream(endpointKind) && (request.Stream || boolValue(body, "stream"))
ensureOpenAIStreamUsage(body, endpointKind, stream)
raw, _ := json.Marshal(body)
upstreamEndpoint := joinURL(openAIBaseURL(endpointKind, request.Candidate), endpoint)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
if err != nil {
return Response{}, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
responseStartedAt := time.Now()
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
correctionScope := newParameterCorrectionScope(request, endpointKind)
correctionEnabled := endpointKind == "chat.completions" || endpointKind == "responses" ||
endpointKind == "images.generations" || endpointKind == "images.edits"
protectedCorrections := callerProtectedCorrectionParameters(request, endpointKind)
if correctionEnabled {
c.Corrections.apply(correctionScope, body, protectedCorrections)
}
provisionalRules := make([]parameterCorrectionRule, 0, 2)
seenCorrectionErrors := make(map[string]struct{})
var resp *http.Response
requestClient := httpClient(request.HTTPClient, c.HTTPClient)
for correctionAttempt := 0; ; correctionAttempt++ {
raw, contentType, payloadErr := openAIRequestPayload(ctx, endpointKind, body, request.Candidate)
if payloadErr != nil {
return Response{}, payloadErr
}
req, requestErr := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
if requestErr != nil {
return Response{}, requestErr
}
req.Header.Set("Content-Type", contentType)
req.Header.Set("Authorization", "Bearer "+apiKey)
applyUpstreamIdempotency(req, request)
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
resp, requestErr = requestClient.Do(req)
if requestErr != nil {
return Response{}, transportClientError(requestErr)
}
if err := notifyResponseReceived(request); err != nil {
_ = resp.Body.Close()
return Response{}, err
}
if !correctionEnabled || (resp.StatusCode != http.StatusBadRequest && resp.StatusCode != http.StatusUnprocessableEntity) {
break
}
_, _, upstreamErr := decodeHTTPResponseForProtocol(resp, openAIWireProtocol(endpointKind))
if upstreamErr == nil {
break
}
fingerprint := fmt.Sprintf("%d:%s", resp.StatusCode, strings.ToLower(upstreamErr.Error()))
if _, repeated := seenCorrectionErrors[fingerprint]; repeated || correctionAttempt >= 2 {
return Response{}, annotateResponseError(upstreamErr, requestIDFromParameterError(upstreamErr), responseStartedAt, time.Now())
}
seenCorrectionErrors[fingerprint] = struct{}{}
rule, safe := deriveParameterCorrection(upstreamErr, body, request.Candidate, protectedCorrections)
if !safe || !applyParameterCorrectionRule(body, rule) {
return Response{}, annotateResponseError(upstreamErr, requestIDFromParameterError(upstreamErr), responseStartedAt, time.Now())
}
c.Corrections.invalidate(correctionScope, rule.Param)
provisionalRules = append(provisionalRules, rule)
}
if len(provisionalRules) > 0 && resp.StatusCode >= 200 && resp.StatusCode < 300 {
c.Corrections.commit(correctionScope, provisionalRules)
}
requestID := requestIDFromHTTPResponse(resp)
var result map[string]any
var wire *WireResponse
var err error
upstreamResponseID := ""
nativeStreamDelta := openAIWireStreamDelta(request.StreamDelta, openAIWireProtocol(endpointKind), resp)
if request.Kind == "responses" && protocol == ProtocolOpenAIResponses && stream {
result, upstreamResponseID, err = decodeNativeResponsesStream(resp, request.StreamDelta)
result, upstreamResponseID, err = decodeNativeResponsesStream(resp, nativeStreamDelta)
wire = &WireResponse{Protocol: openAIWireProtocol(endpointKind), StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
} else {
var streamDelta StreamDelta = request.StreamDelta
var streamDelta StreamDelta = nativeStreamDelta
var adapter *chatResponsesStreamAdapter
if request.Kind == "responses" && protocol == ProtocolOpenAIChatCompletions && stream {
adapter = newChatResponsesStreamAdapter(request.PublicResponseID, request.Model)
adapter = newChatResponsesStreamAdapter(request.PublicResponseID, request.Model, request.Body)
streamDelta = func(event StreamDeltaEvent) error { return adapter.delta(event, request.StreamDelta) }
}
result, err = decodeOpenAIResponse(resp, stream, streamDelta)
if stream {
result, err = decodeOpenAIResponse(resp, true, streamDelta)
wire = &WireResponse{Protocol: openAIWireProtocol(endpointKind), StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
} else {
result, wire, err = decodeHTTPResponseForProtocol(resp, openAIWireProtocol(endpointKind))
}
if err == nil && endpointKind == "chat.completions" {
result = NormalizeChatCompletionResult(result)
}
@@ -103,6 +175,12 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
Progress: providerProgress(request), ResponseStartedAt: responseStartedAt, ResponseFinishedAt: time.Now(),
UpstreamProtocol: protocol, UpstreamEndpoint: endpoint, UpstreamResponseID: upstreamResponseID,
PublicResponseID: request.PublicResponseID, ResponseConverted: true,
Wire: func() *WireResponse {
if wire != nil {
wire.Converted = true
}
return wire
}(),
}, nil
}
}
@@ -138,9 +216,412 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
UpstreamEndpoint: endpoint,
UpstreamResponseID: upstreamResponseID,
PublicResponseID: publicResponseID,
Wire: wire,
}, nil
}
func openAIAdaptationSource(request Request, endpointKind string) map[string]any {
if request.OriginalBody == nil || request.Kind != "responses" || endpointKind != "chat.completions" {
return request.OriginalBody
}
out := cloneBody(request.OriginalBody)
if reasoning, ok := request.OriginalBody["reasoning"].(map[string]any); ok {
if effort, explicit := reasoning["effort"]; explicit {
out["reasoning_effort"] = effort
}
}
return out
}
func normalizeOpenAIImageRequestBody(endpointKind string, body map[string]any, originalBody map[string]any) {
if endpointKind != "images.generations" && endpointKind != "images.edits" {
return
}
normalizeOfficialOpenAIImageSize(body, originalBody)
// OpenAI image endpoints express output geometry through size. The Gateway
// keeps the generic fields for capability validation and billing, but they
// are not valid OpenAI wire parameters.
for _, key := range []string{
"aspect_ratio",
"aspectRatio",
"resolution",
"width",
"height",
"platform_id",
"platformId",
"platform_model_id",
"platformModelId",
} {
delete(body, key)
}
}
func openAIRequestPayload(ctx context.Context, endpointKind string, body map[string]any, candidate store.RuntimeModelCandidate) ([]byte, string, error) {
if endpointKind != "images.edits" {
raw, err := json.Marshal(body)
return raw, "application/json", err
}
if OpenAIImageEditUsesJSONURL(candidate) {
return openAIImageEditJSONPayload(body)
}
var payload bytes.Buffer
writer := multipart.NewWriter(&payload)
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
if len(images) == 0 {
return nil, "", &ClientError{
Code: "invalid_parameter",
Message: "image is required",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
imageFieldName := openAIImageEditFieldName(candidate, len(images))
for index, value := range images {
contentType, image, err := openAIImageEditPayload(ctx, value)
if err != nil {
return nil, "", err
}
if err := writeOpenAIImageEditFile(writer, imageFieldName, fmt.Sprintf("image-%d%s", index+1, openAIImageFileExtension(contentType)), contentType, image); err != nil {
return nil, "", err
}
}
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
contentType, image, err := openAIImageEditPayload(ctx, mask)
if err != nil {
return nil, "", err
}
if err := writeOpenAIImageEditFile(writer, "mask", "mask"+openAIImageFileExtension(contentType), contentType, image); err != nil {
return nil, "", err
}
}
for key, value := range body {
switch key {
case "image", "images", "mask", "mask_image", "maskImage":
continue
}
if strings.HasPrefix(key, "_") || value == nil {
continue
}
fieldValue, err := openAIFormFieldValue(value)
if err != nil {
return nil, "", err
}
if fieldValue == "" {
continue
}
if err := writer.WriteField(key, fieldValue); err != nil {
return nil, "", err
}
}
if err := writer.Close(); err != nil {
return nil, "", err
}
return payload.Bytes(), writer.FormDataContentType(), nil
}
// OpenAIImageEditUsesJSONURL reports whether the platform explicitly opts in to
// JSON URL requests for image edits. Multipart remains the compatibility-safe
// default when the setting is absent or unrecognized.
func OpenAIImageEditUsesJSONURL(candidate store.RuntimeModelCandidate) bool {
format := strings.ToLower(strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditRequestFormat"])))
format = strings.ReplaceAll(format, "-", "_")
format = strings.ReplaceAll(format, " ", "_")
return format == "json" || format == "json_url"
}
func openAIImageEditJSONPayload(body map[string]any) ([]byte, string, error) {
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
if len(images) == 0 {
return nil, "", &ClientError{
Code: "invalid_parameter",
Message: "image is required",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
imageURLs := make([]any, 0, len(images))
for _, image := range images {
imageURL, err := openAIImageEditURLValue(image, "image")
if err != nil {
return nil, "", err
}
imageURLs = append(imageURLs, imageURL)
}
payload := make(map[string]any, len(body))
for key, value := range body {
switch key {
case "image", "images", "mask", "mask_image", "maskImage":
continue
}
if strings.HasPrefix(key, "_") || value == nil {
continue
}
payload[key] = value
}
if len(imageURLs) == 1 {
payload["image"] = imageURLs[0]
} else {
payload["image"] = imageURLs
}
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
maskURL, err := openAIImageEditURLValue(mask, "mask")
if err != nil {
return nil, "", err
}
payload["mask"] = maskURL
}
raw, err := json.Marshal(payload)
return raw, "application/json", err
}
func openAIImageEditURLValue(value any, param string) (string, error) {
switch typed := value.(type) {
case map[string]any:
for _, key := range []string{"url", "image_url", "imageUrl"} {
if nested := typed[key]; nested != nil {
return openAIImageEditURLValue(nested, param)
}
}
case string:
raw := strings.TrimSpace(typed)
parsed, err := url.Parse(raw)
if err == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https")) {
return raw, nil
}
}
return "", &ClientError{
Code: "invalid_parameter",
Message: "OpenAI image edit JSON mode requires " + param + " to be an HTTP(S) URL",
Param: param,
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
func openAIImageEditFieldName(candidate store.RuntimeModelCandidate, imageCount int) string {
if configured := strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditMultipartFieldName"])); configured == "image" || configured == "image[]" {
return configured
}
if imageCount > 1 {
return "image[]"
}
parsed, err := url.Parse(strings.TrimSpace(candidate.BaseURL))
if err == nil && strings.EqualFold(parsed.Hostname(), "api.openai.com") {
return "image[]"
}
return "image"
}
func openAIImageEditValues(value any) []any {
switch typed := value.(type) {
case []any:
return typed
case []string:
out := make([]any, 0, len(typed))
for _, item := range typed {
out = append(out, item)
}
return out
case nil:
return nil
default:
return []any{typed}
}
}
func openAIImageEditPayload(ctx context.Context, value any) (string, []byte, error) {
switch typed := value.(type) {
case map[string]any:
for _, key := range []string{"data", "b64_json", "base64", "url"} {
if nested := typed[key]; nested != nil {
return openAIImageEditPayload(ctx, nested)
}
}
case string:
raw := strings.TrimSpace(typed)
if raw == "" {
break
}
contentType := ""
encoded := raw
if strings.HasPrefix(strings.ToLower(raw), "data:") {
prefix, payload, ok := strings.Cut(raw, ",")
if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
break
}
contentType = strings.TrimSpace(strings.Split(strings.TrimPrefix(prefix, "data:"), ";")[0])
encoded = payload
} else if parsed, parseErr := url.Parse(raw); parseErr == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https")) {
return fetchRemoteMediaInputPayload(ctx, raw, 256<<20)
}
image, err := decodeOpenAIImageEditBase64(encoded)
if err != nil {
break
}
if contentType == "" {
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(image), ";")[0])
}
return contentType, image, nil
case []byte:
if len(typed) > 0 {
return strings.TrimSpace(strings.Split(http.DetectContentType(typed), ";")[0]), typed, nil
}
}
return "", nil, &ClientError{
Code: "invalid_parameter",
Message: "image must be a base64 or data URL payload",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
func fetchRemoteMediaInputPayload(ctx context.Context, sourceURL string, maxBytes int64) (string, []byte, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
if err != nil {
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: false}
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: true}
}
defer resp.Body.Close()
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
return "", nil, &ClientError{
Code: "request_asset_fetch_failed",
Message: resp.Status,
StatusCode: resp.StatusCode,
Retryable: HTTPRetryable(resp.StatusCode),
}
}
if maxBytes <= 0 {
maxBytes = 256 << 20
}
payload, err := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1))
if err != nil {
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: true}
}
if int64(len(payload)) > maxBytes {
return "", nil, &ClientError{
Code: "invalid_parameter",
Message: "remote media input exceeds the download limit",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
if contentType == "" || contentType == "application/octet-stream" {
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(payload), ";")[0])
}
return contentType, payload, nil
}
func decodeOpenAIImageEditBase64(value string) ([]byte, error) {
normalized := strings.Map(func(char rune) rune {
switch char {
case '\n', '\r', '\t', ' ':
return -1
default:
return char
}
}, value)
var lastErr error
for _, encoding := range []*base64.Encoding{
base64.StdEncoding,
base64.RawStdEncoding,
base64.URLEncoding,
base64.RawURLEncoding,
} {
payload, err := encoding.DecodeString(normalized)
if err == nil && len(payload) > 0 {
return payload, nil
}
lastErr = err
}
return nil, lastErr
}
func writeOpenAIImageEditFile(writer *multipart.Writer, fieldName string, fileName string, contentType string, payload []byte) error {
header := make(textproto.MIMEHeader)
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, escapeMultipartFilename(fileName)))
if contentType != "" {
header.Set("Content-Type", contentType)
}
part, err := writer.CreatePart(header)
if err != nil {
return err
}
_, err = io.Copy(part, bytes.NewReader(payload))
return err
}
func openAIImageFileExtension(contentType string) string {
switch strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) {
case "image/jpeg":
return ".jpg"
case "image/webp":
return ".webp"
case "image/gif":
return ".gif"
default:
return ".png"
}
}
func openAIFormFieldValue(value any) (string, error) {
switch typed := value.(type) {
case string:
return typed, nil
case bool:
return fmt.Sprintf("%t", typed), nil
case float64:
return fmt.Sprintf("%v", typed), nil
case float32:
return fmt.Sprintf("%v", typed), nil
case int:
return fmt.Sprintf("%d", typed), nil
case int64:
return fmt.Sprintf("%d", typed), nil
case json.Number:
return typed.String(), nil
default:
encoded, err := json.Marshal(value)
return string(encoded), err
}
}
func openAIWireStreamDelta(next StreamDelta, protocol string, response *http.Response) StreamDelta {
if next == nil {
return nil
}
return func(event StreamDeltaEvent) error {
event.WireProtocol = protocol
if response != nil {
event.WireStatusCode = response.StatusCode
event.WireHeaders = compatibleResponseHeaders(response.Header)
}
return next(event)
}
}
func openAIWireProtocol(kind string) string {
switch kind {
case "chat.completions":
return ProtocolOpenAIChatCompletions
case "responses":
return ProtocolOpenAIResponses
case "embeddings":
return ProtocolOpenAIEmbeddings
case "images.generations", "images.edits":
return ProtocolOpenAIImages
default:
return "openai_" + strings.ReplaceAll(kind, ".", "_")
}
}
func decodeOpenAIResponse(resp *http.Response, stream bool, onDelta StreamDelta) (map[string]any, error) {
if stream {
result, err := decodeOpenAIStreamResponse(resp, onDelta)
@@ -202,12 +683,18 @@ func ensureOpenAIStreamUsage(body map[string]any, kind string, stream bool) {
return
}
streamOptions := map[string]any{}
if existing, ok := body["stream_options"].(map[string]any); ok {
if raw, explicit := body["stream_options"]; explicit {
existing, ok := raw.(map[string]any)
if !ok {
return
}
for key, value := range existing {
streamOptions[key] = value
}
}
streamOptions["include_usage"] = true
if _, explicit := streamOptions["include_usage"]; !explicit {
streamOptions["include_usage"] = true
}
body["stream_options"] = streamOptions
}
@@ -0,0 +1,169 @@
package clients
import (
"net/http"
"net/url"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
var openAIChatAudioFormatByMIME = map[string]string{
"audio/aac": "aac",
"audio/flac": "flac",
"audio/m4a": "m4a",
"audio/mp4": "m4a",
"audio/mpeg": "mp3",
"audio/ogg": "ogg",
"audio/opus": "opus",
"audio/wav": "wav",
"audio/webm": "webm",
"audio/x-m4a": "m4a",
"audio/x-wav": "wav",
}
var openAIChatAudioFormatByExtension = map[string]string{
"aac": "aac",
"flac": "flac",
"m4a": "m4a",
"mp3": "mp3",
"mp4": "m4a",
"oga": "ogg",
"ogg": "ogg",
"opus": "opus",
"wav": "wav",
"webm": "webm",
}
func normalizeOpenAIChatAudioParts(body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
if !isAliyunBailianOpenAI(candidate) {
return body, nil
}
messages, ok := body["messages"].([]any)
if !ok {
return body, nil
}
normalizedMessages := make([]any, 0, len(messages))
for _, rawMessage := range messages {
message, ok := rawMessage.(map[string]any)
if !ok {
normalizedMessages = append(normalizedMessages, rawMessage)
continue
}
content, ok := message["content"].([]any)
if !ok {
normalizedMessages = append(normalizedMessages, rawMessage)
continue
}
normalizedContent := make([]any, 0, len(content))
for _, rawPart := range content {
part, ok := rawPart.(map[string]any)
if !ok {
normalizedContent = append(normalizedContent, rawPart)
continue
}
normalizedPart, err := normalizeAliyunOpenAIChatAudioPart(part)
if err != nil {
return nil, err
}
normalizedContent = append(normalizedContent, normalizedPart)
}
copiedMessage := cloneMapAny(message)
copiedMessage["content"] = normalizedContent
normalizedMessages = append(normalizedMessages, copiedMessage)
}
out := cloneMapAny(body)
out["messages"] = normalizedMessages
return out, nil
}
func normalizeAliyunOpenAIChatAudioPart(part map[string]any) (map[string]any, error) {
switch strings.TrimSpace(stringFromAny(part["type"])) {
case "audio_url":
audioURL, _ := part["audio_url"].(map[string]any)
data := strings.TrimSpace(stringFromAny(audioURL["url"]))
format := resolveOpenAIChatAudioFormat(data, audioURL["format"], audioURL["mime_type"], audioURL["mimeType"])
if data == "" || format == "" {
return nil, invalidOpenAIChatAudioError()
}
return map[string]any{
"type": "input_audio",
"input_audio": map[string]any{
"data": data,
"format": format,
},
}, nil
case "input_audio":
inputAudio, _ := part["input_audio"].(map[string]any)
data := firstNonEmptyString(inputAudio["data"], inputAudio["url"])
format := resolveOpenAIChatAudioFormat(data, inputAudio["format"], inputAudio["mime_type"], inputAudio["mimeType"])
if data == "" || format == "" {
return nil, invalidOpenAIChatAudioError()
}
copiedInput := cloneMapAny(inputAudio)
copiedInput["data"] = data
copiedInput["format"] = format
delete(copiedInput, "url")
copiedPart := cloneMapAny(part)
copiedPart["input_audio"] = copiedInput
return copiedPart, nil
default:
return part, nil
}
}
func resolveOpenAIChatAudioFormat(source string, values ...any) string {
if explicit := strings.ToLower(strings.TrimSpace(firstNonEmptyString(values...))); explicit != "" {
if normalized := openAIChatAudioFormatByExtension[explicit]; normalized != "" {
return normalized
}
if normalized := openAIChatAudioFormatByMIME[strings.TrimSpace(strings.Split(explicit, ";")[0])]; normalized != "" {
return normalized
}
return explicit
}
if dataMIME := openAIChatAudioDataMIME(source); dataMIME != "" {
if normalized := openAIChatAudioFormatByMIME[dataMIME]; normalized != "" {
return normalized
}
}
if extension := openAIChatAudioSourceExtension(source); extension != "" {
return openAIChatAudioFormatByExtension[extension]
}
return ""
}
func openAIChatAudioDataMIME(source string) string {
if !strings.HasPrefix(strings.ToLower(source), "data:") {
return ""
}
header := strings.TrimPrefix(source[:strings.Index(source+",", ",")], "data:")
return strings.ToLower(strings.TrimSpace(strings.Split(header, ";")[0]))
}
func openAIChatAudioSourceExtension(source string) string {
if source == "" || strings.HasPrefix(strings.ToLower(source), "data:") {
return ""
}
if parsed, err := url.Parse(source); err == nil && parsed.Path != "" {
source = parsed.Path
}
source = strings.TrimSuffix(strings.Split(strings.Split(source, "?")[0], "#")[0], "/")
if index := strings.LastIndex(source, "."); index >= 0 && index+1 < len(source) {
return strings.ToLower(source[index+1:])
}
return ""
}
func invalidOpenAIChatAudioError() error {
return &ClientError{
Code: "invalid_parameter",
Message: "input_audio requires data and a resolvable format",
Param: "messages.content.input_audio",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
@@ -0,0 +1,92 @@
package clients
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestOpenAIImageEndpointsRemoveRejectedResponseFormatAndCacheCorrection(t *testing.T) {
for _, test := range []struct {
name string
kind string
body map[string]any
}{
{
name: "generation JSON",
kind: "images.generations",
body: map[string]any{"prompt": "test", "response_format": "url"},
},
{
name: "edit multipart",
kind: "images.edits",
body: map[string]any{"prompt": "test", "image": "aW1hZ2U=", "response_format": "url"},
},
} {
t.Run(test.name, func(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestNo := requests.Add(1)
responseFormat := ""
if test.kind == "images.edits" {
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart image request: %v", err)
}
responseFormat = r.FormValue("response_format")
} else {
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode image request: %v", err)
}
responseFormat, _ = body["response_format"].(string)
}
if requestNo == 1 {
if responseFormat != "url" {
t.Fatalf("first request lost caller response_format: %q", responseFormat)
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
"message": "Unknown parameter: 'response_format'.",
"type": "invalid_request_error", "param": "response_format", "code": "unknown_parameter",
}})
return
}
if responseFormat != "" {
t.Fatalf("corrected request still contains response_format: %q", responseFormat)
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"b64_json": "aW1hZ2U="}}})
}))
defer server.Close()
cache := NewParameterCorrectionCache()
request := Request{
Kind: test.kind, Model: "gpt-image-2", Body: test.body, OriginalBody: test.body,
UpstreamIdempotencyKey: "gateway-task-id",
Candidate: store.RuntimeModelCandidate{
Provider: "openai", BaseURL: server.URL, ProviderModelName: "gpt-image-2",
Credentials: map[string]any{"apiKey": "test-key"},
},
}
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("correct rejected response_format: %v", err)
}
if requests.Load() != 2 {
t.Fatalf("first corrected call requests=%d, want 2", requests.Load())
}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("reuse cached response_format correction: %v", err)
}
if requests.Load() != 3 {
t.Fatalf("cached correction should avoid another 400, requests=%d", requests.Load())
}
})
}
}
@@ -0,0 +1,154 @@
package clients
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"reflect"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestOpenAIChatNativeDeepPassthroughPreservesOfficialAndFutureFields(t *testing.T) {
requestBody := completeChatPassthroughBody()
var captured map[string]any
var idempotencyKey string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
idempotencyKey = r.Header.Get("Idempotency-Key")
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatal(err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "chatcmpl-native", "object": "chat.completion", "model": "provider-chat",
"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "ok"}}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions", Model: "public-chat", Body: requestBody, OriginalBody: requestBody,
UpstreamIdempotencyKey: "gateway-task-id",
Candidate: store.RuntimeModelCandidate{
Provider: "openai", BaseURL: server.URL, ProviderModelName: "provider-chat",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatal(err)
}
want := jsonRoundTripMap(t, requestBody)
delete(want, "request_id")
want["model"] = "provider-chat"
if !reflect.DeepEqual(captured, want) {
t.Fatalf("native Chat request changed\n got: %#v\nwant: %#v", captured, want)
}
if idempotencyKey != "gateway-task-id" {
t.Fatalf("missing Gateway task idempotency key: %q", idempotencyKey)
}
}
func TestOpenAIResponsesNativeDeepPassthroughPreservesOfficialAndFutureFields(t *testing.T) {
requestBody := completeResponsesPassthroughBody()
var captured map[string]any
var idempotencyKey string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
idempotencyKey = r.Header.Get("Idempotency-Key")
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatal(err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "resp-native", "object": "response", "status": "completed", "output": []any{},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "responses", Model: "public-responses", Body: requestBody, OriginalBody: requestBody,
UpstreamProtocol: ProtocolOpenAIResponses, UpstreamPreviousResponseID: "resp_upstream_parent",
UpstreamIdempotencyKey: "gateway-task-id",
Candidate: store.RuntimeModelCandidate{
Provider: "openai", BaseURL: server.URL, ProviderModelName: "provider-responses",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatal(err)
}
want := jsonRoundTripMap(t, requestBody)
delete(want, "request_id")
want["model"] = "provider-responses"
want["previous_response_id"] = "resp_upstream_parent"
if !reflect.DeepEqual(captured, want) {
t.Fatalf("native Responses request changed\n got: %#v\nwant: %#v", captured, want)
}
if idempotencyKey != "gateway-task-id" {
t.Fatalf("missing Gateway task idempotency key: %q", idempotencyKey)
}
}
func completeChatPassthroughBody() map[string]any {
return map[string]any{
"model": "caller-chat", "messages": []any{
map[string]any{"role": "assistant", "content": nil, "function_call": map[string]any{"name": "legacy", "arguments": "{\"x\":1}"}},
map[string]any{"role": "assistant", "content": nil, "tool_calls": []any{
map[string]any{"id": "call_function", "type": "function", "function": map[string]any{"name": "lookup", "arguments": "{\"q\":1}"}},
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
}},
},
"audio": map[string]any{"format": "wav", "voice": "alloy"}, "frequency_penalty": 0.1,
"function_call": map[string]any{"name": "legacy"}, "functions": []any{map[string]any{"name": "legacy", "parameters": map[string]any{"type": "object"}}},
"logit_bias": map[string]any{"1": 2}, "logprobs": true, "max_completion_tokens": 101, "max_tokens": 102,
"metadata": map[string]any{"trace": "1"}, "modalities": []any{"text"}, "moderation": map[string]any{"type": "auto"}, "n": 1,
"parallel_tool_calls": true, "prediction": map[string]any{"type": "content", "content": "answer"}, "presence_penalty": 0.2,
"prompt_cache_key": "cache", "prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "in_memory",
"reasoning_effort": "low", "response_format": map[string]any{"type": "json_object"}, "safety_identifier": "safe", "seed": 7,
"service_tier": "default", "stop": []any{"END"}, "store": false, "stream": false,
"stream_options": map[string]any{"include_usage": true}, "temperature": 0.7,
"tool_choice": map[string]any{"type": "custom", "custom": map[string]any{"name": "shell"}},
"tools": []any{
map[string]any{"type": "function", "function": map[string]any{"name": "lookup", "parameters": map[string]any{"type": "object"}}},
map[string]any{"type": "custom", "custom": map[string]any{"name": "shell", "format": map[string]any{"type": "text"}}},
},
"top_logprobs": 2, "top_p": 0.9, "user": "user-1", "verbosity": "low",
"web_search_options": map[string]any{"search_context_size": "low"},
"future_official_field": map[string]any{"nested": []any{map[string]any{"keep": true}}},
"request_id": "gateway-only",
}
}
func completeResponsesPassthroughBody() map[string]any {
return map[string]any{
"background": false, "context_management": map[string]any{"type": "compaction", "compact_threshold": 2000},
"conversation": "conv_1", "include": []any{"message.output_text.logprobs"},
"input": []any{map[string]any{"type": "message", "role": "user", "content": []any{
map[string]any{"type": "input_text", "text": "hello", "prompt_cache_breakpoint": map[string]any{"type": "ephemeral"}},
}}},
"instructions": "be concise", "max_output_tokens": 200, "max_tool_calls": 3, "metadata": map[string]any{"trace": "1"},
"model": "caller-responses", "moderation": map[string]any{"type": "auto"}, "parallel_tool_calls": true,
"previous_response_id": "resp_caller_parent", "prompt": map[string]any{"id": "pmpt_1", "variables": map[string]any{"x": "y"}},
"prompt_cache_key": "cache", "prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "24h",
"reasoning": map[string]any{"effort": "low", "summary": "auto"}, "safety_identifier": "safe", "service_tier": "default",
"store": true, "stream": false, "stream_options": map[string]any{"include_obfuscation": true}, "temperature": 0.7,
"text": map[string]any{"format": map[string]any{"type": "json_schema", "name": "answer", "schema": map[string]any{"type": "object"}}, "verbosity": "low"},
"tool_choice": map[string]any{"type": "custom", "name": "shell"},
"tools": []any{map[string]any{"type": "custom", "name": "shell", "description": "run", "format": map[string]any{"type": "text"}}},
"top_logprobs": 2, "top_p": 0.9, "truncation": "auto", "user": "user-1",
"future_official_field": map[string]any{"nested": []any{map[string]any{"keep": true}}},
"request_id": "gateway-only",
}
}
func jsonRoundTripMap(t *testing.T, value map[string]any) map[string]any {
t.Helper()
raw, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
t.Fatal(err)
}
return out
}
@@ -1,15 +1,9 @@
package clients
import (
"fmt"
"net/http"
"sort"
)
// Keep these lists aligned with openai-node 6.47.0 and the public OpenAI API
// reference. The Gateway accepts a small, explicit set of routing extensions at
// ingress, but only protocol fields (plus controlled provider adaptations) are
// allowed across the upstream boundary.
// reference. They document and test the currently known surface, but are not an
// upstream allowlist: unknown fields may be introduced by OpenAI after a
// Gateway release and must remain transparent.
var openAIChatRequestParameters = stringSet(
"messages", "model", "audio", "frequency_penalty", "function_call", "functions",
"logit_bias", "logprobs", "max_completion_tokens", "max_tokens", "metadata",
@@ -32,7 +26,7 @@ var openAIResponsesRequestParameters = stringSet(
var gatewayOpenAIRequestExtensions = stringSet(
"runMode", "run_mode", "conversationId", "conversation_id", "sessionId", "session_id",
"requestId", "request_id", "signal", "userMessage", "user_message", "platformId",
"platform_id", "options", "enable_thinking", "thinking_budget_tokens", "enable_web_search",
"platform_id", "options", "enable_thinking", "thinking_budget", "thinking_budget_tokens", "enable_web_search",
"modelType", "model_type", "capability", "capabilityType", "mode", "simulation", "testMode",
"cacheAffinityKey", "cache_affinity_key", "simulationDurationMs", "simulationDurationSeconds",
"simulationMinDurationMs", "simulationMaxDurationMs", "simulationMinDurationSeconds",
@@ -50,56 +44,36 @@ var controlledOpenAIChatProviderParameters = stringSet(
var controlledOpenAIResponsesProviderParameters = stringSet("presence_penalty", "frequency_penalty")
func ValidateOpenAIRequestParameters(kind string, body map[string]any) error {
allowed := openAIChatRequestParameters
if kind == "responses" {
allowed = openAIResponsesRequestParameters
}
unknown := make([]string, 0)
for key := range body {
if _, ok := allowed[key]; ok {
continue
}
if _, ok := gatewayOpenAIRequestExtensions[key]; ok {
continue
}
if kind == "responses" {
if _, ok := gatewayResponsesRequestExtensions[key]; ok {
continue
}
}
unknown = append(unknown, key)
}
if len(unknown) == 0 {
return nil
}
sort.Strings(unknown)
return &ClientError{
Code: "invalid_parameter",
Message: fmt.Sprintf("Unknown parameter: %s", unknown[0]),
Param: unknown[0],
StatusCode: http.StatusBadRequest,
Retryable: false,
}
// OpenAI-compatible public endpoints intentionally accept future official
// fields. Provider validation remains authoritative for fields the selected
// upstream does not support.
_ = kind
_ = body
return nil
}
func FilterOpenAIChatRequestBody(body map[string]any) map[string]any {
return filterOpenAIRequestBody(body, openAIChatRequestParameters, controlledOpenAIChatProviderParameters)
return filterOpenAIRequestBody(body, controlledOpenAIChatProviderParameters, nil)
}
func FilterOpenAIResponsesRequestBody(body map[string]any) map[string]any {
return filterOpenAIRequestBody(body, openAIResponsesRequestParameters, controlledOpenAIResponsesProviderParameters)
return filterOpenAIRequestBody(body, controlledOpenAIResponsesProviderParameters, gatewayResponsesRequestExtensions)
}
func filterOpenAIRequestBody(body map[string]any, allowed map[string]struct{}, extensions map[string]struct{}) map[string]any {
func filterOpenAIRequestBody(body map[string]any, controlled map[string]struct{}, protocolInternal map[string]struct{}) map[string]any {
out := make(map[string]any, len(body))
for key, value := range body {
if _, ok := allowed[key]; ok {
if _, ok := controlled[key]; ok {
out[key] = value
continue
}
if _, ok := extensions[key]; ok {
out[key] = value
if _, internal := gatewayOpenAIRequestExtensions[key]; internal {
continue
}
if _, internal := protocolInternal[key]; internal {
continue
}
out[key] = value
}
return out
}
@@ -1,9 +1,6 @@
package clients
import (
"strings"
"testing"
)
import "testing"
func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
body := map[string]any{}
@@ -11,7 +8,7 @@ func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
body[key] = "sentinel-" + key
}
body["conversationId"] = "internal"
body["unknown"] = "must-not-leak"
body["future_official_field"] = map[string]any{"nested": []any{"must-survive"}}
filtered := FilterOpenAIChatRequestBody(body)
for key := range openAIChatRequestParameters {
@@ -19,7 +16,10 @@ func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
t.Fatalf("official Chat parameter %q was removed", key)
}
}
for _, key := range []string{"conversationId", "unknown"} {
if _, ok := filtered["future_official_field"]; !ok {
t.Fatal("future Chat field was removed at the upstream boundary")
}
for _, key := range []string{"conversationId"} {
if _, ok := filtered[key]; ok {
t.Fatalf("internal/unknown parameter %q leaked upstream", key)
}
@@ -32,7 +32,7 @@ func TestOpenAIResponsesOfficialParametersSurviveBoundary(t *testing.T) {
body[key] = "sentinel-" + key
}
body["request_id"] = "internal"
body["unknown"] = "must-not-leak"
body["future_official_field"] = map[string]any{"nested": []any{"must-survive"}}
filtered := FilterOpenAIResponsesRequestBody(body)
for key := range openAIResponsesRequestParameters {
@@ -40,20 +40,19 @@ func TestOpenAIResponsesOfficialParametersSurviveBoundary(t *testing.T) {
t.Fatalf("official Responses parameter %q was removed", key)
}
}
for _, key := range []string{"request_id", "unknown"} {
if _, ok := filtered["future_official_field"]; !ok {
t.Fatal("future Responses field was removed at the upstream boundary")
}
for _, key := range []string{"request_id"} {
if _, ok := filtered[key]; ok {
t.Fatalf("internal/unknown parameter %q leaked upstream", key)
}
}
}
func TestValidateOpenAIRequestParametersRejectsUnknownTopLevelField(t *testing.T) {
err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "rogue": true})
if err == nil || ErrorCode(err) != "invalid_parameter" || !strings.Contains(err.Error(), "rogue") {
t.Fatalf("expected OpenAI-style invalid_parameter for rogue field, got %v", err)
}
if ErrorParam(err) != "rogue" {
t.Fatalf("expected rogue parameter attribution, got %q", ErrorParam(err))
func TestValidateOpenAIRequestParametersAcceptsFutureTopLevelField(t *testing.T) {
if err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "future_official_field": true}); err != nil {
t.Fatalf("future Responses fields must remain forward compatible, got %v", err)
}
if err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "messages": []any{}, "request_id": "internal"}); err != nil {
t.Fatalf("expected controlled Responses extensions to remain accepted, got %v", err)
@@ -104,3 +103,34 @@ func TestResponsesFallbackMapsEquivalentCurrentParameters(t *testing.T) {
}
}
}
func TestNormalizeChatRequestPreservesCustomAndLegacyFunctionCall(t *testing.T) {
legacy := map[string]any{"name": "legacy", "arguments": "{\"x\":1}"}
custom := map[string]any{
"id": "call_custom", "type": "custom",
"custom": map[string]any{"name": "shell", "input": "pwd"},
}
body := NormalizeChatCompletionRequestBody(map[string]any{"messages": []any{map[string]any{
"role": "assistant", "content": nil, "function_call": legacy, "tool_calls": []any{custom},
}}})
messages, _ := body["messages"].([]any)
message, _ := messages[0].(map[string]any)
if got, ok := message["function_call"].(map[string]any); !ok || got["name"] != "legacy" || got["arguments"] != "{\"x\":1}" {
t.Fatalf("legacy function_call changed: %+v", message)
}
toolCalls, _ := message["tool_calls"].([]any)
got, _ := toolCalls[0].(map[string]any)
gotCustom, _ := got["custom"].(map[string]any)
if got["type"] != "custom" || got["id"] != "call_custom" || gotCustom["name"] != "shell" || gotCustom["input"] != "pwd" || got["function"] != nil {
t.Fatalf("standard custom tool call changed: %+v", got)
}
}
func TestEnsureOpenAIStreamUsageDoesNotOverrideCallerChoice(t *testing.T) {
body := map[string]any{"stream_options": map[string]any{"include_usage": false, "include_obfuscation": false}}
ensureOpenAIStreamUsage(body, "chat.completions", true)
options, _ := body["stream_options"].(map[string]any)
if options["include_usage"] != false || options["include_obfuscation"] != false {
t.Fatalf("caller stream options were changed: %+v", options)
}
}
@@ -0,0 +1,531 @@
package clients
import (
"container/list"
"errors"
"fmt"
"math"
"net/url"
"regexp"
"sort"
"strconv"
"strings"
"sync"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const parameterCorrectionCacheCapacity = 512
var safeUpstreamCorrectionParameters = stringSet(
"reasoning_effort", "reasoning.effort", "enable_thinking", "thinking_budget",
"temperature", "top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "n",
"max_tokens", "max_completion_tokens", "max_output_tokens",
"logprobs", "top_logprobs", "service_tier", "verbosity",
"response_format",
)
var (
anchoredParameterPattern = regexp.MustCompile(`(?i)^(?:unknown|unsupported|invalid) parameter\s*:?\s*[` + "`" + `'\"]?([a-z][a-z0-9_.]*)`)
restrictedParameterPattern = regexp.MustCompile(`(?i)^the value of the ([a-z][a-z0-9_.]*) parameter is restricted to\s+([^\s,.;]+)`)
fixedValuePattern = regexp.MustCompile(`(?i)(?:restricted to|must be(?: set)?(?: to)?|only supports?(?: the)? value(?: of)?)\s*[` + "`" + `'\"]?([^\s,.;` + "`" + `'\"]+)`)
betweenPattern = regexp.MustCompile(`(?i)\bbetween\s+(-?[0-9]+(?:\.[0-9]+)?)\s+and\s+(-?[0-9]+(?:\.[0-9]+)?)\b`)
maximumPattern = regexp.MustCompile(`(?i)(?:less than or equal to|at most|maximum(?: value)?(?: is| of)?|<=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
minimumPattern = regexp.MustCompile(`(?i)(?:greater than or equal to|at least|minimum(?: value)?(?: is| of)?|>=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
)
type parameterCorrectionAction string
const (
parameterCorrectionRemove parameterCorrectionAction = "remove"
parameterCorrectionSet parameterCorrectionAction = "set"
parameterCorrectionClamp parameterCorrectionAction = "clamp"
)
type parameterCorrectionRule struct {
Param string
Action parameterCorrectionAction
Value any
Min *float64
Max *float64
}
type parameterCorrectionScope struct {
Provider string
BaseURL string
Protocol string
Kind string
Model string
}
func newParameterCorrectionScope(request Request, endpointKind string) parameterCorrectionScope {
return parameterCorrectionScope{
Provider: normalizedReasoningString(request.Candidate.Provider),
BaseURL: normalizedCorrectionBaseURL(request.Candidate.BaseURL),
Protocol: strings.TrimSpace(request.UpstreamProtocol),
Kind: strings.TrimSpace(endpointKind),
Model: normalizedReasoningString(upstreamModelName(request.Candidate)),
}
}
func normalizedCorrectionBaseURL(raw string) string {
trimmed := strings.TrimRight(strings.TrimSpace(raw), "/")
parsed, err := url.Parse(trimmed)
if err != nil || parsed.Host == "" {
return strings.ToLower(trimmed)
}
parsed.Scheme = strings.ToLower(parsed.Scheme)
parsed.Host = strings.ToLower(parsed.Host)
parsed.RawQuery = ""
parsed.Fragment = ""
return strings.TrimRight(parsed.String(), "/")
}
func (scope parameterCorrectionScope) prefix() string {
return strings.Join([]string{scope.Provider, scope.BaseURL, scope.Protocol, scope.Kind, scope.Model}, "\x1f") + "\x1f"
}
func (scope parameterCorrectionScope) key(param string) string {
return scope.prefix() + param
}
type parameterCorrectionCacheEntry struct {
key string
rule parameterCorrectionRule
}
type ParameterCorrectionCache struct {
mu sync.Mutex
capacity int
entries map[string]*list.Element
lru *list.List
}
func NewParameterCorrectionCache() *ParameterCorrectionCache {
return &ParameterCorrectionCache{
capacity: parameterCorrectionCacheCapacity,
entries: make(map[string]*list.Element),
lru: list.New(),
}
}
func (cache *ParameterCorrectionCache) apply(scope parameterCorrectionScope, body map[string]any, protected ...map[string]struct{}) []string {
if cache == nil {
return nil
}
cache.mu.Lock()
defer cache.mu.Unlock()
prefix := scope.prefix()
elements := make([]*list.Element, 0)
for key, element := range cache.entries {
if strings.HasPrefix(key, prefix) {
elements = append(elements, element)
}
}
sort.Slice(elements, func(i, j int) bool {
return elements[i].Value.(parameterCorrectionCacheEntry).rule.Param < elements[j].Value.(parameterCorrectionCacheEntry).rule.Param
})
applied := make([]string, 0, len(elements))
for _, element := range elements {
rule := element.Value.(parameterCorrectionCacheEntry).rule
if correctionParamProtected(rule.Param, protected...) {
continue
}
if applyParameterCorrectionRule(body, rule) {
applied = append(applied, rule.Param)
cache.lru.MoveToFront(element)
}
}
return applied
}
func (cache *ParameterCorrectionCache) commit(scope parameterCorrectionScope, rules []parameterCorrectionRule) {
if cache == nil || len(rules) == 0 {
return
}
cache.mu.Lock()
defer cache.mu.Unlock()
for _, rule := range rules {
key := scope.key(rule.Param)
if existing := cache.entries[key]; existing != nil {
existing.Value = parameterCorrectionCacheEntry{key: key, rule: rule}
cache.lru.MoveToFront(existing)
continue
}
element := cache.lru.PushFront(parameterCorrectionCacheEntry{key: key, rule: rule})
cache.entries[key] = element
for cache.lru.Len() > cache.capacity {
oldest := cache.lru.Back()
if oldest == nil {
break
}
delete(cache.entries, oldest.Value.(parameterCorrectionCacheEntry).key)
cache.lru.Remove(oldest)
}
}
}
func (cache *ParameterCorrectionCache) invalidate(scope parameterCorrectionScope, param string) {
if cache == nil {
return
}
cache.mu.Lock()
defer cache.mu.Unlock()
key := scope.key(param)
if element := cache.entries[key]; element != nil {
delete(cache.entries, key)
cache.lru.Remove(element)
}
}
func (cache *ParameterCorrectionCache) size() int {
if cache == nil {
return 0
}
cache.mu.Lock()
defer cache.mu.Unlock()
return cache.lru.Len()
}
func deriveParameterCorrection(err error, body map[string]any, candidate store.RuntimeModelCandidate, protected map[string]struct{}) (parameterCorrectionRule, bool) {
var clientErr *ClientError
if !errors.As(err, &clientErr) || (clientErr.StatusCode != 400 && clientErr.StatusCode != 422) {
return parameterCorrectionRule{}, false
}
param, code, message := structuredParameterError(clientErr)
if param == "" {
param = anchoredParameterFromMessage(message)
}
param = normalizeCorrectionParam(param)
if !isSafeCorrectionParam(param) {
return parameterCorrectionRule{}, false
}
if correctionParamProtected(param, protected) {
return parameterCorrectionRule{}, false
}
lowerMessage := strings.ToLower(message)
lowerCode := strings.ToLower(code)
if conflictRule, ok := deriveConflictCorrection(param, lowerMessage); ok {
if correctionParamProtected(conflictRule.Param, protected) {
return parameterCorrectionRule{}, false
}
return conflictRule, true
}
if strings.Contains(lowerMessage, "unknown parameter") ||
strings.Contains(lowerMessage, "unsupported parameter") ||
strings.Contains(lowerMessage, "not supported") ||
strings.Contains(lowerMessage, "does not support") ||
strings.Contains(lowerCode, "unknown_parameter") ||
strings.Contains(lowerCode, "unsupported_parameter") {
return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
}
if param == "reasoning_effort" || param == "reasoning.effort" {
if effort := nearestErrorSupportedReasoningEffort(message, normalizedReasoningString(valueAtCorrectionPath(body, param)), candidate); effort != "" {
return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: effort}, true
}
}
if match := fixedValuePattern.FindStringSubmatch(message); len(match) == 2 {
if value, ok := parseCorrectionScalar(match[1]); ok {
return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: value}, true
}
}
if min, max, ok := explicitNumericBounds(message); ok {
if isOutputLimitCorrectionParam(param) {
min = nil
}
return parameterCorrectionRule{Param: param, Action: parameterCorrectionClamp, Min: min, Max: max}, true
}
return parameterCorrectionRule{}, false
}
func correctionParamProtected(param string, protected ...map[string]struct{}) bool {
for _, values := range protected {
if _, ok := values[normalizeCorrectionParam(param)]; ok {
return true
}
}
return false
}
func callerProtectedCorrectionParameters(request Request, endpointKind string) map[string]struct{} {
if request.OriginalBody == nil {
return nil
}
protected := make(map[string]struct{})
protect := func(param string) {
param = normalizeCorrectionParam(param)
if isSafeCorrectionParam(param) {
protected[param] = struct{}{}
}
}
for key := range request.OriginalBody {
protect(key)
}
if reasoning, ok := request.OriginalBody["reasoning"].(map[string]any); ok {
if _, explicit := reasoning["effort"]; explicit {
if endpointKind == "chat.completions" {
protect("reasoning_effort")
} else {
protect("reasoning.effort")
}
}
}
if request.Kind == "responses" && endpointKind == "chat.completions" {
if _, explicit := request.OriginalBody["max_output_tokens"]; explicit {
protect("max_completion_tokens")
}
if _, explicit := request.OriginalBody["top_logprobs"]; explicit {
protect("top_logprobs")
protect("logprobs")
}
if responseIncludeContains(request.OriginalBody["include"], "message.output_text.logprobs") {
protect("logprobs")
}
if text, ok := request.OriginalBody["text"].(map[string]any); ok {
if _, explicit := text["verbosity"]; explicit {
protect("verbosity")
}
}
}
if endpointKind == "images.generations" || endpointKind == "images.edits" {
// The Gateway canonicalizes image output independently of the provider's
// URL/Base64 preference. Allow an explicitly supplied response_format to
// be removed only after the selected upstream rejects it as unsupported.
delete(protected, "response_format")
}
return protected
}
func responseIncludeContains(value any, target string) bool {
switch values := value.(type) {
case []any:
for _, value := range values {
if stringFromAny(value) == target {
return true
}
}
case []string:
for _, value := range values {
if value == target {
return true
}
}
}
return false
}
func requestIDFromParameterError(err error) string {
var clientErr *ClientError
if errors.As(err, &clientErr) {
return clientErr.RequestID
}
return ""
}
func structuredParameterError(clientErr *ClientError) (string, string, string) {
param := strings.TrimSpace(clientErr.Param)
code := strings.TrimSpace(clientErr.Code)
message := strings.TrimSpace(clientErr.Message)
if clientErr.Wire == nil {
return param, code, message
}
errorObject, _ := clientErr.Wire.Body["error"].(map[string]any)
if errorObject == nil {
return param, code, message
}
if value := strings.TrimSpace(stringFromAny(errorObject["param"])); value != "" {
param = value
}
if value := strings.TrimSpace(stringFromAny(errorObject["code"])); value != "" {
code = value
}
if value := strings.TrimSpace(stringFromAny(errorObject["message"])); value != "" {
message = value
}
return param, code, message
}
func anchoredParameterFromMessage(message string) string {
trimmed := strings.TrimSpace(message)
for _, pattern := range []*regexp.Regexp{anchoredParameterPattern, restrictedParameterPattern} {
if match := pattern.FindStringSubmatch(trimmed); len(match) >= 2 {
return match[1]
}
}
return ""
}
func normalizeCorrectionParam(param string) string {
return strings.ToLower(strings.Trim(strings.TrimSpace(param), "`'\"[]"))
}
func isSafeCorrectionParam(param string) bool {
_, ok := safeUpstreamCorrectionParameters[param]
return ok
}
func deriveConflictCorrection(param string, message string) (parameterCorrectionRule, bool) {
if !strings.Contains(message, "cannot be used together") &&
!strings.Contains(message, "mutually exclusive") &&
!strings.Contains(message, "not allowed when") {
return parameterCorrectionRule{}, false
}
if strings.Contains(message, "thinking_budget") && strings.Contains(message, "reasoning_effort") {
return parameterCorrectionRule{Param: "reasoning_effort", Action: parameterCorrectionRemove}, true
}
return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
}
func nearestErrorSupportedReasoningEffort(message string, requested string, candidate store.RuntimeModelCandidate) string {
allowed := make(map[string]struct{})
lowerMessage := strings.ToLower(message)
for _, effort := range chatReasoningEffortOrder {
if regexp.MustCompile(`\b` + regexp.QuoteMeta(effort) + `\b`).MatchString(lowerMessage) {
allowed[effort] = struct{}{}
}
}
if len(allowed) == 0 {
if resolved, state := resolveCandidateReasoningEffort(requested, candidate); state == reasoningCapabilitySupported {
return resolved
}
return ""
}
requestedIndex := reasoningEffortIndex(requested)
best := ""
bestDistance := len(chatReasoningEffortOrder) + 1
for index, effort := range chatReasoningEffortOrder {
if _, ok := allowed[effort]; !ok {
continue
}
distance := int(math.Abs(float64(index - requestedIndex)))
if distance < bestDistance {
best = effort
bestDistance = distance
}
}
return best
}
func parseCorrectionScalar(raw string) (any, bool) {
trimmed := strings.Trim(strings.TrimSpace(raw), "`'\"")
switch strings.ToLower(trimmed) {
case "true":
return true, true
case "false":
return false, true
}
if number, err := strconv.ParseFloat(trimmed, 64); err == nil {
if math.Trunc(number) == number {
return int(number), true
}
return number, true
}
if isOpenAIReasoningEffort(strings.ToLower(trimmed)) {
return strings.ToLower(trimmed), true
}
return nil, false
}
func explicitNumericBounds(message string) (*float64, *float64, bool) {
if match := betweenPattern.FindStringSubmatch(message); len(match) == 3 {
minimum, minErr := strconv.ParseFloat(match[1], 64)
maximum, maxErr := strconv.ParseFloat(match[2], 64)
if minErr == nil && maxErr == nil {
return &minimum, &maximum, true
}
}
var minimum *float64
var maximum *float64
if match := minimumPattern.FindStringSubmatch(message); len(match) == 2 {
if value, err := strconv.ParseFloat(match[1], 64); err == nil {
minimum = &value
}
}
if match := maximumPattern.FindStringSubmatch(message); len(match) == 2 {
if value, err := strconv.ParseFloat(match[1], 64); err == nil {
maximum = &value
}
}
return minimum, maximum, minimum != nil || maximum != nil
}
func isOutputLimitCorrectionParam(param string) bool {
return param == "max_tokens" || param == "max_completion_tokens" || param == "max_output_tokens"
}
func applyParameterCorrectionRule(body map[string]any, rule parameterCorrectionRule) bool {
current := valueAtCorrectionPath(body, rule.Param)
switch rule.Action {
case parameterCorrectionRemove:
if current == nil {
return false
}
deleteAtCorrectionPath(body, rule.Param)
return true
case parameterCorrectionSet:
if current == nil {
return false
}
if fmt.Sprint(current) == fmt.Sprint(rule.Value) {
return false
}
setAtCorrectionPath(body, rule.Param, rule.Value)
return true
case parameterCorrectionClamp:
number, ok := finiteFloatFromAny(current)
if !ok {
return false
}
corrected := number
if rule.Min != nil && corrected < *rule.Min {
corrected = *rule.Min
}
if rule.Max != nil && corrected > *rule.Max {
corrected = *rule.Max
}
if corrected == number {
return false
}
if math.Trunc(corrected) == corrected {
setAtCorrectionPath(body, rule.Param, int(corrected))
} else {
setAtCorrectionPath(body, rule.Param, corrected)
}
return true
default:
return false
}
}
func valueAtCorrectionPath(body map[string]any, param string) any {
if param != "reasoning.effort" {
return body[param]
}
reasoning, _ := body["reasoning"].(map[string]any)
return reasoning["effort"]
}
func setAtCorrectionPath(body map[string]any, param string, value any) {
if param != "reasoning.effort" {
body[param] = value
return
}
reasoning, _ := body["reasoning"].(map[string]any)
if reasoning == nil {
reasoning = map[string]any{}
body["reasoning"] = reasoning
}
reasoning["effort"] = value
}
func deleteAtCorrectionPath(body map[string]any, param string) {
if param != "reasoning.effort" {
delete(body, param)
return
}
reasoning, _ := body["reasoning"].(map[string]any)
delete(reasoning, "effort")
if len(reasoning) == 0 {
delete(body, "reasoning")
}
}
@@ -0,0 +1,444 @@
package clients
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"sync"
"sync/atomic"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func correctionTestRequest(server *httptest.Server, body map[string]any) Request {
return Request{
Kind: "chat.completions",
Model: "Qwen3.8-Max-Preview",
Body: body,
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
BaseURL: server.URL,
ProviderModelName: "qwen3.8-max-preview",
Credentials: map[string]any{"apiKey": "test-key"},
},
}
}
func writeCorrectionSuccess(w http.ResponseWriter, model any) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "chatcmpl-corrected", "object": "chat.completion", "model": model,
"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "ok"}}},
"usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
})
}
func writeParameterError(w http.ResponseWriter, param string, message string) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusBadRequest)
_ = json.NewEncoder(w).Encode(map[string]any{
"error": map[string]any{"code": "invalid_parameter", "param": param, "message": message},
})
}
func TestOpenAIParameterCorrectionLearnsOnlyAfterSuccess(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["n"] != float64(1) {
writeParameterError(w, "n", "Parameter n must be set to 1.")
return
}
writeCorrectionSuccess(w, body["model"])
}))
defer server.Close()
cache := NewParameterCorrectionCache()
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
request := correctionTestRequest(server, map[string]any{
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
"n": 2,
})
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("first corrected request failed: %v", err)
}
if got := requests.Load(); got != 2 {
t.Fatalf("first logical request should send twice, got %d", got)
}
if cache.size() != 1 {
t.Fatalf("successful correction should be cached, size=%d", cache.size())
}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("cached request failed: %v", err)
}
if got := requests.Load(); got != 3 {
t.Fatalf("second logical request should hit cache on first send, got total %d", got)
}
}
func TestOpenAIParameterCorrectionLearnsDashScopeThinkingNConstraint(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["n"] != float64(1) {
writeParameterError(w, "n", "The n parameter must be 1 when enable_thinking is true")
return
}
writeCorrectionSuccess(w, body["model"])
}))
defer server.Close()
cache := NewParameterCorrectionCache()
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
"messages": []any{}, "n": 2,
}))
if err != nil {
t.Fatalf("DashScope n constraint should be corrected: %v", err)
}
if cache.size() != 1 {
t.Fatalf("successful DashScope n correction should be cached, size=%d", cache.size())
}
}
func TestOpenAIParameterCorrectionFailureIsNotCached(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
writeParameterError(w, "n", "Parameter n must be set to 1.")
}))
defer server.Close()
cache := NewParameterCorrectionCache()
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
"messages": []any{}, "n": 2,
}))
if err == nil {
t.Fatal("expected upstream error")
}
if cache.size() != 0 {
t.Fatalf("failed correction must not be cached, size=%d", cache.size())
}
}
func TestOpenAIParameterCorrectionStopsAfterTwoRetries(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["n"] != float64(1) {
writeParameterError(w, "n", "Parameter n must be set to 1.")
return
}
if body["temperature"] != float64(1) {
writeParameterError(w, "temperature", "Parameter temperature must be set to 1.")
return
}
writeParameterError(w, "top_p", "Parameter top_p must be set to 1.")
}))
defer server.Close()
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
"messages": []any{}, "n": 2, "temperature": 2, "top_p": 2,
}))
if err == nil {
t.Fatal("expected third parameter error")
}
if got := requests.Load(); got != 3 {
t.Fatalf("expected initial request plus two correction retries, got %d", got)
}
}
func TestOpenAIParameterCorrectionEvictsConflictingRule(t *testing.T) {
cache := NewParameterCorrectionCache()
var seen []float64
var mu sync.Mutex
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
n, _ := body["n"].(float64)
mu.Lock()
seen = append(seen, n)
mu.Unlock()
if n != 2 {
writeParameterError(w, "n", "Parameter n must be set to 2.")
return
}
writeCorrectionSuccess(w, body["model"])
}))
defer server.Close()
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 9})
scope := newParameterCorrectionScope(request, "chat.completions")
cache.commit(scope, []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("conflicting cache correction failed: %v", err)
}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("relearned cache request failed: %v", err)
}
mu.Lock()
defer mu.Unlock()
if fmt.Sprint(seen) != "[1 2 2]" {
t.Fatalf("expected old rule, relearn retry, then new cache hit; got %v", seen)
}
}
func TestOpenAIParameterCorrectionNeverTouchesDangerousFields(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
requests.Add(1)
writeParameterError(w, "model", "Unsupported parameter: model")
}))
defer server.Close()
messages := []any{map[string]any{"role": "user", "content": "secret text"}}
request := correctionTestRequest(server, map[string]any{"messages": messages})
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
_, err := client.Run(context.Background(), request)
if err == nil {
t.Fatal("expected unsafe parameter error")
}
if requests.Load() != 1 {
t.Fatalf("unsafe parameter must not retry, got %d requests", requests.Load())
}
if fmt.Sprint(request.Body["messages"]) != fmt.Sprint(messages) {
t.Fatal("messages were unexpectedly modified")
}
}
func TestOpenAIParameterCorrectionWorksBeforeStreamingOutput(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["n"] != float64(1) {
writeParameterError(w, "n", "Parameter n must be set to 1.")
return
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-stream\",\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
}))
defer server.Close()
var deltas []string
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 2, "stream": true})
request.Stream = true
request.StreamDelta = func(event StreamDeltaEvent) error {
deltas = append(deltas, event.Text)
return nil
}
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("stream correction failed: %v", err)
}
if requests.Load() != 2 || fmt.Sprint(deltas) != "[ok]" {
t.Fatalf("unexpected stream correction requests=%d deltas=%v", requests.Load(), deltas)
}
}
func TestOpenAIResponsesParameterCorrectionLearnsAfterSuccess(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["top_p"] != float64(1) {
writeParameterError(w, "top_p", "Parameter top_p must be set to 1.")
return
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "resp-corrected", "object": "response", "status": "completed",
"output": []any{}, "usage": map[string]any{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
})
}))
defer server.Close()
cache := NewParameterCorrectionCache()
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
request := Request{
Kind: "responses",
Model: "Qwen3.8-Max-Preview",
UpstreamProtocol: ProtocolOpenAIResponses,
Body: map[string]any{"input": "hello", "top_p": 2, "reasoning": map[string]any{"effort": "none"}},
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai", BaseURL: server.URL,
ProviderModelName: "qwen3.8-max-preview", Credentials: map[string]any{"apiKey": "test-key"},
},
}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("first corrected Responses request failed: %v", err)
}
if _, err := client.Run(context.Background(), request); err != nil {
t.Fatalf("cached Responses request failed: %v", err)
}
if requests.Load() != 3 || cache.size() != 1 {
t.Fatalf("expected retry then Responses cache hit, requests=%d cache=%d", requests.Load(), cache.size())
}
}
func TestParameterCorrectionCacheIsConcurrentAndBounded(t *testing.T) {
cache := NewParameterCorrectionCache()
var group sync.WaitGroup
for index := 0; index < 1024; index++ {
group.Add(1)
go func(index int) {
defer group.Done()
scope := parameterCorrectionScope{Provider: "p", BaseURL: "https://example.com", Protocol: "chat", Kind: "chat", Model: fmt.Sprint(index)}
cache.commit(scope, []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
cache.apply(scope, map[string]any{"n": 2})
}(index)
}
group.Wait()
if cache.size() > parameterCorrectionCacheCapacity {
t.Fatalf("cache exceeded capacity: %d", cache.size())
}
}
func TestOpenAIParameterCorrectionNeverChangesCallerExplicitParameterColdOrHot(t *testing.T) {
t.Run("cold", func(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
requests.Add(1)
writeParameterError(w, "n", "Parameter n must be set to 1.")
}))
defer server.Close()
body := map[string]any{"messages": []any{}, "n": 2}
request := correctionTestRequest(server, body)
request.OriginalBody = body
cache := NewParameterCorrectionCache()
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: cache}).Run(context.Background(), request)
if err == nil || requests.Load() != 1 || cache.size() != 0 {
t.Fatalf("explicit cold parameter was corrected: requests=%d cache=%d err=%v", requests.Load(), cache.size(), err)
}
})
t.Run("hot", func(t *testing.T) {
var captured float64
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
captured, _ = body["n"].(float64)
writeCorrectionSuccess(w, body["model"])
}))
defer server.Close()
body := map[string]any{"messages": []any{}, "n": 2}
request := correctionTestRequest(server, body)
request.OriginalBody = body
cache := NewParameterCorrectionCache()
cache.commit(newParameterCorrectionScope(request, "chat.completions"), []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
if _, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: cache}).Run(context.Background(), request); err != nil {
t.Fatal(err)
}
if captured != 2 {
t.Fatalf("cached correction changed explicit n: %v", captured)
}
})
}
func TestOpenAIParameterCorrectionMayChangeGatewayInjectedParameter(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["n"] != float64(1) {
writeParameterError(w, "n", "Parameter n must be set to 1.")
return
}
writeCorrectionSuccess(w, body["model"])
}))
defer server.Close()
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 2})
request.OriginalBody = map[string]any{"messages": []any{}}
if _, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request); err != nil {
t.Fatal(err)
}
if requests.Load() != 2 {
t.Fatalf("Gateway-injected parameter was not corrected: requests=%d", requests.Load())
}
}
func TestOpenAIProviderAdaptationRejectsExplicitSemanticChange(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("semantically incompatible explicit parameter must fail before upstream")
}))
defer server.Close()
body := map[string]any{"messages": []any{}, "reasoning_effort": "none", "temperature": 0.2}
request := correctionTestRequest(server, body)
request.OriginalBody = body
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request)
if ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "reasoning_effort" {
t.Fatalf("expected exact reasoning incompatibility, got %v param=%q", err, ErrorParam(err))
}
}
func TestOpenAIProviderAdaptationRejectsInexactDeepSeekEffort(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
t.Fatal("inexact DeepSeek effort must fail before upstream")
}))
defer server.Close()
body := map[string]any{"messages": []any{}, "reasoning_effort": "low"}
request := Request{
Kind: "chat.completions", Model: "deepseek-v4-pro", Body: body, OriginalBody: body,
Candidate: store.RuntimeModelCandidate{
Provider: "deepseek-openai", BaseURL: server.URL, ProviderModelName: "deepseek-v4-pro", Credentials: map[string]any{"apiKey": "test-key"},
},
}
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request)
if ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "reasoning_effort" {
t.Fatalf("expected exact DeepSeek reasoning incompatibility, got %v param=%q", err, ErrorParam(err))
}
}
func TestResponsesFallbackProtectsExplicitMappedOutputLimitFromCorrection(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["max_completion_tokens"] != float64(128) || body["max_tokens"] != nil {
t.Fatalf("Responses output limit mapped incorrectly: %+v", body)
}
writeParameterError(w, "max_completion_tokens", "Unsupported parameter: max_completion_tokens")
}))
defer server.Close()
body := map[string]any{"input": "hello", "max_output_tokens": 128}
request := Request{
Kind: "responses", Model: "demo", Body: body, OriginalBody: body, UpstreamProtocol: ProtocolOpenAIChatCompletions,
Candidate: store.RuntimeModelCandidate{
Provider: "openai", BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "test-key"},
},
}
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request)
if err == nil || requests.Load() != 1 {
t.Fatalf("explicit mapped output limit was silently corrected: requests=%d err=%v", requests.Load(), err)
}
}
func TestResponsesFallbackProtectsIncludeMappedLogprobsFromCorrection(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
var body map[string]any
_ = json.NewDecoder(r.Body).Decode(&body)
if body["logprobs"] != true {
t.Fatalf("Responses include did not enable Chat logprobs: %+v", body)
}
writeParameterError(w, "logprobs", "Unsupported parameter: logprobs")
}))
defer server.Close()
body := map[string]any{"input": "hello", "include": []any{"message.output_text.logprobs"}}
request := Request{
Kind: "responses", Model: "demo", Body: body, OriginalBody: body, UpstreamProtocol: ProtocolOpenAIChatCompletions,
Candidate: store.RuntimeModelCandidate{
Provider: "openai", BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "test-key"},
},
}
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request)
if err == nil || requests.Load() != 1 {
t.Fatalf("explicit include logprobs was silently corrected: requests=%d err=%v", requests.Load(), err)
}
}
+22 -18
View File
@@ -45,7 +45,15 @@ func (c providerTaskClient) Run(ctx context.Context, request Request) (Response,
requestID := upstreamTaskID
var submitResult map[string]any
if upstreamTaskID == "" {
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
result, id, err := c.submit(ctx, request, payload)
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
return Response{}, notifyErr
}
}
if err != nil {
return Response{}, annotateResponseError(err, id, startedAt, time.Now())
}
@@ -182,7 +190,7 @@ func providerPostJSON(ctx context.Context, client *http.Client, url string, body
applyProviderAuth(req, credentials, auth)
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
@@ -197,7 +205,7 @@ func providerGetJSON(ctx context.Context, client *http.Client, url string, crede
applyProviderAuth(req, credentials, auth)
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
@@ -332,24 +340,20 @@ func hasProviderTaskResult(result map[string]any) bool {
return result["data"] != nil || valueAtPath(result, "data.result") != nil || valueAtPath(result, "data.audio") != nil || valueAtPath(result, "output.image_urls") != nil || valueAtPath(result, "output.video_url") != nil || valueAtPath(result, "Response.ResultVideoUrl") != nil || valueAtPath(result, "Response.ResultImages") != nil || result["audio_url"] != nil || result["urls"] != nil
}
func normalizeProviderTaskResult(request Request, spec providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
out := cloneMapAny(result)
out["status"] = "success"
func normalizeProviderTaskResult(request Request, _ providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
data, ok := result["data"].([]any)
if !ok {
data = providerTaskData(request, result)
}
out := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": data,
}
if upstreamTaskID != "" {
out["upstream_task_id"] = upstreamTaskID
}
if out["created"] == nil {
out["created"] = time.Now().UnixMilli()
}
if out["model"] == nil {
out["model"] = request.Model
}
if _, ok := out["data"].([]any); !ok {
if out["data"] != nil {
out["raw_data"] = out["data"]
}
out["data"] = providerTaskData(request, result)
}
return out
}
@@ -456,7 +460,7 @@ func providerPollInterval(request Request) time.Duration {
}
func providerPollTimeout(request Request) time.Duration {
return durationFromConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
return durationFromConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
}
func durationFromConfig(config map[string]any, fallback time.Duration, keys ...string) time.Duration {
@@ -0,0 +1,312 @@
package clients
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"net/url"
"os"
"strings"
"sync"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type qwen38LiveProxy struct {
server *httptest.Server
upstreamBase *url.URL
mu sync.Mutex
shapes []map[string]any
}
func newQwen38LiveProxy(t *testing.T, upstream string) *qwen38LiveProxy {
t.Helper()
upstreamBase, err := url.Parse(strings.TrimRight(upstream, "/"))
if err != nil {
t.Fatalf("parse Qwen3.8 live upstream: %v", err)
}
proxy := &qwen38LiveProxy{upstreamBase: upstreamBase}
proxy.server = httptest.NewServer(http.HandlerFunc(proxy.forward))
t.Cleanup(proxy.server.Close)
return proxy
}
func (p *qwen38LiveProxy) forward(w http.ResponseWriter, request *http.Request) {
body, err := io.ReadAll(request.Body)
if err != nil {
http.Error(w, "read request", http.StatusBadRequest)
return
}
var decoded map[string]any
if json.Unmarshal(body, &decoded) == nil {
p.mu.Lock()
p.shapes = append(p.shapes, qwen38SafeRequestShape(decoded))
p.mu.Unlock()
}
target := *p.upstreamBase
target.Path = strings.TrimRight(target.Path, "/") + request.URL.Path
upstreamRequest, err := http.NewRequestWithContext(request.Context(), request.Method, target.String(), bytes.NewReader(body))
if err != nil {
http.Error(w, "create upstream request", http.StatusInternalServerError)
return
}
upstreamRequest.Header.Set("Authorization", request.Header.Get("Authorization"))
upstreamRequest.Header.Set("Content-Type", "application/json")
response, err := http.DefaultClient.Do(upstreamRequest)
if err != nil {
http.Error(w, "upstream request failed", http.StatusBadGateway)
return
}
defer response.Body.Close()
responseBody, err := io.ReadAll(response.Body)
if err != nil {
http.Error(w, "read upstream response", http.StatusBadGateway)
return
}
status := response.StatusCode
if status == http.StatusInternalServerError && strings.Contains(string(responseBody), "The n parameter must be 1 when enable_thinking is true") {
// The deployed server-main currently wraps DashScope's 400 as 500. The
// adapter only restores the original parameter-error status and shape so
// this live test can exercise the local client's correction loop.
status = http.StatusBadRequest
responseBody, _ = json.Marshal(map[string]any{
"error": map[string]any{
"code": "invalid_parameter_error",
"param": "n",
"message": "The n parameter must be 1 when enable_thinking is true",
},
})
}
if contentType := response.Header.Get("Content-Type"); contentType != "" {
w.Header().Set("Content-Type", contentType)
}
w.WriteHeader(status)
_, _ = w.Write(responseBody)
}
func (p *qwen38LiveProxy) requestShapes() []map[string]any {
p.mu.Lock()
defer p.mu.Unlock()
out := make([]map[string]any, len(p.shapes))
copy(out, p.shapes)
return out
}
func qwen38SafeRequestShape(body map[string]any) map[string]any {
shape := map[string]any{}
for _, key := range []string{"model", "stream", "reasoning_effort", "enable_thinking", "thinking_budget", "temperature", "n"} {
if value, ok := body[key]; ok {
shape[key] = value
}
}
if reasoning, ok := body["reasoning"].(map[string]any); ok {
safeReasoning := map[string]any{}
for _, key := range []string{"effort", "summary"} {
if value, exists := reasoning[key]; exists {
safeReasoning[key] = value
}
}
shape["reasoning"] = safeReasoning
}
return shape
}
func qwen38LiveRequest(proxy *qwen38LiveProxy, apiKey string, body map[string]any, stream bool) Request {
body = cloneBody(body)
body["stream"] = stream
return Request{
Kind: "chat.completions",
Model: "Qwen3.8-Max-Preview",
Body: body,
Stream: stream,
UpstreamProtocol: ProtocolOpenAIChatCompletions,
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
BaseURL: proxy.server.URL,
ProviderModelName: "Qwen3.8-Max-Preview",
Credentials: map[string]any{"apiKey": apiKey},
Capabilities: map[string]any{
"text_generate": map[string]any{
"supportThinking": true,
"supportThinkingModeSwitch": false,
"thinkingEffortLevels": []any{"low", "medium", "xhigh"},
},
},
},
}
}
func qwen38LiveBody(extra map[string]any) map[string]any {
body := map[string]any{
"messages": []any{map[string]any{"role": "user", "content": "Reply only OK"}},
"max_tokens": 256,
}
for key, value := range extra {
body[key] = value
}
return body
}
func qwen38ResponseHasReasoning(response Response, streamedReasoning string) bool {
if strings.TrimSpace(streamedReasoning) != "" {
return true
}
choices, _ := response.Result["choices"].([]any)
if len(choices) == 0 {
return false
}
choice, _ := choices[0].(map[string]any)
message, _ := choice["message"].(map[string]any)
reasoning, _ := message["reasoning_content"].(string)
return strings.TrimSpace(reasoning) != ""
}
func qwen38LiveConfig(t *testing.T) (string, string) {
t.Helper()
baseURL := strings.TrimSpace(os.Getenv("QWEN38_LIVE_BASE_URL"))
apiKey := strings.TrimSpace(os.Getenv("QWEN38_LIVE_API_KEY"))
if baseURL == "" || apiKey == "" {
t.Skip("set QWEN38_LIVE_BASE_URL and QWEN38_LIVE_API_KEY to run real-model acceptance")
}
return baseURL, apiKey
}
func TestQwen38LiveStaticNormalizationAndReasoning(t *testing.T) {
baseURL, apiKey := qwen38LiveConfig(t)
proxy := newQwen38LiveProxy(t, baseURL)
client := OpenAIClient{Corrections: NewParameterCorrectionCache()}
tests := []struct {
name string
input map[string]any
stream bool
expected map[string]any
absent []string
allowCompatibilityRetry bool
}{
{name: "none", input: map[string]any{"reasoning_effort": "none"}, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
{name: "minimal stream", input: map[string]any{"reasoning_effort": "minimal"}, stream: true, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
{name: "medium", input: map[string]any{"reasoning_effort": "medium"}, expected: map[string]any{"reasoning_effort": "medium", "enable_thinking": true}},
{name: "high stream", input: map[string]any{"reasoning_effort": "high"}, stream: true, expected: map[string]any{"reasoning_effort": "xhigh", "enable_thinking": true}},
{name: "max", input: map[string]any{"reasoning_effort": "max"}, expected: map[string]any{"reasoning_effort": "xhigh", "enable_thinking": true}},
{name: "forced thinking stream", input: map[string]any{"enable_thinking": false}, stream: true, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
{name: "budget wins", input: map[string]any{"reasoning_effort": "high", "thinking_budget": 128}, expected: map[string]any{"thinking_budget": float64(128), "enable_thinking": true}, absent: []string{"reasoning_effort"}, allowCompatibilityRetry: true},
{name: "temperature floor stream", input: map[string]any{"temperature": 0.1}, stream: true, expected: map[string]any{"temperature": 0.6, "enable_thinking": true}},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
before := len(proxy.requestShapes())
streamedReasoning := ""
request := qwen38LiveRequest(proxy, apiKey, qwen38LiveBody(test.input), test.stream)
if test.stream {
request.StreamDelta = func(event StreamDeltaEvent) error {
streamedReasoning += event.ReasoningContent
return nil
}
}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
defer cancel()
response, err := client.Run(ctx, request)
if err != nil {
shapes := proxy.requestShapes()
t.Fatalf("real Qwen3.8 request failed: %v; safe request shapes=%+v", err, shapes[before:])
}
if !qwen38ResponseHasReasoning(response, streamedReasoning) {
t.Fatal("real Qwen3.8 response did not contain reasoning_content")
}
shapes := proxy.requestShapes()
submissions := len(shapes) - before
if submissions != 1 && !(test.allowCompatibilityRetry && submissions == 2) {
t.Fatalf("unexpected submission count, before=%d after=%d safe request shapes=%+v", before, len(shapes), shapes[before:])
}
shape := shapes[before]
for key, expected := range test.expected {
if normalizedJSONScalar(shape[key]) != normalizedJSONScalar(expected) {
t.Fatalf("%s expected %#v, got %#v in shape %+v", key, expected, shape[key], shape)
}
}
for _, key := range test.absent {
if _, ok := shape[key]; ok {
t.Fatalf("%s must be absent in shape %+v", key, shape)
}
}
t.Logf("case=%s response_id=%s correction=%t cache_hit=false", test.name, response.RequestID, submissions > 1)
})
}
}
func TestQwen38LiveCorrectionCacheAndRestart(t *testing.T) {
baseURL, apiKey := qwen38LiveConfig(t)
proxy := newQwen38LiveProxy(t, baseURL)
request := qwen38LiveRequest(proxy, apiKey, qwen38LiveBody(map[string]any{"reasoning_effort": "none", "n": 2}), false)
cache := NewParameterCorrectionCache()
client := OpenAIClient{Corrections: cache}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
defer cancel()
first, err := client.Run(ctx, request)
if err != nil || !qwen38ResponseHasReasoning(first, "") {
t.Fatalf("cold corrected request failed or lacked reasoning: %v", err)
}
second, err := client.Run(ctx, request)
if err != nil || !qwen38ResponseHasReasoning(second, "") {
t.Fatalf("cached corrected request failed or lacked reasoning: %v", err)
}
shapes := proxy.requestShapes()
if len(shapes) != 3 || normalizedJSONScalar(shapes[0]["n"]) != float64(2) || normalizedJSONScalar(shapes[1]["n"]) != float64(1) || normalizedJSONScalar(shapes[2]["n"]) != float64(1) {
t.Fatalf("expected cold n=2 -> corrected n=1 -> cached n=1, got %+v", shapes)
}
restarted := OpenAIClient{Corrections: NewParameterCorrectionCache()}
third, err := restarted.Run(ctx, request)
if err != nil || !qwen38ResponseHasReasoning(third, "") {
t.Fatalf("post-restart corrected request failed or lacked reasoning: %v", err)
}
shapes = proxy.requestShapes()
if len(shapes) != 5 || normalizedJSONScalar(shapes[3]["n"]) != float64(2) || normalizedJSONScalar(shapes[4]["n"]) != float64(1) {
t.Fatalf("new process cache should relearn n rule, got %+v", shapes)
}
t.Logf("response_ids=%s,%s,%s correction=n:set:1 cache_hit_sequence=false,true,false", first.RequestID, second.RequestID, third.RequestID)
}
func TestQwen38LiveNativeResponsesNormalization(t *testing.T) {
baseURL, apiKey := qwen38LiveConfig(t)
proxy := newQwen38LiveProxy(t, baseURL)
client := OpenAIClient{Corrections: NewParameterCorrectionCache()}
request := Request{
Kind: "responses",
Model: "Qwen3.8-Max-Preview",
UpstreamProtocol: ProtocolOpenAIResponses,
Body: map[string]any{
"input": "Reply only OK", "max_output_tokens": 256,
"reasoning": map[string]any{"effort": "none"},
},
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai", BaseURL: proxy.server.URL,
ProviderModelName: "Qwen3.8-Max-Preview", Credentials: map[string]any{"apiKey": apiKey},
Capabilities: map[string]any{"text_generate": map[string]any{
"supportThinking": true, "supportThinkingModeSwitch": false,
"thinkingEffortLevels": []any{"low", "medium", "xhigh"},
}},
},
}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
defer cancel()
response, err := client.Run(ctx, request)
if err != nil {
t.Fatalf("real Qwen3.8 Responses request failed: %v; safe shapes=%+v", err, proxy.requestShapes())
}
shapes := proxy.requestShapes()
if len(shapes) != 1 {
t.Fatalf("expected one Responses submission, got %+v", shapes)
}
reasoning, _ := shapes[0]["reasoning"].(map[string]any)
if reasoning["effort"] != "low" {
t.Fatalf("Responses none was not normalized to low: %+v", shapes)
}
t.Logf("case=responses-none response_id=%s correction=false cache_hit=false", response.RequestID)
}
File diff suppressed because it is too large Load Diff
@@ -74,7 +74,7 @@ func TestOpenAIResponsesChatFallbackPreservesHistoryToolsUsageAndReasoningIntern
if len(messages) != 4 {
t.Fatalf("expected prior user/assistant plus current system/user/tool messages, got %+v", messages)
}
if body["max_tokens"] != float64(128) || body["reasoning_effort"] != "high" {
if body["max_completion_tokens"] != float64(128) || body["reasoning_effort"] != "high" {
t.Fatalf("expected mapped max/reasoning fields: %+v", body)
}
tools, _ := body["tools"].([]any)
@@ -165,6 +165,27 @@ func TestOpenAIResponsesNativeStreamForwardsEventsWithoutChatAggregation(t *test
}
}
func TestOpenAIResponsesNativeStreamAcceptsIncompleteTerminalEvent(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_incomplete\",\"object\":\"response\",\"status\":\"incomplete\",\"output\":[],\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n"))
}))
defer server.Close()
events := make([]StreamDeltaEvent, 0, 1)
response, err := (OpenAIClient{}).Run(context.Background(), Request{
Kind: "responses", Model: "Demo", Body: map[string]any{"input": "hello", "stream": true}, Stream: true,
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "secret"}},
UpstreamProtocol: ProtocolOpenAIResponses,
StreamDelta: func(event StreamDeltaEvent) error { events = append(events, event); return nil },
})
if err != nil {
t.Fatal(err)
}
if len(events) != 1 || events[0].Event["type"] != "response.incomplete" || response.Result["status"] != "incomplete" || response.UpstreamResponseID != "resp_incomplete" {
t.Fatalf("native incomplete terminal event was not preserved: response=%+v events=%+v", response, events)
}
}
func TestOpenAIResponsesChatFallbackStreamsFunctionArgumentFragments(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
@@ -220,6 +241,54 @@ func TestOpenAIResponsesChatFallbackStreamsFunctionArgumentFragments(t *testing.
}
}
func TestOpenAIResponsesChatFallbackStreamsCustomRefusalAndLogprobsEndToEnd(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-custom\",\"object\":\"chat.completion.chunk\",\"model\":\"demo\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hi\",\"refusal\":\"no\",\"annotations\":[{\"type\":\"url_citation\"}],\"tool_calls\":[{\"index\":0,\"id\":\"call_custom\",\"type\":\"custom\",\"custom\":{\"name\":\"shell\",\"input\":\"pw\"}}]},\"logprobs\":{\"content\":[{\"token\":\"hi\",\"logprob\":-0.1}]},\"finish_reason\":null}]}\n\n"))
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-custom\",\"object\":\"chat.completion.chunk\",\"model\":\"demo\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"type\":\"custom\",\"custom\":{\"input\":\"d\"}}]},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}}\n\n"))
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer server.Close()
events := make([]StreamDeltaEvent, 0)
response, err := (OpenAIClient{}).Run(context.Background(), Request{
Kind: "responses", Model: "Demo", Body: map[string]any{
"input": "call it", "stream": true, "include": []any{"message.output_text.logprobs"},
"tools": []any{map[string]any{"type": "custom", "name": "shell", "format": map[string]any{"type": "text"}}},
}, Stream: true,
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "secret"}},
UpstreamProtocol: ProtocolOpenAIChatCompletions, PublicResponseID: "resp_12345678901234567890123456789012",
StreamDelta: func(event StreamDeltaEvent) error { events = append(events, event); return nil },
})
if err != nil {
t.Fatal(err)
}
output, _ := response.Result["output"].([]any)
if len(output) != 2 {
t.Fatalf("expected message and custom tool call: %+v", response.Result)
}
message, _ := output[0].(map[string]any)
parts, _ := message["content"].([]any)
text, _ := parts[0].(map[string]any)
refusal, _ := parts[1].(map[string]any)
custom, _ := output[1].(map[string]any)
if custom["type"] != "custom_tool_call" || custom["call_id"] != "call_custom" || custom["input"] != "pwd" {
t.Fatalf("custom tool stream was not aggregated: %+v", output)
}
if refusal["refusal"] != "no" || len(text["logprobs"].([]any)) != 1 || len(text["annotations"].([]any)) != 1 {
t.Fatalf("refusal/logprobs were not aggregated: %+v", message)
}
types := make([]string, 0, len(events))
for _, event := range events {
types = append(types, stringFromAny(event.Event["type"]))
}
for _, required := range []string{"response.custom_tool_call_input.delta", "response.custom_tool_call_input.done", "response.refusal.delta", "response.refusal.done", "response.completed"} {
if !containsTestString(types, required) {
t.Fatalf("missing %s in converted stream: %v", required, types)
}
}
}
func TestChatResultToResponseMapsIncompleteFinishReason(t *testing.T) {
response := ChatResultToResponse(map[string]any{
"choices": []any{map[string]any{"finish_reason": "length", "message": map[string]any{"role": "assistant", "content": "partial"}}},
@@ -233,6 +302,23 @@ func TestChatResultToResponseMapsIncompleteFinishReason(t *testing.T) {
}
}
func TestChatResultToResponseConvertsLegacyFunctionCall(t *testing.T) {
response := ChatResultToResponse(map[string]any{
"choices": []any{map[string]any{"finish_reason": "function_call", "message": map[string]any{
"role": "assistant", "content": nil,
"function_call": map[string]any{"name": "legacy_lookup", "arguments": "{\"city\":\"Paris\"}"},
}}},
}, "resp_12345678901234567890123456789012", "demo", map[string]any{})
output, _ := response["output"].([]any)
if len(output) != 1 {
t.Fatalf("legacy function_call was lost: %+v", response)
}
call, _ := output[0].(map[string]any)
if call["type"] != "function_call" || call["name"] != "legacy_lookup" || call["arguments"] != "{\"city\":\"Paris\"}" {
t.Fatalf("legacy function_call was not converted: %+v", call)
}
}
func TestNativeResponsesStreamSupportsMultilineDataFrames(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
@@ -277,6 +363,7 @@ func TestChatResponsesStreamAdapterEmitsTextItemLifecycle(t *testing.T) {
want := []string{
"response.created", "response.in_progress", "response.output_item.added", "response.content_part.added",
"response.output_text.delta", "response.output_text.done", "response.content_part.done", "response.output_item.done",
"response.completed",
}
if strings.Join(events, ",") != strings.Join(want, ",") {
t.Fatalf("unexpected text event lifecycle got=%v want=%v", events, want)
@@ -285,8 +372,8 @@ func TestChatResponsesStreamAdapterEmitsTextItemLifecycle(t *testing.T) {
func TestResponsesChatFallbackRejectsBuiltInToolsAndUnknownParameters(t *testing.T) {
_, err := ResponsesRequestToChat(map[string]any{"input": "hello", "tools": []any{map[string]any{"type": "web_search_preview"}}}, nil)
if ErrorCode(err) != "unsupported_response_tool" {
t.Fatalf("expected unsupported_response_tool, got %v", err)
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "tools[0].type" {
t.Fatalf("expected precise unsupported_response_parameter, got %v param=%q", err, ErrorParam(err))
}
_, err = ResponsesRequestToChat(map[string]any{"input": "hello", "conversation": "conv_1"}, nil)
if ErrorCode(err) != "unsupported_response_parameter" {
@@ -317,6 +404,223 @@ func TestResponsesChatFallbackKeepsClientManagedStateAuthoritative(t *testing.T)
}
}
func TestResponsesChatFallbackMapsMultimodalCustomToolsAndCallHistory(t *testing.T) {
body, err := ResponsesRequestToChat(map[string]any{
"input": []any{
map[string]any{"type": "message", "role": "user", "content": []any{
map[string]any{"type": "input_text", "text": "hello", "prompt_cache_breakpoint": map[string]any{"type": "ephemeral"}},
map[string]any{"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high"},
map[string]any{"type": "input_file", "file_id": "file_1", "filename": "a.pdf"},
map[string]any{"type": "input_audio", "input_audio": map[string]any{"data": "AAAA", "format": "wav"}},
}},
map[string]any{"type": "function_call", "call_id": "call_fn", "name": "lookup", "arguments": "{\"x\":1}"},
map[string]any{"type": "function_call_output", "call_id": "call_fn", "output": map[string]any{"ok": true}},
map[string]any{"type": "custom_tool_call", "call_id": "call_custom", "name": "shell", "input": "pwd"},
map[string]any{"type": "custom_tool_call_output", "call_id": "call_custom", "output": "done"},
},
"tools": []any{
map[string]any{"type": "function", "name": "lookup", "parameters": map[string]any{"type": "object"}, "strict": true},
map[string]any{"type": "custom", "name": "shell", "description": "run", "format": map[string]any{"type": "text"}},
},
"tool_choice": map[string]any{"type": "custom", "name": "shell"},
"include": []any{"message.output_text.logprobs"}, "max_output_tokens": 256,
}, nil)
if err != nil {
t.Fatal(err)
}
if body["max_completion_tokens"] != 256 || body["logprobs"] != true {
t.Fatalf("token/logprobs mapping is incomplete: %+v", body)
}
messages, _ := body["messages"].([]any)
if len(messages) != 5 {
t.Fatalf("call history conversion changed item ownership: %+v", messages)
}
first, _ := messages[0].(map[string]any)
parts, _ := first["content"].([]any)
if len(parts) != 4 {
t.Fatalf("multimodal content was dropped: %+v", parts)
}
text, _ := parts[0].(map[string]any)
image, _ := parts[1].(map[string]any)
imageURL, _ := image["image_url"].(map[string]any)
file, _ := parts[2].(map[string]any)
fileValue, _ := file["file"].(map[string]any)
if text["prompt_cache_breakpoint"] == nil || imageURL["detail"] != "high" || fileValue["file_id"] != "file_1" || fileValue["filename"] != "a.pdf" {
t.Fatalf("multimodal nested fields changed: %+v", parts)
}
customMessage, _ := messages[3].(map[string]any)
customCalls, _ := customMessage["tool_calls"].([]any)
customCall, _ := customCalls[0].(map[string]any)
custom, _ := customCall["custom"].(map[string]any)
if customCall["type"] != "custom" || custom["name"] != "shell" || custom["input"] != "pwd" {
t.Fatalf("custom call history changed: %+v", customMessage)
}
tools, _ := body["tools"].([]any)
customTool, _ := tools[1].(map[string]any)
if customTool["type"] != "custom" || customTool["custom"] == nil {
t.Fatalf("custom tool definition changed: %+v", tools)
}
choice, _ := body["tool_choice"].(map[string]any)
if choice["type"] != "custom" || choice["custom"] == nil {
t.Fatalf("custom tool choice changed: %+v", choice)
}
}
func TestResponsesChatFallbackMapsTextToolOutputAndRejectsNativeOnlyToolOutputContent(t *testing.T) {
body, err := ResponsesRequestToChat(map[string]any{"input": []any{map[string]any{
"type": "function_call_output", "call_id": "call_1", "output": []any{map[string]any{
"type": "input_text", "text": "done", "prompt_cache_breakpoint": map[string]any{"mode": "explicit"},
}},
}}}, nil)
if err != nil {
t.Fatal(err)
}
messages, _ := body["messages"].([]any)
message, _ := messages[0].(map[string]any)
parts, _ := message["content"].([]any)
part, _ := parts[0].(map[string]any)
if part["type"] != "text" || part["text"] != "done" || part["prompt_cache_breakpoint"] == nil {
t.Fatalf("text tool output was not preserved: %+v", message)
}
err = ValidateResponsesChatFallback(map[string]any{"input": []any{map[string]any{
"type": "function_call_output", "call_id": "call_1", "output": []any{map[string]any{
"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high",
}},
}}})
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "input[0].output[0].type" {
t.Fatalf("native-only tool output must be rejected precisely, got %v param=%q", err, ErrorParam(err))
}
}
func TestResponsesChatFallbackAcceptsOnlyNoopNativeFields(t *testing.T) {
if err := ValidateResponsesChatFallback(map[string]any{
"input": "hello", "background": false, "context_management": map[string]any{}, "conversation": "",
"prompt": nil, "truncation": "disabled", "reasoning": map[string]any{"summary": nil},
}); err != nil {
t.Fatalf("no-op native fields should be accepted: %v", err)
}
for param, value := range map[string]any{
"background": true, "conversation": "conv_1", "max_tool_calls": 2, "prompt": map[string]any{"id": "pmpt_1"}, "truncation": "auto",
} {
err := ValidateResponsesChatFallback(map[string]any{"input": "hello", param: value})
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != param {
t.Fatalf("expected precise rejection for %s, got %v param=%q", param, err, ErrorParam(err))
}
}
err := ValidateResponsesChatFallback(map[string]any{"input": []any{map[string]any{
"type": "custom_tool_call", "call_id": "call_1", "name": "tool", "input": "x", "namespace": "native-only",
}}})
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "input[0].namespace" {
t.Fatalf("custom call namespace must be rejected precisely, got %v param=%q", err, ErrorParam(err))
}
}
func TestChatResultToResponseBuildsCompleteResponseWithRefusalLogprobsAndCustomTool(t *testing.T) {
request := map[string]any{
"background": false, "conversation": nil, "instructions": "be concise", "max_output_tokens": 300,
"max_tool_calls": nil, "metadata": map[string]any{"trace": "1"}, "moderation": map[string]any{"type": "auto"},
"parallel_tool_calls": true, "previous_response_id": "resp_parent", "prompt": nil, "prompt_cache_key": "cache",
"prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "in_memory",
"reasoning": map[string]any{"effort": "low"}, "safety_identifier": "safe", "service_tier": "default",
"temperature": 0.4, "text": map[string]any{"verbosity": "low"}, "tool_choice": "auto", "tools": []any{},
"top_logprobs": 2, "top_p": 0.9, "truncation": "disabled", "user": "user-1",
}
response := ChatResultToResponse(map[string]any{
"created": 1710000000, "service_tier": "priority", "moderation": map[string]any{"flagged": false},
"choices": []any{map[string]any{
"finish_reason": "stop", "logprobs": map[string]any{"content": []any{map[string]any{"token": "hello", "logprob": -0.1}}},
"message": map[string]any{
"role": "assistant", "content": "hello", "refusal": "cannot continue", "annotations": []any{map[string]any{"type": "url_citation"}},
"tool_calls": []any{
map[string]any{"id": "call_fn", "type": "function", "function": map[string]any{"name": "lookup", "arguments": "{\"x\":1}"}},
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
},
},
}},
"usage": map[string]any{"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
}, "resp_12345678901234567890123456789012", "demo", request)
for _, key := range []string{
"id", "object", "created_at", "status", "completed_at", "error", "incomplete_details", "instructions", "metadata", "model", "output",
"parallel_tool_calls", "temperature", "tool_choice", "tools", "top_p", "background", "conversation", "max_output_tokens", "max_tool_calls",
"moderation", "output_text", "previous_response_id", "prompt", "prompt_cache_key", "prompt_cache_options", "prompt_cache_retention", "reasoning",
"safety_identifier", "service_tier", "text", "top_logprobs", "truncation", "usage", "user",
} {
if _, ok := response[key]; !ok {
t.Fatalf("complete fallback Response omitted %q: %+v", key, response)
}
}
output, _ := response["output"].([]any)
if len(output) != 3 {
t.Fatalf("expected message, function, and custom output items: %+v", output)
}
message, _ := output[0].(map[string]any)
parts, _ := message["content"].([]any)
text, _ := parts[0].(map[string]any)
refusal, _ := parts[1].(map[string]any)
custom, _ := output[2].(map[string]any)
if len(text["logprobs"].([]any)) != 1 || len(text["annotations"].([]any)) != 1 || refusal["type"] != "refusal" || custom["type"] != "custom_tool_call" || custom["input"] != "pwd" {
t.Fatalf("Response content/tool details changed: %+v", output)
}
if response["service_tier"] != "priority" {
t.Fatalf("upstream response service tier was not preserved: %+v", response)
}
}
func TestChatResponsesStreamAdapterEmitsRefusalLogprobsAndCustomToolLifecycle(t *testing.T) {
adapter := newChatResponsesStreamAdapter("resp_12345678901234567890123456789012", "demo")
events := make([]map[string]any, 0)
onDelta := func(event StreamDeltaEvent) error {
events = append(events, event.Event)
return nil
}
if err := adapter.delta(StreamDeltaEvent{Event: map[string]any{"choices": []any{map[string]any{
"logprobs": map[string]any{"content": []any{map[string]any{"token": "hi", "logprob": -0.1}}},
"delta": map[string]any{
"content": "hi", "refusal": "no",
"tool_calls": []any{map[string]any{"index": 0, "id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pw"}}},
},
}}}}, onDelta); err != nil {
t.Fatal(err)
}
if err := adapter.delta(StreamDeltaEvent{Event: map[string]any{"choices": []any{map[string]any{
"delta": map[string]any{"tool_calls": []any{map[string]any{"index": 0, "type": "custom", "custom": map[string]any{"input": "d"}}}},
}}}}, onDelta); err != nil {
t.Fatal(err)
}
result := ChatResultToResponse(map[string]any{"choices": []any{map[string]any{
"message": map[string]any{"role": "assistant", "content": "hi", "refusal": "no", "tool_calls": []any{
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
}},
}}}, adapter.publicID, "demo", map[string]any{})
if err := adapter.done(result, onDelta); err != nil {
t.Fatal(err)
}
types := make([]string, 0, len(events))
var sawLogprobs bool
for index, event := range events {
types = append(types, stringFromAny(event["type"]))
if intFromAny(event["sequence_number"]) != index {
t.Fatalf("non-contiguous sequence at %d: %+v", index, event)
}
if event["type"] == "response.output_text.delta" {
logprobs, _ := event["logprobs"].([]any)
sawLogprobs = len(logprobs) == 1
}
}
for _, required := range []string{
"response.refusal.delta", "response.refusal.done", "response.custom_tool_call_input.delta",
"response.custom_tool_call_input.done", "response.output_item.done", "response.completed",
} {
if !containsTestString(types, required) {
t.Fatalf("missing %s lifecycle event: %v", required, types)
}
}
if !sawLogprobs {
t.Fatalf("streamed output_text delta lost logprobs: %+v", events)
}
}
func mustJSON(t *testing.T, value any) []byte {
t.Helper()
data, err := json.Marshal(value)
@@ -0,0 +1,51 @@
package clients
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestLocalUnsupportedKindDoesNotMarkUpstreamSubmissionStarted(t *testing.T) {
started, responseReceived := false, false
_, err := (VolcesClient{}).Run(context.Background(), Request{
Kind: "speech.generations",
Candidate: store.RuntimeModelCandidate{
Provider: "volces", Credentials: map[string]any{"apiKey": "test"},
},
OnUpstreamSubmissionStarted: func() error { started = true; return nil },
OnUpstreamResponseReceived: func() error { responseReceived = true; return nil },
})
if err == nil || ErrorCode(err) != "unsupported_kind" {
t.Fatalf("expected local unsupported_kind, got %v", err)
}
if started || responseReceived {
t.Fatalf("local validation must remain not_submitted: started=%v response=%v", started, responseReceived)
}
}
func TestNetworkDisconnectMarksSubmittingWithoutResponseReceived(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
baseURL := upstream.URL
upstream.Close()
started, responseReceived := false, false
_, err := (VolcesClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
Kind: "images.generations",
Body: map[string]any{"prompt": "test"},
Candidate: store.RuntimeModelCandidate{
Provider: "volces", BaseURL: baseURL, ModelName: "seedream-test",
Credentials: map[string]any{"apiKey": "test"},
},
OnUpstreamSubmissionStarted: func() error { started = true; return nil },
OnUpstreamResponseReceived: func() error { responseReceived = true; return nil },
})
if err == nil || ErrorCode(err) != "network" {
t.Fatalf("expected network error, got %v", err)
}
if !started || responseReceived {
t.Fatalf("network disconnect must remain submitting: started=%v response=%v", started, responseReceived)
}
}
@@ -0,0 +1,10 @@
[
{ "name": "none-to-low", "input": { "reasoning_effort": "none" }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
{ "name": "minimal-to-low", "input": { "reasoning_effort": "minimal" }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
{ "name": "medium", "input": { "reasoning_effort": "medium" }, "expected": { "enable_thinking": true, "reasoning_effort": "medium" } },
{ "name": "high-to-xhigh", "input": { "reasoning_effort": "high" }, "expected": { "enable_thinking": true, "reasoning_effort": "xhigh" } },
{ "name": "max-to-xhigh", "input": { "reasoning_effort": "max" }, "expected": { "enable_thinking": true, "reasoning_effort": "xhigh" } },
{ "name": "force-thinking", "input": { "enable_thinking": false }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
{ "name": "budget-wins", "input": { "reasoning_effort": "high", "thinking_budget": 300000 }, "expected": { "enable_thinking": true, "thinking_budget": 262144 } },
{ "name": "temperature-floor", "input": { "temperature": 0.1 }, "expected": { "enable_thinking": true, "temperature": 0.6 } }
]
+617
View File
@@ -0,0 +1,617 @@
package clients
import (
"bytes"
"context"
"crypto/md5"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"time"
)
const (
topazDefaultPollInterval = 15 * time.Second
topazDefaultPollTimeout = 60 * time.Minute
topazDefaultMaxInput = int64(2 << 30)
)
type TopazClient struct {
HTTPClient *http.Client
Probe func(context.Context, string) (TopazVideoMetadata, error)
}
type TopazVideoMetadata struct {
Width int
Height int
Duration float64
FrameRate float64
FrameCount int
HasAudio bool
}
type topazSource struct {
Path string
Container string
Size int64
MD5 string
Resolution map[string]any
Duration float64
FrameRate float64
FrameCount int
}
func (c TopazClient) Run(ctx context.Context, request Request) (Response, error) {
startedAt := time.Now()
apiKey := credential(request.Candidate.Credentials, "apiKey", "api_key", "key", "token")
if apiKey == "" {
return Response{}, &ClientError{Code: "missing_credentials", Message: "Topaz API key is required", Retryable: false}
}
requestID := strings.TrimSpace(request.RemoteTaskID)
payload := cloneMapAny(request.RemoteTaskPayload)
var source *topazSource
var target map[string]int
if requestID == "" || strings.TrimSpace(firstNonEmptyString(payload["phase"])) != "uploaded" {
videoURL := strings.TrimSpace(firstNonEmptyString(request.Body["video_url"], request.Body["videoUrl"]))
if videoURL == "" {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "video_url is required", Param: "video_url", StatusCode: http.StatusBadRequest}
}
prepared, err := c.prepareSource(ctx, request, videoURL)
if err != nil {
return Response{}, err
}
source = &prepared
defer os.Remove(prepared.Path)
target = topazTargetResolution(request.Body, prepared.Resolution)
if requestID == "" {
created, err := c.createRequest(ctx, request, apiKey, prepared, target)
if err != nil {
return Response{}, annotateResponseError(err, "", startedAt, time.Now())
}
requestID = strings.TrimSpace(firstNonEmptyString(created["requestId"], created["request_id"], created["id"]))
if requestID == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz create response is missing requestId", Retryable: false}
}
payload = map[string]any{"phase": "created", "targetResolution": target}
if request.OnRemoteTaskSubmitted != nil {
if err := request.OnRemoteTaskSubmitted(requestID, payload); err != nil {
return Response{}, err
}
}
}
if err := c.uploadSource(ctx, request, apiKey, requestID, prepared); err != nil {
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
}
payload = map[string]any{"phase": "uploaded", "targetResolution": target}
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(requestID, payload); err != nil {
return Response{}, err
}
}
}
if target == nil {
target = topazTargetFromPayload(payload, request.Body)
}
completed, err := c.poll(ctx, request, apiKey, requestID, target)
if err != nil {
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
}
outputURL := topazDownloadURL(completed)
if outputURL == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz completed without a download URL", RequestID: requestID, Retryable: false}
}
metadata, err := c.probe(ctx, outputURL)
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: "cannot validate Topaz output metadata: " + err.Error(), RequestID: requestID, Retryable: true}
}
if target["width"] > 0 && target["height"] > 0 && (metadata.Width+4 < target["width"] || metadata.Height+4 < target["height"]) {
return Response{}, &ClientError{Code: "invalid_response", Message: fmt.Sprintf("Topaz output resolution %dx%d does not reach target %dx%d", metadata.Width, metadata.Height, target["width"], target["height"]), RequestID: requestID, Retryable: false}
}
finishedAt := time.Now()
resultItem := map[string]any{
"type": "video",
"url": outputURL,
"video_url": outputURL,
"width": metadata.Width,
"height": metadata.Height,
"target_resolution": fmt.Sprintf("%dx%d", target["width"], target["height"]),
"duration": metadata.Duration,
"target_frame_rate": metadata.FrameRate,
"slow_motion_rate": numericValue(request.Body["slow_motion_rate"], 1),
}
if source != nil {
resultItem["source_resolution"] = fmt.Sprintf("%vx%v", source.Resolution["width"], source.Resolution["height"])
resultItem["source_frame_rate"] = source.FrameRate
resultItem["duration"] = source.Duration
}
return Response{
Result: map[string]any{
"status": "success",
"model": request.Model,
"task_id": requestID,
"upstream_task_id": requestID,
"data": []any{resultItem},
},
RequestID: requestID,
Progress: append(providerProgress(request), Progress{Phase: "polling", Progress: 0.9, Message: "Topaz upscale completed", Payload: map[string]any{"upstreamTaskId": requestID}}),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
func (c TopazClient) prepareSource(ctx context.Context, request Request, rawURL string) (topazSource, error) {
parsed, err := url.Parse(rawURL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url must be an http(s) URL", Param: "video_url", StatusCode: http.StatusBadRequest}
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return topazSource{}, err
}
downloadClient := topazSourceHTTPClient(httpClient(request.HTTPClient, c.HTTPClient), boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false))
resp, err := downloadClient.Do(req)
if err != nil {
return topazSource{}, transportClientError(err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url download failed: " + resp.Status, Param: "video_url", StatusCode: http.StatusBadRequest}
}
limit := int64(numericValue(firstPresent(request.Candidate.PlatformConfig["maxInputBytes"], request.Candidate.PlatformConfig["max_input_bytes"]), float64(topazDefaultMaxInput)))
if limit <= 0 {
limit = topazDefaultMaxInput
}
tmp, err := os.CreateTemp("", "easyai-topaz-*"+topazContainerExtension(parsed.Path))
if err != nil {
return topazSource{}, err
}
path := tmp.Name()
cleanup := func() {
_ = tmp.Close()
_ = os.Remove(path)
}
hash := md5.New()
written, err := io.Copy(io.MultiWriter(tmp, hash), io.LimitReader(resp.Body, limit+1))
closeErr := tmp.Close()
if err != nil || closeErr != nil || written <= 0 || written > limit {
cleanup()
if written > limit {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video source exceeds configured size limit", Param: "video_url", StatusCode: http.StatusBadRequest}
}
return topazSource{}, firstError(err, closeErr)
}
metadata, err := c.probe(ctx, path)
if err != nil || metadata.Width <= 0 || metadata.Height <= 0 {
cleanup()
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "cannot probe source video metadata", Param: "video_url", StatusCode: http.StatusBadRequest}
}
frameRate := metadata.FrameRate
if frameRate <= 0 {
frameRate = numericValue(request.Candidate.PlatformConfig["frameRate"], 30)
}
duration := math.Max(1, metadata.Duration)
frameCount := metadata.FrameCount
if frameCount <= 0 {
frameCount = int(math.Round(duration * frameRate))
}
return topazSource{
Path: path, Container: strings.TrimPrefix(strings.ToLower(filepath.Ext(path)), "."), Size: written,
MD5: hex.EncodeToString(hash.Sum(nil)), Resolution: map[string]any{"width": metadata.Width, "height": metadata.Height},
Duration: duration, FrameRate: frameRate, FrameCount: frameCount,
}, nil
}
func (c TopazClient) createRequest(ctx context.Context, request Request, apiKey string, source topazSource, target map[string]int) (map[string]any, error) {
model := topazModelName(firstNonEmptyString(request.Candidate.ProviderModelName, request.Model))
scale := math.Max(float64(target["width"])/numericValue(source.Resolution["width"], 1), float64(target["height"])/numericValue(source.Resolution["height"], 1))
filter := map[string]any{"model": model}
if scale > 1 {
key := "scale"
if model == "slf-2" || model == "slhq-1" || model == "slm-1" || model == "slp-2.5" {
key = "upscaling_factor"
}
filter[key] = math.Round(scale*1000) / 1000
}
preserveAudio := boolishDefault(request.Body["preserve_audio"], true)
filters := []any{filter}
targetFrameRate := numericValue(firstPresent(request.Body["target_frame_rate"], request.Body["output_frame_rate"]), source.FrameRate)
slowMotionRate := numericValue(request.Body["slow_motion_rate"], 1)
if targetFrameRate <= 0 {
targetFrameRate = source.FrameRate
}
if slowMotionRate <= 0 {
slowMotionRate = 1
}
if targetFrameRate != source.FrameRate || slowMotionRate > 1 {
interpolationModel := strings.TrimSpace(firstNonEmptyString(request.Body["frame_interpolation_model"]))
if interpolationModel == "" {
return nil, &ClientError{Code: "invalid_parameter", Message: "frame_interpolation_model is required when output frame rate or slow motion changes", Param: "frame_interpolation_model", StatusCode: http.StatusBadRequest}
}
interpolation := map[string]any{"model": topazModelName(interpolationModel)}
if targetFrameRate > 0 {
interpolation["fps"] = targetFrameRate
}
if slowMotionRate > 1 {
interpolation["slowmo"] = slowMotionRate
}
filters = append(filters, interpolation)
}
body := map[string]any{
"source": map[string]any{"container": source.Container, "size": source.Size, "duration": source.Duration, "frameCount": source.FrameCount, "frameRate": source.FrameRate, "resolution": source.Resolution},
"filters": filters,
"output": map[string]any{"resolution": target, "frameRate": targetFrameRate, "audioCodec": "AAC", "audioTransfer": map[bool]string{true: "Copy", false: "None"}[preserveAudio], "dynamicCompressionLevel": "High", "videoEncoder": "H265", "videoProfile": "Main", "container": "mp4"},
}
return c.topazJSON(ctx, request, apiKey, http.MethodPost, "/video/", body)
}
func (c TopazClient) uploadSource(ctx context.Context, request Request, apiKey, requestID string, source topazSource) error {
accepted, err := c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/accept", nil)
if err != nil {
return err
}
urls := stringList(accepted["urls"])
if len(urls) == 0 {
return &ClientError{Code: "invalid_response", Message: "Topaz accept response has no upload URLs", RequestID: requestID}
}
file, err := os.Open(source.Path)
if err != nil {
return err
}
defer file.Close()
chunkSize := (source.Size + int64(len(urls)) - 1) / int64(len(urls))
results := make([]any, 0, len(urls))
for index, uploadURL := range urls {
remaining := source.Size - int64(index)*chunkSize
if remaining <= 0 {
break
}
length := minInt64(chunkSize, remaining)
section := io.NewSectionReader(file, int64(index)*chunkSize, length)
req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, section)
if err != nil {
return err
}
req.ContentLength = length
req.Header.Set("Content-Type", topazContainerContentType(source.Container))
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return transportClientError(err)
}
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20))
resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return &ClientError{Code: "provider_failed", Message: "Topaz part upload failed: " + resp.Status, StatusCode: resp.StatusCode, Retryable: resp.StatusCode >= 500}
}
etag := strings.Trim(resp.Header.Get("ETag"), `"`)
if etag == "" {
return &ClientError{Code: "invalid_response", Message: "Topaz part upload is missing ETag", Retryable: false}
}
results = append(results, map[string]any{"partNum": index + 1, "eTag": etag})
}
_, err = c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/complete-upload/", map[string]any{"md5Hash": source.MD5, "uploadResults": results})
return err
}
func (c TopazClient) poll(ctx context.Context, request Request, apiKey, requestID string, target map[string]int) (map[string]any, error) {
interval := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollInterval, "pollIntervalMs", "poll_interval_ms")
timeout := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollTimeout, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
deadline := time.NewTimer(timeout)
defer deadline.Stop()
for {
status, err := c.topazJSON(ctx, request, apiKey, http.MethodGet, "/video/"+url.PathEscape(requestID)+"/status", nil)
if err != nil {
return nil, err
}
state := strings.ToLower(strings.TrimSpace(firstNonEmptyString(status["status"], status["state"])))
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(requestID, map[string]any{"phase": "uploaded", "status": state, "targetResolution": target}); err != nil {
return nil, err
}
}
switch state {
case "complete", "completed", "success", "succeeded":
return status, nil
case "failed", "failure", "cancelled", "canceled", "error":
return nil, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(status["message"], "Topaz task failed"), RequestID: requestID, Retryable: false}
}
select {
case <-ctx.Done():
return nil, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
case <-deadline.C:
return nil, &ClientError{Code: "timeout", Message: "Topaz task polling timed out", RequestID: requestID, Retryable: true}
case <-time.After(interval):
}
}
}
func (c TopazClient) topazJSON(ctx context.Context, request Request, apiKey, method, path string, body map[string]any) (map[string]any, error) {
var reader io.Reader
if body != nil {
raw, _ := json.Marshal(body)
reader = bytes.NewReader(raw)
}
req, err := http.NewRequestWithContext(ctx, method, providerURL(request.Candidate.BaseURL, path), reader)
if err != nil {
return nil, err
}
req.Header.Set("X-API-Key", apiKey)
req.Header.Set("Accept", "application/json")
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, transportClientError(err)
}
result, decodeErr := decodeHTTPResponse(resp)
if decodeErr != nil {
return nil, decodeErr
}
return result, nil
}
func (c TopazClient) probe(ctx context.Context, target string) (TopazVideoMetadata, error) {
if c.Probe != nil {
return c.Probe(ctx, target)
}
return probeTopazVideo(ctx, target)
}
func probeTopazVideo(ctx context.Context, target string) (TopazVideoMetadata, error) {
if _, err := exec.LookPath("ffprobe"); err != nil {
return TopazVideoMetadata{}, err
}
probeCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
output, err := exec.CommandContext(probeCtx, "ffprobe", "-v", "error", "-show_entries", "format=duration:stream=codec_type,width,height,avg_frame_rate,nb_frames", "-of", "json", target).Output()
if err != nil {
return TopazVideoMetadata{}, err
}
var decoded struct {
Format struct {
Duration string `json:"duration"`
} `json:"format"`
Streams []struct {
CodecType string `json:"codec_type"`
Width int `json:"width"`
Height int `json:"height"`
AverageRate string `json:"avg_frame_rate"`
FrameCount string `json:"nb_frames"`
} `json:"streams"`
}
if err := json.Unmarshal(output, &decoded); err != nil {
return TopazVideoMetadata{}, err
}
metadata := TopazVideoMetadata{}
metadata.Duration, _ = strconv.ParseFloat(decoded.Format.Duration, 64)
for _, stream := range decoded.Streams {
if stream.CodecType == "audio" {
metadata.HasAudio = true
}
if stream.CodecType != "video" || stream.Width <= 0 || stream.Height <= 0 {
continue
}
metadata.Width, metadata.Height = stream.Width, stream.Height
metadata.FrameRate = parseTopazRate(stream.AverageRate)
metadata.FrameCount, _ = strconv.Atoi(stream.FrameCount)
}
if metadata.Width <= 0 || metadata.Height <= 0 {
return TopazVideoMetadata{}, errors.New("video stream metadata is missing")
}
return metadata, nil
}
func topazTargetResolution(body map[string]any, source map[string]any) map[string]int {
width := int(numericValue(body["output_width"], 0))
height := int(numericValue(body["output_height"], 0))
if width > 0 && height > 0 {
return map[string]int{"width": width, "height": height}
}
value := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["target_resolution"], body["output_resolution"], "1080p")))
if parsedWidth, parsedHeight, ok := parseTopazSize(value); ok {
return map[string]int{"width": parsedWidth, "height": parsedHeight}
}
base := map[string][2]int{"480p": {854, 480}, "720p": {1280, 720}, "1080p": {1920, 1080}, "1440p": {2560, 1440}, "2k": {2560, 1440}, "2160p": {3840, 2160}, "4k": {3840, 2160}}
value = strings.TrimSuffix(value, "_upscale")
target, ok := base[value]
if !ok {
target = base["1080p"]
}
sourceWidth := numericValue(source["width"], 1)
sourceHeight := numericValue(source["height"], 1)
if sourceWidth == sourceHeight {
side := minInt(target[0], target[1])
return map[string]int{"width": side, "height": side}
}
if sourceWidth > sourceHeight {
return map[string]int{"width": int(math.Round(sourceWidth / sourceHeight * float64(target[1]))), "height": target[1]}
}
return map[string]int{"width": target[1], "height": int(math.Round(sourceHeight / sourceWidth * float64(target[1])))}
}
func topazTargetFromPayload(payload map[string]any, body map[string]any) map[string]int {
if raw, ok := payload["targetResolution"].(map[string]any); ok {
return map[string]int{"width": int(numericValue(raw["width"], 0)), "height": int(numericValue(raw["height"], 0))}
}
return topazTargetResolution(body, map[string]any{"width": 16, "height": 9})
}
func topazDownloadURL(result map[string]any) string {
if download, ok := result["download"].(map[string]any); ok {
return firstNonEmptyString(download["url"], download["downloadUrl"])
}
return firstNonEmptyString(result["download_url"], result["downloadUrl"], result["url"])
}
func topazModelName(model string) string {
switch strings.TrimSpace(model) {
case "easy-proteus-standard-4":
return "prob-4"
case "easy-starlight-fast-2":
return "slf-2"
case "easy-starlight-hq-1":
return "slhq-1"
case "easy-starlight-mini-1":
return "slm-1"
default:
return strings.TrimSpace(model)
}
}
func topazContainerExtension(path string) string {
ext := strings.ToLower(filepath.Ext(strings.Split(path, "?")[0]))
switch ext {
case ".mov", ".mkv", ".webm", ".mp4":
return ext
default:
return ".mp4"
}
}
func topazContainerContentType(container string) string {
switch strings.ToLower(container) {
case "mov":
return "video/quicktime"
case "mkv":
return "video/x-matroska"
case "webm":
return "video/webm"
default:
return "video/mp4"
}
}
func parseTopazRate(value string) float64 {
parts := strings.Split(value, "/")
if len(parts) == 2 {
numerator, _ := strconv.ParseFloat(parts[0], 64)
denominator, _ := strconv.ParseFloat(parts[1], 64)
if denominator > 0 {
return numerator / denominator
}
}
parsed, _ := strconv.ParseFloat(value, 64)
return parsed
}
func parseTopazSize(value string) (int, int, bool) {
parts := strings.Split(strings.ReplaceAll(value, " ", ""), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, errWidth := strconv.Atoi(parts[0])
height, errHeight := strconv.Atoi(parts[1])
return width, height, errWidth == nil && errHeight == nil && width > 0 && height > 0
}
func stringList(value any) []string {
items, ok := value.([]any)
if !ok {
if values, ok := value.([]string); ok {
return values
}
return nil
}
out := make([]string, 0, len(items))
for _, item := range items {
if text := strings.TrimSpace(fmt.Sprint(item)); text != "" {
out = append(out, text)
}
}
return out
}
func boolishDefault(value any, fallback bool) bool {
if value == nil {
return fallback
}
switch typed := value.(type) {
case bool:
return typed
case string:
parsed, err := strconv.ParseBool(strings.TrimSpace(typed))
if err == nil {
return parsed
}
}
return fallback
}
func firstError(values ...error) error {
for _, value := range values {
if value != nil {
return value
}
}
return nil
}
func minInt64(left, right int64) int64 {
if left < right {
return left
}
return right
}
func minInt(left, right int) int {
if left < right {
return left
}
return right
}
func topazSourceHTTPClient(base *http.Client, allowPrivate bool) *http.Client {
client := *base
client.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
if client.Timeout <= 0 {
client.Timeout = 10 * time.Minute
}
if allowPrivate {
return &client
}
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil || len(addresses) == 0 {
return nil, errors.New("video source DNS resolution failed")
}
for _, address := range addresses {
if topazBlockedAddress(address.IP) {
return nil, errors.New("video source resolved to a blocked network")
}
}
dialer := &net.Dialer{Timeout: 10 * time.Second}
var attempts []error
for _, resolved := range addresses {
connection, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
if dialErr == nil {
return connection, nil
}
attempts = append(attempts, dialErr)
}
return nil, errors.Join(attempts...)
}
client.Transport = transport
return &client
}
func topazBlockedAddress(ip net.IP) bool {
return ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast()
}
+110 -17
View File
@@ -3,6 +3,7 @@ package clients
import (
"context"
"errors"
"net"
"net/http"
"strings"
"time"
@@ -10,24 +11,85 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const (
defaultProviderRequestTimeout = 10 * time.Minute
imageProviderRequestTimeout = 20 * time.Minute
videoProviderRequestTimeout = 30 * time.Minute
)
// ProviderRequestTimeout returns the default end-to-end provider request or
// polling budget for a task kind. Explicit platform polling configuration may
// still override these defaults for providers with a documented requirement.
func ProviderRequestTimeout(kind string) time.Duration {
switch {
case strings.HasPrefix(strings.TrimSpace(kind), "images."):
return imageProviderRequestTimeout
case strings.HasPrefix(strings.TrimSpace(kind), "videos."):
return videoProviderRequestTimeout
default:
return defaultProviderRequestTimeout
}
}
func transportClientError(err error) *ClientError {
if err == nil {
return nil
}
var netErr net.Error
if errors.Is(err, context.DeadlineExceeded) || (errors.As(err, &netErr) && netErr.Timeout()) {
return &ClientError{
Code: "timeout",
Message: "upstream request timed out: " + err.Error(),
Retryable: false,
}
}
return &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
type Request struct {
Kind string
ModelType string
Model string
Body map[string]any
Candidate store.RuntimeModelCandidate
HTTPClient *http.Client
RemoteTaskID string
RemoteTaskPayload map[string]any
OnRemoteTaskSubmitted func(remoteTaskID string, payload map[string]any) error
OnRemoteTaskPolled func(remoteTaskID string, payload map[string]any) error
Stream bool
StreamDelta StreamDelta
UpstreamProtocol string
PublicResponseID string
PublicPreviousResponseID string
UpstreamPreviousResponseID string
PreviousResponseTurns []ResponseTurn
Kind string
ModelType string
Model string
Body map[string]any
OriginalBody map[string]any
Candidate store.RuntimeModelCandidate
HTTPClient *http.Client
RemoteTaskID string
RemoteTaskPayload map[string]any
UpstreamIdempotencyKey string
OnRemoteTaskSubmitted func(remoteTaskID string, payload map[string]any) error
OnRemoteTaskPolled func(remoteTaskID string, payload map[string]any) error
OnUpstreamSubmissionStarted func() error
OnUpstreamResponseReceived func() error
OnUpstreamWireResponse func(*WireResponse) error
Stream bool
StreamDelta StreamDelta
UpstreamProtocol string
PublicResponseID string
PublicPreviousResponseID string
UpstreamPreviousResponseID string
PreviousResponseTurns []ResponseTurn
}
func applyUpstreamIdempotency(request *http.Request, input Request) {
if request == nil {
return
}
if key := strings.TrimSpace(input.UpstreamIdempotencyKey); key != "" {
request.Header.Set("Idempotency-Key", key)
}
}
// WireResponse preserves the provider-facing response independently from the
// canonical result used by billing, retries, and cross-protocol conversion.
// Headers are filtered before this value leaves the clients package.
type WireResponse struct {
Protocol string `json:"protocol,omitempty"`
StatusCode int `json:"statusCode,omitempty"`
Headers map[string][]string `json:"headers,omitempty"`
Body map[string]any `json:"body,omitempty"`
RawJSON []byte `json:"-"`
Converted bool `json:"converted,omitempty"`
}
type ResponseTurn struct {
@@ -54,6 +116,7 @@ type Response struct {
ResponseChainDepth int
ResponseConverted bool
InternalResult map[string]any
Wire *WireResponse
}
type Usage struct {
@@ -75,6 +138,10 @@ type StreamDeltaEvent struct {
Text string
ReasoningContent string
Event map[string]any
WireProtocol string
WireConverted bool
WireStatusCode int
WireHeaders map[string][]string
}
type StreamDelta func(event StreamDeltaEvent) error
@@ -106,12 +173,22 @@ type ClientError struct {
Code string
Message string
Param string
Details map[string]any
StatusCode int
RequestID string
ResponseStartedAt time.Time
ResponseFinishedAt time.Time
ResponseDurationMS int64
Retryable bool
Wire *WireResponse
}
func ErrorWireResponse(err error) *WireResponse {
var clientErr *ClientError
if errors.As(err, &clientErr) {
return clientErr.Wire
}
return nil
}
func ErrorParam(err error) string {
@@ -122,6 +199,18 @@ func ErrorParam(err error) string {
return ""
}
func ErrorDetails(err error) map[string]any {
var clientErr *ClientError
if !errors.As(err, &clientErr) || len(clientErr.Details) == 0 {
return nil
}
details := make(map[string]any, len(clientErr.Details))
for key, value := range clientErr.Details {
details[key] = value
}
return details
}
func (e *ClientError) Error() string {
if e.Message != "" {
return e.Message
@@ -142,6 +231,10 @@ func upstreamModelName(candidate store.RuntimeModelCandidate) string {
}
func ErrorCode(err error) string {
var coded interface{ ErrorCode() string }
if errors.As(err, &coded) && strings.TrimSpace(coded.ErrorCode()) != "" {
return strings.TrimSpace(coded.ErrorCode())
}
var clientErr *ClientError
if errors.As(err, &clientErr) && clientErr.Code != "" {
return clientErr.Code
+12 -4
View File
@@ -40,7 +40,15 @@ func (c UniversalClient) Run(ctx context.Context, request Request) (Response, er
if err != nil {
return Response{}, err
}
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
submitResult, submitRequestID, err = c.universalSubmit(ctx, executor, request, modelType, payload)
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
return Response{}, notifyErr
}
}
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, startedAt, time.Now())
}
@@ -145,7 +153,7 @@ func (c UniversalClient) universalSubmit(ctx context.Context, executor *scripten
func (c UniversalClient) universalPollUntilDone(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string, upstreamTaskID string, payload map[string]any, requestID string, startedAt time.Time) (map[string]any, string, error) {
interval := universalDurationConfig(request.Candidate.PlatformConfig, 2*time.Second, "pollIntervalMs", "poll_interval_ms")
timeout := universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
timeout := universalDurationConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
deadline := time.NewTimer(timeout)
defer deadline.Stop()
ticker := time.NewTicker(interval)
@@ -226,7 +234,7 @@ func universalScriptContext(request Request, modelType string, payload map[strin
"platformModelId": request.Candidate.PlatformModelID,
"canonicalModelKey": request.Candidate.CanonicalModelKey,
"modelType": modelType,
"timeout": universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms").Milliseconds(),
"timeout": universalDurationConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms").Milliseconds(),
},
"env": cloneMapAny(request.Candidate.PlatformConfig),
"candidate": universalCandidateSnapshot(request),
@@ -310,7 +318,7 @@ func universalPostJSON(ctx context.Context, client *http.Client, baseURL string,
}
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
@@ -327,7 +335,7 @@ func universalGetJSON(ctx context.Context, client *http.Client, url string, cred
}
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
+271
View File
@@ -0,0 +1,271 @@
package clients
import (
"bytes"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"fmt"
"io"
"mime/multipart"
"net"
"net/http"
"net/textproto"
"net/url"
"strings"
"time"
)
const vectorizerMaxResponseBytes = 128 << 20
// VectorizerClient implements the Vectorizer.AI binary vectorization API.
// Gateway async semantics are supplied by the outer River-backed task runner.
type VectorizerClient struct {
HTTPClient *http.Client
LookupIP func(context.Context, string) ([]net.IPAddr, error)
}
func (c VectorizerClient) Run(ctx context.Context, request Request) (Response, error) {
startedAt := time.Now()
format := strings.ToLower(strings.TrimSpace(firstNonEmptyString(request.Body["format"], request.Body["output_format"])))
if format == "" {
format = "svg"
}
if !vectorizerFormatAllowed(format) {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer format must be svg, eps, pdf, dxf, or png", Param: "format", StatusCode: http.StatusBadRequest}
}
imageToken := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_image_token"]))
receipt := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_receipt"]))
endpoint := "vectorize"
fields := map[string]string{"output.file_format": format}
if imageToken != "" {
endpoint = "download"
fields["image.token"] = imageToken
if receipt != "" {
fields["receipt"] = receipt
}
} else {
imageURL := vectorizerImageURL(request.Body)
if imageURL == "" {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer source image URL is required", Param: "source.url", StatusCode: http.StatusBadRequest}
}
if err := c.validateSourceURL(ctx, imageURL, boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false)); err != nil {
return Response{}, err
}
fields["image.url"] = imageURL
fields["mode"] = firstNonEmptyString(request.Candidate.PlatformConfig["mode"], "production")
fields["policy.retention_days"] = firstNonEmptyString(request.Candidate.PlatformConfig["retentionDays"], request.Candidate.PlatformConfig["retention_days"], "7")
fields["processing.shapes.min_area_px"] = vectorizerCleanupMinArea(request.Body)
if maxColors := vectorizerMaxColors(request.Body); maxColors > 0 {
fields["processing.max_colors"] = fmt.Sprint(maxColors)
}
}
appendVectorizerOutputFields(fields, format)
payload, contentType, err := vectorizerMultipartBody(fields)
if err != nil {
return Response{}, err
}
url := providerURL(request.Candidate.BaseURL, endpoint)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, payload)
if err != nil {
return Response{}, err
}
req.Header.Set("Content-Type", contentType)
applyVectorizerAuth(req, request.Candidate.Credentials)
client := httpClient(request.HTTPClient, c.HTTPClient)
resp, err := client.Do(req)
if err != nil {
return Response{}, transportClientError(err)
}
defer resp.Body.Close()
requestID := requestIDFromHTTPResponse(resp)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
result, decodeErr := decodeHTTPResponse(resp)
if decodeErr != nil {
return Response{}, annotateResponseError(decodeErr, requestID, startedAt, time.Now())
}
return Response{}, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(result["message"], result["error"], resp.Status), RequestID: requestID, StatusCode: resp.StatusCode, Retryable: resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500}
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, vectorizerMaxResponseBytes+1))
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: err.Error(), RequestID: requestID, Retryable: true}
}
if len(raw) == 0 || len(raw) > vectorizerMaxResponseBytes {
return Response{}, &ClientError{Code: "invalid_response", Message: "vectorizer response is empty or too large", RequestID: requestID, Retryable: false}
}
returnedToken := strings.TrimSpace(resp.Header.Get("X-Image-Token"))
returnedReceipt := strings.TrimSpace(resp.Header.Get("X-Receipt"))
if request.OnRemoteTaskSubmitted != nil && returnedToken != "" {
digest := sha256.Sum256([]byte(returnedToken))
if err := request.OnRemoteTaskSubmitted("vectorizer-"+hex.EncodeToString(digest[:8]), map[string]any{
"imageToken": returnedToken,
"receipt": returnedReceipt,
}); err != nil {
return Response{}, err
}
}
finishedAt := time.Now()
mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
if mimeType == "" || mimeType == "application/octet-stream" {
mimeType = vectorizerContentType(format)
}
return Response{
Result: map[string]any{
"status": "success",
"model": request.Model,
"data": []any{map[string]any{
"type": vectorizerOutputKind(format),
"b64_json": base64.StdEncoding.EncodeToString(raw),
"mime_type": mimeType,
"format": format,
}},
"vectorizer": map[string]any{
"format": format,
"creditsCharged": numericHeader(resp.Header.Get("X-Credits-Charged")),
"creditsCalculated": numericHeader(resp.Header.Get("X-Credits-Calculated")),
},
},
RequestID: requestID,
Progress: append(providerProgress(request), Progress{Phase: "uploading", Progress: 0.9, Message: "vectorizer result received"}),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
func (c VectorizerClient) validateSourceURL(ctx context.Context, rawURL string, allowPrivate bool) error {
parsed, err := url.Parse(rawURL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || strings.TrimSpace(parsed.Hostname()) == "" || parsed.User != nil {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source URL must be a public http(s) URL without userinfo", Param: "source.url", StatusCode: http.StatusBadRequest}
}
if allowPrivate {
return nil
}
lookup := c.LookupIP
if lookup == nil {
lookup = net.DefaultResolver.LookupIPAddr
}
addresses, err := lookup(ctx, parsed.Hostname())
if err != nil || len(addresses) == 0 {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source DNS resolution failed", Param: "source.url", StatusCode: http.StatusBadRequest}
}
for _, address := range addresses {
if topazBlockedAddress(address.IP) {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source resolved to a blocked network", Param: "source.url", StatusCode: http.StatusBadRequest}
}
}
return nil
}
func vectorizerImageURL(body map[string]any) string {
if source, ok := body["source"].(map[string]any); ok {
return firstNonEmptyString(source["url"], source["image_url"], source["imageUrl"])
}
return firstNonEmptyString(body["image_url"], body["imageUrl"])
}
func vectorizerMultipartBody(fields map[string]string) (*bytes.Buffer, string, error) {
var payload bytes.Buffer
writer := multipart.NewWriter(&payload)
for key, value := range fields {
header := make(textproto.MIMEHeader)
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"`, strings.ReplaceAll(key, `"`, `\"`)))
part, err := writer.CreatePart(header)
if err != nil {
return nil, "", err
}
if _, err := io.WriteString(part, value); err != nil {
return nil, "", err
}
}
if err := writer.Close(); err != nil {
return nil, "", err
}
return &payload, writer.FormDataContentType(), nil
}
func applyVectorizerAuth(req *http.Request, credentials map[string]any) {
username := credential(credentials, "username", "accessKey", "access_key", "apiId", "api_id", "id")
password := credential(credentials, "password", "secretKey", "secret_key", "apiSecret", "api_secret", "secret")
if username != "" || password != "" {
req.SetBasicAuth(username, password)
return
}
if apiKey := credential(credentials, "apiKey", "api_key", "token"); apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
}
func vectorizerFormatAllowed(format string) bool {
switch format {
case "svg", "eps", "pdf", "dxf", "png":
return true
default:
return false
}
}
func vectorizerCleanupMinArea(body map[string]any) string {
cleanup := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["cleanupLevel"], body["cleanup_level"])))
switch cleanup {
case "low":
return "0"
case "strong":
return "1"
default:
return "0.125"
}
}
func vectorizerMaxColors(body map[string]any) int {
for _, key := range []string{"maxColors", "max_colors"} {
if value := int(numericValue(body[key], 0)); value == 0 || value == 2 || value == 4 || value == 8 || value == 16 || value == 32 {
return value
}
}
return 0
}
func appendVectorizerOutputFields(fields map[string]string, format string) {
if format == "svg" {
fields["output.svg.version"] = "svg_1_1"
fields["output.svg.fixed_size"] = "false"
fields["output.svg.adobe_compatibility_mode"] = "true"
}
if format == "dxf" {
fields["output.dxf.compatibility_level"] = "lines_and_arcs"
}
}
func vectorizerContentType(format string) string {
switch format {
case "svg":
return "image/svg+xml"
case "eps":
return "application/postscript"
case "pdf":
return "application/pdf"
case "dxf":
return "application/dxf"
default:
return "image/png"
}
}
func vectorizerOutputKind(format string) string {
if format == "svg" || format == "png" {
return "image"
}
return "file"
}
func numericHeader(value string) any {
value = strings.TrimSpace(value)
if value == "" {
return nil
}
return numericValue(value, 0)
}
+102 -31
View File
@@ -44,14 +44,21 @@ func (c VolcesClient) runImage(ctx context.Context, request Request, apiKey stri
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
applyUpstreamIdempotency(req, request)
responseStartedAt := time.Now()
if err := notifySubmissionStarted(request); err != nil {
return Response{}, err
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return Response{}, transportClientError(err)
}
if err := notifyResponseReceived(request); err != nil {
return Response{}, err
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolOpenAIImages)
responseFinishedAt := time.Now()
if err != nil {
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
@@ -67,6 +74,8 @@ func (c VolcesClient) runImage(ctx context.Context, request Request, apiKey stri
ResponseStartedAt: responseStartedAt,
ResponseFinishedAt: responseFinishedAt,
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
UpstreamProtocol: ProtocolOpenAIImages,
Wire: wire,
}, nil
}
@@ -80,7 +89,7 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
if err := validateVolcesVideoTaskBody(body); err != nil {
return Response{}, err
}
submitResult, requestID, err := c.postJSON(ctx, request, request.Candidate.BaseURL, taskPath, apiKey, body)
submitResult, requestID, _, err := c.postJSON(ctx, request, request.Candidate.BaseURL, taskPath, apiKey, body)
submitRequestID = requestID
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
@@ -119,7 +128,7 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
}
case <-nextPoll.C:
pollStartedAt := time.Now()
pollResult, pollRequestID, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
pollResult, pollRequestID, pollWire, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
pollFinishedAt := time.Now()
requestID := firstNonEmpty(pollRequestID, submitRequestID, upstreamTaskID)
lastRequestID = requestID
@@ -151,6 +160,8 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
ResponseStartedAt: submitStartedAt,
ResponseFinishedAt: pollFinishedAt,
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
UpstreamProtocol: ProtocolVolcesContents,
Wire: pollWire,
}, nil
case "failed", "cancelled":
return Response{}, &ClientError{
@@ -188,7 +199,7 @@ func (c VolcesClient) DeleteVideoTask(ctx context.Context, request Request) (map
req.Header.Set("Authorization", "Bearer "+apiKey)
response, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", transportClientError(err)
}
requestID := requestIDFromHTTPResponse(response)
result, err := decodeHTTPResponse(response)
@@ -214,44 +225,74 @@ func volcesVideoTaskPath(request Request) string {
return strings.TrimRight(path, "/")
}
func (c VolcesClient) postJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string, body map[string]any) (map[string]any, string, error) {
func (c VolcesClient) postJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string, body map[string]any) (map[string]any, string, *WireResponse, error) {
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(baseURL, path), bytes.NewReader(raw))
if err != nil {
return nil, "", err
return nil, "", nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+apiKey)
applyUpstreamIdempotency(req, request)
if err := notifySubmissionStarted(request); err != nil {
return nil, "", nil, err
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", nil, transportClientError(err)
}
if err := notifyResponseReceived(request); err != nil {
return nil, "", nil, err
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolVolcesContents)
if err != nil {
return result, requestID, err
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
return result, requestID, wire, notifyErr
}
return result, requestID, wire, err
}
original := result
result, envelopeRequestID, err := normalizeVolcesCompatibleResult(result)
return result, firstNonEmpty(requestID, envelopeRequestID), err
if wire != nil {
wire.Converted = volcesCompatibleResultConverted(original)
}
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
return result, firstNonEmpty(requestID, envelopeRequestID), wire, notifyErr
}
return result, firstNonEmpty(requestID, envelopeRequestID), wire, err
}
func (c VolcesClient) getJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string) (map[string]any, string, error) {
func (c VolcesClient) getJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string) (map[string]any, string, *WireResponse, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(baseURL, path), nil)
if err != nil {
return nil, "", err
return nil, "", nil, err
}
req.Header.Set("Authorization", "Bearer "+apiKey)
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return nil, "", nil, transportClientError(err)
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolVolcesContents)
if err != nil {
return result, requestID, err
return result, requestID, wire, err
}
original := result
result, envelopeRequestID, err := normalizeVolcesCompatibleResult(result)
return result, firstNonEmpty(requestID, envelopeRequestID), err
if wire != nil {
wire.Converted = volcesCompatibleResultConverted(original)
}
return result, firstNonEmpty(requestID, envelopeRequestID), wire, err
}
func volcesCompatibleResultConverted(result map[string]any) bool {
if _, ok := result["error"].(map[string]any); ok {
return true
}
_, hasCode := result["code"]
_, hasData := result["data"].(map[string]any)
return hasCode && hasData
}
func normalizeVolcesCompatibleResult(result map[string]any) (map[string]any, string, error) {
@@ -318,16 +359,33 @@ func volcesImageBody(request Request) map[string]any {
body["seed"] = -1
}
}
if resolution := strings.TrimSpace(stringFromAny(body["resolution"])); resolution != "" {
resolution := normalizedProviderImageResolution(body["resolution"])
size := widthHeightSize(body)
if volcesImageUsesResolutionSize(request) && resolution != "" {
body["size"] = resolution
} else if size != "" {
body["size"] = size
} else if resolution != "" {
body["size"] = resolution
}
if size := widthHeightSize(body); size != "" {
body["size"] = size
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio", "resolution", "width", "height"} {
delete(body, key)
}
normalizeVolcesSequentialImageGeneration(body, request)
return body
}
func volcesImageUsesResolutionSize(request Request) bool {
modelType := firstNonEmpty(request.ModelType, request.Candidate.ModelType)
if capability, ok := request.Candidate.Capabilities[modelType].(map[string]any); ok {
if strings.EqualFold(strings.TrimSpace(stringFromAny(capability["size_param_format"])), "resolution") {
return true
}
}
model := strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate)))
return strings.HasPrefix(model, "doubao-seedream-5-0")
}
func volcesVideoBody(request Request) map[string]any {
body := cleanProviderBody(request.Body)
body["model"] = upstreamModelName(request.Candidate)
@@ -354,6 +412,16 @@ func cleanProviderBody(body map[string]any) map[string]any {
"poll_interval_ms",
"pollTimeoutSeconds",
"poll_timeout_seconds",
"_gateway_compatibility",
"_gateway_target_protocol",
"_compat_provider",
"_kling_compat_version",
"platform_id",
"platformId",
"platform_model_id",
"platformModelId",
"modelType",
"model_type",
} {
delete(out, key)
}
@@ -1057,9 +1125,11 @@ func volcesTaskErrorMessage(result map[string]any) string {
}
func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[string]any) map[string]any {
result := cloneMapAny(raw)
if result == nil {
result = map[string]any{}
result := map[string]any{}
for _, key := range []string{"seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if raw[key] != nil {
result[key] = raw[key]
}
}
content, _ := raw["content"].(map[string]any)
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
@@ -1069,18 +1139,18 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
}
data := []any{}
if videoURL != "" {
data = append(data, map[string]any{"url": videoURL, "type": "video"})
item := map[string]any{"url": videoURL, "type": "video"}
if lastFrameURL := strings.TrimSpace(stringFromAny(content["last_frame_url"])); lastFrameURL != "" {
item["last_frame_url"] = lastFrameURL
}
data = append(data, item)
}
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
result["model"] = upstreamModelName(request.Candidate)
}
result["model"] = firstNonEmpty(stringFromAny(raw["model"]), upstreamModelName(request.Candidate))
result["status"] = "succeeded"
result["object"] = "video.generation"
result["created"] = created
result["upstream_task_id"] = upstreamTaskID
result["data"] = data
result["raw"] = cloneMapAny(raw)
return result
}
@@ -1114,9 +1184,10 @@ func volcesPollInterval(request Request) time.Duration {
}
func volcesPollTimeout(request Request) time.Duration {
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["volcesPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), 600)
fallbackSeconds := ProviderRequestTimeout(request.Kind).Seconds()
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["volcesPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), fallbackSeconds)
if seconds < 1 {
seconds = 600
seconds = fallbackSeconds
}
return time.Duration(seconds) * time.Second
}
+1 -1
View File
@@ -106,7 +106,7 @@ func (c VolcesAssetClient) call(ctx context.Context, credentials VolcesAssetCred
response, err := httpClient(nil, c.HTTPClient).Do(req)
if err != nil {
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
return "", transportClientError(err)
}
defer response.Body.Close()
var envelope struct {
@@ -79,7 +79,7 @@ func TestVolcesClientSupportsDeyunEnvelope(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "deyun-task-1" || item["url"] != "https://example.com/deyun.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/deyun.mp4" {
t.Fatalf("unexpected response: %+v", response.Result)
}
}
@@ -0,0 +1,125 @@
package clients
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) {
response := &http.Response{
StatusCode: http.StatusTooManyRequests,
Status: "429 Too Many Requests",
Header: http.Header{
"Content-Type": {"application/json"},
"X-Request-Id": {"req_1"},
"X-Ratelimit-Reset-Requests": {"1s"},
"Set-Cookie": {"secret=must-not-leak"},
"Connection": {"keep-alive"},
},
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"slow down","future_field":true}}`)),
}
_, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
if err == nil {
t.Fatal("expected upstream error")
}
if wire == nil || wire.StatusCode != http.StatusTooManyRequests || wire.Protocol != ProtocolOpenAIResponses {
t.Fatalf("unexpected wire response: %+v", wire)
}
if wire.Headers["X-Request-Id"][0] != "req_1" || wire.Headers["X-Ratelimit-Reset-Requests"][0] != "1s" {
t.Fatalf("official response headers were lost: %+v", wire.Headers)
}
if _, ok := wire.Headers["Set-Cookie"]; ok {
t.Fatalf("sensitive header leaked: %+v", wire.Headers)
}
if _, ok := wire.Headers["Connection"]; ok {
t.Fatalf("hop-by-hop header leaked: %+v", wire.Headers)
}
if ErrorWireResponse(err) != wire {
t.Fatal("client error did not retain its wire response")
}
}
func TestDecodeHTTPResponseForProtocolAllowsLargeImagePayload(t *testing.T) {
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
response := &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: http.Header{"Content-Type": {"application/json"}},
Body: io.NopCloser(strings.NewReader(payload)),
}
result, wire, err := decodeHTTPResponseForProtocol(response, ProtocolGeminiGenerateContent)
if err != nil {
t.Fatalf("large Gemini image response failed: %v", err)
}
if got, _ := result["data"].(string); len(got) != int(defaultMaxJSONResponseBytes) {
t.Fatalf("large Gemini image payload length = %d", len(got))
}
if wire == nil || len(wire.RawJSON) != len(payload) {
t.Fatalf("large Gemini wire payload was truncated: %+v", wire)
}
}
func TestDecodeHTTPResponseForProtocolRejectsOversizedDefaultPayload(t *testing.T) {
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
response := &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: http.Header{"Content-Type": {"application/json"}},
Body: io.NopCloser(strings.NewReader(payload)),
}
_, _, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
if err == nil {
t.Fatal("expected oversized default response to fail")
}
clientErr, ok := err.(*ClientError)
if !ok || clientErr.Code != "response_too_large" {
t.Fatalf("unexpected oversized response error: %T %v", err, err)
}
}
func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) {
var requestedPath string
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestedPath = r.URL.RequestURI()
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("X-Goog-Request-Id", "goog-stream-1")
_, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}],\"futureOfficialField\":true}\n\n")
}))
defer upstream.Close()
var events []StreamDeltaEvent
response, err := (GeminiClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
Kind: "chat.completions", ModelType: "text", Model: "gemini-test", Stream: true,
Body: map[string]any{"prompt": "hello"},
Candidate: store.RuntimeModelCandidate{
Provider: "gemini", BaseURL: upstream.URL, ProviderModelName: "gemini-test",
Credentials: map[string]any{"apiKey": "test-key"},
},
StreamDelta: func(event StreamDeltaEvent) error {
events = append(events, event)
return nil
},
})
if err != nil {
t.Fatalf("Gemini stream failed: %v", err)
}
if !strings.Contains(requestedPath, ":streamGenerateContent") || !strings.Contains(requestedPath, "alt=sse") {
t.Fatalf("unexpected Gemini stream endpoint: %s", requestedPath)
}
if len(events) != 1 || events[0].WireProtocol != ProtocolGeminiGenerateContent || events[0].Event["futureOfficialField"] != true {
raw, _ := json.Marshal(events)
t.Fatalf("official Gemini event was not preserved: %s", raw)
}
if response.Wire == nil || response.Wire.Headers["X-Goog-Request-Id"][0] != "goog-stream-1" {
t.Fatalf("Gemini stream wire metadata was lost: %+v", response.Wire)
}
}
+535 -4
View File
@@ -1,6 +1,7 @@
package config
import (
"encoding/json"
"errors"
"log/slog"
"net/url"
@@ -39,15 +40,81 @@ type Config struct {
LocalGeneratedStorageDir string
LocalUploadedStorageDir string
LocalTempAssetTTLHours int
LocalResultTTLHours int
LocalResultMinFreeBytes int64
LocalResultMaxBytes int64
LocalResultMaxTaskBytes int64
TaskProgressCallbackEnabled bool
TaskProgressCallbackURL string
TaskProgressCallbackTimeoutMS string
TaskProgressCallbackMaxAttempts string
TaskProgressCallbackTimeoutMS int
TaskProgressCallbackMaxAttempts int
TaskCleanupEnabled bool
TaskRetentionDays int
TaskAnalysisRetentionDays int
TaskCleanupIntervalSeconds int
TaskCleanupBatchSize int
CORSAllowedOrigin string
GlobalHTTPProxy string
GlobalHTTPProxySource string
PlatformProxyBypassIDs string
LogLevel slog.Level
BillingEngineMode string
ProcessRole string
DatabaseMaxConns int
DatabaseCriticalMaxConns int
DatabaseRiverMaxConns int
DatabaseMinIdleConns int
DatabaseMaxConnIdleSeconds int
DatabaseIdleInTransactionTimeoutSeconds int
DatabaseLockTimeoutSeconds int
MediaRequestConcurrency int
MediaMaterializationConcurrency int
MediaImageNormalizationConcurrency int
MediaOSSDirectEnabled bool
MediaOSSEndpoint string
MediaOSSBucket string
MediaOSSAccessKeyID string
MediaOSSAccessKeySecret string
MediaOSSPublicBaseURL string
MediaOSSObjectPrefix string
AsyncQueueWorkerEnabled bool
AsyncWorkerHardLimit int
AsyncWorkerInstanceHardLimit int
AsyncWorkerRefreshIntervalSeconds int
AsyncWorkerLoadMode string
AsyncAdmissionMicrobatchSize int
AsyncAdmissionDispatcherEnabled bool
AsyncAdmissionDispatcherConfigured bool
RoutingMode string
ExecutionPoolID string
ExecutionPoolLabels string
WorkerID string
WorkerAdvertiseEndpoint string
WorkerOrchestratorInstanceRef string
WorkerEndpointAllowedSuffixes string
WorkerEndpointAllowPrivate bool
WorkerExecutionSecret string
WorkerExecutionCAFile string
WorkerExecutionCertFile string
WorkerExecutionKeyFile string
RouteProbeEnabled bool
RouteProbeHotIntervalSeconds int
RouteProbeColdIntervalSeconds int
RouteProbeTimeoutMS int
RouteProbeConcurrency int
WorkerAutoscalingEnabled bool
CapacityOrchestratorAdapter string
CapacityPoolsJSON string
WorkerTargetOutstandingPerReplica int
WorkerScaleUpWindowSeconds int
WorkerScaleDownStabilizationSeconds int
WorkerDrainTimeoutSeconds int
NodeMemoryTargetPercent int
NodeMemoryHardPercent int
NodeCPUTargetPercent int
PostgresConnectionBudget int
PostgresNonWorkerConnectionBudget int
WorkerDatabaseMaxConns int
}
func Load() Config {
@@ -81,26 +148,291 @@ func Load() Config {
LocalGeneratedStorageDir: env("AI_GATEWAY_GENERATED_STORAGE_DIR", env("LOCAL_GENERATED_STORAGE_DIR", env("AI_GATEWAY_STATIC_STORAGE_DIR", DefaultLocalGeneratedStorageDir))),
LocalUploadedStorageDir: env("AI_GATEWAY_UPLOADED_STORAGE_DIR", env("LOCAL_UPLOADED_STORAGE_DIR", DefaultLocalUploadedStorageDir)),
LocalTempAssetTTLHours: envInt("AI_GATEWAY_LOCAL_TEMP_ASSET_TTL_HOURS", 24),
LocalResultTTLHours: envIntValidated("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS", 24),
LocalResultMinFreeBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES", 10*1024*1024*1024),
LocalResultMaxBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES", 256*1024*1024),
LocalResultMaxTaskBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES", 512*1024*1024),
TaskProgressCallbackEnabled: env("TASK_PROGRESS_CALLBACK_ENABLED", "true") == "true",
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
),
TaskProgressCallbackTimeoutMS: env("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", "5000"),
TaskProgressCallbackMaxAttempts: env("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", "10"),
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
PlatformProxyBypassIDs: env("AI_GATEWAY_PLATFORM_PROXY_BYPASS_IDS", ""),
LogLevel: logLevel(env("LOG_LEVEL", "info")),
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))),
DatabaseMaxConns: envInt("AI_GATEWAY_DATABASE_MAX_CONNS", 0),
DatabaseCriticalMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS", 0),
DatabaseRiverMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_RIVER_MAX_CONNS", 0),
DatabaseMinIdleConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", 0),
DatabaseMaxConnIdleSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS", 0),
DatabaseIdleInTransactionTimeoutSeconds: envOptionalIntValidated(
"AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS",
60,
),
DatabaseLockTimeoutSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", 30),
MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16),
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
MediaImageNormalizationConcurrency: envIntValidated(
"AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY",
2,
),
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated(
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT",
envOptionalIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
),
AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32),
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
AsyncWorkerLoadMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_WORKER_LOAD_MODE", "adaptive"))),
AsyncAdmissionMicrobatchSize: envIntValidated("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", 8),
AsyncAdmissionDispatcherEnabled: envValue("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED") == "true",
AsyncAdmissionDispatcherConfigured: envValue(
"AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED",
) != "",
RoutingMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_ROUTING_MODE", "legacy"))),
ExecutionPoolID: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_ID", "legacy-default")),
ExecutionPoolLabels: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_LABELS", "{}")),
WorkerID: strings.TrimSpace(env("AI_GATEWAY_WORKER_ID", "")),
WorkerAdvertiseEndpoint: workerAdvertiseEndpoint(),
WorkerOrchestratorInstanceRef: strings.TrimSpace(env("AI_GATEWAY_ORCHESTRATOR_INSTANCE_REF", "")),
WorkerEndpointAllowedSuffixes: strings.TrimSpace(env("AI_GATEWAY_WORKER_ENDPOINT_ALLOWED_SUFFIXES", "svc,cluster.local")),
WorkerEndpointAllowPrivate: env("AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE", "true") == "true",
WorkerExecutionSecret: env("AI_GATEWAY_WORKER_EXECUTION_SECRET", env("SERVER_MAIN_INTERNAL_SECRET", env("SERVER_MAIN_INTERNAL_TOKEN", ""))),
WorkerExecutionCAFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CA_FILE", "")),
WorkerExecutionCertFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE", "")),
WorkerExecutionKeyFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_KEY_FILE", "")),
RouteProbeEnabled: env("AI_GATEWAY_ROUTE_PROBE_ENABLED", "true") == "true",
RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15),
RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60),
RouteProbeTimeoutMS: envIntValidated("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS", 3000),
RouteProbeConcurrency: envIntValidated("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY", 8),
WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true",
CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))),
CapacityPoolsJSON: strings.TrimSpace(env("AI_GATEWAY_CAPACITY_POOLS", "")),
WorkerTargetOutstandingPerReplica: envOptionalIntValidated("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA", 0),
WorkerScaleUpWindowSeconds: envIntValidated("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS", 20),
WorkerScaleDownStabilizationSeconds: envIntValidated(
"AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS",
600,
),
WorkerDrainTimeoutSeconds: envIntValidated("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS", 600),
NodeMemoryTargetPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT", 75),
NodeMemoryHardPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_HARD_PERCENT", 85),
NodeCPUTargetPercent: envIntValidated("AI_GATEWAY_NODE_CPU_TARGET_PERCENT", 70),
PostgresConnectionBudget: envIntValidated("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET", 150),
PostgresNonWorkerConnectionBudget: envIntValidated(
"AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET",
72,
),
WorkerDatabaseMaxConns: envIntValidated("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS", 32),
}
}
func (c Config) Validate() error {
switch strings.ToLower(strings.TrimSpace(c.ProcessRole)) {
case "", "all", "api", "worker", "capacity-controller":
default:
return errors.New("AI_GATEWAY_PROCESS_ROLE must be all, api, worker, or capacity-controller")
}
if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 {
return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured")
}
if c.DatabaseMinIdleConns < 0 || c.DatabaseMinIdleConns > 1000 ||
(c.DatabaseMaxConns > 0 && c.DatabaseMinIdleConns > c.DatabaseMaxConns) {
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must be between 0 and AI_GATEWAY_DATABASE_MAX_CONNS")
}
if c.DatabaseCriticalMaxConns < 0 || c.DatabaseCriticalMaxConns > 64 ||
(c.DatabaseMaxConns > 0 && c.DatabaseCriticalMaxConns >= c.DatabaseMaxConns) {
return errors.New("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS must be between 0 and less than AI_GATEWAY_DATABASE_MAX_CONNS")
}
if c.DatabaseRiverMaxConns < 0 || c.DatabaseRiverMaxConns > 256 ||
(c.DatabaseMaxConns > 0 &&
c.DatabaseCriticalMaxConns+c.DatabaseRiverMaxConns >= c.DatabaseMaxConns) {
return errors.New("critical and River database pools must leave at least one execution connection")
}
if c.DatabaseMaxConns > 0 &&
c.DatabaseMinIdleConns > c.DatabaseMaxConns-c.DatabaseCriticalMaxConns-c.DatabaseRiverMaxConns {
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must not exceed the execution pool after reserving critical and River connections")
}
if c.DatabaseMaxConnIdleSeconds < 0 || c.DatabaseMaxConnIdleSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS must be between 0 and 3600")
}
if c.DatabaseIdleInTransactionTimeoutSeconds < 0 || c.DatabaseIdleInTransactionTimeoutSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS must be between 0 and 3600")
}
if c.DatabaseLockTimeoutSeconds < 0 || c.DatabaseLockTimeoutSeconds > 3600 {
return errors.New("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS must be between 0 and 3600")
}
if c.MediaMaterializationConcurrency != 0 && (c.MediaMaterializationConcurrency < 1 || c.MediaMaterializationConcurrency > 256) {
return errors.New("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY must be between 1 and 256")
}
if c.MediaRequestConcurrency != 0 && (c.MediaRequestConcurrency < 1 || c.MediaRequestConcurrency > 1024) {
return errors.New("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY must be between 1 and 1024")
}
if c.MediaImageNormalizationConcurrency != 0 && (c.MediaImageNormalizationConcurrency < 1 || c.MediaImageNormalizationConcurrency > 64) {
return errors.New("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY must be between 1 and 64")
}
if c.MediaOSSDirectEnabled && !c.RunsCapacityController() {
if strings.TrimSpace(c.MediaOSSAccessKeyID) == "" || strings.TrimSpace(c.MediaOSSAccessKeySecret) == "" {
return errors.New("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID and AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET are required when direct media OSS is enabled")
}
if !validMediaOSSBucket(c.MediaOSSBucket) {
return errors.New("AI_GATEWAY_MEDIA_OSS_BUCKET is invalid")
}
if !validMediaOSSURL(c.MediaOSSEndpoint, c.AppEnv) {
return errors.New("AI_GATEWAY_MEDIA_OSS_ENDPOINT must be an HTTPS URL")
}
if !validMediaOSSURL(c.MediaOSSPublicBaseURL, c.AppEnv) {
return errors.New("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL must be an HTTPS URL")
}
if !validMediaOSSObjectPrefix(c.MediaOSSObjectPrefix) {
return errors.New("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX is invalid")
}
}
switch strings.ToLower(strings.TrimSpace(c.BillingEngineMode)) {
case "", "observe", "enforce", "hold":
default:
return errors.New("BILLING_ENGINE_MODE must be observe, enforce, or hold")
}
if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerInstanceHardLimit < 1 || c.AsyncWorkerInstanceHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
}
switch strings.ToLower(strings.TrimSpace(c.AsyncWorkerLoadMode)) {
case "", "adaptive", "legacy":
default:
return errors.New("AI_GATEWAY_WORKER_LOAD_MODE must be adaptive or legacy")
}
switch strings.ToLower(strings.TrimSpace(c.RoutingMode)) {
case "", "legacy", "shadow", "enforced":
default:
return errors.New("AI_GATEWAY_ROUTING_MODE must be legacy, shadow, or enforced")
}
if strings.TrimSpace(c.ExecutionPoolID) != "" && !validExecutionPoolID(c.ExecutionPoolID) {
return errors.New("AI_GATEWAY_EXECUTION_POOL_ID is invalid")
}
if strings.TrimSpace(c.ExecutionPoolLabels) != "" && !jsonObject(c.ExecutionPoolLabels) {
return errors.New("AI_GATEWAY_EXECUTION_POOL_LABELS must be a JSON object")
}
if c.RouteProbeHotIntervalSeconds != 0 && (c.RouteProbeHotIntervalSeconds < 5 || c.RouteProbeHotIntervalSeconds > 300) {
return errors.New("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS must be between 5 and 300")
}
if c.RouteProbeColdIntervalSeconds != 0 && (c.RouteProbeColdIntervalSeconds < c.RouteProbeHotIntervalSeconds || c.RouteProbeColdIntervalSeconds > 3600) {
return errors.New("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS must be between the hot interval and 3600")
}
if c.RouteProbeTimeoutMS != 0 && (c.RouteProbeTimeoutMS < 250 || c.RouteProbeTimeoutMS > 10000) {
return errors.New("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS must be between 250 and 10000")
}
if c.RouteProbeConcurrency != 0 && (c.RouteProbeConcurrency < 1 || c.RouteProbeConcurrency > 64) {
return errors.New("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY must be between 1 and 64")
}
if strings.EqualFold(c.RoutingMode, "enforced") && len(c.WorkerExecutionSecret) < 32 {
return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode")
}
if (c.WorkerExecutionCertFile == "") != (c.WorkerExecutionKeyFile == "") {
return errors.New("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE and AI_GATEWAY_WORKER_EXECUTION_KEY_FILE must be configured together")
}
if c.AsyncAdmissionMicrobatchSize < 1 || c.AsyncAdmissionMicrobatchSize > 32 {
return errors.New("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE must be between 1 and 32")
}
pools, err := c.CapacityPools()
if err != nil {
return err
}
if c.RunsCapacityController() && len(pools) == 0 {
return errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool for the capacity-controller role")
}
switch strings.ToLower(strings.TrimSpace(c.CapacityOrchestratorAdapter)) {
case "", "kubernetes", "static":
default:
return errors.New("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER must be kubernetes or static")
}
if c.WorkerTargetOutstandingPerReplica < 0 || c.WorkerTargetOutstandingPerReplica > 100000 {
return errors.New("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA must be between 0 and 100000")
}
if c.WorkerScaleUpWindowSeconds != 0 && (c.WorkerScaleUpWindowSeconds < 5 || c.WorkerScaleUpWindowSeconds > 600) {
return errors.New("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS must be between 5 and 600")
}
if c.WorkerScaleDownStabilizationSeconds != 0 &&
(c.WorkerScaleDownStabilizationSeconds < 60 || c.WorkerScaleDownStabilizationSeconds > 86400) {
return errors.New("AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS must be between 60 and 86400")
}
if c.WorkerDrainTimeoutSeconds != 0 && (c.WorkerDrainTimeoutSeconds < 60 || c.WorkerDrainTimeoutSeconds > 86400) {
return errors.New("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS must be between 60 and 86400")
}
if (c.NodeMemoryTargetPercent != 0 || c.NodeMemoryHardPercent != 0) &&
(c.NodeMemoryTargetPercent < 50 || c.NodeMemoryTargetPercent > 90 ||
c.NodeMemoryHardPercent <= c.NodeMemoryTargetPercent || c.NodeMemoryHardPercent > 95) {
return errors.New("Worker node memory target/hard percentages must be ordered within 50..95")
}
if c.NodeCPUTargetPercent != 0 && (c.NodeCPUTargetPercent < 40 || c.NodeCPUTargetPercent > 90) {
return errors.New("AI_GATEWAY_NODE_CPU_TARGET_PERCENT must be between 40 and 90")
}
if c.PostgresConnectionBudget < 0 || c.PostgresConnectionBudget > 1000 {
return errors.New("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET must be between 1 and 1000")
}
if c.PostgresNonWorkerConnectionBudget < 0 ||
(c.PostgresConnectionBudget > 0 &&
c.PostgresNonWorkerConnectionBudget >= c.PostgresConnectionBudget) {
return errors.New("AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET must be below the PostgreSQL connection budget")
}
if c.WorkerDatabaseMaxConns < 0 || c.WorkerDatabaseMaxConns > 256 {
return errors.New("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS must be between 1 and 256")
}
if c.TaskProgressCallbackTimeoutMS != 0 && (c.TaskProgressCallbackTimeoutMS < 100 || c.TaskProgressCallbackTimeoutMS > 60000) {
return errors.New("TASK_PROGRESS_CALLBACK_TIMEOUT_MS must be between 100 and 60000")
}
if c.TaskProgressCallbackMaxAttempts != 0 && (c.TaskProgressCallbackMaxAttempts < 1 || c.TaskProgressCallbackMaxAttempts > 100) {
return errors.New("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS must be between 1 and 100")
}
if c.TaskRetentionDays != 0 && (c.TaskRetentionDays < 1 || c.TaskRetentionDays > 3650) {
return errors.New("AI_GATEWAY_TASK_RETENTION_DAYS must be between 1 and 3650")
}
if c.TaskAnalysisRetentionDays != 0 && (c.TaskAnalysisRetentionDays < 1 || (c.TaskRetentionDays != 0 && c.TaskAnalysisRetentionDays > c.TaskRetentionDays)) {
return errors.New("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS must be between 1 and AI_GATEWAY_TASK_RETENTION_DAYS")
}
if c.TaskCleanupIntervalSeconds != 0 && c.TaskCleanupIntervalSeconds < 60 {
return errors.New("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS must be at least 60")
}
if c.TaskCleanupBatchSize != 0 && (c.TaskCleanupBatchSize < 100 || c.TaskCleanupBatchSize > 5000) {
return errors.New("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE must be between 100 and 5000")
}
if c.LocalResultTTLHours != 0 && (c.LocalResultTTLHours < 1 || c.LocalResultTTLHours > 24*30) {
return errors.New("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS must be between 1 and 720")
}
if c.LocalResultMinFreeBytes < 0 {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES must not be negative")
}
if c.LocalResultMaxBytes != 0 && c.LocalResultMaxBytes < 1024 {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES must be at least 1024")
}
if c.LocalResultMaxTaskBytes != 0 && c.LocalResultMaxTaskBytes < c.LocalResultMaxBytes {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES must be at least AI_GATEWAY_LOCAL_RESULT_MAX_BYTES")
}
switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) {
case "":
case "file":
@@ -124,11 +456,174 @@ func (c Config) Validate() error {
return nil
}
func validMediaOSSBucket(value string) bool {
value = strings.TrimSpace(value)
if len(value) < 3 || len(value) > 63 {
return false
}
for index, item := range value {
if (item >= 'a' && item <= 'z') || (item >= '0' && item <= '9') || (item == '-' && index > 0 && index < len(value)-1) {
continue
}
return false
}
return true
}
func validMediaOSSURL(value string, appEnv string) bool {
parsed, err := url.Parse(strings.TrimSpace(value))
if err != nil || parsed.Host == "" || parsed.RawQuery != "" || parsed.Fragment != "" || parsed.User != nil {
return false
}
if parsed.Scheme == "https" {
return true
}
host := strings.ToLower(parsed.Hostname())
return !strings.EqualFold(strings.TrimSpace(appEnv), "production") &&
parsed.Scheme == "http" && (host == "127.0.0.1" || host == "localhost" || host == "::1")
}
func validMediaOSSObjectPrefix(value string) bool {
value = strings.Trim(value, "/")
if value == "" || len(value) > 256 || strings.Contains(value, "..") {
return false
}
for _, item := range value {
if (item >= 'a' && item <= 'z') || (item >= 'A' && item <= 'Z') || (item >= '0' && item <= '9') ||
item == '-' || item == '_' || item == '/' {
continue
}
return false
}
return true
}
func (c Config) EffectiveProcessRole() string {
role := strings.ToLower(strings.TrimSpace(c.ProcessRole))
if role == "" {
return "all"
}
return role
}
func (c Config) RunsPublicHTTP() bool {
switch c.EffectiveProcessRole() {
case "worker", "capacity-controller":
return false
default:
return true
}
}
func (c Config) RunsAsyncExecutionWorker() bool {
switch c.EffectiveProcessRole() {
case "api", "capacity-controller":
return false
case "worker":
return true
default:
return c.AsyncQueueWorkerEnabled
}
}
// RunsAsyncAdmissionDispatcher keeps the legacy all-in-one/Worker behaviour
// unless a deployment explicitly separates admission coordination from task
// execution. Production enables this on both API sites so the instance nearest
// the current PostgreSQL primary wins the existing database scope locks, while
// remote Worker replicas only execute River jobs.
func (c Config) RunsAsyncAdmissionDispatcher() bool {
if c.AsyncAdmissionDispatcherConfigured {
return c.AsyncAdmissionDispatcherEnabled
}
return c.RunsAsyncExecutionWorker()
}
func (c Config) RunsBackgroundWorkers() bool {
switch c.EffectiveProcessRole() {
case "api", "capacity-controller":
return false
default:
return true
}
}
func (c Config) RunsCapacityController() bool {
return c.EffectiveProcessRole() == "capacity-controller"
}
func (c Config) RunsRouteProber() bool {
return c.RunsAsyncExecutionWorker() && c.RouteProbeEnabled && !strings.EqualFold(c.RoutingMode, "legacy")
}
func workerAdvertiseEndpoint() string {
return strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_WORKER_ADVERTISE_ENDPOINT")), "/")
}
func validExecutionPoolID(value string) bool {
value = strings.TrimSpace(value)
if value == "" || len(value) > 128 {
return false
}
for index, character := range value {
valid := (character >= 'a' && character <= 'z') ||
(character >= 'A' && character <= 'Z') ||
(character >= '0' && character <= '9') ||
character == '.' || character == '_' || character == ':' || character == '/' || character == '-'
if !valid || (index == 0 && !((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || (character >= '0' && character <= '9'))) {
return false
}
}
return true
}
func jsonObject(value string) bool {
var object map[string]any
return json.Unmarshal([]byte(value), &object) == nil && object != nil
}
type GlobalHTTPProxyStatus struct {
HTTPProxy string
Source string
}
type ExecutionPoolCapacityConfig struct {
ID string `json:"id"`
AdapterRef string `json:"adapterRef"`
BootstrapReplicas int `json:"bootstrapReplicas"`
MinReplicas int `json:"minReplicas"`
MaxReplicas int `json:"maxReplicas"`
}
func (c Config) CapacityPools() ([]ExecutionPoolCapacityConfig, error) {
raw := strings.TrimSpace(c.CapacityPoolsJSON)
if raw == "" {
return nil, nil
}
var pools []ExecutionPoolCapacityConfig
if err := json.Unmarshal([]byte(raw), &pools); err != nil {
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must be a JSON array")
}
if len(pools) == 0 {
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool")
}
seen := make(map[string]struct{}, len(pools))
for index := range pools {
pools[index].ID = strings.TrimSpace(pools[index].ID)
pools[index].AdapterRef = strings.TrimSpace(pools[index].AdapterRef)
if pools[index].AdapterRef == "" {
pools[index].AdapterRef = pools[index].ID
}
if !validExecutionPoolID(pools[index].ID) || pools[index].MinReplicas < 0 || pools[index].MaxReplicas < pools[index].MinReplicas || pools[index].MaxReplicas > 64 || pools[index].BootstrapReplicas < pools[index].MinReplicas || pools[index].BootstrapReplicas > pools[index].MaxReplicas {
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains an invalid pool definition")
}
if _, duplicate := seen[pools[index].ID]; duplicate {
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains duplicate pool IDs")
}
seen[pools[index].ID] = struct{}{}
}
return pools, nil
}
func LoadGlobalHTTPProxyStatus() GlobalHTTPProxyStatus {
for _, key := range []string{
"AI_GATEWAY_GLOBAL_HTTP_PROXY",
@@ -210,6 +705,42 @@ func envInt(key string, fallback int) int {
return parsed
}
func envIntValidated(key string, fallback int) int {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil {
return 0
}
return parsed
}
func envOptionalIntValidated(key string, fallback int) int {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil {
return -1
}
return parsed
}
func envInt64Validated(key string, fallback int64) int64 {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.ParseInt(value, 10, 64)
if err != nil {
return 0
}
return parsed
}
func logLevel(value string) slog.Level {
switch strings.ToLower(value) {
case "debug":
+348 -2
View File
@@ -26,7 +26,13 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
}
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
cfg := Config{IdentitySecretStore: "file"}
cfg := Config{
IdentitySecretStore: "file",
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
AsyncAdmissionMicrobatchSize: 8,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
}
@@ -37,7 +43,14 @@ func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
}
func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
cfg := Config{IdentitySecretStore: "kubernetes", IdentityKubernetesSecretName: "easyai-gateway-identity"}
cfg := Config{
IdentitySecretStore: "kubernetes",
IdentityKubernetesSecretName: "easyai-gateway-identity",
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
AsyncAdmissionMicrobatchSize: 8,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
t.Fatalf("Validate() error = %v, want missing namespace", err)
}
@@ -52,8 +65,341 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
IdentitySecurityEventHeartbeatIntervalSeconds: 60,
IdentitySecurityEventStaleAfterSeconds: 60,
IdentitySecurityEventClockSkewSeconds: 60,
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
AsyncAdmissionMicrobatchSize: 8,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
t.Fatalf("Validate() error = %v, want invalid stale threshold", err)
}
}
func TestValidateAsyncWorkerSettings(t *testing.T) {
cfg := Config{
AsyncWorkerHardLimit: 10001,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
AsyncAdmissionMicrobatchSize: 8,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid hard limit", err)
}
cfg.AsyncWorkerHardLimit = 2048
cfg.AsyncWorkerInstanceHardLimit = 10001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid instance hard limit", err)
}
cfg.AsyncWorkerInstanceHardLimit = 32
cfg.AsyncWorkerRefreshIntervalSeconds = 0
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
}
cfg.AsyncWorkerRefreshIntervalSeconds = 5
cfg.AsyncWorkerLoadMode = "invalid"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOAD_MODE") {
t.Fatalf("Validate() error = %v, want invalid worker load mode", err)
}
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "not-an-integer")
loaded := Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
}
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "2048")
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "not-an-integer")
loaded = Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer instance hard limit", err)
}
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "32")
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", "33")
loaded = Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "MICROBATCH_SIZE") {
t.Fatalf("Validate() error = %v, want invalid admission microbatch size", err)
}
}
func TestLoadAsyncAdmissionMicrobatchSize(t *testing.T) {
cfg := Load()
if cfg.AsyncAdmissionMicrobatchSize != 8 {
t.Fatalf("admission microbatch size=%d, want 8", cfg.AsyncAdmissionMicrobatchSize)
}
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", "4")
cfg = Load()
if err := cfg.Validate(); err != nil {
t.Fatalf("valid admission microbatch size was rejected: %v", err)
}
if cfg.AsyncAdmissionMicrobatchSize != 4 {
t.Fatalf("admission microbatch size=%d, want 4", cfg.AsyncAdmissionMicrobatchSize)
}
}
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
cfg := Load()
if !cfg.AsyncQueueWorkerEnabled {
t.Fatal("async queue worker must be enabled by default")
}
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
cfg = Load()
if cfg.AsyncQueueWorkerEnabled {
t.Fatal("async queue worker remained enabled after explicit disable")
}
}
func TestProcessRolePrecedenceAndCompatibility(t *testing.T) {
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
cfg := Load()
if got := cfg.EffectiveProcessRole(); got != "all" {
t.Fatalf("legacy configuration changed process role to %q", got)
}
if cfg.RunsAsyncExecutionWorker() {
t.Fatal("legacy async worker disable was ignored")
}
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "worker")
cfg = Load()
if got := cfg.EffectiveProcessRole(); got != "worker" {
t.Fatalf("effective process role = %q, want worker", got)
}
if cfg.RunsPublicHTTP() || !cfg.RunsAsyncExecutionWorker() || !cfg.RunsBackgroundWorkers() {
t.Fatalf("worker role capabilities are inconsistent: %+v", cfg)
}
if !cfg.RunsAsyncAdmissionDispatcher() {
t.Fatal("legacy Worker role no longer runs the admission dispatcher")
}
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "api")
cfg = Load()
if !cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() {
t.Fatalf("api role capabilities are inconsistent: %+v", cfg)
}
if cfg.RunsAsyncAdmissionDispatcher() {
t.Fatal("API role unexpectedly runs the admission dispatcher without explicit configuration")
}
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED", "true")
cfg = Load()
if !cfg.RunsAsyncAdmissionDispatcher() || cfg.RunsAsyncExecutionWorker() {
t.Fatalf("API admission-only role is inconsistent: %+v", cfg)
}
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED", "false")
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "worker")
cfg = Load()
if cfg.RunsAsyncAdmissionDispatcher() || !cfg.RunsAsyncExecutionWorker() {
t.Fatalf("Worker execution-only role is inconsistent: %+v", cfg)
}
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "capacity-controller")
cfg = Load()
if cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() || !cfg.RunsCapacityController() {
t.Fatalf("capacity controller role capabilities are inconsistent: %+v", cfg)
}
}
func TestWorkerAutoscalingConfiguration(t *testing.T) {
t.Setenv("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "true")
t.Setenv("AI_GATEWAY_CAPACITY_POOLS", `[{"id":"pool-a","bootstrapReplicas":2,"minReplicas":1,"maxReplicas":3}]`)
cfg := Load()
if err := cfg.Validate(); err != nil {
t.Fatalf("valid autoscaling configuration was rejected: %v", err)
}
pools, err := cfg.CapacityPools()
if err != nil || !cfg.WorkerAutoscalingEnabled || len(pools) != 1 || pools[0].BootstrapReplicas != 2 || pools[0].MaxReplicas != 3 {
t.Fatalf("autoscaling configuration=%+v", cfg)
}
cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":4,"minReplicas":1,"maxReplicas":3}]`
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "invalid pool") {
t.Fatalf("Validate() error=%v, want bootstrap bound failure", err)
}
}
func TestWorkerTopologyHasNoBuiltInSites(t *testing.T) {
t.Setenv("AI_GATEWAY_CAPACITY_POOLS", "")
cfg := Load()
pools, err := cfg.CapacityPools()
if err != nil || len(pools) != 0 {
t.Fatalf("default pools=%+v error=%v, want no built-in sites", pools, err)
}
}
func TestLoadPrefersWorkerGlobalHardLimitAlias(t *testing.T) {
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "48")
t.Setenv("AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT", "96")
cfg := Load()
if cfg.AsyncWorkerHardLimit != 96 {
t.Fatalf("global worker hard limit=%d, want 96", cfg.AsyncWorkerHardLimit)
}
}
func TestValidatePostgresReplicaBudgets(t *testing.T) {
cfg := Load()
if cfg.PostgresConnectionBudget != 150 || cfg.PostgresNonWorkerConnectionBudget != 72 {
t.Fatalf(
"PostgreSQL budgets=%d/%d, want 150/72",
cfg.PostgresConnectionBudget,
cfg.PostgresNonWorkerConnectionBudget,
)
}
cfg.PostgresNonWorkerConnectionBudget = cfg.PostgresConnectionBudget
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "NON_WORKER") {
t.Fatalf("Validate() error=%v, want non-Worker budget failure", err)
}
}
func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
cfg := Load()
if cfg.DatabaseMinIdleConns != 0 ||
cfg.DatabaseMaxConnIdleSeconds != 0 ||
cfg.DatabaseIdleInTransactionTimeoutSeconds != 60 ||
cfg.DatabaseLockTimeoutSeconds != 30 {
t.Fatalf(
"database pool minimum/max-idle/transaction timeouts = %d/%d/%d/%d, want 0/0/60/30",
cfg.DatabaseMinIdleConns,
cfg.DatabaseMaxConnIdleSeconds,
cfg.DatabaseIdleInTransactionTimeoutSeconds,
cfg.DatabaseLockTimeoutSeconds,
)
}
cfg.ProcessRole = "invalid"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "PROCESS_ROLE") {
t.Fatalf("Validate() error = %v, want invalid process role", err)
}
cfg.ProcessRole = "api"
cfg.DatabaseMaxConns = -1
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONNS") {
t.Fatalf("Validate() error = %v, want invalid database max conns", err)
}
cfg.DatabaseMaxConns = 16
cfg.DatabaseMinIdleConns = 17
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MIN_IDLE_CONNS") {
t.Fatalf("Validate() error = %v, want invalid database min idle conns", err)
}
cfg.DatabaseMinIdleConns = 4
cfg.DatabaseCriticalMaxConns = 4
cfg.DatabaseRiverMaxConns = 12
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "River") {
t.Fatalf("Validate() error = %v, want database pool budget failure", err)
}
cfg.DatabaseRiverMaxConns = 4
t.Setenv("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", "not-an-integer")
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MIN_IDLE_CONNS") {
t.Fatalf("Validate() error = %v, want invalid non-integer database min idle conns", err)
}
t.Setenv("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", "4")
cfg.DatabaseMaxConnIdleSeconds = 3601
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONN_IDLE_SECONDS") {
t.Fatalf("Validate() error = %v, want invalid database max connection idle seconds", err)
}
cfg.DatabaseMaxConnIdleSeconds = 30
cfg.DatabaseIdleInTransactionTimeoutSeconds = 3601
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDLE_IN_TRANSACTION") {
t.Fatalf("Validate() error = %v, want invalid idle transaction timeout", err)
}
cfg.DatabaseIdleInTransactionTimeoutSeconds = 60
cfg.DatabaseLockTimeoutSeconds = 3601
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
t.Fatalf("Validate() error = %v, want invalid lock timeout", err)
}
cfg.DatabaseLockTimeoutSeconds = 30
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "not-an-integer")
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
t.Fatalf("Validate() error = %v, want invalid non-integer lock timeout", err)
}
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "30")
cfg.MediaMaterializationConcurrency = 257
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_MATERIALIZATION_CONCURRENCY") {
t.Fatalf("Validate() error = %v, want invalid media materialization concurrency", err)
}
t.Setenv("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", "12")
if got := Load().MediaMaterializationConcurrency; got != 12 {
t.Fatalf("media materialization concurrency = %d, want 12", got)
}
cfg = Load()
cfg.MediaRequestConcurrency = 1025
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_REQUEST_CONCURRENCY") {
t.Fatalf("Validate() error = %v, want invalid media request concurrency", err)
}
t.Setenv("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", "96")
if got := Load().MediaRequestConcurrency; got != 96 {
t.Fatalf("media request concurrency = %d, want 96", got)
}
cfg = Load()
cfg.MediaImageNormalizationConcurrency = 65
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_IMAGE_NORMALIZATION_CONCURRENCY") {
t.Fatalf("Validate() error = %v, want invalid image normalization concurrency", err)
}
t.Setenv("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY", "3")
if got := Load().MediaImageNormalizationConcurrency; got != 3 {
t.Fatalf("image normalization concurrency = %d, want 3", got)
}
}
func TestValidateTaskHistorySettings(t *testing.T) {
cfg := Load()
cfg.TaskRetentionDays = 30
cfg.TaskAnalysisRetentionDays = 31
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ANALYSIS_RETENTION") {
t.Fatalf("Validate() error = %v, want invalid analysis retention", err)
}
cfg.TaskAnalysisRetentionDays = 7
cfg.TaskCleanupIntervalSeconds = 59
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid cleanup interval", err)
}
cfg.TaskCleanupIntervalSeconds = 300
cfg.TaskCleanupBatchSize = 5001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_BATCH") {
t.Fatalf("Validate() error = %v, want invalid cleanup batch", err)
}
}
func TestLoadAndValidateDirectMediaOSS(t *testing.T) {
t.Setenv("APP_ENV", "production")
t.Setenv("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "true")
t.Setenv("AI_GATEWAY_MEDIA_OSS_ENDPOINT", "https://media-bucket.oss-cn-shanghai.aliyuncs.com")
t.Setenv("AI_GATEWAY_MEDIA_OSS_BUCKET", "media-bucket")
t.Setenv("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "test-access-key")
t.Setenv("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "test-access-secret")
t.Setenv("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", "https://cdn.example.com")
t.Setenv("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/production/media")
cfg := Load()
if !cfg.MediaOSSDirectEnabled || cfg.MediaOSSBucket != "media-bucket" ||
cfg.MediaOSSObjectPrefix != "easyai-ai-gateway/production/media" {
t.Fatalf("unexpected direct media OSS config: %+v", cfg)
}
if err := cfg.Validate(); err != nil {
t.Fatalf("validate direct media OSS: %v", err)
}
cfg.MediaOSSAccessKeySecret = ""
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ACCESS_KEY_SECRET") {
t.Fatalf("Validate() error = %v, want missing direct OSS secret", err)
}
cfg.MediaOSSAccessKeySecret = "test-access-secret"
cfg.MediaOSSEndpoint = "http://media-bucket.example.com"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HTTPS URL") {
t.Fatalf("Validate() error = %v, want insecure production endpoint", err)
}
cfg.MediaOSSEndpoint = "https://media-bucket.example.com"
cfg.MediaOSSObjectPrefix = "../outside"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "OBJECT_PREFIX") {
t.Fatalf("Validate() error = %v, want invalid object prefix", err)
}
}
func TestCapacityControllerDoesNotRequireMediaOSSCredentials(t *testing.T) {
cfg := Load()
cfg.ProcessRole = "capacity-controller"
cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}]`
cfg.MediaOSSDirectEnabled = true
cfg.MediaOSSAccessKeyID = ""
cfg.MediaOSSAccessKeySecret = ""
if err := cfg.Validate(); err != nil {
t.Fatalf("capacity controller rejected unused media OSS credentials: %v", err)
}
cfg.ProcessRole = "worker"
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ACCESS_KEY") {
t.Fatalf("worker Validate() error = %v, want missing direct OSS credentials", err)
}
}
@@ -0,0 +1,27 @@
package config
import "testing"
func TestCapacityPoolsAreDataDriven(t *testing.T) {
cfg := Config{CapacityPoolsJSON: `[
{"id":"pool-a","adapterRef":"deployment-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":3},
{"id":"pool-b","adapterRef":"deployment-b","bootstrapReplicas":0,"minReplicas":0,"maxReplicas":2}
]`}
pools, err := cfg.CapacityPools()
if err != nil {
t.Fatal(err)
}
if len(pools) != 2 || pools[0].ID != "pool-a" || pools[1].AdapterRef != "deployment-b" {
t.Fatalf("pools=%+v", pools)
}
}
func TestCapacityPoolsRejectDuplicates(t *testing.T) {
cfg := Config{CapacityPoolsJSON: `[
{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1},
{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}
]`}
if _, err := cfg.CapacityPools(); err == nil {
t.Fatal("expected duplicate pool rejection")
}
}
+137
View File
@@ -0,0 +1,137 @@
package executionpool
import (
"crypto/hmac"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"net"
"net/url"
"strings"
"time"
)
func QueueName(poolID string) string {
poolID = strings.TrimSpace(poolID)
if poolID == "" {
return "gateway_tasks"
}
digest := sha256.Sum256([]byte(poolID))
return "gateway_pool_" + fmt.Sprintf("%x", digest[:8])
}
func RouteProfileKey(provider, protocol, endpointHost, proxyMode, configRevision string) string {
value := strings.Join([]string{
strings.ToLower(strings.TrimSpace(provider)),
strings.ToLower(strings.TrimSpace(protocol)),
strings.ToLower(strings.TrimSpace(endpointHost)),
strings.ToLower(strings.TrimSpace(proxyMode)),
strings.TrimSpace(configRevision),
}, "\x00")
digest := sha256.Sum256([]byte(value))
return fmt.Sprintf("route_%x", digest[:16])
}
func EndpointHost(rawURL string) string {
parsed, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil {
return ""
}
return strings.ToLower(parsed.Hostname())
}
func ValidateAdvertisedEndpoint(rawURL string, allowedSuffixes []string, allowPrivate bool) error {
parsed, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil || parsed.Host == "" || parsed.User != nil {
return errors.New("worker endpoint must be an absolute URL without user info")
}
if parsed.Scheme != "https" && parsed.Scheme != "http" {
return errors.New("worker endpoint scheme must be http or https")
}
host := strings.ToLower(parsed.Hostname())
if host == "" {
return errors.New("worker endpoint host is required")
}
if ip := net.ParseIP(host); ip != nil {
if allowPrivate && (ip.IsPrivate() || ip.IsLoopback()) {
return nil
}
return errors.New("worker endpoint IP is outside the configured trust boundary")
}
for _, suffix := range allowedSuffixes {
suffix = strings.ToLower(strings.TrimSpace(suffix))
if suffix != "" && (host == strings.TrimPrefix(suffix, ".") || strings.HasSuffix(host, "."+strings.TrimPrefix(suffix, "."))) {
return nil
}
}
return errors.New("worker endpoint host is outside the configured trust boundary")
}
type ExecutionClaims struct {
Audience string `json:"aud"`
TaskID string `json:"task_id"`
PoolID string `json:"pool_id"`
WorkerID string `json:"worker_id"`
Nonce string `json:"nonce"`
ExpiresAt int64 `json:"exp"`
}
type TokenSigner struct {
Secret []byte
Now func() time.Time
}
func (s TokenSigner) Sign(claims ExecutionClaims) (string, error) {
if len(s.Secret) < 32 {
return "", errors.New("execution token secret must be at least 32 bytes")
}
if claims.Audience == "" || claims.TaskID == "" || claims.PoolID == "" || claims.WorkerID == "" || claims.Nonce == "" || claims.ExpiresAt <= 0 {
return "", errors.New("execution token claims are incomplete")
}
payload, err := json.Marshal(claims)
if err != nil {
return "", err
}
encoded := base64.RawURLEncoding.EncodeToString(payload)
mac := hmac.New(sha256.New, s.Secret)
_, _ = mac.Write([]byte(encoded))
signature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
return encoded + "." + signature, nil
}
func (s TokenSigner) Verify(token, audience string) (ExecutionClaims, error) {
var claims ExecutionClaims
if len(s.Secret) < 32 {
return claims, errors.New("execution token secret must be at least 32 bytes")
}
parts := strings.Split(token, ".")
if len(parts) != 2 {
return claims, errors.New("invalid execution token")
}
signature, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return claims, errors.New("invalid execution token signature")
}
mac := hmac.New(sha256.New, s.Secret)
_, _ = mac.Write([]byte(parts[0]))
if !hmac.Equal(signature, mac.Sum(nil)) {
return claims, errors.New("invalid execution token signature")
}
payload, err := base64.RawURLEncoding.DecodeString(parts[0])
if err != nil || json.Unmarshal(payload, &claims) != nil {
return ExecutionClaims{}, errors.New("invalid execution token payload")
}
now := time.Now()
if s.Now != nil {
now = s.Now()
}
if claims.ExpiresAt <= now.Unix() {
return ExecutionClaims{}, errors.New("execution token expired")
}
if claims.Audience != audience {
return ExecutionClaims{}, errors.New("execution token audience mismatch")
}
return claims, nil
}
@@ -0,0 +1,44 @@
package executionpool
import (
"strings"
"testing"
"time"
)
func TestQueueNameDoesNotExposePoolID(t *testing.T) {
name := QueueName("region-sensitive-name")
if strings.Contains(name, "sensitive") || name == QueueName("other") {
t.Fatalf("unexpected queue name %q", name)
}
}
func TestTokenSigner(t *testing.T) {
now := time.Unix(1000, 0)
signer := TokenSigner{Secret: []byte("01234567890123456789012345678901"), Now: func() time.Time { return now }}
token, err := signer.Sign(ExecutionClaims{
Audience: "worker", TaskID: "task", PoolID: "pool", WorkerID: "worker", Nonce: "nonce", ExpiresAt: now.Add(30 * time.Second).Unix(),
})
if err != nil {
t.Fatal(err)
}
claims, err := signer.Verify(token, "worker")
if err != nil || claims.TaskID != "task" {
t.Fatalf("claims=%+v err=%v", claims, err)
}
if _, err := signer.Verify(token, "other"); err == nil {
t.Fatal("expected audience mismatch")
}
}
func TestValidateAdvertisedEndpoint(t *testing.T) {
if err := ValidateAdvertisedEndpoint("http://10.0.0.2:8088", nil, true); err != nil {
t.Fatal(err)
}
if err := ValidateAdvertisedEndpoint("https://worker.internal.example", []string{"internal.example"}, false); err != nil {
t.Fatal(err)
}
if err := ValidateAdvertisedEndpoint("https://public.example", []string{"internal.example"}, false); err == nil {
t.Fatal("expected untrusted endpoint rejection")
}
}
@@ -0,0 +1,63 @@
package executionpool
import (
"strings"
"time"
)
type RoutePreference struct {
RouteProfileKey string
CurrentPoolID string
CurrentSince time.Time
ChallengerPoolID string
ChallengerWins int
UpdatedAt time.Time
}
func AdvanceRoutePreference(
previous RoutePreference,
proposedPoolID string,
scores map[string]float64,
now time.Time,
) RoutePreference {
if now.IsZero() {
now = time.Now()
}
proposedPoolID = strings.TrimSpace(proposedPoolID)
next := previous
next.UpdatedAt = now
if proposedPoolID == "" {
return next
}
if _, eligible := scores[next.CurrentPoolID]; next.CurrentPoolID == "" || !eligible {
next.CurrentPoolID = proposedPoolID
next.CurrentSince = now
next.ChallengerPoolID = ""
next.ChallengerWins = 0
return next
}
if proposedPoolID == next.CurrentPoolID {
next.ChallengerPoolID = ""
next.ChallengerWins = 0
return next
}
if relativeImprovement(scores[proposedPoolID], scores[next.CurrentPoolID]) < defaultSwitchImprovement {
next.ChallengerPoolID = ""
next.ChallengerWins = 0
return next
}
if next.ChallengerPoolID == proposedPoolID {
next.ChallengerWins++
} else {
next.ChallengerPoolID = proposedPoolID
next.ChallengerWins = 1
}
if next.ChallengerWins >= defaultBetterWindowCount &&
!next.CurrentSince.IsZero() && now.Sub(next.CurrentSince) >= defaultMinimumDwell {
next.CurrentPoolID = proposedPoolID
next.CurrentSince = now
next.ChallengerPoolID = ""
next.ChallengerWins = 0
}
return next
}
@@ -0,0 +1,32 @@
package executionpool
import (
"testing"
"time"
)
func TestAdvanceRoutePreferenceRequiresThreeWinsAndMinimumDwell(t *testing.T) {
now := time.Date(2026, 8, 5, 12, 0, 0, 0, time.UTC)
state := RoutePreference{CurrentPoolID: "slow", CurrentSince: now}
scores := map[string]float64{"slow": 0.5, "fast": 0.7}
for window := 1; window <= 3; window++ {
state = AdvanceRoutePreference(state, "fast", scores, now.Add(time.Duration(window)*10*time.Second))
if state.CurrentPoolID != "slow" {
t.Fatalf("switched before minimum dwell in window %d", window)
}
}
state = AdvanceRoutePreference(state, "fast", scores, now.Add(2*time.Minute))
if state.CurrentPoolID != "fast" {
t.Fatalf("current pool=%q, want fast", state.CurrentPoolID)
}
}
func TestAdvanceRoutePreferenceImmediatelyLeavesIneligibleCurrentPool(t *testing.T) {
now := time.Now()
state := AdvanceRoutePreference(RoutePreference{
CurrentPoolID: "offline", CurrentSince: now.Add(-time.Minute),
}, "healthy", map[string]float64{"healthy": 0.6}, now)
if state.CurrentPoolID != "healthy" || state.ChallengerWins != 0 {
t.Fatalf("unexpected preference: %#v", state)
}
}
+101
View File
@@ -0,0 +1,101 @@
package executionpool
import (
"context"
"crypto/tls"
"net/http"
"net/http/httptrace"
"time"
)
type ProbeTarget struct {
RouteProfile RouteProfile
URL string
Method string
Timeout time.Duration
}
type ProbeResult struct {
Reachable bool
StatusCode int
DNSDuration time.Duration
TCPDuration time.Duration
TLSDuration time.Duration
FirstByte time.Duration
SampledAt time.Time
ErrorClass string
}
type NetworkProbe interface {
Probe(context.Context, ProbeTarget) ProbeResult
}
type HTTPProbe struct {
Client *http.Client
Now func() time.Time
}
func (p HTTPProbe) Probe(ctx context.Context, target ProbeTarget) ProbeResult {
result := ProbeResult{}
now := time.Now
if p.Now != nil {
now = p.Now
}
result.SampledAt = now()
method := target.Method
if method == "" {
method = http.MethodHead
}
timeout := target.Timeout
if timeout <= 0 {
timeout = 3 * time.Second
}
probeCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
var dnsStart, tcpStart, tlsStart time.Time
requestStart := now()
trace := &httptrace.ClientTrace{
DNSStart: func(httptrace.DNSStartInfo) { dnsStart = now() },
DNSDone: func(httptrace.DNSDoneInfo) {
if !dnsStart.IsZero() {
result.DNSDuration = now().Sub(dnsStart)
}
},
ConnectStart: func(_, _ string) { tcpStart = now() },
ConnectDone: func(_, _ string, _ error) {
if !tcpStart.IsZero() {
result.TCPDuration = now().Sub(tcpStart)
}
},
TLSHandshakeStart: func() { tlsStart = now() },
TLSHandshakeDone: func(tls.ConnectionState, error) {
if !tlsStart.IsZero() {
result.TLSDuration = now().Sub(tlsStart)
}
},
GotFirstResponseByte: func() { result.FirstByte = now().Sub(requestStart) },
}
request, err := http.NewRequestWithContext(httptrace.WithClientTrace(probeCtx, trace), method, target.URL, nil)
if err != nil {
result.ErrorClass = "invalid_target"
return result
}
client := p.Client
if client == nil {
client = http.DefaultClient
}
response, err := client.Do(request)
if err != nil {
if probeCtx.Err() != nil {
result.ErrorClass = "timeout"
} else {
result.ErrorClass = "network"
}
return result
}
defer response.Body.Close()
result.StatusCode = response.StatusCode
// Any HTTP response, including 401/403/404/405, proves transport reachability.
result.Reachable = true
return result
}
+202
View File
@@ -0,0 +1,202 @@
package executionpool
import (
"errors"
"math"
"sort"
"strings"
"time"
)
const (
defaultHealthMaxAge = 45 * time.Second
defaultMinimumDwell = 2 * time.Minute
defaultSwitchImprovement = 0.20
defaultBetterWindowCount = 3
)
var ErrNoEligiblePool = errors.New("no eligible execution pool")
type SelectionCandidate struct {
Pool ExecutionPool
Health RouteHealth
Capacity CapacitySnapshot
CapabilityMatched bool
BetterWindows int
}
type SelectionRequest struct {
Candidates []SelectionCandidate
CurrentPoolID string
CurrentPoolSince time.Time
Now time.Time
}
type Decision struct {
PoolID string
Score float64
Reason string
Components map[string]float64
Rejected map[string]string
}
type Selector struct {
HealthMaxAge time.Duration
MinimumDwell time.Duration
SwitchImprovement float64
BetterWindows int
}
func NewSelector() Selector {
return Selector{
HealthMaxAge: defaultHealthMaxAge, MinimumDwell: defaultMinimumDwell,
SwitchImprovement: defaultSwitchImprovement, BetterWindows: defaultBetterWindowCount,
}
}
type scoredCandidate struct {
candidate SelectionCandidate
score float64
components map[string]float64
}
func (s Selector) Select(request SelectionRequest) (Decision, error) {
now := request.Now
if now.IsZero() {
now = time.Now()
}
if s.HealthMaxAge <= 0 {
s.HealthMaxAge = defaultHealthMaxAge
}
if s.MinimumDwell <= 0 {
s.MinimumDwell = defaultMinimumDwell
}
if s.SwitchImprovement <= 0 {
s.SwitchImprovement = defaultSwitchImprovement
}
if s.BetterWindows <= 0 {
s.BetterWindows = defaultBetterWindowCount
}
rejected := make(map[string]string)
scored := make([]scoredCandidate, 0, len(request.Candidates))
for _, candidate := range request.Candidates {
if reason := rejectionReason(candidate, now, s.HealthMaxAge); reason != "" {
rejected[candidate.Pool.ID] = reason
continue
}
network := routeQuality(candidate.Health)
queue := queueQuality(candidate.Capacity)
resources := clamp01(candidate.Capacity.ResourceHeadroom)
stability := clamp01(candidate.Capacity.Stability)
components := map[string]float64{
"network": network, "queue": queue, "resources": resources, "stability": stability,
}
scored = append(scored, scoredCandidate{
candidate: candidate,
score: 0.55*network + 0.25*queue + 0.15*resources + 0.05*stability,
components: components,
})
}
if len(scored) == 0 {
return Decision{Rejected: rejected}, ErrNoEligiblePool
}
sort.SliceStable(scored, func(i, j int) bool {
if math.Abs(scored[i].score-scored[j].score) > 0.000001 {
return scored[i].score > scored[j].score
}
return scored[i].candidate.Pool.ID < scored[j].candidate.Pool.ID
})
selected := scored[0]
if currentIndex := indexOfPool(scored, strings.TrimSpace(request.CurrentPoolID)); currentIndex >= 0 && currentIndex != 0 {
current := scored[currentIndex]
withinDwell := !request.CurrentPoolSince.IsZero() && now.Sub(request.CurrentPoolSince) < s.MinimumDwell
improvement := relativeImprovement(selected.score, current.score)
if withinDwell || improvement < s.SwitchImprovement || selected.candidate.BetterWindows < s.BetterWindows {
selected = current
}
}
return Decision{
PoolID: selected.candidate.Pool.ID,
Score: selected.score,
Reason: "network_quality_then_safe_capacity",
Components: selected.components,
Rejected: rejected,
}, nil
}
func rejectionReason(candidate SelectionCandidate, now time.Time, maxAge time.Duration) string {
if strings.TrimSpace(candidate.Pool.ID) == "" {
return "invalid_pool"
}
if candidate.Pool.State != PoolActive {
return "pool_not_active"
}
if !candidate.CapabilityMatched {
return "capability_mismatch"
}
if candidate.Health.State == RouteUnreachable {
return "route_unreachable"
}
if candidate.Health.State == RouteUnknown || candidate.Health.State == "" {
return "route_unknown"
}
if candidate.Health.SampledAt.IsZero() || now.Sub(candidate.Health.SampledAt) > maxAge || (!candidate.Health.ExpiresAt.IsZero() && !now.Before(candidate.Health.ExpiresAt)) {
return "route_stale"
}
if candidate.Capacity.Critical {
return "resource_critical"
}
if candidate.Capacity.WorkerCount < 1 {
return "no_worker"
}
if candidate.Capacity.SafeCapacity-candidate.Capacity.ActiveTasks < 1 {
return "capacity_exhausted"
}
return ""
}
func routeQuality(health RouteHealth) float64 {
success := clamp01(health.SuccessRate)
connect := latencyQuality(health.ConnectTLSP95, 100*time.Millisecond)
firstByte := latencyQuality(health.FirstByteP95, 250*time.Millisecond)
throughput := clamp01(health.UploadBytesPerSecond / (10 * 1024 * 1024))
jitter := latencyQuality(health.JitterP95, 100*time.Millisecond)
return 0.35*success + 0.25*connect + 0.25*firstByte + 0.10*throughput + 0.05*jitter
}
func queueQuality(capacity CapacitySnapshot) float64 {
available := capacity.SafeCapacity - capacity.ActiveTasks
if available <= 0 {
return 0
}
headroom := float64(available) / float64(max(capacity.SafeCapacity, 1))
waitPenalty := 1 / (1 + capacity.EstimatedQueueWait.Seconds())
return clamp01(0.6*headroom + 0.4*waitPenalty)
}
func latencyQuality(value, target time.Duration) float64 {
if value <= 0 {
return 0
}
return clamp01(1 / (1 + float64(value)/float64(target)))
}
func relativeImprovement(candidate, current float64) float64 {
if current <= 0 {
return 1
}
return (candidate - current) / current
}
func indexOfPool(candidates []scoredCandidate, poolID string) int {
for index := range candidates {
if candidates[index].candidate.Pool.ID == poolID {
return index
}
}
return -1
}
func clamp01(value float64) float64 {
return math.Max(0, math.Min(1, value))
}
@@ -0,0 +1,86 @@
package executionpool
import (
"errors"
"testing"
"time"
)
func TestSelectorRejectsUnreachableAndPrefersNetwork(t *testing.T) {
now := time.Now()
selector := NewSelector()
decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
candidate("slow", RouteHealthy, now, 120*time.Millisecond, 8, 1),
candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1),
candidate("down", RouteUnreachable, now, time.Millisecond, 32, 0),
}})
if err != nil {
t.Fatal(err)
}
if decision.PoolID != "fast" {
t.Fatalf("pool=%s, want fast", decision.PoolID)
}
if decision.Rejected["down"] != "route_unreachable" {
t.Fatalf("down rejection=%q", decision.Rejected["down"])
}
}
func TestSelectorUsesCapacityAsHardGate(t *testing.T) {
now := time.Now()
selector := NewSelector()
decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
candidate("fast-full", RouteHealthy, now, 10*time.Millisecond, 2, 2),
candidate("slower-free", RouteHealthy, now, 80*time.Millisecond, 8, 1),
}})
if err != nil {
t.Fatal(err)
}
if decision.PoolID != "slower-free" {
t.Fatalf("pool=%s, want slower-free", decision.PoolID)
}
}
func TestSelectorRejectsUnknownAndStale(t *testing.T) {
now := time.Now()
selector := NewSelector()
_, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
candidate("unknown", RouteUnknown, now, time.Millisecond, 8, 0),
candidate("stale", RouteHealthy, now.Add(-time.Minute), time.Millisecond, 8, 0),
}})
if !errors.Is(err, ErrNoEligiblePool) {
t.Fatalf("err=%v, want ErrNoEligiblePool", err)
}
}
func TestSelectorHonorsHysteresis(t *testing.T) {
now := time.Now()
selector := NewSelector()
fast := candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1)
fast.BetterWindows = 3
current := candidate("current", RouteHealthy, now, 80*time.Millisecond, 8, 1)
decision, err := selector.Select(SelectionRequest{
Now: now, Candidates: []SelectionCandidate{fast, current},
CurrentPoolID: "current", CurrentPoolSince: now.Add(-time.Minute),
})
if err != nil {
t.Fatal(err)
}
if decision.PoolID != "current" {
t.Fatalf("pool=%s, want current during dwell", decision.PoolID)
}
}
func candidate(id string, state RouteState, sampled time.Time, latency time.Duration, capacity, active int) SelectionCandidate {
return SelectionCandidate{
Pool: ExecutionPool{ID: id, State: PoolActive}, CapabilityMatched: true,
Health: RouteHealth{
PoolID: id, State: state, SuccessRate: 1, ConnectTLSP95: latency,
FirstByteP95: latency, UploadBytesPerSecond: 10 * 1024 * 1024,
JitterP95: latency / 10, SampledAt: sampled, ExpiresAt: sampled.Add(45 * time.Second),
},
Capacity: CapacitySnapshot{
PoolID: id, WorkerCount: 1, SafeCapacity: capacity, ActiveTasks: active,
ResourceHeadroom: 0.8, Stability: 1,
},
}
}
+169
View File
@@ -0,0 +1,169 @@
package executionpool
import (
"context"
"io"
"time"
)
const ProtocolVersion = "v1"
type PoolState string
const (
PoolActive PoolState = "active"
PoolDraining PoolState = "draining"
PoolDisabled PoolState = "disabled"
)
type RouteState string
const (
RouteHealthy RouteState = "healthy"
RouteDegraded RouteState = "degraded"
RouteUnreachable RouteState = "unreachable"
RouteUnknown RouteState = "unknown"
)
type ExecutionPool struct {
ID string
Labels map[string]string
Capabilities map[string]any
State PoolState
UpdatedAt time.Time
}
type WorkerDescriptor struct {
WorkerID string
InstanceID string
PoolID string
Endpoint string
ProtocolVersion string
Revision string
Capabilities map[string]any
Allocated int
SafeCapacity int
HeavyCapacity int
ActiveTasks int
PressureState string
HeartbeatAt time.Time
LoadSampledAt time.Time
}
func (w WorkerDescriptor) AvailableCapacity() int {
limit := w.Allocated
if w.SafeCapacity < limit {
limit = w.SafeCapacity
}
if limit < 0 {
return 0
}
available := limit - w.ActiveTasks
if available < 0 {
return 0
}
return available
}
type RouteProfile struct {
Key string
Provider string
Protocol string
EndpointHost string
ProxyMode string
ConfigRevision string
}
type RouteHealth struct {
PoolID string
RouteProfileKey string
State RouteState
SuccessRate float64
ConnectTLSP95 time.Duration
FirstByteP95 time.Duration
UploadBytesPerSecond float64
JitterP95 time.Duration
ConsecutiveFailures int
ConsecutiveSuccesses int
SampleCount int
SampledAt time.Time
ExpiresAt time.Time
}
// RouteObservation contains only transport-phase facts. Model generation time
// is deliberately excluded from this structure and therefore from routing
// quality updates.
type RouteObservation struct {
PoolID string
RouteProfileKey string
SampleCount int
SuccessCount int
ConnectTLSP95 time.Duration
UploadBytesPerSecond float64
}
type CapacitySnapshot struct {
PoolID string
WorkerCount int
Allocated int
SafeCapacity int
ActiveTasks int
HeavyCapacity int
HeavyTasks int
QueueDepth int
EstimatedQueueWait time.Duration
ResourceHeadroom float64
Stability float64
Critical bool
SampledAt time.Time
}
type DesiredCapacity struct {
PoolID string
Desired int
Reason string
ValidUntil time.Time
}
// WorkerDirectory is the platform-neutral discovery boundary. Implementations
// may use PostgreSQL, Consul, etcd, or another registry.
type WorkerDirectory interface {
ListWorkers(context.Context, time.Time) ([]WorkerDescriptor, error)
}
// ExecutionBroker hides the durable queue implementation from routing code.
type ExecutionBroker interface {
Publish(context.Context, string, string, time.Time) (int64, error)
}
// ExecutionTransport hides HTTP, gRPC, or service-mesh transport details.
type ExecutionTransport interface {
Execute(context.Context, WorkerDescriptor, ExecutionRequest) (ExecutionResponse, error)
}
type ExecutionRequest struct {
TaskID string
PoolID string
WorkerID string
LeaseID string
AuthorizationToken string
Deadline time.Time
Stream bool
}
type ExecutionResponse struct {
StatusCode int
Headers map[string][]string
Body io.ReadCloser
}
type RouteHealthRepository interface {
ListRouteHealth(context.Context, string, time.Time) ([]RouteHealth, error)
RecordRouteHealth(context.Context, RouteHealth) error
RecordRouteObservation(context.Context, RouteObservation) error
}
type CapacityProvider interface {
ListCapacity(context.Context, time.Time) ([]CapacitySnapshot, error)
PublishDesiredCapacity(context.Context, DesiredCapacity) error
}
@@ -0,0 +1,481 @@
package httpapi
import (
"context"
"crypto/sha256"
"encoding/json"
"errors"
"net/http"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const (
acceptanceRunHeader = "X-EasyAI-Acceptance-Run"
acceptanceTokenHeader = "X-EasyAI-Acceptance-Token"
acceptanceUpstreamHeader = "X-EasyAI-Acceptance-Upstream"
taskTrafficAuthorizationCacheTTL = 5 * time.Second
)
type taskTrafficAuthorizationCacheKey struct {
runID string
apiKeyID string
userID string
tokenHash [sha256.Size]byte
}
type taskTrafficAuthorizationCacheEntry struct {
runID string
expiresAt time.Time
}
type taskTrafficAuthorizationCall struct {
done chan struct{}
runID string
err error
}
type taskTrafficAdmission struct {
RunMode string
AcceptanceRunID string
}
type taskTrafficError struct {
Status int
Code string
Message string
Err error
}
func (s *Server) acceptanceStore() *store.Store {
if s.coordinationStore != nil {
return s.coordinationStore
}
return s.store
}
func (e *taskTrafficError) Error() string {
return e.Message
}
func (e *taskTrafficError) ErrorCode() string {
return e.Code
}
func (s *Server) authorizeTaskTraffic(r *http.Request, user *auth.User) (taskTrafficAdmission, error) {
runID, err := s.authorizeAcceptanceTask(r.Context(), r.Header.Get(acceptanceRunHeader), r.Header.Get(acceptanceTokenHeader), user)
if err != nil {
switch {
case errors.Is(err, store.ErrProductionTrafficPaused):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "validation_in_progress",
Message: "new production tasks are paused while validation is running", Err: err,
}
case errors.Is(err, store.ErrAcceptanceNotAuthorized):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
Message: "acceptance credentials do not match the active run", Err: err,
}
case errors.Is(err, store.ErrAcceptanceRunNotActive):
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusConflict, Code: "acceptance_run_not_active",
Message: "acceptance headers are not valid while live traffic is enabled", Err: err,
}
default:
return taskTrafficAdmission{}, &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
Message: "task traffic gate is unavailable", Err: err,
}
}
}
if runID != "" {
runMode := "acceptance"
if strings.EqualFold(strings.TrimSpace(r.Header.Get(acceptanceUpstreamHeader)), "real") {
runMode = "acceptance_canary"
}
return taskTrafficAdmission{RunMode: runMode, AcceptanceRunID: runID}, nil
}
return taskTrafficAdmission{RunMode: "production"}, nil
}
// authorizeAcceptanceTask collapses concurrent validation requests for the same
// participant into one critical-pool lookup and briefly caches only successful
// validation grants. Live requests and incomplete acceptance credentials always
// consult PostgreSQL so a validation transition remains fail-closed.
func (s *Server) authorizeAcceptanceTask(ctx context.Context, runID string, token string, user *auth.User) (string, error) {
runID = strings.TrimSpace(runID)
token = strings.TrimSpace(token)
if runID == "" || token == "" || user == nil ||
strings.TrimSpace(user.APIKeyID) == "" || strings.TrimSpace(user.ID) == "" {
return s.acceptanceStore().AuthorizeAcceptanceTask(ctx, runID, token, user)
}
key := taskTrafficAuthorizationCacheKey{
runID: runID, apiKeyID: strings.TrimSpace(user.APIKeyID), userID: strings.TrimSpace(user.ID),
tokenHash: sha256.Sum256([]byte(token)),
}
return s.cachedTaskTrafficAuthorization(ctx, key, taskTrafficAuthorizationCacheTTL, func() (string, error) {
return s.acceptanceStore().AuthorizeAcceptanceTask(ctx, runID, token, user)
})
}
func (s *Server) cachedTaskTrafficAuthorization(
ctx context.Context,
key taskTrafficAuthorizationCacheKey,
ttl time.Duration,
load func() (string, error),
) (string, error) {
now := time.Now()
s.taskTrafficCacheMu.Lock()
if entry, ok := s.taskTrafficCache[key]; ok {
if now.Before(entry.expiresAt) {
s.taskTrafficCacheMu.Unlock()
return entry.runID, nil
}
delete(s.taskTrafficCache, key)
}
if call, ok := s.taskTrafficInflight[key]; ok {
s.taskTrafficCacheMu.Unlock()
select {
case <-call.done:
return call.runID, call.err
case <-ctx.Done():
return "", ctx.Err()
}
}
if s.taskTrafficInflight == nil {
s.taskTrafficInflight = make(map[taskTrafficAuthorizationCacheKey]*taskTrafficAuthorizationCall)
}
call := &taskTrafficAuthorizationCall{done: make(chan struct{})}
s.taskTrafficInflight[key] = call
s.taskTrafficCacheMu.Unlock()
call.runID, call.err = load()
s.taskTrafficCacheMu.Lock()
delete(s.taskTrafficInflight, key)
if call.err == nil && call.runID != "" {
if s.taskTrafficCache == nil {
s.taskTrafficCache = make(map[taskTrafficAuthorizationCacheKey]taskTrafficAuthorizationCacheEntry)
}
callTTL := ttl
if callTTL <= 0 {
callTTL = taskTrafficAuthorizationCacheTTL
}
s.taskTrafficCache[key] = taskTrafficAuthorizationCacheEntry{
runID: call.runID, expiresAt: time.Now().Add(callTTL),
}
}
close(call.done)
s.taskTrafficCacheMu.Unlock()
return call.runID, call.err
}
func (s *Server) admittedTaskRunMode(admission taskTrafficAdmission, body map[string]any) (string, error) {
if admission.RunMode == "acceptance" || admission.RunMode == "acceptance_canary" {
return admission.RunMode, nil
}
requested := runModeFromRequest(body)
requestedMode := strings.ToLower(strings.TrimSpace(requested))
if requestedMode == "acceptance" || requestedMode == "acceptance_canary" {
return "", &taskTrafficError{
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
Message: "acceptance mode is available only through an active acceptance run",
}
}
if strings.EqualFold(strings.TrimSpace(s.cfg.AppEnv), "production") && requestedMode == "simulation" {
return "", &taskTrafficError{
Status: http.StatusForbidden, Code: "simulation_not_authorized",
Message: "simulation mode is available only through an active acceptance run",
}
}
return requested, nil
}
func writeTaskTrafficError(w http.ResponseWriter, err error, protocol string) {
var trafficErr *taskTrafficError
if !errors.As(err, &trafficErr) {
trafficErr = &taskTrafficError{
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
Message: "task traffic gate is unavailable", Err: err,
}
}
if protocol != "" {
writeProtocolError(w, protocol, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
return
}
writeErrorWithDetails(w, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
}
// getGatewayTrafficMode godoc
// @Summary 获取 Gateway 流量模式
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/traffic-mode [get]
func (s *Server) getGatewayTrafficMode(w http.ResponseWriter, r *http.Request) {
mode, err := s.acceptanceStore().GetGatewayTrafficMode(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "get gateway traffic mode failed")
return
}
writeJSON(w, http.StatusOK, mode)
}
// pauseGatewayTraffic godoc
// @Summary 上线后硬门禁异常时通过 CAS 暂停正式新任务
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param body body store.PauseGatewayTrafficInput true "当前线上 release 与 revision CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/traffic-mode/pause [post]
func (s *Server) pauseGatewayTraffic(w http.ResponseWriter, r *http.Request) {
var input store.PauseGatewayTrafficInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
mode, err := s.acceptanceStore().PauseGatewayTraffic(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// listCapacityProfiles godoc
// @Summary 获取经生产同构验收认证的容量配置
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Success 200 {array} store.CapacityProfile
// @Router /api/admin/system/acceptance/capacity-profiles [get]
func (s *Server) listCapacityProfiles(w http.ResponseWriter, r *http.Request) {
profiles, err := s.acceptanceStore().ListCapacityProfiles(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "list capacity profiles failed")
return
}
writeJSON(w, http.StatusOK, profiles)
}
// getWorkerClusterRuntime godoc
// @Summary 获取集群 Worker 实时负载与领取容量
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Success 200 {object} store.WorkerClusterRuntime
// @Router /api/admin/runtime/workers [get]
func (s *Server) getWorkerClusterRuntime(w http.ResponseWriter, r *http.Request) {
runtime, err := s.acceptanceStore().GetWorkerClusterRuntime(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, "get worker cluster runtime failed")
return
}
writeJSON(w, http.StatusOK, runtime)
}
// createAcceptanceRun godoc
// @Summary 创建生产同构验收 Run
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param body body store.CreateAcceptanceRunInput true "验收 Run"
// @Success 201 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs [post]
func (s *Server) createAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.CreateAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
run, err := s.acceptanceStore().CreateAcceptanceRun(r.Context(), input)
if err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
}
writeJSON(w, http.StatusCreated, run)
}
// getAcceptanceRun godoc
// @Summary 获取生产同构验收 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID} [get]
func (s *Server) getAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.acceptanceStore().GetAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "acceptance run not found")
return
}
writeError(w, http.StatusInternalServerError, "get acceptance run failed")
return
}
writeJSON(w, http.StatusOK, run)
}
// activateAcceptanceRun godoc
// @Summary 切换到 validation 并激活 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/activate [post]
func (s *Server) activateAcceptanceRun(w http.ResponseWriter, r *http.Request) {
mode, err := s.acceptanceStore().ActivateAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// stageAcceptanceCapacityProfiles godoc
// @Summary 为当前 validation Run 暂存 80% 容量门禁
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.StageAcceptanceCapacityProfilesInput true "待验证容量配置"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/capacity-profiles [post]
func (s *Server) stageAcceptanceCapacityProfiles(w http.ResponseWriter, r *http.Request) {
var input store.StageAcceptanceCapacityProfilesInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
run, err := s.acceptanceStore().StageAcceptanceCapacityProfiles(r.Context(), input)
if err != nil {
switch {
case store.IsNotFound(err):
writeError(w, http.StatusConflict, "acceptance run is not running")
default:
writeError(w, http.StatusBadRequest, err.Error())
}
return
}
writeJSON(w, http.StatusOK, run)
}
// finishAcceptanceRun godoc
// @Summary 记录验收门禁结果
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.FinishAcceptanceRunInput true "门禁结果"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/finish [post]
func (s *Server) finishAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.FinishAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
run, err := s.acceptanceStore().FinishAcceptanceRun(r.Context(), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusConflict, "acceptance run is not running")
return
}
writeError(w, http.StatusInternalServerError, "finish acceptance run failed")
return
}
writeJSON(w, http.StatusOK, run)
}
// retryAcceptanceRun godoc
// @Summary 在 validation 关闭正式流量的状态下重试失败 Run
// @Tags acceptance
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Success 200 {object} store.AcceptanceRun
// @Router /api/admin/system/acceptance/runs/{runID}/retry [post]
func (s *Server) retryAcceptanceRun(w http.ResponseWriter, r *http.Request) {
run, err := s.acceptanceStore().RetryAcceptanceRun(r.Context(), r.PathValue("runID"))
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, run)
}
// promoteAcceptanceRun godoc
// @Summary 通过 CAS 切回 live
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/promote [post]
func (s *Server) promoteAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.PromoteAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
mode, err := s.acceptanceStore().PromoteAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
// abortAcceptanceRun godoc
// @Summary 人工 CAS 中止验收并切回 live
// @Tags acceptance
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param runID path string true "验收 Run ID"
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
// @Success 200 {object} store.GatewayTrafficMode
// @Router /api/admin/system/acceptance/runs/{runID}/abort [post]
func (s *Server) abortAcceptanceRun(w http.ResponseWriter, r *http.Request) {
var input store.PromoteAcceptanceRunInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
input.RunID = r.PathValue("runID")
mode, err := s.acceptanceStore().AbortAcceptanceRun(r.Context(), input)
if err != nil {
writeAcceptanceMutationError(w, err)
return
}
writeJSON(w, http.StatusOK, mode)
}
func writeAcceptanceMutationError(w http.ResponseWriter, err error) {
switch {
case errors.Is(err, store.ErrAcceptanceStateConflict):
writeError(w, http.StatusConflict, err.Error())
case errors.Is(err, store.ErrAcceptancePromotionGates):
writeError(w, http.StatusPreconditionFailed, err.Error())
case store.IsNotFound(err):
writeError(w, http.StatusNotFound, "acceptance run not found")
default:
writeError(w, http.StatusInternalServerError, "acceptance state update failed")
}
}
@@ -0,0 +1,103 @@
package httpapi
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
)
func TestAdmittedTaskRunModeReservesAcceptanceAndProductionSimulation(t *testing.T) {
server := &Server{cfg: config.Config{AppEnv: "production"}}
for _, runMode := range []string{"simulation", "SIMULATION", "acceptance", "acceptance_canary"} {
_, err := server.admittedTaskRunMode(taskTrafficAdmission{}, map[string]any{"runMode": runMode})
var trafficErr *taskTrafficError
if !errors.As(err, &trafficErr) || trafficErr.Status != 403 {
t.Fatalf("run mode %q error=%v", runMode, err)
}
}
for _, runMode := range []string{"acceptance", "acceptance_canary"} {
got, err := server.admittedTaskRunMode(taskTrafficAdmission{RunMode: runMode}, map[string]any{
"runMode": "production",
})
if err != nil || got != runMode {
t.Fatalf("admitted mode=%q got=%q err=%v", runMode, got, err)
}
}
}
func TestCachedTaskTrafficAuthorizationCollapsesConcurrentLoads(t *testing.T) {
server := &Server{}
key := taskTrafficAuthorizationCacheKey{runID: "run", apiKeyID: "key", userID: "user"}
var loads atomic.Int32
start := make(chan struct{})
const callers = 64
var wg sync.WaitGroup
errCh := make(chan error, callers)
for range callers {
wg.Add(1)
go func() {
defer wg.Done()
<-start
runID, err := server.cachedTaskTrafficAuthorization(
context.Background(), key, time.Minute,
func() (string, error) {
loads.Add(1)
time.Sleep(10 * time.Millisecond)
return "run", nil
},
)
if err != nil || runID != "run" {
errCh <- errors.New("unexpected cached authorization result")
}
}()
}
close(start)
wg.Wait()
close(errCh)
for err := range errCh {
t.Fatal(err)
}
if got := loads.Load(); got != 1 {
t.Fatalf("authorization loads=%d, want 1", got)
}
if _, err := server.cachedTaskTrafficAuthorization(
context.Background(), key, time.Minute,
func() (string, error) {
loads.Add(1)
return "run", nil
},
); err != nil {
t.Fatal(err)
}
if got := loads.Load(); got != 1 {
t.Fatalf("cached authorization loads=%d, want 1", got)
}
}
func TestCachedTaskTrafficAuthorizationDoesNotCacheFailures(t *testing.T) {
server := &Server{}
key := taskTrafficAuthorizationCacheKey{runID: "run", apiKeyID: "key", userID: "user"}
wantErr := errors.New("database unavailable")
loads := 0
for range 2 {
_, err := server.cachedTaskTrafficAuthorization(
context.Background(), key, time.Minute,
func() (string, error) {
loads++
return "", wantErr
},
)
if !errors.Is(err, wantErr) {
t.Fatalf("authorization error=%v, want %v", err, wantErr)
}
}
if loads != 2 {
t.Fatalf("failed authorization loads=%d, want 2", loads)
}
}
@@ -0,0 +1,165 @@
package httpapi
import (
"context"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
type acceptanceStatusSnapshot struct {
ID string
Status string
}
type acceptanceUserGroupSnapshot struct {
ID string
RateLimitPolicy []byte
}
// isolateHTTPAcceptanceState lets a stress test disable unrelated runtime
// resources without leaking those mutations into later integration tests that
// intentionally share the same dedicated test database.
func isolateHTTPAcceptanceState(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
includeStorage bool,
) func() {
t.Helper()
startedAt := time.Now()
platforms := readAcceptanceStatusSnapshots(t, ctx, pool, `
SELECT id::text, status
FROM integration_platforms
WHERE deleted_at IS NULL`)
userGroups := make([]acceptanceUserGroupSnapshot, 0)
rows, err := pool.Query(ctx, `
SELECT id::text, rate_limit_policy
FROM gateway_user_groups`)
if err != nil {
t.Fatalf("snapshot acceptance user groups: %v", err)
}
for rows.Next() {
var item acceptanceUserGroupSnapshot
if err := rows.Scan(&item.ID, &item.RateLimitPolicy); err != nil {
rows.Close()
t.Fatalf("scan acceptance user group snapshot: %v", err)
}
userGroups = append(userGroups, item)
}
if err := rows.Err(); err != nil {
rows.Close()
t.Fatalf("read acceptance user group snapshots: %v", err)
}
rows.Close()
storageChannels := []acceptanceStatusSnapshot(nil)
if includeStorage {
storageChannels = readAcceptanceStatusSnapshots(t, ctx, pool, `
SELECT id::text, status
FROM file_storage_channels
WHERE deleted_at IS NULL`)
}
if _, err := pool.Exec(ctx, `
UPDATE integration_platforms
SET status = 'disabled'
WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
if includeStorage {
if _, err := pool.Exec(ctx, `
UPDATE file_storage_channels
SET status = 'disabled'
WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance storage channels: %v", err)
}
}
return func() {
restoreCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
tx, err := pool.Begin(restoreCtx)
if err != nil {
t.Errorf("begin acceptance state restore: %v", err)
return
}
defer func() {
_ = tx.Rollback(restoreCtx)
}()
for _, item := range platforms {
if _, err := tx.Exec(restoreCtx, `
UPDATE integration_platforms
SET status = $2
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
t.Errorf("restore acceptance platform %s: %v", item.ID, err)
return
}
}
if _, err := tx.Exec(restoreCtx, `
UPDATE integration_platforms
SET status = 'disabled'
WHERE created_at >= $1
AND deleted_at IS NULL`, startedAt); err != nil {
t.Errorf("disable acceptance-created platforms: %v", err)
return
}
for _, item := range userGroups {
if _, err := tx.Exec(restoreCtx, `
UPDATE gateway_user_groups
SET rate_limit_policy = $2::jsonb
WHERE id = $1::uuid`, item.ID, string(item.RateLimitPolicy)); err != nil {
t.Errorf("restore acceptance user group %s: %v", item.ID, err)
return
}
}
if includeStorage {
for _, item := range storageChannels {
if _, err := tx.Exec(restoreCtx, `
UPDATE file_storage_channels
SET status = $2
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
t.Errorf("restore acceptance storage channel %s: %v", item.ID, err)
return
}
}
if _, err := tx.Exec(restoreCtx, `
UPDATE file_storage_channels
SET status = 'disabled'
WHERE created_at >= $1
AND deleted_at IS NULL`, startedAt); err != nil {
t.Errorf("disable acceptance-created storage channels: %v", err)
return
}
}
if err := tx.Commit(restoreCtx); err != nil {
t.Errorf("commit acceptance state restore: %v", err)
}
}
}
func readAcceptanceStatusSnapshots(
t *testing.T,
ctx context.Context,
pool *pgxpool.Pool,
query string,
) []acceptanceStatusSnapshot {
t.Helper()
rows, err := pool.Query(ctx, query)
if err != nil {
t.Fatalf("snapshot acceptance statuses: %v", err)
}
defer rows.Close()
items := make([]acceptanceStatusSnapshot, 0)
for rows.Next() {
var item acceptanceStatusSnapshot
if err := rows.Scan(&item.ID, &item.Status); err != nil {
t.Fatalf("scan acceptance status snapshot: %v", err)
}
items = append(items, item)
}
if err := rows.Err(); err != nil {
t.Fatalf("read acceptance status snapshots: %v", err)
}
return items
}
@@ -12,7 +12,7 @@ import (
// listAccessRules godoc
// @Summary 列出访问规则
// @Description 管理端返回用户组、租户、用户或 API Key 到平台、平台模型、基础模型的访问规则。
// @Description 管理端返回用户组、租户、用户或 API Key 到平台、平台模型、基础模型的分层访问规则。主体当前层无 allow 时继承上级,存在 allow 时仅允许白名单,deny 始终优先。
// @Tags access-rules
// @Produce json
// @Security BearerAuth
@@ -33,7 +33,7 @@ func (s *Server) listAccessRules(w http.ResponseWriter, r *http.Request) {
// listAPIKeyAccessRules godoc
// @Summary 列出 API Key 访问规则
// @Description 返回当前本地用户可管理的 API Key 访问规则。
// @Description 返回当前本地用户拥有的 API Key 访问规则;不会混入其他用户或其他 API Key 的规则
// @Tags api-keys
// @Produce json
// @Security BearerAuth
@@ -60,16 +60,18 @@ func (s *Server) listAPIKeyAccessRules(w http.ResponseWriter, r *http.Request) {
// listAPIKeyAssignableModels godoc
// @Summary 列出 API Key 可分配模型
// @Description 按当前用户自身的户、租户和用户组权限返回可分配给 API Key 的启用模型,不任何 API Key 权限规则影响
// @Description 按当前用户自身的户、用户组和用户分层白名单返回可分配给 API Key 的启用模型,不应用任何 API Key 层规则
// @Tags api-keys
// @Produce json
// @Security BearerAuth
// @Deprecated
// @Success 200 {object} PlatformModelListResponse
// @Failure 401 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/v1/api-keys/assignable-models [get]
func (s *Server) listAPIKeyAssignableModels(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Deprecation", "true")
user, _ := auth.UserFromContext(r.Context())
models, err := s.store.ListAPIKeyAssignablePlatformModels(r.Context(), user)
if err != nil {
@@ -84,9 +86,44 @@ func (s *Server) listAPIKeyAssignableModels(w http.ResponseWriter, r *http.Reque
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
// listAPIKeyAssignableModelsForKey godoc
// @Summary 列出指定 API Key 可分配模型
// @Description 返回全局启用、命中指定 API Key 的租户/用户组/用户基线且符合 Key scope 的可分配平台来源;当前 Key 的 allow/deny 不缩减候选,仅作为已有规则有效性诊断返回。
// @Tags api-keys
// @Produce json
// @Security BearerAuth
// @Param apiKeyID path string true "API Key ID"
// @Success 200 {object} APIKeyAssignableModelsResponse
// @Failure 401 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/v1/api-keys/{apiKeyID}/assignable-models [get]
func (s *Server) listAPIKeyAssignableModelsForKey(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context())
models, diagnostics, err := s.store.ListAPIKeyAssignablePlatformModelsForKey(r.Context(), user, r.PathValue("apiKeyID"))
if err != nil {
if errors.Is(err, store.ErrLocalUserRequired) {
writeLocalUserRequired(w)
return
}
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "api key not found")
return
}
s.logger.Error("list api key assignable models failed", "error", err)
writeError(w, http.StatusInternalServerError, "list api key assignable models failed")
return
}
writeJSON(w, http.StatusOK, APIKeyAssignableModelsResponse{
Items: s.platformModelResponses(r.Context(), models),
RuleDiagnostics: diagnostics,
})
}
// createAccessRule godoc
// @Summary 创建访问规则
// @Description 管理端创建一条访问控制规则。
// @Description 管理端创建一条访问控制规则;同一主体层存在任意有效 allow 后该层启用白名单,deny 始终优先
// @Tags access-rules
// @Accept json
// @Produce json
@@ -124,7 +161,7 @@ func (s *Server) createAccessRule(w http.ResponseWriter, r *http.Request) {
// batchAccessRules godoc
// @Summary 批量写入访问规则
// @Description 管理端为同一主体批量新增、更新或删除资源访问规则。
// @Description 管理端为同一主体批量新增、更新或删除资源访问规则。清空该主体全部 allow 会恢复上级继承。
// @Tags access-rules
// @Accept json
// @Produce json
@@ -157,7 +194,7 @@ func (s *Server) batchAccessRules(w http.ResponseWriter, r *http.Request) {
// batchAPIKeyAccessRules godoc
// @Summary 批量写入 API Key 访问规则
// @Description 当前本地用户为自己的 API Key 批量新增、更新或删除可访问资源
// @Description 当前本地用户为自己的 API Key 批量新增、更新或删除白名单/拒绝资源;Key 无 allow 时继承父级范围,存在 allow 后仅允许命中项
// @Tags api-keys
// @Accept json
// @Produce json
@@ -192,7 +229,7 @@ func (s *Server) batchAPIKeyAccessRules(w http.ResponseWriter, r *http.Request)
return
}
if errors.Is(err, store.ErrAccessRuleResourceDenied) {
writeError(w, http.StatusForbidden, "resource is not available for current user group")
writeError(w, http.StatusForbidden, "resource is not available for current user group or API key scope")
return
}
s.logger.Error("batch api key access rules failed", "error", err)
@@ -0,0 +1,143 @@
package httpapi
import (
"net/http"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
)
type AdminExecutionPoolWorker struct {
WorkerID string `json:"workerId"`
InstanceID string `json:"instanceId"`
Revision string `json:"revision,omitempty"`
ProtocolVersion string `json:"protocolVersion"`
Capabilities map[string]any `json:"capabilities"`
Allocated int `json:"allocated"`
SafeCapacity int `json:"safeCapacity"`
ActiveTasks int `json:"activeTasks"`
PressureState string `json:"pressureState"`
HeartbeatAt time.Time `json:"heartbeatAt"`
}
type AdminExecutionPool struct {
PoolID string `json:"poolId"`
Labels map[string]string `json:"labels"`
Capabilities map[string]any `json:"capabilities"`
State string `json:"state"`
Capacity *AdminExecutionPoolCapacity `json:"capacity,omitempty"`
Workers []AdminExecutionPoolWorker `json:"workers"`
}
type AdminExecutionPoolCapacity struct {
WorkerCount int `json:"workerCount"`
Allocated int `json:"allocated"`
SafeCapacity int `json:"safeCapacity"`
ActiveTasks int `json:"activeTasks"`
HeavyCapacity int `json:"heavyCapacity"`
HeavyTasks int `json:"heavyTasks"`
QueueDepth int `json:"queueDepth"`
EstimatedQueueWaitMilli int64 `json:"estimatedQueueWaitMs"`
ResourceHeadroom float64 `json:"resourceHeadroom"`
Stability float64 `json:"stability"`
Critical bool `json:"critical"`
SampledAt time.Time `json:"sampledAt"`
}
type AdminRouteHealth struct {
PoolID string `json:"poolId"`
RouteProfileKey string `json:"routeProfileKey"`
State string `json:"state"`
SuccessRate float64 `json:"successRate"`
ConnectTLSP95Milli int64 `json:"connectTlsP95Ms"`
FirstByteP95Milli int64 `json:"firstByteP95Ms"`
UploadBytesPerSecond float64 `json:"uploadBytesPerSecond"`
JitterP95Milli int64 `json:"jitterP95Ms"`
ConsecutiveFailures int `json:"consecutiveFailures"`
ConsecutiveSuccesses int `json:"consecutiveSuccesses"`
SampleCount int `json:"sampleCount"`
SampledAt time.Time `json:"sampledAt"`
ExpiresAt time.Time `json:"expiresAt"`
}
type AdminExecutionPoolResponse struct {
Pools []AdminExecutionPool `json:"pools"`
RouteHealth []AdminRouteHealth `json:"routeHealth"`
CapturedAt time.Time `json:"capturedAt"`
}
// listExecutionPools godoc
// @Summary 查询逻辑执行池、Worker 容量和上游路由健康状态
// @Tags Admin Runtime
// @Produce json
// @Success 200 {object} AdminExecutionPoolResponse
// @Failure 500 {object} ErrorEnvelope
// @Security BearerAuth
// @Router /api/admin/runtime/execution-pools [get]
func (s *Server) listExecutionPools(w http.ResponseWriter, r *http.Request) {
now := time.Now()
pools, err := s.coordinationStore.ListExecutionPools(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, err.Error())
return
}
workers, err := s.coordinationStore.ListWorkers(r.Context(), now)
if err != nil {
writeError(w, http.StatusInternalServerError, err.Error())
return
}
capacities, err := s.coordinationStore.ListCapacity(r.Context(), now)
if err != nil {
writeError(w, http.StatusInternalServerError, err.Error())
return
}
health, err := s.coordinationStore.ListAllRouteHealth(r.Context())
if err != nil {
writeError(w, http.StatusInternalServerError, err.Error())
return
}
capacityByPool := make(map[string]executionpool.CapacitySnapshot, len(capacities))
for _, capacity := range capacities {
capacityByPool[capacity.PoolID] = capacity
}
workersByPool := make(map[string][]AdminExecutionPoolWorker)
for _, worker := range workers {
workersByPool[worker.PoolID] = append(workersByPool[worker.PoolID], AdminExecutionPoolWorker{
WorkerID: worker.WorkerID, InstanceID: worker.InstanceID, Revision: worker.Revision,
ProtocolVersion: worker.ProtocolVersion, Capabilities: worker.Capabilities,
Allocated: worker.Allocated, SafeCapacity: worker.SafeCapacity,
ActiveTasks: worker.ActiveTasks, PressureState: worker.PressureState,
HeartbeatAt: worker.HeartbeatAt,
})
}
response := AdminExecutionPoolResponse{
Pools: make([]AdminExecutionPool, 0, len(pools)), RouteHealth: make([]AdminRouteHealth, 0, len(health)), CapturedAt: now,
}
for _, item := range health {
response.RouteHealth = append(response.RouteHealth, AdminRouteHealth{
PoolID: item.PoolID, RouteProfileKey: item.RouteProfileKey, State: string(item.State),
SuccessRate: item.SuccessRate, ConnectTLSP95Milli: item.ConnectTLSP95.Milliseconds(),
FirstByteP95Milli: item.FirstByteP95.Milliseconds(), UploadBytesPerSecond: item.UploadBytesPerSecond,
JitterP95Milli: item.JitterP95.Milliseconds(), ConsecutiveFailures: item.ConsecutiveFailures,
ConsecutiveSuccesses: item.ConsecutiveSuccesses, SampleCount: item.SampleCount,
SampledAt: item.SampledAt, ExpiresAt: item.ExpiresAt,
})
}
for _, pool := range pools {
item := AdminExecutionPool{
PoolID: pool.ID, Labels: pool.Labels, Capabilities: pool.Capabilities,
State: string(pool.State), Workers: workersByPool[pool.ID],
}
if capacity, ok := capacityByPool[pool.ID]; ok {
item.Capacity = &AdminExecutionPoolCapacity{
WorkerCount: capacity.WorkerCount, Allocated: capacity.Allocated, SafeCapacity: capacity.SafeCapacity,
ActiveTasks: capacity.ActiveTasks, HeavyCapacity: capacity.HeavyCapacity, HeavyTasks: capacity.HeavyTasks,
QueueDepth: capacity.QueueDepth, EstimatedQueueWaitMilli: capacity.EstimatedQueueWait.Milliseconds(),
ResourceHeadroom: capacity.ResourceHeadroom, Stability: capacity.Stability,
Critical: capacity.Critical, SampledAt: capacity.SampledAt,
}
}
response.Pools = append(response.Pools, item)
}
writeJSON(w, http.StatusOK, response)
}
@@ -0,0 +1,213 @@
package httpapi
import (
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
)
// listAdminTasks godoc
// @Summary 管理员列出任务
// @Description 跨租户分页查询任务;请求、结果和执行快照中的常见敏感字段默认脱敏。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param q query string false "关键词,匹配任务、用户、租户、模型、平台和 API Key"
// @Param tenantId query string false "网关租户 ID"
// @Param userId query string false "网关用户 ID"
// @Param userGroupId query string false "用户组 ID"
// @Param status query string false "任务状态"
// @Param platformId query string false "执行平台 ID"
// @Param model query string false "调用或实际模型名称"
// @Param modelType query string false "模型类型"
// @Param runMode query string false "运行模式"
// @Param billingStatus query string false "计费状态"
// @Param apiKey query string false "API Key ID、名称或前缀"
// @Param createdFrom query string false "创建时间起点,支持 RFC3339 或日期格式"
// @Param createdTo query string false "创建时间终点,支持 RFC3339 或日期格式"
// @Param page query int false "页码" default(1)
// @Param pageSize query int false "每页数量" default(50)
// @Success 200 {object} AdminTaskListResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks [get]
func (s *Server) listAdminTasks(w http.ResponseWriter, r *http.Request) {
query := r.URL.Query()
page, err := positiveQueryInt(query.Get("page"), 1)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid page")
return
}
pageSize, err := positiveQueryInt(query.Get("pageSize"), 50)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid pageSize")
return
}
createdFrom, err := parseTaskListTime(query.Get("createdFrom"), query.Get("from"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdFrom")
return
}
createdTo, err := parseTaskListTime(query.Get("createdTo"), query.Get("to"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdTo")
return
}
for name, value := range map[string]string{
"tenantId": query.Get("tenantId"),
"userId": query.Get("userId"),
"userGroupId": query.Get("userGroupId"),
"platformId": query.Get("platformId"),
} {
if value != "" && !validUUID(value) {
writeError(w, http.StatusBadRequest, "invalid "+name)
return
}
}
result, err := s.store.ListAdminTasks(r.Context(), store.AdminTaskListFilter{
Query: firstNonEmpty(query.Get("q"), query.Get("query")),
GatewayTenant: query.Get("tenantId"),
GatewayUser: query.Get("userId"),
UserGroup: query.Get("userGroupId"),
Status: query.Get("status"),
Platform: query.Get("platformId"),
Model: query.Get("model"),
ModelType: query.Get("modelType"),
RunMode: query.Get("runMode"),
BillingStatus: query.Get("billingStatus"),
APIKey: query.Get("apiKey"),
CreatedFrom: createdFrom,
CreatedTo: createdTo,
Page: page,
PageSize: pageSize,
})
if err != nil {
s.logger.Error("list admin tasks failed", "error", err)
writeError(w, http.StatusInternalServerError, "list admin tasks failed")
return
}
for index := range result.Items {
result.Items[index] = store.MaskAdminGatewayTask(result.Items[index])
}
writeJSON(w, http.StatusOK, map[string]any{
"items": result.Items,
"total": result.Total,
"page": result.Page,
"pageSize": result.PageSize,
})
}
// getAdminTask godoc
// @Summary 管理员获取任务详情
// @Description 返回任意租户的任务详情;默认脱敏,sensitive=full 时返回完整存储内容。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
// @Success 200 {object} store.AdminGatewayTask
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks/{taskID} [get]
func (s *Server) getAdminTask(w http.ResponseWriter, r *http.Request) {
taskID := r.PathValue("taskID")
if !validUUID(taskID) {
writeError(w, http.StatusBadRequest, "invalid taskID")
return
}
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
if !ok {
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
return
}
task, err := s.store.GetAdminTask(r.Context(), taskID)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "task not found")
return
}
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "get admin task failed")
return
}
task.GatewayTask, err = s.hydrateTaskResult(r.Context(), task.GatewayTask)
if err != nil {
writeStoredBinaryResultError(w, err)
return
}
if !full {
task = store.MaskAdminGatewayTask(task)
}
writeJSON(w, http.StatusOK, task)
}
// adminTaskParamPreprocessing godoc
// @Summary 管理员获取任务参数预处理日志
// @Description 返回任意租户任务的参数改写、校验或模板处理日志;默认脱敏。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
// @Success 200 {object} TaskParamPreprocessingLogListResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks/{taskID}/param-preprocessing [get]
func (s *Server) adminTaskParamPreprocessing(w http.ResponseWriter, r *http.Request) {
taskID := r.PathValue("taskID")
if !validUUID(taskID) {
writeError(w, http.StatusBadRequest, "invalid taskID")
return
}
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
if !ok {
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
return
}
if _, err := s.store.GetAdminTask(r.Context(), taskID); err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "task not found")
return
}
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "get admin task failed")
return
}
items, err := s.store.ListTaskParamPreprocessingLogs(r.Context(), taskID)
if err != nil {
s.logger.Error("list admin task parameter preprocessing logs failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "list admin task parameter preprocessing logs failed")
return
}
if !full {
items = store.MaskTaskParamPreprocessingLogs(items)
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func parseAdminTaskSensitiveMode(value string) (bool, bool) {
switch strings.ToLower(strings.TrimSpace(value)) {
case "", "masked":
return false, true
case "full":
return true, true
default:
return false, false
}
}
func validUUID(value string) bool {
_, err := uuid.Parse(strings.TrimSpace(value))
return err == nil
}
@@ -0,0 +1,22 @@
package httpapi
import "testing"
func TestParseAdminTaskSensitiveMode(t *testing.T) {
tests := []struct {
value string
full bool
accepted bool
}{
{value: "", full: false, accepted: true},
{value: "masked", full: false, accepted: true},
{value: "FULL", full: true, accepted: true},
{value: "invalid", full: false, accepted: false},
}
for _, test := range tests {
full, accepted := parseAdminTaskSensitiveMode(test.value)
if full != test.full || accepted != test.accepted {
t.Fatalf("parseAdminTaskSensitiveMode(%q)=(%v,%v), want (%v,%v)", test.value, full, accepted, test.full, test.accepted)
}
}
}
@@ -0,0 +1,321 @@
package httpapi
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/url"
"os"
"strconv"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdminTasksCrossTenantQueryAndRedaction(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
}
ctx := context.Background()
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
createTenant := func(key, name string) store.GatewayTenant {
t.Helper()
tenant, err := db.CreateTenant(ctx, store.GatewayTenantInput{
TenantKey: key + "-" + suffix,
Name: name,
Source: "gateway",
Status: "active",
})
if err != nil {
t.Fatalf("create tenant %s: %v", name, err)
}
return tenant
}
createUser := func(tenant store.GatewayTenant, prefix string, roles []string) store.GatewayUser {
t.Helper()
username := prefix + "_" + suffix
user, err := db.CreateGatewayUser(ctx, store.GatewayUserInput{
UserKey: "gateway:" + username,
Username: username,
DisplayName: prefix + " display",
Email: username + "@example.com",
Source: "gateway",
GatewayTenantID: tenant.ID,
TenantKey: tenant.TenantKey,
Roles: roles,
Status: "active",
})
if err != nil {
t.Fatalf("create gateway user %s: %v", prefix, err)
}
return user
}
authUser := func(user store.GatewayUser, apiKeyName string) *auth.User {
return &auth.User{
ID: user.ID,
Username: user.Username,
Roles: user.Roles,
Source: user.Source,
GatewayUserID: user.ID,
GatewayTenantID: user.GatewayTenantID,
TenantKey: user.TenantKey,
APIKeyID: "key-" + apiKeyName + "-" + suffix,
APIKeyName: apiKeyName,
APIKeyPrefix: "sk-" + apiKeyName,
}
}
createTask := func(user *auth.User, model, secret string) store.GatewayTask {
t.Helper()
task, err := db.CreateTask(ctx, store.CreateTaskInput{
Kind: "chat.completions",
Model: model,
RunMode: "simulation",
Request: map[string]any{
"model": model,
"diagnostic": map[string]any{
"apiKey": secret,
},
},
}, user)
if err != nil {
t.Fatalf("create task for %s: %v", user.Username, err)
}
return task
}
tenantA := createTenant("admin-task-a", "Admin Task Tenant A")
tenantB := createTenant("admin-task-b", "Admin Task Tenant B")
operator := createUser(tenantA, "admin_task_operator", []string{"operator"})
userA := createUser(tenantA, "admin_task_user_a", []string{"user"})
userB := createUser(tenantB, "admin_task_user_b", []string{"user"})
userAAuth := authUser(userA, "alpha-key")
userBAuth := authUser(userB, "bravo-key")
modelA := "admin-task-model-a-" + suffix
modelB := "admin-task-model-b-" + suffix
taskA := createTask(userAAuth, modelA, "alpha-secret-"+suffix)
taskBSecret := "bravo-secret-" + suffix
taskB := createTask(userBAuth, modelB, taskBSecret)
platform, err := db.CreatePlatform(ctx, store.CreatePlatformInput{
Provider: "openai",
PlatformKey: "admin-task-platform-" + suffix,
Name: "Admin Task Platform " + suffix,
BaseURL: "https://example.invalid/v1",
AuthType: "bearer",
Credentials: map[string]any{"apiKey": "platform-secret-" + suffix},
Status: "enabled",
})
if err != nil {
t.Fatalf("create platform: %v", err)
}
attemptID, err := db.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
TaskID: taskB.ID,
AttemptNo: 1,
PlatformID: platform.ID,
QueueKey: "admin-task-integration-" + suffix,
Status: "succeeded",
Simulated: true,
})
if err != nil {
t.Fatalf("create task attempt: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
UPDATE gateway_tasks
SET status = 'succeeded',
model_type = 'text_generate',
requested_model = $2,
resolved_model = $2,
request_id = $3,
billing_status = 'settled',
result = $4::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
taskB.ID,
modelB,
"request-admin-task-"+suffix,
`{"authorization":"Bearer result-secret","output":"ok"}`,
); err != nil {
t.Fatalf("enrich task: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
UPDATE gateway_task_attempts
SET request_snapshot = $2::jsonb,
response_snapshot = $3::jsonb,
finished_at = now()
WHERE id = $1::uuid`,
attemptID,
`{"headers":{"authorization":"Bearer attempt-secret"}}`,
`{"cookie":"attempt-cookie","status":"ok"}`,
); err != nil {
t.Fatalf("enrich task attempt: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
INSERT INTO gateway_task_param_preprocessing_logs (
task_id, attempt_id, attempt_no, model_type, platform_id, client_id,
changed, change_count, actual_input, converted_output, changes, model_snapshot
)
VALUES (
$1::uuid, $2::uuid, 1, 'text_generate', $3::uuid, $4,
true, 1, $5::jsonb, $6::jsonb, $7::jsonb, '{}'::jsonb
)`,
taskB.ID,
attemptID,
platform.ID,
"admin-task-client-"+suffix,
`{"password":"input-secret","prompt":"hello"}`,
`{"token":"converted-secret","prompt":"hello"}`,
`[{"path":"credentials.secret","before":"old-secret","after":"new-secret"}]`,
); err != nil {
t.Fatalf("create parameter preprocessing log: %v", err)
}
const jwtSecret = "admin-task-integration-jwt-secret"
handlerCtx, cancelHandler := context.WithCancel(ctx)
defer cancelHandler()
server := httptest.NewServer(NewServerWithContext(handlerCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: jwtSecret,
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
authenticator := auth.New(jwtSecret, "", "")
signToken := func(user *auth.User) string {
t.Helper()
token, err := authenticator.SignJWT(user, time.Hour)
if err != nil {
t.Fatalf("sign JWT for %s: %v", user.Username, err)
}
return token
}
operatorToken := signToken(&auth.User{
ID: operator.ID,
Username: operator.Username,
Roles: operator.Roles,
Source: operator.Source,
GatewayUserID: operator.ID,
GatewayTenantID: operator.GatewayTenantID,
TenantKey: operator.TenantKey,
})
userAToken := signToken(userAAuth)
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", userAToken, nil, http.StatusForbidden, nil)
query := url.Values{
"q": {taskB.ID},
"tenantId": {tenantB.ID},
"userId": {userB.ID},
"status": {"succeeded"},
"platformId": {platform.ID},
"model": {modelB},
"modelType": {"text_generate"},
"runMode": {"simulation"},
"billingStatus": {"settled"},
"apiKey": {"bravo-key"},
"page": {"1"},
"pageSize": {"10"},
}
var listResponse struct {
Items []store.AdminGatewayTask `json:"items"`
Total int `json:"total"`
Page int `json:"page"`
PageSize int `json:"pageSize"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks?"+query.Encode(), operatorToken, nil, http.StatusOK, &listResponse)
if len(listResponse.Items) != 1 || listResponse.Total != 1 || listResponse.Page != 1 || listResponse.PageSize != 10 {
t.Fatalf("unexpected filtered admin task list: %+v", listResponse)
}
listed := listResponse.Items[0]
if listed.ID != taskB.ID || listed.AdminContext.User == nil || listed.AdminContext.User.ID != userB.ID ||
listed.AdminContext.Tenant == nil || listed.AdminContext.Tenant.ID != tenantB.ID ||
listed.AdminContext.LatestPlatform == nil || listed.AdminContext.LatestPlatform.ID != platform.ID {
t.Fatalf("admin task list should expose cross-tenant identity and platform summaries: %+v", listed)
}
if nestedString(listed.Request, "diagnostic", "apiKey") != "***" ||
nestedString(listed.Result, "authorization") != "***" ||
nestedString(listed.Attempts[0].RequestSnapshot, "headers", "authorization") != "***" {
t.Fatalf("admin task list should mask task and attempt secrets: %+v", listed)
}
var maskedDetail store.AdminGatewayTask
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID, operatorToken, nil, http.StatusOK, &maskedDetail)
if nestedString(maskedDetail.Request, "diagnostic", "apiKey") != "***" {
t.Fatalf("admin task detail should be masked by default: %+v", maskedDetail.Request)
}
var fullDetail store.AdminGatewayTask
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"?sensitive=full", operatorToken, nil, http.StatusOK, &fullDetail)
if nestedString(fullDetail.Request, "diagnostic", "apiKey") != taskBSecret {
t.Fatalf("explicit full task detail should expose stored content: %+v", fullDetail.Request)
}
var maskedLogs struct {
Items []store.TaskParamPreprocessingLog `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing", operatorToken, nil, http.StatusOK, &maskedLogs)
if len(maskedLogs.Items) != 1 ||
nestedString(maskedLogs.Items[0].ActualInput, "password") != "***" ||
nestedString(maskedLogs.Items[0].ConvertedOutput, "token") != "***" {
t.Fatalf("admin task preprocessing logs should be masked by default: %+v", maskedLogs.Items)
}
var fullLogs struct {
Items []store.TaskParamPreprocessingLog `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing?sensitive=full", operatorToken, nil, http.StatusOK, &fullLogs)
if len(fullLogs.Items) != 1 ||
nestedString(fullLogs.Items[0].ActualInput, "password") != "input-secret" ||
nestedString(fullLogs.Items[0].ConvertedOutput, "token") != "converted-secret" {
t.Fatalf("explicit full preprocessing logs should expose stored content: %+v", fullLogs.Items)
}
var ownTasks struct {
Items []store.GatewayTask `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks?pageSize=100", userAToken, nil, http.StatusOK, &ownTasks)
if !gatewayTaskListHasID(ownTasks.Items, taskA.ID) || gatewayTaskListHasID(ownTasks.Items, taskB.ID) {
t.Fatalf("workspace task list must remain scoped to the current user: %+v", ownTasks.Items)
}
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID, userAToken, nil, http.StatusNotFound, nil)
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID+"/param-preprocessing", userAToken, nil, http.StatusNotFound, nil)
}
func nestedString(value map[string]any, path ...string) string {
var current any = value
for _, key := range path {
object, ok := current.(map[string]any)
if !ok {
return ""
}
current = object[key]
}
result, _ := current.(string)
return result
}
func gatewayTaskListHasID(items []store.GatewayTask, taskID string) bool {
for _, item := range items {
if item.ID == taskID {
return true
}
}
return false
}
@@ -0,0 +1,62 @@
package httpapi
import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdvancedMediaScopeAliases(t *testing.T) {
tests := []struct {
kind string
scope string
}{
{kind: "images.vectorize", scope: "image_vectorize"},
{kind: "images.vectorize", scope: "image"},
{kind: "images.vectorize", scope: "vectorize"},
{kind: "videos.upscales", scope: "video_enhance"},
{kind: "videos.upscales", scope: "video"},
{kind: "videos.upscales", scope: "video_upscale"},
}
for _, test := range tests {
user := &auth.User{APIKeyID: "key", APIKeyScopes: []string{test.scope}}
if !apiKeyScopeAllowed(user, test.kind) {
t.Fatalf("scope %q should allow %q", test.scope, test.kind)
}
}
if apiKeyScopeAllowed(&auth.User{APIKeyID: "key", APIKeyScopes: []string{"chat"}}, "videos.upscales") {
t.Fatal("chat scope must not allow video upscale")
}
}
func TestFillDailyTokenUsageDaysKeepsGapsAndStreaks(t *testing.T) {
location, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
to := time.Date(2026, 7, 24, 0, 0, 0, 0, location)
items, summary := fillDailyTokenUsageDays(map[string]store.DailyTokenUsage{
"2026-07-20": {Date: "2026-07-20", TotalTokens: 10, TaskCount: 1},
"2026-07-22": {Date: "2026-07-22", TotalTokens: 20, TaskCount: 1},
"2026-07-23": {Date: "2026-07-23", TotalTokens: 30, TaskCount: 2},
"2026-07-24": {Date: "2026-07-24", TotalTokens: 40, TaskCount: 1},
}, from, to)
if len(items) != 5 || items[1].Date != "2026-07-21" || items[1].TaskCount != 0 {
t.Fatalf("daily usage should contain continuous zero days: %+v", items)
}
if summary.CumulativeTokens != 100 || summary.PeakDailyTokens != 40 || summary.CurrentStreakDays != 3 || summary.LongestStreakDays != 3 {
t.Fatalf("unexpected daily usage summary: %+v", summary)
}
}
func TestAdvancedMediaDefaultModels(t *testing.T) {
if got := canonicalTaskModelName("images.vectorize", ""); got != "easy-image-vectorizer-1" {
t.Fatalf("vectorizer default model=%q", got)
}
if got := canonicalTaskModelName("videos.upscales", ""); got != "easy-proteus-standard-4" {
t.Fatalf("Topaz default model=%q", got)
}
}

Some files were not shown because too many files have changed in this diff Show More