github-actions[bot]
commited on
Commit
·
6fefda3
1
Parent(s):
daebe81
Update from GitHub Actions
Browse files- .env.example +24 -0
- CONFIGURATION_SYSTEM.md +274 -0
- Dockerfile +25 -0
- Dockerfile.ci +9 -0
- MULTI_JWT_README.md +284 -0
- docker-compose.env.example +11 -0
- docker-compose.yml +26 -0
- examples/.env.example +24 -0
- examples/complete_example.md +348 -0
- examples/demo_balancer.go +180 -0
- examples/start_with_multiple_jwt.sh +71 -0
- go.mod +33 -0
- go.sum +77 -0
- img.png +0 -0
- internal/apiserver/router.go +83 -0
- internal/balancer/health_checker.go +216 -0
- internal/balancer/jwt_balancer.go +175 -0
- internal/balancer/jwt_balancer_test.go +226 -0
- internal/config/config.go +430 -0
- internal/config/discovery.go +293 -0
- internal/jetbrains/client.go +188 -0
- internal/jetbrains/sse.go +331 -0
- internal/middleware/auth.go +31 -0
- internal/types/jetbrains.go +131 -0
- internal/utils/req_client.go +25 -0
- internal/utils/string.go +18 -0
- internal/utils/token.go +27 -0
- main.go +278 -0
- start.sh +200 -0
.env.example
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# JetBrains AI Proxy 多JWT配置示例
|
2 |
+
|
3 |
+
# 方式1: 使用多个JWT tokens(推荐)
|
4 |
+
# 多个tokens用逗号分隔,系统会自动进行负载均衡
|
5 |
+
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
6 |
+
|
7 |
+
# 方式2: 使用单个JWT token(向后兼容)
|
8 |
+
# JWT_TOKEN=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
9 |
+
|
10 |
+
# Bearer Token(必需)
|
11 |
+
BEARER_TOKEN=your_bearer_token_here
|
12 |
+
|
13 |
+
# 负载均衡策略(可选)
|
14 |
+
# round_robin: 轮询策略(默认)
|
15 |
+
# random: 随机策略
|
16 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
17 |
+
|
18 |
+
# 使用说明:
|
19 |
+
# 1. 复制此文件为 .env
|
20 |
+
# 2. 替换上面的示例值为真实的tokens
|
21 |
+
# 3. 启动服务: ./jetbrains-ai-proxy
|
22 |
+
#
|
23 |
+
# 或者使用命令行参数:
|
24 |
+
# ./jetbrains-ai-proxy -c "jwt1,jwt2,jwt3" -k "bearer_token" -s "random" -p 8080
|
CONFIGURATION_SYSTEM.md
ADDED
@@ -0,0 +1,274 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# 统一配置系统说明
|
2 |
+
|
3 |
+
JetBrains AI Proxy现在采用了全新的统一配置系统,实现了自动配置发现、多种配置方式支持和配置热重载等功能。
|
4 |
+
|
5 |
+
## 🏗️ 系统架构
|
6 |
+
|
7 |
+
### 核心组件
|
8 |
+
|
9 |
+
1. **ConfigManager** (`internal/config/config.go`)
|
10 |
+
- 统一的配置管理器
|
11 |
+
- 支持多种配置源的合并
|
12 |
+
- 线程安全的配置访问
|
13 |
+
- 配置验证和默认值处理
|
14 |
+
|
15 |
+
2. **ConfigDiscovery** (`internal/config/discovery.go`)
|
16 |
+
- 自动配置文件发现
|
17 |
+
- 配置文件格式验证
|
18 |
+
- 配置文件监控和热重载
|
19 |
+
- 示例配置生成
|
20 |
+
|
21 |
+
3. **JWTBalancer** (`internal/balancer/jwt_balancer.go`)
|
22 |
+
- JWT负载均衡器
|
23 |
+
- 支持轮询和随机策略
|
24 |
+
- 并发安全的token管理
|
25 |
+
- 动态token更新
|
26 |
+
|
27 |
+
4. **HealthChecker** (`internal/balancer/health_checker.go`)
|
28 |
+
- JWT健康检查器
|
29 |
+
- 自动故障检测和恢复
|
30 |
+
- 可配置的检查间隔
|
31 |
+
- 并发健康检查
|
32 |
+
|
33 |
+
## 📁 文件结构
|
34 |
+
|
35 |
+
```
|
36 |
+
jetbrains-ai-proxy/
|
37 |
+
├── main.go # 主程序,支持新配置系统
|
38 |
+
├── start.sh # 智能启动脚本
|
39 |
+
├── CONFIGURATION_SYSTEM.md # 配置系统说明(本文件)
|
40 |
+
├── MULTI_JWT_README.md # 多JWT功能说明
|
41 |
+
├── internal/
|
42 |
+
│ ├── config/
|
43 |
+
│ │ ├── config.go # 配置管理器
|
44 |
+
│ │ └── discovery.go # 配置发现器
|
45 |
+
│ ├── balancer/
|
46 |
+
│ │ ├── jwt_balancer.go # JWT负载均衡器
|
47 |
+
│ │ ├── health_checker.go # 健康检查器
|
48 |
+
│ │ └── jwt_balancer_test.go # 测试用例
|
49 |
+
│ └── jetbrains/
|
50 |
+
│ └── client.go # 集成配置系统的客户端
|
51 |
+
└── examples/
|
52 |
+
├── complete_example.md # 完整使用示例
|
53 |
+
├── start_with_multiple_jwt.sh # 多JWT启动脚本
|
54 |
+
└── .env.example # 环境变量示例
|
55 |
+
```
|
56 |
+
|
57 |
+
## ⚙️ 配置优先级
|
58 |
+
|
59 |
+
系统按以下优先级加载和合并配置:
|
60 |
+
|
61 |
+
1. **命令行参数** (最高优先级)
|
62 |
+
2. **环境变量**
|
63 |
+
3. **配置文件**
|
64 |
+
4. **默认值** (最低优先级)
|
65 |
+
|
66 |
+
## 🔍 配置发现机制
|
67 |
+
|
68 |
+
系统会按以下顺序搜索配置文件:
|
69 |
+
|
70 |
+
1. `CONFIG_FILE` 环境变量指定的路径
|
71 |
+
2. 当前目录:`config.json`, `jetbrains-ai-proxy.json`
|
72 |
+
3. config目录:`config/config.json`, `configs/config.json`
|
73 |
+
4. 隐藏目录:`.config/config.json`
|
74 |
+
5. 用户主目录:`$HOME/.config/jetbrains-ai-proxy/config.json`
|
75 |
+
6. 系统目录:`/etc/jetbrains-ai-proxy/config.json`
|
76 |
+
|
77 |
+
如果没有找到配置文件,系统会自动生成示例配置。
|
78 |
+
|
79 |
+
## 🚀 使用方式
|
80 |
+
|
81 |
+
### 1. 自动配置(推荐)
|
82 |
+
|
83 |
+
```bash
|
84 |
+
# 生成示例配置
|
85 |
+
./jetbrains-ai-proxy --generate-config
|
86 |
+
|
87 |
+
# 编辑配置文件
|
88 |
+
vim config/config.json
|
89 |
+
|
90 |
+
# 启动服务(自动发现配置)
|
91 |
+
./jetbrains-ai-proxy
|
92 |
+
```
|
93 |
+
|
94 |
+
### 2. 使用启动脚本
|
95 |
+
|
96 |
+
```bash
|
97 |
+
# 智能启动(自动检查配置)
|
98 |
+
./start.sh
|
99 |
+
|
100 |
+
# 生成配置
|
101 |
+
./start.sh --generate
|
102 |
+
|
103 |
+
# 查看配置
|
104 |
+
./start.sh --config
|
105 |
+
```
|
106 |
+
|
107 |
+
### 3. 指定配置文件
|
108 |
+
|
109 |
+
```bash
|
110 |
+
# 使用特定配置文件
|
111 |
+
./jetbrains-ai-proxy --config /path/to/config.json
|
112 |
+
|
113 |
+
# 或设置环境变量
|
114 |
+
export CONFIG_FILE=/path/to/config.json
|
115 |
+
./jetbrains-ai-proxy
|
116 |
+
```
|
117 |
+
|
118 |
+
## 📋 配置文件格式
|
119 |
+
|
120 |
+
### JSON配置文件示例
|
121 |
+
|
122 |
+
```json
|
123 |
+
{
|
124 |
+
"jetbrains_tokens": [
|
125 |
+
{
|
126 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
127 |
+
"name": "Primary_JWT",
|
128 |
+
"description": "Primary JWT token for JetBrains AI",
|
129 |
+
"priority": 1,
|
130 |
+
"metadata": {
|
131 |
+
"environment": "production",
|
132 |
+
"region": "us-east-1"
|
133 |
+
}
|
134 |
+
}
|
135 |
+
],
|
136 |
+
"bearer_token": "your_bearer_token_here",
|
137 |
+
"load_balance_strategy": "round_robin",
|
138 |
+
"health_check_interval": "30s",
|
139 |
+
"server_port": 8080,
|
140 |
+
"server_host": "0.0.0.0"
|
141 |
+
}
|
142 |
+
```
|
143 |
+
|
144 |
+
### 环境变量配置
|
145 |
+
|
146 |
+
```bash
|
147 |
+
# JWT Tokens(逗号分隔)
|
148 |
+
JWT_TOKENS=token1,token2,token3
|
149 |
+
|
150 |
+
# Bearer Token
|
151 |
+
BEARER_TOKEN=your_bearer_token
|
152 |
+
|
153 |
+
# 负载均衡策略
|
154 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
155 |
+
|
156 |
+
# 服务器配置
|
157 |
+
SERVER_HOST=0.0.0.0
|
158 |
+
SERVER_PORT=8080
|
159 |
+
|
160 |
+
# 配置文件路径(可选)
|
161 |
+
CONFIG_FILE=config/config.json
|
162 |
+
```
|
163 |
+
|
164 |
+
## 🔄 配置热重载
|
165 |
+
|
166 |
+
系统支持运行时配置重载,无需重启服务:
|
167 |
+
|
168 |
+
### 自动重载
|
169 |
+
|
170 |
+
配置文件监控器会自动检测配置文件变化并重新加载。
|
171 |
+
|
172 |
+
### 手动重载
|
173 |
+
|
174 |
+
```bash
|
175 |
+
# 通过API端点重载
|
176 |
+
curl -X POST http://localhost:8080/reload
|
177 |
+
|
178 |
+
# 响应
|
179 |
+
{
|
180 |
+
"message": "Configuration reloaded successfully"
|
181 |
+
}
|
182 |
+
```
|
183 |
+
|
184 |
+
## 🛠️ 管理端点
|
185 |
+
|
186 |
+
系统提供了丰富的管理端点:
|
187 |
+
|
188 |
+
| 端点 | 方法 | 描述 |
|
189 |
+
|------|------|------|
|
190 |
+
| `/health` | GET | 健康检查和负载均衡状态 |
|
191 |
+
| `/config` | GET | 当前配置信息(隐藏敏感数据) |
|
192 |
+
| `/stats` | GET | 详细统计信息 |
|
193 |
+
| `/reload` | POST | 重新加载配置 |
|
194 |
+
|
195 |
+
## 🔧 高级功能
|
196 |
+
|
197 |
+
### 1. JWT Token元数据
|
198 |
+
|
199 |
+
支持为每个JWT token配置元数据:
|
200 |
+
|
201 |
+
```json
|
202 |
+
{
|
203 |
+
"token": "jwt_token_here",
|
204 |
+
"name": "Production_Primary",
|
205 |
+
"description": "Primary production JWT token",
|
206 |
+
"priority": 1,
|
207 |
+
"metadata": {
|
208 |
+
"environment": "production",
|
209 |
+
"region": "us-east-1",
|
210 |
+
"tier": "primary",
|
211 |
+
"max_requests_per_minute": "1000"
|
212 |
+
}
|
213 |
+
}
|
214 |
+
```
|
215 |
+
|
216 |
+
### 2. 配置验证
|
217 |
+
|
218 |
+
系统会自动验证配置的有效性:
|
219 |
+
- JWT token格式检查
|
220 |
+
- 必需字段验证
|
221 |
+
- 数值范围检查
|
222 |
+
- 策略有效性验证
|
223 |
+
|
224 |
+
### 3. 配置合并策略
|
225 |
+
|
226 |
+
多个配置源的合并规则:
|
227 |
+
- 数组类型:高优先级完全覆盖低优先级
|
228 |
+
- 对象类型:递归合并,高优先级字段覆盖低优先级
|
229 |
+
- 基本类型:高优先级直接覆盖低优先级
|
230 |
+
|
231 |
+
## 🚨 故障排除
|
232 |
+
|
233 |
+
### 配置问题诊断
|
234 |
+
|
235 |
+
```bash
|
236 |
+
# 查看当前配置
|
237 |
+
./jetbrains-ai-proxy --print-config
|
238 |
+
|
239 |
+
# 验证配置文件
|
240 |
+
./jetbrains-ai-proxy --config config.json --print-config
|
241 |
+
|
242 |
+
# 生成新的示例配置
|
243 |
+
./jetbrains-ai-proxy --generate-config
|
244 |
+
```
|
245 |
+
|
246 |
+
### 常见问题
|
247 |
+
|
248 |
+
1. **配置文件未找到**
|
249 |
+
- 检查文件路径和权限
|
250 |
+
- 使用 `--generate-config` 生成示例配置
|
251 |
+
|
252 |
+
2. **JWT tokens无效**
|
253 |
+
- 检查token格式和有效性
|
254 |
+
- 查看健康检查日志
|
255 |
+
|
256 |
+
3. **配置合并问题**
|
257 |
+
- 使用 `--print-config` 查看最终配置
|
258 |
+
- 检查配置优先级
|
259 |
+
|
260 |
+
## 📈 性能考虑
|
261 |
+
|
262 |
+
1. **配置缓存**: 配置在内存中缓存,避免重复读取
|
263 |
+
2. **并发安全**: 使用读写锁保护配置访问
|
264 |
+
3. **懒加载**: 配置发现器按需加载配置文件
|
265 |
+
4. **监控优化**: 配置文件监控使用高效的文件系统事件
|
266 |
+
|
267 |
+
## 🔮 未来扩展
|
268 |
+
|
269 |
+
系统设计支持以下扩展:
|
270 |
+
- 远程配置中心集成(如Consul、etcd)
|
271 |
+
- 配置加密和安全存储
|
272 |
+
- 配置版本管理和回滚
|
273 |
+
- 更多负载均衡策略
|
274 |
+
- 动态配置更新API
|
Dockerfile
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Builder Stage - 使用明确的版本并优化缓存
|
2 |
+
FROM golang:1.24-alpine as builder
|
3 |
+
|
4 |
+
# 设置工作目录
|
5 |
+
WORKDIR /app
|
6 |
+
|
7 |
+
# 1. 仅复制依赖描述文件
|
8 |
+
COPY go.mod go.sum ./
|
9 |
+
|
10 |
+
# 2. 下载依赖项。这一步会被缓存,只有在 go.mod/go.sum 变化时才会重新运行
|
11 |
+
RUN go mod download
|
12 |
+
|
13 |
+
# 3. 复制项目源码
|
14 |
+
COPY . .
|
15 |
+
|
16 |
+
# 4. 编译应用。现在此步骤将使用缓存的依赖
|
17 |
+
RUN go build -o jetbrains-ai-proxy
|
18 |
+
|
19 |
+
# Final Stage - 保持不变
|
20 |
+
FROM alpine
|
21 |
+
LABEL maintainer="zouyq <[email protected]>"
|
22 |
+
|
23 |
+
COPY --from=builder /app/jetbrains-ai-proxy /usr/local/bin/
|
24 |
+
|
25 |
+
ENTRYPOINT ["jetbrains-ai-proxy"]
|
Dockerfile.ci
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
FROM alpine
|
2 |
+
LABEL maintainer="zouyq <[email protected]>"
|
3 |
+
|
4 |
+
ARG TARGETOS
|
5 |
+
ARG TARGETARCH
|
6 |
+
|
7 |
+
COPY dist/jetbrains-ai-proxy-${TARGETOS}-${TARGETARCH} /usr/local/bin/jetbrains-ai-proxy
|
8 |
+
|
9 |
+
ENTRYPOINT ["jetbrains-ai-proxy"]
|
MULTI_JWT_README.md
ADDED
@@ -0,0 +1,284 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# JetBrains AI Proxy - 多JWT负载均衡系统
|
2 |
+
|
3 |
+
本项目提供了一个功能完整的多JWT负载均衡系统,支持自动配置发现、健康检查和故障转移。
|
4 |
+
|
5 |
+
## 🎯 核心特性
|
6 |
+
|
7 |
+
- ✅ **智能配置管理**: 自动发现和加载配置文件,支持多种配置方式
|
8 |
+
- ✅ **多JWT支持**: 支持配置多个JWT tokens进行负载均衡
|
9 |
+
- ✅ **负载均衡策略**: 支持轮询(round_robin)和随机(random)两种策略
|
10 |
+
- ✅ **健康检查**: 自动检测失效的tokens并从负载均衡池中移除
|
11 |
+
- ✅ **故障转移**: 当某个token失效时自动切换到其他健康的token
|
12 |
+
- ✅ **配置热重载**: 支持运行时重新加载配置
|
13 |
+
- ✅ **并发安全**: 支持高并发环境下的安全使用
|
14 |
+
- ✅ **管理端点**: 提供健康检查、配置查看、统计信息等管理接口
|
15 |
+
- ✅ **优雅关闭**: 支持优雅关闭和资源清理
|
16 |
+
|
17 |
+
## 🚀 快速开始
|
18 |
+
|
19 |
+
### 方式1: 使用启动脚本(推荐)
|
20 |
+
|
21 |
+
```bash
|
22 |
+
# 1. 生成示例配置
|
23 |
+
./start.sh --generate
|
24 |
+
|
25 |
+
# 2. 编辑配置文件
|
26 |
+
vim config/config.json
|
27 |
+
# 或编辑环境变量文件
|
28 |
+
cp .env.example .env && vim .env
|
29 |
+
|
30 |
+
# 3. 启动服务
|
31 |
+
./start.sh
|
32 |
+
```
|
33 |
+
|
34 |
+
### 方式2: 直接使用可执行文件
|
35 |
+
|
36 |
+
```bash
|
37 |
+
# 生成示例配置
|
38 |
+
./jetbrains-ai-proxy --generate-config
|
39 |
+
|
40 |
+
# 查看当前配置
|
41 |
+
./jetbrains-ai-proxy --print-config
|
42 |
+
|
43 |
+
# 启动服务
|
44 |
+
./jetbrains-ai-proxy
|
45 |
+
```
|
46 |
+
|
47 |
+
## ⚙️ 配置方式
|
48 |
+
|
49 |
+
系统支持多种配置方式,优先级从高到低:
|
50 |
+
|
51 |
+
1. **命令行参数** (最高优先级)
|
52 |
+
2. **环境变量**
|
53 |
+
3. **配置文件**
|
54 |
+
4. **默认值** (最低优先级)
|
55 |
+
|
56 |
+
### 1. 配置文件方式(推荐)
|
57 |
+
|
58 |
+
系统会自动搜索以下路径的配置文件:
|
59 |
+
|
60 |
+
- `config.json`
|
61 |
+
- `config/config.json`
|
62 |
+
- `configs/config.json`
|
63 |
+
- `.config/jetbrains-ai-proxy.json`
|
64 |
+
- `$HOME/.config/jetbrains-ai-proxy/config.json`
|
65 |
+
|
66 |
+
配置文件示例:
|
67 |
+
|
68 |
+
```json
|
69 |
+
{
|
70 |
+
"jetbrains_tokens": [
|
71 |
+
{
|
72 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
73 |
+
"name": "Primary_JWT",
|
74 |
+
"description": "Primary JWT token for JetBrains AI",
|
75 |
+
"priority": 1,
|
76 |
+
"metadata": {
|
77 |
+
"environment": "production",
|
78 |
+
"region": "us-east-1"
|
79 |
+
}
|
80 |
+
},
|
81 |
+
{
|
82 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
83 |
+
"name": "Secondary_JWT",
|
84 |
+
"description": "Secondary JWT token for load balancing",
|
85 |
+
"priority": 2,
|
86 |
+
"metadata": {
|
87 |
+
"environment": "production",
|
88 |
+
"region": "us-west-2"
|
89 |
+
}
|
90 |
+
}
|
91 |
+
],
|
92 |
+
"bearer_token": "your_bearer_token_here",
|
93 |
+
"load_balance_strategy": "round_robin",
|
94 |
+
"health_check_interval": "30s",
|
95 |
+
"server_port": 8080,
|
96 |
+
"server_host": "0.0.0.0"
|
97 |
+
}
|
98 |
+
```
|
99 |
+
|
100 |
+
### 2. 环境变量配置
|
101 |
+
|
102 |
+
```bash
|
103 |
+
# 多个JWT tokens,用逗号分隔
|
104 |
+
export JWT_TOKENS="jwt_token_1,jwt_token_2,jwt_token_3"
|
105 |
+
|
106 |
+
# 或者使用旧的单token配置(向后兼容)
|
107 |
+
export JWT_TOKEN="single_jwt_token"
|
108 |
+
|
109 |
+
# Bearer token
|
110 |
+
export BEARER_TOKEN="your_bearer_token"
|
111 |
+
|
112 |
+
# 负载均衡策略(可选,默认为round_robin)
|
113 |
+
export LOAD_BALANCE_STRATEGY="random"
|
114 |
+
|
115 |
+
# 服务器配置
|
116 |
+
export SERVER_HOST="0.0.0.0"
|
117 |
+
export SERVER_PORT="8080"
|
118 |
+
|
119 |
+
# 指定配置文件路径(可选)
|
120 |
+
export CONFIG_FILE="path/to/config.json"
|
121 |
+
```
|
122 |
+
|
123 |
+
### 3. 命令行参数配置
|
124 |
+
|
125 |
+
```bash
|
126 |
+
# 查看所有选项
|
127 |
+
./jetbrains-ai-proxy --help
|
128 |
+
|
129 |
+
# 使用命令行参数启动
|
130 |
+
./jetbrains-ai-proxy \
|
131 |
+
-c "jwt1,jwt2,jwt3" \
|
132 |
+
-k "bearer_token" \
|
133 |
+
-s "round_robin" \
|
134 |
+
-p 8080 \
|
135 |
+
-h "0.0.0.0"
|
136 |
+
|
137 |
+
# 指定配置文件
|
138 |
+
./jetbrains-ai-proxy --config config/my-config.json
|
139 |
+
```
|
140 |
+
|
141 |
+
## 负载均衡策略
|
142 |
+
|
143 |
+
### 轮询策略 (round_robin)
|
144 |
+
|
145 |
+
- **默认策略**
|
146 |
+
- 按顺序依次使用每个健康的JWT token
|
147 |
+
- 确保负载均匀分布
|
148 |
+
- 适合大多数场景
|
149 |
+
|
150 |
+
### 随机策略 (random)
|
151 |
+
|
152 |
+
- 随机选择一个健康的JWT token
|
153 |
+
- 避免可预测的请求模式
|
154 |
+
- 适合需要随机分布的场景
|
155 |
+
|
156 |
+
## 健康检查机制
|
157 |
+
|
158 |
+
系统会自动进行JWT token健康检查:
|
159 |
+
|
160 |
+
- **检查间隔**: 每30秒检查一次
|
161 |
+
- **检查方式**: 发送测试请求到JetBrains AI API
|
162 |
+
- **故障处理**: 自动标记失效的tokens为不健康状态
|
163 |
+
- **恢复机制**: 定期重新检查不健康的tokens
|
164 |
+
|
165 |
+
## 监控端点
|
166 |
+
|
167 |
+
访问 `/health` 端点可以查看负载均衡器状态:
|
168 |
+
|
169 |
+
```bash
|
170 |
+
curl http://localhost:8080/health
|
171 |
+
```
|
172 |
+
|
173 |
+
响应示例:
|
174 |
+
|
175 |
+
```json
|
176 |
+
{
|
177 |
+
"status": "ok",
|
178 |
+
"healthy_tokens": 2,
|
179 |
+
"total_tokens": 3,
|
180 |
+
"strategy": "round_robin"
|
181 |
+
}
|
182 |
+
```
|
183 |
+
|
184 |
+
## 使用示例
|
185 |
+
|
186 |
+
### 启动服务
|
187 |
+
|
188 |
+
```bash
|
189 |
+
# 使用3个JWT tokens,轮询策略
|
190 |
+
./jetbrains-ai-proxy \
|
191 |
+
-p 8080 \
|
192 |
+
-c "eyJ0eXAiOiJKV1QiLCJhbGc...,eyJ0eXAiOiJKV1QiLCJhbGc...,eyJ0eXAiOiJKV1QiLCJhbGc..." \
|
193 |
+
-k "your_bearer_token" \
|
194 |
+
-s "round_robin"
|
195 |
+
```
|
196 |
+
|
197 |
+
### 发送请求
|
198 |
+
|
199 |
+
```bash
|
200 |
+
curl -X POST http://localhost:8080/v1/chat/completions \
|
201 |
+
-H "Authorization: Bearer your_bearer_token" \
|
202 |
+
-H "Content-Type: application/json" \
|
203 |
+
-d '{
|
204 |
+
"model": "gpt-4o",
|
205 |
+
"messages": [
|
206 |
+
{"role": "user", "content": "Hello, world!"}
|
207 |
+
],
|
208 |
+
"stream": false
|
209 |
+
}'
|
210 |
+
```
|
211 |
+
|
212 |
+
## 日志输出
|
213 |
+
|
214 |
+
启动时会显示负载均衡器配置信息:
|
215 |
+
|
216 |
+
```
|
217 |
+
2024/01/01 12:00:00 JWT balancer initialized with 3 tokens, strategy: round_robin
|
218 |
+
2024/01/01 12:00:00 JWT health checker started
|
219 |
+
2024/01/01 12:00:00 Server starting on 0.0.0.0:8080
|
220 |
+
2024/01/01 12:00:00 JWT tokens configured: 3
|
221 |
+
2024/01/01 12:00:00 Load balance strategy: round_robin
|
222 |
+
```
|
223 |
+
|
224 |
+
运行时会显示健康检查和token使用情况:
|
225 |
+
|
226 |
+
```
|
227 |
+
2024/01/01 12:01:00 Performing JWT health check...
|
228 |
+
2024/01/01 12:01:01 Health check completed: 3/3 tokens healthy
|
229 |
+
2024/01/01 12:01:30 JWT token marked as unhealthy: eyJ0eXAiOi... (errors: 1)
|
230 |
+
2024/01/01 12:02:00 JWT token marked as healthy: eyJ0eXAiOi...
|
231 |
+
```
|
232 |
+
|
233 |
+
## 故障排除
|
234 |
+
|
235 |
+
### 1. 所有tokens都不健康
|
236 |
+
|
237 |
+
如果所有JWT tokens都被标记为不健康,请求会返回错误:
|
238 |
+
|
239 |
+
```json
|
240 |
+
{
|
241 |
+
"error": "no available JWT tokens: no healthy JWT tokens available"
|
242 |
+
}
|
243 |
+
```
|
244 |
+
|
245 |
+
**解决方案**:
|
246 |
+
|
247 |
+
- 检查JWT tokens是否有效
|
248 |
+
- 检查网络连接
|
249 |
+
- 查看健康检查日志
|
250 |
+
|
251 |
+
### 2. 部分tokens不健康
|
252 |
+
|
253 |
+
系统会自动使用健康的tokens,但建议:
|
254 |
+
|
255 |
+
- 检查不健康tokens的有效性
|
256 |
+
- 考虑更换失效的tokens
|
257 |
+
- 监控健康检查日志
|
258 |
+
|
259 |
+
### 3. 性能优化
|
260 |
+
|
261 |
+
- 根据实际负载调整JWT tokens数量
|
262 |
+
- 选择合适的负载均衡策略
|
263 |
+
- 监控 `/health` 端点的响应
|
264 |
+
|
265 |
+
## 技术实现
|
266 |
+
|
267 |
+
### 核心组件
|
268 |
+
|
269 |
+
1. **JWTBalancer**: 负载均衡器接口和实现
|
270 |
+
2. **HealthChecker**: JWT健康检查器
|
271 |
+
3. **Config**: 配置管理,支持多JWT配置
|
272 |
+
4. **Client**: 集成负载均衡器的HTTP客户端
|
273 |
+
|
274 |
+
### 并发安全
|
275 |
+
|
276 |
+
- 使用读写锁保护共享状态
|
277 |
+
- 原子操作处理计数器
|
278 |
+
- 线程安全的随机数生成器
|
279 |
+
|
280 |
+
### 错误处理
|
281 |
+
|
282 |
+
- 自动重试机制
|
283 |
+
- 优雅的错误降级
|
284 |
+
- 详细的错误日志记录
|
docker-compose.env.example
ADDED
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# JetBrains AI Proxy Docker Compose 配置
|
2 |
+
|
3 |
+
# JWT Tokens (必需) - 多个tokens用逗号分隔
|
4 |
+
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
5 |
+
|
6 |
+
# Bearer Token (必需)
|
7 |
+
BEARER_TOKEN=your_bearer_token_here
|
8 |
+
|
9 |
+
# 负载均衡策略 (可选,默认: round_robin)
|
10 |
+
# 可选值: round_robin, random
|
11 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
docker-compose.yml
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
version: '3.8'
|
2 |
+
|
3 |
+
services:
|
4 |
+
jetbrains-ai-proxy:
|
5 |
+
build:
|
6 |
+
context: ./jetbrains-ai-proxy
|
7 |
+
dockerfile: Dockerfile
|
8 |
+
container_name: jetbrains-ai-proxy
|
9 |
+
ports:
|
10 |
+
- "8080:8080"
|
11 |
+
environment:
|
12 |
+
- JWT_TOKENS=${JWT_TOKENS}
|
13 |
+
- BEARER_TOKEN=${BEARER_TOKEN}
|
14 |
+
- LOAD_BALANCE_STRATEGY=${LOAD_BALANCE_STRATEGY:-round_robin}
|
15 |
+
- SERVER_HOST=0.0.0.0
|
16 |
+
- SERVER_PORT=8080
|
17 |
+
volumes:
|
18 |
+
- ./jetbrains-ai-proxy/config:/app/config:ro
|
19 |
+
- ./jetbrains-ai-proxy/logs:/app/logs
|
20 |
+
restart: unless-stopped
|
21 |
+
healthcheck:
|
22 |
+
test: ["CMD", "wget", "--no-verbose", "--tries=1", "--spider", "http://localhost:8080/health"]
|
23 |
+
interval: 1000s
|
24 |
+
timeout: 30s
|
25 |
+
retries: 3
|
26 |
+
start_period: 40s
|
examples/.env.example
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# JetBrains AI Proxy 多JWT配置示例
|
2 |
+
|
3 |
+
# 方式1: 使用多个JWT tokens(推荐)
|
4 |
+
# 多个tokens用逗号分隔,系统会自动进行负载均衡
|
5 |
+
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
6 |
+
|
7 |
+
# 方式2: 使用单个JWT token(向后兼容)
|
8 |
+
# JWT_TOKEN=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
9 |
+
|
10 |
+
# Bearer Token(必需)
|
11 |
+
BEARER_TOKEN=your_bearer_token_here
|
12 |
+
|
13 |
+
# 负载均衡策略(可选)
|
14 |
+
# round_robin: 轮询策略(默认)
|
15 |
+
# random: 随机策略
|
16 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
17 |
+
|
18 |
+
# 使用说明:
|
19 |
+
# 1. 复制此文件为 .env
|
20 |
+
# 2. 替换上面的示例值为真实的tokens
|
21 |
+
# 3. 启动服务: ./jetbrains-ai-proxy
|
22 |
+
#
|
23 |
+
# 或者使用命令行参数:
|
24 |
+
# ./jetbrains-ai-proxy -c "jwt1,jwt2,jwt3" -k "bearer_token" -s "random" -p 8080
|
examples/complete_example.md
ADDED
@@ -0,0 +1,348 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# 完整使用示例
|
2 |
+
|
3 |
+
本文档提供了JetBrains AI Proxy多JWT负载均衡系统的完整使用示例。
|
4 |
+
|
5 |
+
## 📋 准备工作
|
6 |
+
|
7 |
+
### 1. 编译项目
|
8 |
+
|
9 |
+
```bash
|
10 |
+
# 进入项目目录
|
11 |
+
cd jetbrains-ai-proxy
|
12 |
+
|
13 |
+
# 编译项目
|
14 |
+
go build -o jetbrains-ai-proxy
|
15 |
+
|
16 |
+
# 或者使用交叉编译
|
17 |
+
GOOS=linux GOARCH=amd64 go build -o jetbrains-ai-proxy-linux
|
18 |
+
```
|
19 |
+
|
20 |
+
### 2. 准备JWT Tokens
|
21 |
+
|
22 |
+
确保你有有效的JetBrains AI JWT tokens。你可以从以下途径获取:
|
23 |
+
- JetBrains AI服务控制台
|
24 |
+
- 现有的JetBrains IDE配置
|
25 |
+
- API密钥管理界面
|
26 |
+
|
27 |
+
## 🎯 使用场景示例
|
28 |
+
|
29 |
+
### 场景1: 开发环境快速启动
|
30 |
+
|
31 |
+
```bash
|
32 |
+
# 1. 生成示例配置
|
33 |
+
./jetbrains-ai-proxy --generate-config
|
34 |
+
|
35 |
+
# 2. 编辑配置文件
|
36 |
+
vim config/config.json
|
37 |
+
|
38 |
+
# 3. 启动服务
|
39 |
+
./jetbrains-ai-proxy
|
40 |
+
```
|
41 |
+
|
42 |
+
配置文件内容:
|
43 |
+
```json
|
44 |
+
{
|
45 |
+
"jetbrains_tokens": [
|
46 |
+
{
|
47 |
+
"token": "your_jwt_token_here",
|
48 |
+
"name": "Dev_Token",
|
49 |
+
"description": "Development JWT token"
|
50 |
+
}
|
51 |
+
],
|
52 |
+
"bearer_token": "your_bearer_token_here",
|
53 |
+
"load_balance_strategy": "round_robin",
|
54 |
+
"server_port": 8080,
|
55 |
+
"server_host": "127.0.0.1"
|
56 |
+
}
|
57 |
+
```
|
58 |
+
|
59 |
+
### 场景2: 生产环境多JWT负载均衡
|
60 |
+
|
61 |
+
```bash
|
62 |
+
# 1. 创建生产配置目录
|
63 |
+
mkdir -p /etc/jetbrains-ai-proxy
|
64 |
+
|
65 |
+
# 2. 创建生产配置文件
|
66 |
+
cat > /etc/jetbrains-ai-proxy/config.json << 'EOF'
|
67 |
+
{
|
68 |
+
"jetbrains_tokens": [
|
69 |
+
{
|
70 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
71 |
+
"name": "Primary_Production",
|
72 |
+
"description": "Primary production JWT token",
|
73 |
+
"priority": 1,
|
74 |
+
"metadata": {
|
75 |
+
"environment": "production",
|
76 |
+
"region": "us-east-1",
|
77 |
+
"tier": "primary"
|
78 |
+
}
|
79 |
+
},
|
80 |
+
{
|
81 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
82 |
+
"name": "Secondary_Production",
|
83 |
+
"description": "Secondary production JWT token",
|
84 |
+
"priority": 2,
|
85 |
+
"metadata": {
|
86 |
+
"environment": "production",
|
87 |
+
"region": "us-west-2",
|
88 |
+
"tier": "secondary"
|
89 |
+
}
|
90 |
+
},
|
91 |
+
{
|
92 |
+
"token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
93 |
+
"name": "Backup_Production",
|
94 |
+
"description": "Backup production JWT token",
|
95 |
+
"priority": 3,
|
96 |
+
"metadata": {
|
97 |
+
"environment": "production",
|
98 |
+
"region": "eu-west-1",
|
99 |
+
"tier": "backup"
|
100 |
+
}
|
101 |
+
}
|
102 |
+
],
|
103 |
+
"bearer_token": "prod_bearer_token_here",
|
104 |
+
"load_balance_strategy": "random",
|
105 |
+
"health_check_interval": "30s",
|
106 |
+
"server_port": 8080,
|
107 |
+
"server_host": "0.0.0.0"
|
108 |
+
}
|
109 |
+
EOF
|
110 |
+
|
111 |
+
# 3. 启动服务
|
112 |
+
./jetbrains-ai-proxy --config /etc/jetbrains-ai-proxy/config.json
|
113 |
+
```
|
114 |
+
|
115 |
+
### 场景3: 使用环境变量配置
|
116 |
+
|
117 |
+
```bash
|
118 |
+
# 1. 设置环境变量
|
119 |
+
export JWT_TOKENS="jwt1,jwt2,jwt3"
|
120 |
+
export BEARER_TOKEN="your_bearer_token"
|
121 |
+
export LOAD_BALANCE_STRATEGY="random"
|
122 |
+
export SERVER_PORT="9090"
|
123 |
+
|
124 |
+
# 2. 启动服务
|
125 |
+
./jetbrains-ai-proxy
|
126 |
+
|
127 |
+
# 或者使用.env文件
|
128 |
+
cat > .env << 'EOF'
|
129 |
+
JWT_TOKENS=jwt_token_1,jwt_token_2,jwt_token_3
|
130 |
+
BEARER_TOKEN=your_bearer_token_here
|
131 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
132 |
+
SERVER_HOST=0.0.0.0
|
133 |
+
SERVER_PORT=8080
|
134 |
+
EOF
|
135 |
+
|
136 |
+
./jetbrains-ai-proxy
|
137 |
+
```
|
138 |
+
|
139 |
+
### 场景4: Docker容器部署
|
140 |
+
|
141 |
+
```bash
|
142 |
+
# 1. 创建Dockerfile(如果不存在)
|
143 |
+
cat > Dockerfile << 'EOF'
|
144 |
+
FROM golang:1.21-alpine AS builder
|
145 |
+
WORKDIR /app
|
146 |
+
COPY go.mod go.sum ./
|
147 |
+
RUN go mod download
|
148 |
+
COPY . .
|
149 |
+
RUN go build -o jetbrains-ai-proxy
|
150 |
+
|
151 |
+
FROM alpine:latest
|
152 |
+
RUN apk --no-cache add ca-certificates
|
153 |
+
WORKDIR /root/
|
154 |
+
COPY --from=builder /app/jetbrains-ai-proxy .
|
155 |
+
EXPOSE 8080
|
156 |
+
CMD ["./jetbrains-ai-proxy"]
|
157 |
+
EOF
|
158 |
+
|
159 |
+
# 2. 构建镜像
|
160 |
+
docker build -t jetbrains-ai-proxy .
|
161 |
+
|
162 |
+
# 3. 运行容器
|
163 |
+
docker run -d \
|
164 |
+
--name jetbrains-ai-proxy \
|
165 |
+
-p 8080:8080 \
|
166 |
+
-e JWT_TOKENS="jwt1,jwt2,jwt3" \
|
167 |
+
-e BEARER_TOKEN="your_bearer_token" \
|
168 |
+
-e LOAD_BALANCE_STRATEGY="random" \
|
169 |
+
jetbrains-ai-proxy
|
170 |
+
|
171 |
+
# 4. 或者使用配置文件挂载
|
172 |
+
docker run -d \
|
173 |
+
--name jetbrains-ai-proxy \
|
174 |
+
-p 8080:8080 \
|
175 |
+
-v $(pwd)/config:/app/config \
|
176 |
+
jetbrains-ai-proxy
|
177 |
+
```
|
178 |
+
|
179 |
+
## 🔧 管理和监控
|
180 |
+
|
181 |
+
### 健康检查
|
182 |
+
|
183 |
+
```bash
|
184 |
+
# 检查服务状态
|
185 |
+
curl http://localhost:8080/health
|
186 |
+
|
187 |
+
# 响应示例
|
188 |
+
{
|
189 |
+
"status": "ok",
|
190 |
+
"healthy_tokens": 3,
|
191 |
+
"total_tokens": 3,
|
192 |
+
"strategy": "round_robin",
|
193 |
+
"server_info": {
|
194 |
+
"host": "0.0.0.0",
|
195 |
+
"port": 8080
|
196 |
+
}
|
197 |
+
}
|
198 |
+
```
|
199 |
+
|
200 |
+
### 查看配置信息
|
201 |
+
|
202 |
+
```bash
|
203 |
+
# 查看当前配置
|
204 |
+
curl http://localhost:8080/config
|
205 |
+
|
206 |
+
# 响应示例
|
207 |
+
{
|
208 |
+
"jwt_tokens_count": 3,
|
209 |
+
"jwt_tokens": [
|
210 |
+
{
|
211 |
+
"name": "Primary_JWT",
|
212 |
+
"description": "Primary JWT token",
|
213 |
+
"priority": 1,
|
214 |
+
"token_preview": "eyJ0eXAiOiJKV1QiLCJhbGc..."
|
215 |
+
}
|
216 |
+
],
|
217 |
+
"bearer_token_set": true,
|
218 |
+
"load_balance_strategy": "round_robin",
|
219 |
+
"health_check_interval": "30s",
|
220 |
+
"server_host": "0.0.0.0",
|
221 |
+
"server_port": 8080
|
222 |
+
}
|
223 |
+
```
|
224 |
+
|
225 |
+
### 查看统计信息
|
226 |
+
|
227 |
+
```bash
|
228 |
+
# 查看负载均衡统计
|
229 |
+
curl http://localhost:8080/stats
|
230 |
+
|
231 |
+
# 响应示例
|
232 |
+
{
|
233 |
+
"balancer": {
|
234 |
+
"healthy_tokens": 3,
|
235 |
+
"total_tokens": 3,
|
236 |
+
"strategy": "round_robin"
|
237 |
+
},
|
238 |
+
"config": {
|
239 |
+
"health_check_interval": "30s",
|
240 |
+
"server_host": "0.0.0.0",
|
241 |
+
"server_port": 8080
|
242 |
+
}
|
243 |
+
}
|
244 |
+
```
|
245 |
+
|
246 |
+
### 重载配置
|
247 |
+
|
248 |
+
```bash
|
249 |
+
# 重新加载配置(无需重启服务)
|
250 |
+
curl -X POST http://localhost:8080/reload
|
251 |
+
|
252 |
+
# 响应示例
|
253 |
+
{
|
254 |
+
"message": "Configuration reloaded successfully"
|
255 |
+
}
|
256 |
+
```
|
257 |
+
|
258 |
+
## 🧪 测试API
|
259 |
+
|
260 |
+
### 发送聊天请求
|
261 |
+
|
262 |
+
```bash
|
263 |
+
# 发送非流式请求
|
264 |
+
curl -X POST http://localhost:8080/v1/chat/completions \
|
265 |
+
-H "Authorization: Bearer your_bearer_token" \
|
266 |
+
-H "Content-Type: application/json" \
|
267 |
+
-d '{
|
268 |
+
"model": "gpt-4o",
|
269 |
+
"messages": [
|
270 |
+
{"role": "user", "content": "Hello, how are you?"}
|
271 |
+
],
|
272 |
+
"stream": false
|
273 |
+
}'
|
274 |
+
|
275 |
+
# 发送流式请求
|
276 |
+
curl -X POST http://localhost:8080/v1/chat/completions \
|
277 |
+
-H "Authorization: Bearer your_bearer_token" \
|
278 |
+
-H "Content-Type: application/json" \
|
279 |
+
-d '{
|
280 |
+
"model": "gpt-4o",
|
281 |
+
"messages": [
|
282 |
+
{"role": "user", "content": "Tell me a story"}
|
283 |
+
],
|
284 |
+
"stream": true
|
285 |
+
}'
|
286 |
+
```
|
287 |
+
|
288 |
+
### 获取支持的模型
|
289 |
+
|
290 |
+
```bash
|
291 |
+
# 获取模型列表
|
292 |
+
curl http://localhost:8080/v1/models \
|
293 |
+
-H "Authorization: Bearer your_bearer_token"
|
294 |
+
```
|
295 |
+
|
296 |
+
## 🚨 故障排除
|
297 |
+
|
298 |
+
### 常见问题
|
299 |
+
|
300 |
+
1. **所有JWT tokens都不健康**
|
301 |
+
```bash
|
302 |
+
# 检查token有效性
|
303 |
+
curl http://localhost:8080/health
|
304 |
+
|
305 |
+
# 查看日志
|
306 |
+
tail -f /var/log/jetbrains-ai-proxy.log
|
307 |
+
```
|
308 |
+
|
309 |
+
2. **配置文件未找到**
|
310 |
+
```bash
|
311 |
+
# 生成示例配置
|
312 |
+
./jetbrains-ai-proxy --generate-config
|
313 |
+
|
314 |
+
# 检查配置文件路径
|
315 |
+
./jetbrains-ai-proxy --print-config
|
316 |
+
```
|
317 |
+
|
318 |
+
3. **端口被占用**
|
319 |
+
```bash
|
320 |
+
# 检查端口使用情况
|
321 |
+
lsof -i :8080
|
322 |
+
|
323 |
+
# 使用不同端口启动
|
324 |
+
./jetbrains-ai-proxy -p 9090
|
325 |
+
```
|
326 |
+
|
327 |
+
### 日志分析
|
328 |
+
|
329 |
+
```bash
|
330 |
+
# 查看实时日志
|
331 |
+
tail -f jetbrains-ai-proxy.log
|
332 |
+
|
333 |
+
# 过滤健康检查日志
|
334 |
+
grep "health check" jetbrains-ai-proxy.log
|
335 |
+
|
336 |
+
# 过滤错误日志
|
337 |
+
grep -i error jetbrains-ai-proxy.log
|
338 |
+
```
|
339 |
+
|
340 |
+
## 📈 性能优化建议
|
341 |
+
|
342 |
+
1. **JWT Token数量**: 建议配置3-5个JWT tokens以获得最佳负载均衡效果
|
343 |
+
2. **健康检查间隔**: 生产环境建议设置为30-60秒
|
344 |
+
3. **负载均衡策略**:
|
345 |
+
- 使用`round_robin`获得均匀分布
|
346 |
+
- 使用`random`避免可预测的请求模式
|
347 |
+
4. **监控**: 定期检查`/health`和`/stats`端点
|
348 |
+
5. **日志**: 配置适当的日志级别和轮转策略
|
examples/demo_balancer.go
ADDED
@@ -0,0 +1,180 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package main
|
2 |
+
|
3 |
+
import (
|
4 |
+
"fmt"
|
5 |
+
"jetbrains-ai-proxy/internal/balancer"
|
6 |
+
"jetbrains-ai-proxy/internal/config"
|
7 |
+
"log"
|
8 |
+
"sync"
|
9 |
+
"time"
|
10 |
+
)
|
11 |
+
|
12 |
+
func main() {
|
13 |
+
fmt.Println("=== JWT负载均衡器演示 ===")
|
14 |
+
|
15 |
+
// 演示轮询策略
|
16 |
+
fmt.Println("\n1. 轮询策略演示:")
|
17 |
+
demoRoundRobin()
|
18 |
+
|
19 |
+
// 演示随机策略
|
20 |
+
fmt.Println("\n2. 随机策略演示:")
|
21 |
+
demoRandom()
|
22 |
+
|
23 |
+
// 演示健康检查
|
24 |
+
fmt.Println("\n3. 健康检查演示:")
|
25 |
+
demoHealthCheck()
|
26 |
+
|
27 |
+
// 演示并发访问
|
28 |
+
fmt.Println("\n4. 并发访问演示:")
|
29 |
+
demoConcurrent()
|
30 |
+
|
31 |
+
fmt.Println("\n=== 演示完成 ===")
|
32 |
+
}
|
33 |
+
|
34 |
+
func demoRoundRobin() {
|
35 |
+
tokens := []string{"JWT_TOKEN_1", "JWT_TOKEN_2", "JWT_TOKEN_3"}
|
36 |
+
balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)
|
37 |
+
|
38 |
+
fmt.Printf("配置了 %d 个JWT tokens,使用轮询策略\n", len(tokens))
|
39 |
+
fmt.Println("获取token顺序:")
|
40 |
+
|
41 |
+
for i := 0; i < 9; i++ {
|
42 |
+
token, err := balancer.GetToken()
|
43 |
+
if err != nil {
|
44 |
+
log.Printf("错误: %v", err)
|
45 |
+
continue
|
46 |
+
}
|
47 |
+
fmt.Printf(" 第%d次: %s\n", i+1, token)
|
48 |
+
}
|
49 |
+
}
|
50 |
+
|
51 |
+
func demoRandom() {
|
52 |
+
tokens := []string{"JWT_TOKEN_A", "JWT_TOKEN_B", "JWT_TOKEN_C"}
|
53 |
+
balancer := balancer.NewJWTBalancer(tokens, config.Random)
|
54 |
+
|
55 |
+
fmt.Printf("配置了 %d 个JWT tokens,使用随机策略\n", len(tokens))
|
56 |
+
fmt.Println("获取token顺序:")
|
57 |
+
|
58 |
+
tokenCounts := make(map[string]int)
|
59 |
+
for i := 0; i < 12; i++ {
|
60 |
+
token, err := balancer.GetToken()
|
61 |
+
if err != nil {
|
62 |
+
log.Printf("错误: %v", err)
|
63 |
+
continue
|
64 |
+
}
|
65 |
+
tokenCounts[token]++
|
66 |
+
fmt.Printf(" 第%d次: %s\n", i+1, token)
|
67 |
+
}
|
68 |
+
|
69 |
+
fmt.Println("使用统计:")
|
70 |
+
for token, count := range tokenCounts {
|
71 |
+
fmt.Printf(" %s: %d次\n", token, count)
|
72 |
+
}
|
73 |
+
}
|
74 |
+
|
75 |
+
func demoHealthCheck() {
|
76 |
+
tokens := []string{"JWT_HEALTHY_1", "JWT_HEALTHY_2", "JWT_UNHEALTHY_3"}
|
77 |
+
balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)
|
78 |
+
|
79 |
+
fmt.Printf("初始状态: %d/%d tokens健康\n",
|
80 |
+
balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())
|
81 |
+
|
82 |
+
// 标记一个token为不健康
|
83 |
+
fmt.Println("标记 JWT_UNHEALTHY_3 为不健康...")
|
84 |
+
balancer.MarkTokenUnhealthy("JWT_UNHEALTHY_3")
|
85 |
+
|
86 |
+
fmt.Printf("更新后状态: %d/%d tokens健康\n",
|
87 |
+
balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())
|
88 |
+
|
89 |
+
fmt.Println("获取token(应该只返回健康的tokens):")
|
90 |
+
for i := 0; i < 6; i++ {
|
91 |
+
token, err := balancer.GetToken()
|
92 |
+
if err != nil {
|
93 |
+
log.Printf("错误: %v", err)
|
94 |
+
continue
|
95 |
+
}
|
96 |
+
fmt.Printf(" 第%d次: %s\n", i+1, token)
|
97 |
+
}
|
98 |
+
|
99 |
+
// 恢复token健康状态
|
100 |
+
fmt.Println("恢复 JWT_UNHEALTHY_3 为健康...")
|
101 |
+
balancer.MarkTokenHealthy("JWT_UNHEALTHY_3")
|
102 |
+
|
103 |
+
fmt.Printf("恢复后状态: %d/%d tokens健康\n",
|
104 |
+
balancer.GetHealthyTokenCount(), balancer.GetTotalTokenCount())
|
105 |
+
}
|
106 |
+
|
107 |
+
func demoConcurrent() {
|
108 |
+
tokens := []string{"JWT_CONCURRENT_1", "JWT_CONCURRENT_2", "JWT_CONCURRENT_3", "JWT_CONCURRENT_4"}
|
109 |
+
balancer := balancer.NewJWTBalancer(tokens, config.RoundRobin)
|
110 |
+
|
111 |
+
fmt.Printf("使用 %d 个JWT tokens进行并发测试\n", len(tokens))
|
112 |
+
|
113 |
+
var wg sync.WaitGroup
|
114 |
+
numGoroutines := 5
|
115 |
+
requestsPerGoroutine := 10
|
116 |
+
|
117 |
+
tokenCounts := make(map[string]int)
|
118 |
+
var mutex sync.Mutex
|
119 |
+
|
120 |
+
startTime := time.Now()
|
121 |
+
|
122 |
+
for i := 0; i < numGoroutines; i++ {
|
123 |
+
wg.Add(1)
|
124 |
+
go func(goroutineID int) {
|
125 |
+
defer wg.Done()
|
126 |
+
|
127 |
+
for j := 0; j < requestsPerGoroutine; j++ {
|
128 |
+
token, err := balancer.GetToken()
|
129 |
+
if err != nil {
|
130 |
+
log.Printf("Goroutine %d 错误: %v", goroutineID, err)
|
131 |
+
continue
|
132 |
+
}
|
133 |
+
|
134 |
+
mutex.Lock()
|
135 |
+
tokenCounts[token]++
|
136 |
+
mutex.Unlock()
|
137 |
+
|
138 |
+
// 模拟一些处理时间
|
139 |
+
time.Sleep(time.Millisecond * 10)
|
140 |
+
}
|
141 |
+
}(i)
|
142 |
+
}
|
143 |
+
|
144 |
+
wg.Wait()
|
145 |
+
duration := time.Since(startTime)
|
146 |
+
|
147 |
+
fmt.Printf("并发测试完成,耗时: %v\n", duration)
|
148 |
+
fmt.Printf("总请求数: %d\n", numGoroutines*requestsPerGoroutine)
|
149 |
+
fmt.Println("Token使用分布:")
|
150 |
+
|
151 |
+
for token, count := range tokenCounts {
|
152 |
+
percentage := float64(count) / float64(numGoroutines*requestsPerGoroutine) * 100
|
153 |
+
fmt.Printf(" %s: %d次 (%.1f%%)\n", token, count, percentage)
|
154 |
+
}
|
155 |
+
}
|
156 |
+
|
157 |
+
// 演示配置加载
|
158 |
+
//func demoConfigLoading() {
|
159 |
+
// fmt.Println("\n5. 配置加载演示:")
|
160 |
+
//
|
161 |
+
// // 模拟环境变量配置
|
162 |
+
// fmt.Println("模拟配置加载...")
|
163 |
+
//
|
164 |
+
// cfg := &config.Config{}
|
165 |
+
//
|
166 |
+
// // 设置多个JWT tokens
|
167 |
+
// cfg.SetJetbrainsTokens("token1,token2,token3")
|
168 |
+
// cfg.BearerToken = "bearer_token_example"
|
169 |
+
// cfg.LoadBalanceStrategy = config.RoundRobin
|
170 |
+
//
|
171 |
+
// fmt.Printf("JWT Tokens数量: %d\n", len(cfg.GetJetbrainsTokens()))
|
172 |
+
// fmt.Printf("Bearer Token: %s\n", cfg.BearerToken)
|
173 |
+
// fmt.Printf("负载均衡策略: %s\n", cfg.LoadBalanceStrategy)
|
174 |
+
// fmt.Printf("是否有JWT Tokens: %v\n", cfg.HasJetbrainsTokens())
|
175 |
+
//
|
176 |
+
// fmt.Println("JWT Tokens列表:")
|
177 |
+
// for i, token := range cfg.GetJetbrainsTokens() {
|
178 |
+
// fmt.Printf(" %d: %s\n", i+1, token)
|
179 |
+
// }
|
180 |
+
//}
|
examples/start_with_multiple_jwt.sh
ADDED
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/bin/bash
|
2 |
+
|
3 |
+
# 多JWT负载均衡启动示例脚本
|
4 |
+
|
5 |
+
echo "=== JetBrains AI Proxy 多JWT负载均衡启动示例 ==="
|
6 |
+
|
7 |
+
# 检查是否提供了JWT tokens
|
8 |
+
if [ -z "$1" ]; then
|
9 |
+
echo "用法: $0 \"jwt_token1,jwt_token2,jwt_token3\" [bearer_token] [strategy] [port]"
|
10 |
+
echo ""
|
11 |
+
echo "参数说明:"
|
12 |
+
echo " jwt_tokens - 多个JWT tokens,用逗号分隔(必需)"
|
13 |
+
echo " bearer_token - Bearer token(可选,默认从环境变量读取)"
|
14 |
+
echo " strategy - 负载均衡策略:round_robin 或 random(可选,默认round_robin)"
|
15 |
+
echo " port - 监听端口(可选,默认8080)"
|
16 |
+
echo ""
|
17 |
+
echo "示例:"
|
18 |
+
echo " $0 \"jwt1,jwt2,jwt3\""
|
19 |
+
echo " $0 \"jwt1,jwt2,jwt3\" \"bearer123\" \"random\" 9090"
|
20 |
+
echo ""
|
21 |
+
echo "环境变量配置示例:"
|
22 |
+
echo " export JWT_TOKENS=\"jwt1,jwt2,jwt3\""
|
23 |
+
echo " export BEARER_TOKEN=\"your_bearer_token\""
|
24 |
+
echo " export LOAD_BALANCE_STRATEGY=\"random\""
|
25 |
+
echo " ./jetbrains-ai-proxy"
|
26 |
+
exit 1
|
27 |
+
fi
|
28 |
+
|
29 |
+
# 参数设置
|
30 |
+
JWT_TOKENS="$1"
|
31 |
+
BEARER_TOKEN="${2:-$BEARER_TOKEN}"
|
32 |
+
STRATEGY="${3:-round_robin}"
|
33 |
+
PORT="${4:-8080}"
|
34 |
+
|
35 |
+
# 检查Bearer token
|
36 |
+
if [ -z "$BEARER_TOKEN" ]; then
|
37 |
+
echo "错误: 需要提供Bearer token"
|
38 |
+
echo "请通过参数提供或设置环境变量 BEARER_TOKEN"
|
39 |
+
exit 1
|
40 |
+
fi
|
41 |
+
|
42 |
+
# 检查策略有效性
|
43 |
+
if [ "$STRATEGY" != "round_robin" ] && [ "$STRATEGY" != "random" ]; then
|
44 |
+
echo "警告: 无效的负载均衡策略 '$STRATEGY',使用默认策略 'round_robin'"
|
45 |
+
STRATEGY="round_robin"
|
46 |
+
fi
|
47 |
+
|
48 |
+
# 计算JWT tokens数量
|
49 |
+
TOKEN_COUNT=$(echo "$JWT_TOKENS" | tr ',' '\n' | wc -l | tr -d ' ')
|
50 |
+
|
51 |
+
echo "配置信息:"
|
52 |
+
echo " JWT Tokens数量: $TOKEN_COUNT"
|
53 |
+
echo " 负载均衡策略: $STRATEGY"
|
54 |
+
echo " 监听端口: $PORT"
|
55 |
+
echo " Bearer Token: ${BEARER_TOKEN:0:10}..."
|
56 |
+
echo ""
|
57 |
+
|
58 |
+
# 启动服务
|
59 |
+
echo "启动 JetBrains AI Proxy..."
|
60 |
+
echo "命令: ./jetbrains-ai-proxy -p $PORT -c \"$JWT_TOKENS\" -k \"$BEARER_TOKEN\" -s \"$STRATEGY\""
|
61 |
+
echo ""
|
62 |
+
|
63 |
+
# 检查可执行文件是否存在
|
64 |
+
if [ ! -f "./jetbrains-ai-proxy" ]; then
|
65 |
+
echo "错误: 找不到可执行文件 './jetbrains-ai-proxy'"
|
66 |
+
echo "请先编译项目: go build -o jetbrains-ai-proxy"
|
67 |
+
exit 1
|
68 |
+
fi
|
69 |
+
|
70 |
+
# 启动服务
|
71 |
+
exec ./jetbrains-ai-proxy -p "$PORT" -c "$JWT_TOKENS" -k "$BEARER_TOKEN" -s "$STRATEGY"
|
go.mod
ADDED
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
module jetbrains-ai-proxy
|
2 |
+
|
3 |
+
go 1.24
|
4 |
+
|
5 |
+
require (
|
6 |
+
github.com/bytedance/sonic v1.13.3 // indirect
|
7 |
+
github.com/bytedance/sonic/loader v0.2.4 // indirect
|
8 |
+
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
9 |
+
github.com/cloudwego/base64x v0.1.5 // indirect
|
10 |
+
github.com/dgrijalva/jwt-go v3.2.0+incompatible // indirect
|
11 |
+
github.com/dlclark/regexp2 v1.10.0 // indirect
|
12 |
+
github.com/go-resty/resty/v2 v2.16.5 // indirect
|
13 |
+
github.com/google/uuid v1.6.0 // indirect
|
14 |
+
github.com/joho/godotenv v1.5.1 // indirect
|
15 |
+
github.com/klauspost/cpuid/v2 v2.0.9 // indirect
|
16 |
+
github.com/labstack/echo v3.3.10+incompatible // indirect
|
17 |
+
github.com/labstack/echo/v4 v4.13.4 // indirect
|
18 |
+
github.com/labstack/gommon v0.4.2 // indirect
|
19 |
+
github.com/mattn/go-colorable v0.1.14 // indirect
|
20 |
+
github.com/mattn/go-isatty v0.0.20 // indirect
|
21 |
+
github.com/pkoukk/tiktoken-go v0.1.7 // indirect
|
22 |
+
github.com/samber/lo v1.51.0 // indirect
|
23 |
+
github.com/sashabaranov/go-openai v1.40.3 // indirect
|
24 |
+
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
|
25 |
+
github.com/valyala/bytebufferpool v1.0.0 // indirect
|
26 |
+
github.com/valyala/fasttemplate v1.2.2 // indirect
|
27 |
+
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
|
28 |
+
golang.org/x/crypto v0.39.0 // indirect
|
29 |
+
golang.org/x/net v0.41.0 // indirect
|
30 |
+
golang.org/x/sys v0.33.0 // indirect
|
31 |
+
golang.org/x/text v0.26.0 // indirect
|
32 |
+
golang.org/x/time v0.11.0 // indirect
|
33 |
+
)
|
go.sum
ADDED
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
github.com/bytedance/sonic v1.13.3 h1:MS8gmaH16Gtirygw7jV91pDCN33NyMrPbN7qiYhEsF0=
|
2 |
+
github.com/bytedance/sonic v1.13.3/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4=
|
3 |
+
github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
|
4 |
+
github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY=
|
5 |
+
github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
|
6 |
+
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
7 |
+
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
8 |
+
github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
|
9 |
+
github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
|
10 |
+
github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
|
11 |
+
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
12 |
+
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
13 |
+
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
14 |
+
github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM=
|
15 |
+
github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
|
16 |
+
github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
|
17 |
+
github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
|
18 |
+
github.com/go-resty/resty/v2 v2.16.5 h1:hBKqmWrr7uRc3euHVqmh1HTHcKn99Smr7o5spptdhTM=
|
19 |
+
github.com/go-resty/resty/v2 v2.16.5/go.mod h1:hkJtXbA2iKHzJheXYvQ8snQES5ZLGKMwQ07xAwp/fiA=
|
20 |
+
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
21 |
+
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
22 |
+
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
|
23 |
+
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
|
24 |
+
github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
|
25 |
+
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
|
26 |
+
github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M=
|
27 |
+
github.com/labstack/echo v3.3.10+incompatible h1:pGRcYk231ExFAyoAjAfD85kQzRJCRI8bbnE7CX5OEgg=
|
28 |
+
github.com/labstack/echo v3.3.10+incompatible/go.mod h1:0INS7j/VjnFxD4E2wkz67b8cVwCLbBmJyDaka6Cmk1s=
|
29 |
+
github.com/labstack/echo/v4 v4.13.4 h1:oTZZW+T3s9gAu5L8vmzihV7/lkXGZuITzTQkTEhcXEA=
|
30 |
+
github.com/labstack/echo/v4 v4.13.4/go.mod h1:g63b33BZ5vZzcIUF8AtRH40DrTlXnx4UMC8rBdndmjQ=
|
31 |
+
github.com/labstack/gommon v0.4.2 h1:F8qTUNXgG1+6WQmqoUWnz8WiEU60mXVVw0P4ht1WRA0=
|
32 |
+
github.com/labstack/gommon v0.4.2/go.mod h1:QlUFxVM+SNXhDL/Z7YhocGIBYOiwB0mXm1+1bAPHPyU=
|
33 |
+
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
34 |
+
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
|
35 |
+
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
36 |
+
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
37 |
+
github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw=
|
38 |
+
github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg=
|
39 |
+
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
40 |
+
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
41 |
+
github.com/samber/lo v1.51.0 h1:kysRYLbHy/MB7kQZf5DSN50JHmMsNEdeY24VzJFu7wI=
|
42 |
+
github.com/samber/lo v1.51.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0=
|
43 |
+
github.com/sashabaranov/go-openai v1.40.3 h1:PkOw0SK34wrvYVOuXF1HZzuTBRh992qRZHil4kG3eYE=
|
44 |
+
github.com/sashabaranov/go-openai v1.40.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
|
45 |
+
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
46 |
+
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
|
47 |
+
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
|
48 |
+
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
49 |
+
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
50 |
+
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
|
51 |
+
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
52 |
+
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
53 |
+
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
54 |
+
github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
|
55 |
+
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
|
56 |
+
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
|
57 |
+
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
|
58 |
+
github.com/valyala/fasttemplate v1.2.2 h1:lxLXG0uE3Qnshl9QyaK6XJxMXlQZELvChBOCmQD0Loo=
|
59 |
+
github.com/valyala/fasttemplate v1.2.2/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
|
60 |
+
golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU=
|
61 |
+
golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
|
62 |
+
golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
|
63 |
+
golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
|
64 |
+
golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
|
65 |
+
golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
|
66 |
+
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
67 |
+
golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
|
68 |
+
golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
69 |
+
golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
|
70 |
+
golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
|
71 |
+
golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0=
|
72 |
+
golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
|
73 |
+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
74 |
+
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
75 |
+
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
76 |
+
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
77 |
+
nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50=
|
img.png
ADDED
![]() |
internal/apiserver/router.go
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package apiserver
|
2 |
+
|
3 |
+
import (
|
4 |
+
"fmt"
|
5 |
+
"github.com/labstack/echo"
|
6 |
+
"jetbrains-ai-proxy/internal/jetbrains"
|
7 |
+
"jetbrains-ai-proxy/internal/middleware"
|
8 |
+
"jetbrains-ai-proxy/internal/types"
|
9 |
+
"jetbrains-ai-proxy/internal/utils"
|
10 |
+
"net/http"
|
11 |
+
|
12 |
+
"github.com/sashabaranov/go-openai"
|
13 |
+
)
|
14 |
+
|
15 |
+
func RegisterRoutes(e *echo.Echo) {
|
16 |
+
e.Use(middleware.BearerAuth())
|
17 |
+
e.POST("/v1/chat/completions", handleChatCompletion)
|
18 |
+
e.GET("/v1/models", handleListModels)
|
19 |
+
}
|
20 |
+
|
21 |
+
func handleChatCompletion(c echo.Context) error {
|
22 |
+
var req openai.ChatCompletionRequest
|
23 |
+
|
24 |
+
if err := c.Bind(&req); err != nil {
|
25 |
+
return c.JSON(http.StatusBadRequest, map[string]interface{}{
|
26 |
+
"error": "Invalid request payload",
|
27 |
+
})
|
28 |
+
}
|
29 |
+
|
30 |
+
_, err := types.GetModelByName(req.Model)
|
31 |
+
if err != nil {
|
32 |
+
return c.JSON(http.StatusBadRequest, map[string]interface{}{
|
33 |
+
"error": fmt.Sprintf("Model '%s' not supported", req.Model),
|
34 |
+
})
|
35 |
+
}
|
36 |
+
|
37 |
+
if len(req.Messages) == 0 {
|
38 |
+
return c.JSON(http.StatusBadRequest, map[string]interface{}{
|
39 |
+
"error": "No messages found",
|
40 |
+
})
|
41 |
+
}
|
42 |
+
|
43 |
+
jetbrainsReq, err := types.ChatGPTToJetbrainsAI(req)
|
44 |
+
if err != nil {
|
45 |
+
return c.JSON(http.StatusInternalServerError, map[string]interface{}{
|
46 |
+
"error": err.Error(),
|
47 |
+
})
|
48 |
+
}
|
49 |
+
|
50 |
+
stream, err := jetbrains.SendJetbrainsRequest(c.Request().Context(), jetbrainsReq)
|
51 |
+
if err != nil {
|
52 |
+
return c.JSON(http.StatusInternalServerError, map[string]interface{}{
|
53 |
+
"error": err.Error(),
|
54 |
+
})
|
55 |
+
}
|
56 |
+
defer stream.RawBody().Close()
|
57 |
+
|
58 |
+
// 根据请求的 stream 参数决定使用哪种处理方式
|
59 |
+
fingerprint := utils.RandStringUsingMathRand(10)
|
60 |
+
if req.Stream {
|
61 |
+
// 流式处理
|
62 |
+
c.Response().Header().Set(echo.HeaderContentType, "text/event-stream")
|
63 |
+
c.Response().Header().Set("Cache-Control", "no-cache")
|
64 |
+
c.Response().Header().Set("Transfer-Encoding", "chunked")
|
65 |
+
c.Response().WriteHeader(http.StatusOK)
|
66 |
+
|
67 |
+
return jetbrains.StreamJetbrainsAISSEToClient(c.Request().Context(), req, c.Response().Writer, stream.RawBody(), fingerprint)
|
68 |
+
} else {
|
69 |
+
// 非流式处理
|
70 |
+
response, err := jetbrains.ResponseJetbrainsAIToClient(c.Request().Context(), req, stream.RawBody(), fingerprint)
|
71 |
+
if err != nil {
|
72 |
+
return c.JSON(http.StatusInternalServerError, map[string]interface{}{
|
73 |
+
"error": err.Error(),
|
74 |
+
})
|
75 |
+
}
|
76 |
+
return c.JSON(http.StatusOK, response)
|
77 |
+
}
|
78 |
+
}
|
79 |
+
|
80 |
+
func handleListModels(c echo.Context) error {
|
81 |
+
models := types.GetSupportedModels()
|
82 |
+
return c.JSON(http.StatusOK, models)
|
83 |
+
}
|
internal/balancer/health_checker.go
ADDED
@@ -0,0 +1,216 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package balancer
|
2 |
+
|
3 |
+
import (
|
4 |
+
"context"
|
5 |
+
"github.com/go-resty/resty/v2"
|
6 |
+
"jetbrains-ai-proxy/internal/types"
|
7 |
+
"log"
|
8 |
+
"sync"
|
9 |
+
"time"
|
10 |
+
)
|
11 |
+
|
12 |
+
// HealthChecker JWT健康检查器
|
13 |
+
type HealthChecker struct {
|
14 |
+
balancer JWTBalancer
|
15 |
+
client *resty.Client
|
16 |
+
checkInterval time.Duration
|
17 |
+
timeout time.Duration
|
18 |
+
maxRetries int
|
19 |
+
stopChan chan struct{}
|
20 |
+
wg sync.WaitGroup
|
21 |
+
running bool
|
22 |
+
mutex sync.RWMutex
|
23 |
+
}
|
24 |
+
|
25 |
+
// NewHealthChecker 创建健康检查器
|
26 |
+
func NewHealthChecker(balancer JWTBalancer) *HealthChecker {
|
27 |
+
client := resty.New().
|
28 |
+
SetTimeout(10 * time.Second).
|
29 |
+
SetHeaders(map[string]string{
|
30 |
+
"Content-Type": "application/json",
|
31 |
+
})
|
32 |
+
|
33 |
+
return &HealthChecker{
|
34 |
+
balancer: balancer,
|
35 |
+
client: client,
|
36 |
+
checkInterval: 30 * time.Second, // 每30秒检查一次
|
37 |
+
timeout: 10 * time.Second,
|
38 |
+
maxRetries: 3,
|
39 |
+
stopChan: make(chan struct{}),
|
40 |
+
}
|
41 |
+
}
|
42 |
+
|
43 |
+
// Start 启动健康检查
|
44 |
+
func (hc *HealthChecker) Start() {
|
45 |
+
hc.mutex.Lock()
|
46 |
+
defer hc.mutex.Unlock()
|
47 |
+
|
48 |
+
if hc.running {
|
49 |
+
return
|
50 |
+
}
|
51 |
+
|
52 |
+
hc.running = true
|
53 |
+
hc.wg.Add(1)
|
54 |
+
|
55 |
+
go hc.healthCheckLoop()
|
56 |
+
log.Println("JWT health checker started")
|
57 |
+
}
|
58 |
+
|
59 |
+
// Stop 停止健康检查
|
60 |
+
func (hc *HealthChecker) Stop() {
|
61 |
+
hc.mutex.Lock()
|
62 |
+
defer hc.mutex.Unlock()
|
63 |
+
|
64 |
+
if !hc.running {
|
65 |
+
return
|
66 |
+
}
|
67 |
+
|
68 |
+
hc.running = false
|
69 |
+
close(hc.stopChan)
|
70 |
+
hc.wg.Wait()
|
71 |
+
log.Println("JWT health checker stopped")
|
72 |
+
}
|
73 |
+
|
74 |
+
// healthCheckLoop 健康检查循环
|
75 |
+
func (hc *HealthChecker) healthCheckLoop() {
|
76 |
+
defer hc.wg.Done()
|
77 |
+
|
78 |
+
ticker := time.NewTicker(hc.checkInterval)
|
79 |
+
defer ticker.Stop()
|
80 |
+
|
81 |
+
// 启动时立即执行一次检查
|
82 |
+
hc.performHealthCheck()
|
83 |
+
|
84 |
+
for {
|
85 |
+
select {
|
86 |
+
case <-ticker.C:
|
87 |
+
hc.performHealthCheck()
|
88 |
+
case <-hc.stopChan:
|
89 |
+
return
|
90 |
+
}
|
91 |
+
}
|
92 |
+
}
|
93 |
+
|
94 |
+
// performHealthCheck 执行健康检查
|
95 |
+
func (hc *HealthChecker) performHealthCheck() {
|
96 |
+
log.Println("Performing JWT health check...")
|
97 |
+
|
98 |
+
// 获取所有tokens进行检查
|
99 |
+
baseBalancer, ok := hc.balancer.(*BaseBalancer)
|
100 |
+
if !ok {
|
101 |
+
log.Println("Warning: Cannot access tokens for health check")
|
102 |
+
return
|
103 |
+
}
|
104 |
+
|
105 |
+
baseBalancer.mutex.RLock()
|
106 |
+
tokens := make([]string, 0, len(baseBalancer.tokens))
|
107 |
+
for token := range baseBalancer.tokens {
|
108 |
+
tokens = append(tokens, token)
|
109 |
+
}
|
110 |
+
baseBalancer.mutex.RUnlock()
|
111 |
+
|
112 |
+
// 并发检查所有tokens
|
113 |
+
var wg sync.WaitGroup
|
114 |
+
for _, token := range tokens {
|
115 |
+
wg.Add(1)
|
116 |
+
go func(t string) {
|
117 |
+
defer wg.Done()
|
118 |
+
hc.checkTokenHealth(t)
|
119 |
+
}(token)
|
120 |
+
}
|
121 |
+
wg.Wait()
|
122 |
+
|
123 |
+
healthyCount := hc.balancer.GetHealthyTokenCount()
|
124 |
+
totalCount := hc.balancer.GetTotalTokenCount()
|
125 |
+
log.Printf("Health check completed: %d/%d tokens healthy", healthyCount, totalCount)
|
126 |
+
}
|
127 |
+
|
128 |
+
// checkTokenHealth 检查单个token的健康状态
|
129 |
+
func (hc *HealthChecker) checkTokenHealth(token string) {
|
130 |
+
ctx, cancel := context.WithTimeout(context.Background(), hc.timeout)
|
131 |
+
defer cancel()
|
132 |
+
|
133 |
+
// 创建一个简单的测试请求
|
134 |
+
testRequest := &types.JetbrainsRequest{
|
135 |
+
Prompt: types.PROMPT,
|
136 |
+
Profile: "openai-gpt-4o", // 使用一个通用的profile进行测试
|
137 |
+
Chat: types.ChatField{
|
138 |
+
MessageField: []types.MessageField{
|
139 |
+
{
|
140 |
+
Type: "user_message",
|
141 |
+
Content: "test", // 简单的测试消息
|
142 |
+
},
|
143 |
+
},
|
144 |
+
},
|
145 |
+
}
|
146 |
+
|
147 |
+
success := false
|
148 |
+
for retry := 0; retry < hc.maxRetries; retry++ {
|
149 |
+
if hc.testTokenRequest(ctx, token, testRequest) {
|
150 |
+
success = true
|
151 |
+
break
|
152 |
+
}
|
153 |
+
|
154 |
+
// 重试前等待一小段时间
|
155 |
+
if retry < hc.maxRetries-1 {
|
156 |
+
time.Sleep(time.Second)
|
157 |
+
}
|
158 |
+
}
|
159 |
+
|
160 |
+
if success {
|
161 |
+
hc.balancer.MarkTokenHealthy(token)
|
162 |
+
} else {
|
163 |
+
hc.balancer.MarkTokenUnhealthy(token)
|
164 |
+
log.Printf("JWT token health check failed: %s...", token[:min(len(token), 10)])
|
165 |
+
}
|
166 |
+
}
|
167 |
+
|
168 |
+
// testTokenRequest 测试token请求
|
169 |
+
func (hc *HealthChecker) testTokenRequest(ctx context.Context, token string, req *types.JetbrainsRequest) bool {
|
170 |
+
resp, err := hc.client.R().
|
171 |
+
SetContext(ctx).
|
172 |
+
SetHeader(types.JwtTokenKey, token).
|
173 |
+
SetBody(req).
|
174 |
+
Post(types.ChatStreamV7)
|
175 |
+
|
176 |
+
if err != nil {
|
177 |
+
log.Printf("Health check request error for token %s...: %v", token[:min(len(token), 10)], err)
|
178 |
+
return false
|
179 |
+
}
|
180 |
+
|
181 |
+
// 检查响应状态码
|
182 |
+
if resp.StatusCode() == 200 {
|
183 |
+
return true
|
184 |
+
}
|
185 |
+
|
186 |
+
// 401表示token无效,403可能表示配额用完但token有效
|
187 |
+
if resp.StatusCode() == 403 {
|
188 |
+
// 配额用完但token有效,仍然标记为健康
|
189 |
+
return true
|
190 |
+
}
|
191 |
+
|
192 |
+
log.Printf("Health check failed for token %s...: status %d",
|
193 |
+
token[:min(len(token), 10)], resp.StatusCode())
|
194 |
+
return false
|
195 |
+
}
|
196 |
+
|
197 |
+
// SetCheckInterval 设置检查间隔
|
198 |
+
func (hc *HealthChecker) SetCheckInterval(interval time.Duration) {
|
199 |
+
hc.mutex.Lock()
|
200 |
+
defer hc.mutex.Unlock()
|
201 |
+
hc.checkInterval = interval
|
202 |
+
}
|
203 |
+
|
204 |
+
// SetTimeout 设置请求超时
|
205 |
+
func (hc *HealthChecker) SetTimeout(timeout time.Duration) {
|
206 |
+
hc.mutex.Lock()
|
207 |
+
defer hc.mutex.Unlock()
|
208 |
+
hc.timeout = timeout
|
209 |
+
}
|
210 |
+
|
211 |
+
// SetMaxRetries 设置最大重试次数
|
212 |
+
func (hc *HealthChecker) SetMaxRetries(retries int) {
|
213 |
+
hc.mutex.Lock()
|
214 |
+
defer hc.mutex.Unlock()
|
215 |
+
hc.maxRetries = retries
|
216 |
+
}
|
internal/balancer/jwt_balancer.go
ADDED
@@ -0,0 +1,175 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package balancer
|
2 |
+
|
3 |
+
import (
|
4 |
+
"fmt"
|
5 |
+
"jetbrains-ai-proxy/internal/config"
|
6 |
+
"math/rand"
|
7 |
+
"sync"
|
8 |
+
"sync/atomic"
|
9 |
+
"time"
|
10 |
+
)
|
11 |
+
|
12 |
+
// JWTBalancer JWT负载均衡器接口
|
13 |
+
type JWTBalancer interface {
|
14 |
+
GetToken() (string, error)
|
15 |
+
MarkTokenUnhealthy(token string)
|
16 |
+
MarkTokenHealthy(token string)
|
17 |
+
GetHealthyTokenCount() int
|
18 |
+
GetTotalTokenCount() int
|
19 |
+
RefreshTokens(tokens []string)
|
20 |
+
}
|
21 |
+
|
22 |
+
// TokenStatus token状态
|
23 |
+
type TokenStatus struct {
|
24 |
+
Token string
|
25 |
+
Healthy bool
|
26 |
+
LastUsed time.Time
|
27 |
+
ErrorCount int64
|
28 |
+
}
|
29 |
+
|
30 |
+
// BaseBalancer 基础负载均衡器
|
31 |
+
type BaseBalancer struct {
|
32 |
+
tokens map[string]*TokenStatus
|
33 |
+
strategy config.LoadBalanceStrategy
|
34 |
+
mutex sync.RWMutex
|
35 |
+
counter int64 // 用于轮询计数
|
36 |
+
rand *rand.Rand
|
37 |
+
}
|
38 |
+
|
39 |
+
// NewJWTBalancer 创建JWT负载均衡器
|
40 |
+
func NewJWTBalancer(tokens []string, strategy config.LoadBalanceStrategy) JWTBalancer {
|
41 |
+
balancer := &BaseBalancer{
|
42 |
+
tokens: make(map[string]*TokenStatus),
|
43 |
+
strategy: strategy,
|
44 |
+
rand: rand.New(rand.NewSource(time.Now().UnixNano())),
|
45 |
+
}
|
46 |
+
|
47 |
+
// 初始化tokens
|
48 |
+
for _, token := range tokens {
|
49 |
+
balancer.tokens[token] = &TokenStatus{
|
50 |
+
Token: token,
|
51 |
+
Healthy: true,
|
52 |
+
LastUsed: time.Now(),
|
53 |
+
ErrorCount: 0,
|
54 |
+
}
|
55 |
+
}
|
56 |
+
|
57 |
+
return balancer
|
58 |
+
}
|
59 |
+
|
60 |
+
// GetToken 获取一个可用的token
|
61 |
+
func (b *BaseBalancer) GetToken() (string, error) {
|
62 |
+
b.mutex.RLock()
|
63 |
+
defer b.mutex.RUnlock()
|
64 |
+
|
65 |
+
// 获取所有健康的tokens
|
66 |
+
healthyTokens := make([]*TokenStatus, 0)
|
67 |
+
for _, status := range b.tokens {
|
68 |
+
if status.Healthy {
|
69 |
+
healthyTokens = append(healthyTokens, status)
|
70 |
+
}
|
71 |
+
}
|
72 |
+
|
73 |
+
if len(healthyTokens) == 0 {
|
74 |
+
return "", fmt.Errorf("no healthy JWT tokens available")
|
75 |
+
}
|
76 |
+
|
77 |
+
var selectedToken *TokenStatus
|
78 |
+
|
79 |
+
switch b.strategy {
|
80 |
+
case config.RoundRobin:
|
81 |
+
// 轮询策略
|
82 |
+
index := atomic.AddInt64(&b.counter, 1) % int64(len(healthyTokens))
|
83 |
+
selectedToken = healthyTokens[index]
|
84 |
+
case config.Random:
|
85 |
+
// 随机策略
|
86 |
+
index := b.rand.Intn(len(healthyTokens))
|
87 |
+
selectedToken = healthyTokens[index]
|
88 |
+
default:
|
89 |
+
// 默认使用轮询
|
90 |
+
index := atomic.AddInt64(&b.counter, 1) % int64(len(healthyTokens))
|
91 |
+
selectedToken = healthyTokens[index]
|
92 |
+
}
|
93 |
+
|
94 |
+
// 更新最后使用时间
|
95 |
+
selectedToken.LastUsed = time.Now()
|
96 |
+
|
97 |
+
return selectedToken.Token, nil
|
98 |
+
}
|
99 |
+
|
100 |
+
// MarkTokenUnhealthy 标记token为不健康
|
101 |
+
func (b *BaseBalancer) MarkTokenUnhealthy(token string) {
|
102 |
+
b.mutex.Lock()
|
103 |
+
defer b.mutex.Unlock()
|
104 |
+
|
105 |
+
if status, exists := b.tokens[token]; exists {
|
106 |
+
status.Healthy = false
|
107 |
+
atomic.AddInt64(&status.ErrorCount, 1)
|
108 |
+
fmt.Printf("JWT token marked as unhealthy: %s (errors: %d)\n",
|
109 |
+
token[:min(len(token), 10)]+"...", status.ErrorCount)
|
110 |
+
}
|
111 |
+
}
|
112 |
+
|
113 |
+
// MarkTokenHealthy 标记token为健康
|
114 |
+
func (b *BaseBalancer) MarkTokenHealthy(token string) {
|
115 |
+
b.mutex.Lock()
|
116 |
+
defer b.mutex.Unlock()
|
117 |
+
|
118 |
+
if status, exists := b.tokens[token]; exists {
|
119 |
+
status.Healthy = true
|
120 |
+
atomic.StoreInt64(&status.ErrorCount, 0)
|
121 |
+
fmt.Printf("JWT token marked as healthy: %s\n",
|
122 |
+
token[:min(len(token), 10)]+"...")
|
123 |
+
}
|
124 |
+
}
|
125 |
+
|
126 |
+
// GetHealthyTokenCount 获取健康token数量
|
127 |
+
func (b *BaseBalancer) GetHealthyTokenCount() int {
|
128 |
+
b.mutex.RLock()
|
129 |
+
defer b.mutex.RUnlock()
|
130 |
+
|
131 |
+
count := 0
|
132 |
+
for _, status := range b.tokens {
|
133 |
+
if status.Healthy {
|
134 |
+
count++
|
135 |
+
}
|
136 |
+
}
|
137 |
+
return count
|
138 |
+
}
|
139 |
+
|
140 |
+
// GetTotalTokenCount 获取总token数量
|
141 |
+
func (b *BaseBalancer) GetTotalTokenCount() int {
|
142 |
+
b.mutex.RLock()
|
143 |
+
defer b.mutex.RUnlock()
|
144 |
+
|
145 |
+
return len(b.tokens)
|
146 |
+
}
|
147 |
+
|
148 |
+
// RefreshTokens 刷新token列表
|
149 |
+
func (b *BaseBalancer) RefreshTokens(tokens []string) {
|
150 |
+
b.mutex.Lock()
|
151 |
+
defer b.mutex.Unlock()
|
152 |
+
|
153 |
+
// 清空现有tokens
|
154 |
+
b.tokens = make(map[string]*TokenStatus)
|
155 |
+
|
156 |
+
// 添加新tokens
|
157 |
+
for _, token := range tokens {
|
158 |
+
b.tokens[token] = &TokenStatus{
|
159 |
+
Token: token,
|
160 |
+
Healthy: true,
|
161 |
+
LastUsed: time.Now(),
|
162 |
+
ErrorCount: 0,
|
163 |
+
}
|
164 |
+
}
|
165 |
+
|
166 |
+
fmt.Printf("JWT tokens refreshed, total: %d\n", len(tokens))
|
167 |
+
}
|
168 |
+
|
169 |
+
// min 辅助函数
|
170 |
+
func min(a, b int) int {
|
171 |
+
if a < b {
|
172 |
+
return a
|
173 |
+
}
|
174 |
+
return b
|
175 |
+
}
|
internal/balancer/jwt_balancer_test.go
ADDED
@@ -0,0 +1,226 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package balancer
|
2 |
+
|
3 |
+
import (
|
4 |
+
"jetbrains-ai-proxy/internal/config"
|
5 |
+
"sync"
|
6 |
+
"testing"
|
7 |
+
"time"
|
8 |
+
)
|
9 |
+
|
10 |
+
func TestNewJWTBalancer(t *testing.T) {
|
11 |
+
tokens := []string{"token1", "token2", "token3"}
|
12 |
+
|
13 |
+
// 测试轮询策略
|
14 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
15 |
+
if balancer == nil {
|
16 |
+
t.Fatal("Expected balancer to be created")
|
17 |
+
}
|
18 |
+
|
19 |
+
if balancer.GetTotalTokenCount() != 3 {
|
20 |
+
t.Errorf("Expected 3 tokens, got %d", balancer.GetTotalTokenCount())
|
21 |
+
}
|
22 |
+
|
23 |
+
if balancer.GetHealthyTokenCount() != 3 {
|
24 |
+
t.Errorf("Expected 3 healthy tokens, got %d", balancer.GetHealthyTokenCount())
|
25 |
+
}
|
26 |
+
}
|
27 |
+
|
28 |
+
func TestRoundRobinStrategy(t *testing.T) {
|
29 |
+
tokens := []string{"token1", "token2", "token3"}
|
30 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
31 |
+
|
32 |
+
// 测试轮询顺序
|
33 |
+
expectedOrder := []string{"token1", "token2", "token3", "token1", "token2", "token3"}
|
34 |
+
|
35 |
+
for i, expected := range expectedOrder {
|
36 |
+
token, err := balancer.GetToken()
|
37 |
+
if err != nil {
|
38 |
+
t.Fatalf("Unexpected error at iteration %d: %v", i, err)
|
39 |
+
}
|
40 |
+
|
41 |
+
if token != expected {
|
42 |
+
t.Errorf("At iteration %d, expected %s, got %s", i, expected, token)
|
43 |
+
}
|
44 |
+
}
|
45 |
+
}
|
46 |
+
|
47 |
+
func TestRandomStrategy(t *testing.T) {
|
48 |
+
tokens := []string{"token1", "token2", "token3"}
|
49 |
+
balancer := NewJWTBalancer(tokens, config.Random)
|
50 |
+
|
51 |
+
// 测试随机策略 - 多次获取token,确保都是有效的
|
52 |
+
tokenCounts := make(map[string]int)
|
53 |
+
iterations := 100
|
54 |
+
|
55 |
+
for i := 0; i < iterations; i++ {
|
56 |
+
token, err := balancer.GetToken()
|
57 |
+
if err != nil {
|
58 |
+
t.Fatalf("Unexpected error at iteration %d: %v", i, err)
|
59 |
+
}
|
60 |
+
|
61 |
+
// 检查token是否在预期列表中
|
62 |
+
found := false
|
63 |
+
for _, expectedToken := range tokens {
|
64 |
+
if token == expectedToken {
|
65 |
+
found = true
|
66 |
+
break
|
67 |
+
}
|
68 |
+
}
|
69 |
+
|
70 |
+
if !found {
|
71 |
+
t.Errorf("Got unexpected token: %s", token)
|
72 |
+
}
|
73 |
+
|
74 |
+
tokenCounts[token]++
|
75 |
+
}
|
76 |
+
|
77 |
+
// 确保所有token都被使用过(随机策略下应该都有机会被选中)
|
78 |
+
for _, token := range tokens {
|
79 |
+
if tokenCounts[token] == 0 {
|
80 |
+
t.Errorf("Token %s was never selected", token)
|
81 |
+
}
|
82 |
+
}
|
83 |
+
}
|
84 |
+
|
85 |
+
func TestMarkTokenUnhealthy(t *testing.T) {
|
86 |
+
tokens := []string{"token1", "token2", "token3"}
|
87 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
88 |
+
|
89 |
+
// 标记一个token为不健康
|
90 |
+
balancer.MarkTokenUnhealthy("token2")
|
91 |
+
|
92 |
+
if balancer.GetHealthyTokenCount() != 2 {
|
93 |
+
t.Errorf("Expected 2 healthy tokens, got %d", balancer.GetHealthyTokenCount())
|
94 |
+
}
|
95 |
+
|
96 |
+
// 获取token,应该只返回健康的token
|
97 |
+
for i := 0; i < 10; i++ {
|
98 |
+
token, err := balancer.GetToken()
|
99 |
+
if err != nil {
|
100 |
+
t.Fatalf("Unexpected error: %v", err)
|
101 |
+
}
|
102 |
+
|
103 |
+
if token == "token2" {
|
104 |
+
t.Errorf("Got unhealthy token: %s", token)
|
105 |
+
}
|
106 |
+
}
|
107 |
+
}
|
108 |
+
|
109 |
+
func TestMarkTokenHealthy(t *testing.T) {
|
110 |
+
tokens := []string{"token1", "token2", "token3"}
|
111 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
112 |
+
|
113 |
+
// 先标记为不健康,再标记为健康
|
114 |
+
balancer.MarkTokenUnhealthy("token2")
|
115 |
+
if balancer.GetHealthyTokenCount() != 2 {
|
116 |
+
t.Errorf("Expected 2 healthy tokens after marking unhealthy, got %d", balancer.GetHealthyTokenCount())
|
117 |
+
}
|
118 |
+
|
119 |
+
balancer.MarkTokenHealthy("token2")
|
120 |
+
if balancer.GetHealthyTokenCount() != 3 {
|
121 |
+
t.Errorf("Expected 3 healthy tokens after marking healthy, got %d", balancer.GetHealthyTokenCount())
|
122 |
+
}
|
123 |
+
}
|
124 |
+
|
125 |
+
func TestNoHealthyTokens(t *testing.T) {
|
126 |
+
tokens := []string{"token1", "token2"}
|
127 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
128 |
+
|
129 |
+
// 标记所有token为不健康
|
130 |
+
balancer.MarkTokenUnhealthy("token1")
|
131 |
+
balancer.MarkTokenUnhealthy("token2")
|
132 |
+
|
133 |
+
// 尝试获取token应该返回错误
|
134 |
+
_, err := balancer.GetToken()
|
135 |
+
if err == nil {
|
136 |
+
t.Error("Expected error when no healthy tokens available")
|
137 |
+
}
|
138 |
+
}
|
139 |
+
|
140 |
+
func TestConcurrentAccess(t *testing.T) {
|
141 |
+
tokens := []string{"token1", "token2", "token3", "token4", "token5"}
|
142 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
143 |
+
|
144 |
+
var wg sync.WaitGroup
|
145 |
+
numGoroutines := 10
|
146 |
+
tokensPerGoroutine := 100
|
147 |
+
|
148 |
+
// 并发获取tokens
|
149 |
+
for i := 0; i < numGoroutines; i++ {
|
150 |
+
wg.Add(1)
|
151 |
+
go func() {
|
152 |
+
defer wg.Done()
|
153 |
+
for j := 0; j < tokensPerGoroutine; j++ {
|
154 |
+
_, err := balancer.GetToken()
|
155 |
+
if err != nil {
|
156 |
+
t.Errorf("Unexpected error in concurrent access: %v", err)
|
157 |
+
}
|
158 |
+
}
|
159 |
+
}()
|
160 |
+
}
|
161 |
+
|
162 |
+
// 并发标记tokens健康状态
|
163 |
+
for i := 0; i < numGoroutines; i++ {
|
164 |
+
wg.Add(1)
|
165 |
+
go func(index int) {
|
166 |
+
defer wg.Done()
|
167 |
+
token := tokens[index%len(tokens)]
|
168 |
+
for j := 0; j < 10; j++ {
|
169 |
+
if j%2 == 0 {
|
170 |
+
balancer.MarkTokenUnhealthy(token)
|
171 |
+
} else {
|
172 |
+
balancer.MarkTokenHealthy(token)
|
173 |
+
}
|
174 |
+
time.Sleep(time.Millisecond)
|
175 |
+
}
|
176 |
+
}(i)
|
177 |
+
}
|
178 |
+
|
179 |
+
wg.Wait()
|
180 |
+
|
181 |
+
// 确保最终状态正常
|
182 |
+
if balancer.GetTotalTokenCount() != len(tokens) {
|
183 |
+
t.Errorf("Expected %d total tokens, got %d", len(tokens), balancer.GetTotalTokenCount())
|
184 |
+
}
|
185 |
+
}
|
186 |
+
|
187 |
+
func TestRefreshTokens(t *testing.T) {
|
188 |
+
tokens := []string{"token1", "token2"}
|
189 |
+
balancer := NewJWTBalancer(tokens, config.RoundRobin)
|
190 |
+
|
191 |
+
if balancer.GetTotalTokenCount() != 2 {
|
192 |
+
t.Errorf("Expected 2 tokens initially, got %d", balancer.GetTotalTokenCount())
|
193 |
+
}
|
194 |
+
|
195 |
+
// 刷新tokens
|
196 |
+
newTokens := []string{"token3", "token4", "token5"}
|
197 |
+
balancer.RefreshTokens(newTokens)
|
198 |
+
|
199 |
+
if balancer.GetTotalTokenCount() != 3 {
|
200 |
+
t.Errorf("Expected 3 tokens after refresh, got %d", balancer.GetTotalTokenCount())
|
201 |
+
}
|
202 |
+
|
203 |
+
if balancer.GetHealthyTokenCount() != 3 {
|
204 |
+
t.Errorf("Expected 3 healthy tokens after refresh, got %d", balancer.GetHealthyTokenCount())
|
205 |
+
}
|
206 |
+
|
207 |
+
// 验证新tokens可以被获取
|
208 |
+
for i := 0; i < 6; i++ { // 两轮完整轮询
|
209 |
+
token, err := balancer.GetToken()
|
210 |
+
if err != nil {
|
211 |
+
t.Fatalf("Unexpected error: %v", err)
|
212 |
+
}
|
213 |
+
|
214 |
+
found := false
|
215 |
+
for _, newToken := range newTokens {
|
216 |
+
if token == newToken {
|
217 |
+
found = true
|
218 |
+
break
|
219 |
+
}
|
220 |
+
}
|
221 |
+
|
222 |
+
if !found {
|
223 |
+
t.Errorf("Got unexpected token after refresh: %s", token)
|
224 |
+
}
|
225 |
+
}
|
226 |
+
}
|
internal/config/config.go
ADDED
@@ -0,0 +1,430 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package config
|
2 |
+
|
3 |
+
import (
|
4 |
+
"encoding/json"
|
5 |
+
"fmt"
|
6 |
+
"io/ioutil"
|
7 |
+
"log"
|
8 |
+
"os"
|
9 |
+
"path/filepath"
|
10 |
+
"strings"
|
11 |
+
"sync"
|
12 |
+
"time"
|
13 |
+
|
14 |
+
"github.com/joho/godotenv"
|
15 |
+
)
|
16 |
+
|
17 |
+
var (
|
18 |
+
GlobalConfig *Manager
|
19 |
+
once sync.Once
|
20 |
+
)
|
21 |
+
|
22 |
+
// LoadBalanceStrategy 负载均衡策略
|
23 |
+
type LoadBalanceStrategy string
|
24 |
+
|
25 |
+
const (
|
26 |
+
RoundRobin LoadBalanceStrategy = "round_robin"
|
27 |
+
Random LoadBalanceStrategy = "random"
|
28 |
+
)
|
29 |
+
|
30 |
+
// JWTTokenConfig JWT token配置
|
31 |
+
type JWTTokenConfig struct {
|
32 |
+
Token string `json:"token"`
|
33 |
+
Name string `json:"name,omitempty"`
|
34 |
+
Description string `json:"description,omitempty"`
|
35 |
+
Priority int `json:"priority,omitempty"`
|
36 |
+
Metadata map[string]string `json:"metadata,omitempty"`
|
37 |
+
}
|
38 |
+
|
39 |
+
// Config 应用配置
|
40 |
+
type Config struct {
|
41 |
+
JetbrainsTokens []JWTTokenConfig `json:"jetbrains_tokens"`
|
42 |
+
BearerToken string `json:"bearer_token"`
|
43 |
+
LoadBalanceStrategy LoadBalanceStrategy `json:"load_balance_strategy"`
|
44 |
+
HealthCheckInterval time.Duration `json:"health_check_interval"`
|
45 |
+
ServerPort int `json:"server_port"`
|
46 |
+
ServerHost string `json:"server_host"`
|
47 |
+
}
|
48 |
+
|
49 |
+
// Manager 配置管理器
|
50 |
+
type Manager struct {
|
51 |
+
config *Config
|
52 |
+
configPath string
|
53 |
+
mutex sync.RWMutex
|
54 |
+
}
|
55 |
+
|
56 |
+
// GetGlobalConfig 获取全局配置管理器(单例)
|
57 |
+
func GetGlobalConfig() *Manager {
|
58 |
+
once.Do(func() {
|
59 |
+
GlobalConfig = NewManager()
|
60 |
+
})
|
61 |
+
return GlobalConfig
|
62 |
+
}
|
63 |
+
|
64 |
+
// NewManager 创建新的配置管理器
|
65 |
+
func NewManager() *Manager {
|
66 |
+
return &Manager{
|
67 |
+
config: &Config{
|
68 |
+
LoadBalanceStrategy: RoundRobin,
|
69 |
+
HealthCheckInterval: 30 * time.Second,
|
70 |
+
ServerPort: 8080,
|
71 |
+
ServerHost: "0.0.0.0",
|
72 |
+
},
|
73 |
+
}
|
74 |
+
}
|
75 |
+
|
76 |
+
// LoadConfig 加载配置
|
77 |
+
func (m *Manager) LoadConfig() error {
|
78 |
+
m.mutex.Lock()
|
79 |
+
defer m.mutex.Unlock()
|
80 |
+
|
81 |
+
// 1. 首先尝试加载 .env 文件
|
82 |
+
_ = godotenv.Load()
|
83 |
+
|
84 |
+
// 2. 自动发现并加载配置文件
|
85 |
+
if err := m.loadConfigFile(); err != nil {
|
86 |
+
log.Printf("Warning: Failed to load config file: %v", err)
|
87 |
+
}
|
88 |
+
|
89 |
+
// 3. 从环境变量加载配置
|
90 |
+
m.loadFromEnv()
|
91 |
+
|
92 |
+
// 4. 验证配置
|
93 |
+
return m.validateConfig()
|
94 |
+
}
|
95 |
+
|
96 |
+
// loadConfigFile 自动发现并加载配置文件
|
97 |
+
func (m *Manager) loadConfigFile() error {
|
98 |
+
// 配置文件搜索路径
|
99 |
+
searchPaths := []string{
|
100 |
+
"config.json",
|
101 |
+
"config/config.json",
|
102 |
+
"configs/config.json",
|
103 |
+
".config/jetbrains-ai-proxy.json",
|
104 |
+
os.ExpandEnv("$HOME/.config/jetbrains-ai-proxy/config.json"),
|
105 |
+
}
|
106 |
+
|
107 |
+
for _, path := range searchPaths {
|
108 |
+
if _, err := os.Stat(path); err == nil {
|
109 |
+
log.Printf("Found config file: %s", path)
|
110 |
+
return m.loadFromFile(path)
|
111 |
+
}
|
112 |
+
}
|
113 |
+
|
114 |
+
return fmt.Errorf("no config file found in search paths")
|
115 |
+
}
|
116 |
+
|
117 |
+
// loadFromFile 从文件加载配置
|
118 |
+
func (m *Manager) loadFromFile(path string) error {
|
119 |
+
data, err := ioutil.ReadFile(path)
|
120 |
+
if err != nil {
|
121 |
+
return fmt.Errorf("failed to read config file %s: %v", path, err)
|
122 |
+
}
|
123 |
+
|
124 |
+
var fileConfig Config
|
125 |
+
if err := json.Unmarshal(data, &fileConfig); err != nil {
|
126 |
+
return fmt.Errorf("failed to parse config file %s: %v", path, err)
|
127 |
+
}
|
128 |
+
|
129 |
+
// 合并配置
|
130 |
+
m.mergeConfig(&fileConfig)
|
131 |
+
m.configPath = path
|
132 |
+
log.Printf("Loaded config from file: %s", path)
|
133 |
+
return nil
|
134 |
+
}
|
135 |
+
|
136 |
+
// loadFromEnv 从环境变量加载配置
|
137 |
+
func (m *Manager) loadFromEnv() {
|
138 |
+
// JWT Tokens
|
139 |
+
jwtTokensStr := os.Getenv("JWT_TOKENS")
|
140 |
+
if jwtTokensStr == "" {
|
141 |
+
// 兼容旧的单token配置
|
142 |
+
jwtTokensStr = os.Getenv("JWT_TOKEN")
|
143 |
+
}
|
144 |
+
|
145 |
+
if jwtTokensStr != "" {
|
146 |
+
tokens := m.parseJWTTokens(jwtTokensStr)
|
147 |
+
if len(tokens) > 0 {
|
148 |
+
m.config.JetbrainsTokens = tokens
|
149 |
+
}
|
150 |
+
}
|
151 |
+
|
152 |
+
// Bearer Token
|
153 |
+
if bearerToken := os.Getenv("BEARER_TOKEN"); bearerToken != "" {
|
154 |
+
m.config.BearerToken = bearerToken
|
155 |
+
}
|
156 |
+
|
157 |
+
// Load Balance Strategy
|
158 |
+
if strategy := os.Getenv("LOAD_BALANCE_STRATEGY"); strategy != "" {
|
159 |
+
if strategy == string(RoundRobin) || strategy == string(Random) {
|
160 |
+
m.config.LoadBalanceStrategy = LoadBalanceStrategy(strategy)
|
161 |
+
}
|
162 |
+
}
|
163 |
+
|
164 |
+
// Server configuration
|
165 |
+
if port := os.Getenv("SERVER_PORT"); port != "" {
|
166 |
+
if p, err := parsePort(port); err == nil {
|
167 |
+
m.config.ServerPort = p
|
168 |
+
}
|
169 |
+
}
|
170 |
+
|
171 |
+
if host := os.Getenv("SERVER_HOST"); host != "" {
|
172 |
+
m.config.ServerHost = host
|
173 |
+
}
|
174 |
+
}
|
175 |
+
|
176 |
+
// parseJWTTokens 解析JWT tokens字符串
|
177 |
+
func (m *Manager) parseJWTTokens(tokensStr string) []JWTTokenConfig {
|
178 |
+
var tokens []JWTTokenConfig
|
179 |
+
tokenList := strings.Split(tokensStr, ",")
|
180 |
+
|
181 |
+
for i, token := range tokenList {
|
182 |
+
token = strings.TrimSpace(token)
|
183 |
+
if token != "" {
|
184 |
+
tokens = append(tokens, JWTTokenConfig{
|
185 |
+
Token: token,
|
186 |
+
Name: fmt.Sprintf("JWT_%d", i+1),
|
187 |
+
Priority: 1,
|
188 |
+
})
|
189 |
+
}
|
190 |
+
}
|
191 |
+
|
192 |
+
return tokens
|
193 |
+
}
|
194 |
+
|
195 |
+
// mergeConfig 合并配置
|
196 |
+
func (m *Manager) mergeConfig(other *Config) {
|
197 |
+
if len(other.JetbrainsTokens) > 0 {
|
198 |
+
m.config.JetbrainsTokens = other.JetbrainsTokens
|
199 |
+
}
|
200 |
+
if other.BearerToken != "" {
|
201 |
+
m.config.BearerToken = other.BearerToken
|
202 |
+
}
|
203 |
+
if other.LoadBalanceStrategy != "" {
|
204 |
+
m.config.LoadBalanceStrategy = other.LoadBalanceStrategy
|
205 |
+
}
|
206 |
+
if other.HealthCheckInterval > 0 {
|
207 |
+
m.config.HealthCheckInterval = other.HealthCheckInterval
|
208 |
+
}
|
209 |
+
if other.ServerPort > 0 {
|
210 |
+
m.config.ServerPort = other.ServerPort
|
211 |
+
}
|
212 |
+
if other.ServerHost != "" {
|
213 |
+
m.config.ServerHost = other.ServerHost
|
214 |
+
}
|
215 |
+
}
|
216 |
+
|
217 |
+
// validateConfig 验证配置
|
218 |
+
func (m *Manager) validateConfig() error {
|
219 |
+
if len(m.config.JetbrainsTokens) == 0 {
|
220 |
+
return fmt.Errorf("no JWT tokens configured")
|
221 |
+
}
|
222 |
+
|
223 |
+
if m.config.BearerToken == "" {
|
224 |
+
return fmt.Errorf("bearer token is required")
|
225 |
+
}
|
226 |
+
|
227 |
+
if m.config.ServerPort <= 0 || m.config.ServerPort > 65535 {
|
228 |
+
return fmt.Errorf("invalid server port: %d", m.config.ServerPort)
|
229 |
+
}
|
230 |
+
|
231 |
+
return nil
|
232 |
+
}
|
233 |
+
|
234 |
+
// GetConfig 获取当前配置
|
235 |
+
func (m *Manager) GetConfig() *Config {
|
236 |
+
m.mutex.RLock()
|
237 |
+
defer m.mutex.RUnlock()
|
238 |
+
|
239 |
+
// 返回配置的副本
|
240 |
+
configCopy := *m.config
|
241 |
+
return &configCopy
|
242 |
+
}
|
243 |
+
|
244 |
+
// GetJWTTokens 获取JWT tokens字符串列表
|
245 |
+
func (m *Manager) GetJWTTokens() []string {
|
246 |
+
m.mutex.RLock()
|
247 |
+
defer m.mutex.RUnlock()
|
248 |
+
|
249 |
+
tokens := make([]string, len(m.config.JetbrainsTokens))
|
250 |
+
for i, tokenConfig := range m.config.JetbrainsTokens {
|
251 |
+
tokens[i] = tokenConfig.Token
|
252 |
+
}
|
253 |
+
return tokens
|
254 |
+
}
|
255 |
+
|
256 |
+
// GetJWTTokenConfigs 获取JWT token配置列表
|
257 |
+
func (m *Manager) GetJWTTokenConfigs() []JWTTokenConfig {
|
258 |
+
m.mutex.RLock()
|
259 |
+
defer m.mutex.RUnlock()
|
260 |
+
|
261 |
+
configs := make([]JWTTokenConfig, len(m.config.JetbrainsTokens))
|
262 |
+
copy(configs, m.config.JetbrainsTokens)
|
263 |
+
return configs
|
264 |
+
}
|
265 |
+
|
266 |
+
// SetJWTTokens 设置JWT tokens(用于命令行参数)
|
267 |
+
func (m *Manager) SetJWTTokens(tokensStr string) {
|
268 |
+
m.mutex.Lock()
|
269 |
+
defer m.mutex.Unlock()
|
270 |
+
|
271 |
+
if tokensStr != "" {
|
272 |
+
m.config.JetbrainsTokens = m.parseJWTTokens(tokensStr)
|
273 |
+
}
|
274 |
+
}
|
275 |
+
|
276 |
+
// SetBearerToken 设置Bearer token
|
277 |
+
func (m *Manager) SetBearerToken(token string) {
|
278 |
+
m.mutex.Lock()
|
279 |
+
defer m.mutex.Unlock()
|
280 |
+
|
281 |
+
m.config.BearerToken = token
|
282 |
+
}
|
283 |
+
|
284 |
+
// SetLoadBalanceStrategy 设置负载均衡策略
|
285 |
+
func (m *Manager) SetLoadBalanceStrategy(strategy string) {
|
286 |
+
m.mutex.Lock()
|
287 |
+
defer m.mutex.Unlock()
|
288 |
+
|
289 |
+
if strategy == string(RoundRobin) || strategy == string(Random) {
|
290 |
+
m.config.LoadBalanceStrategy = LoadBalanceStrategy(strategy)
|
291 |
+
}
|
292 |
+
}
|
293 |
+
|
294 |
+
// HasJWTTokens 检查是否有可用的JWT tokens
|
295 |
+
func (m *Manager) HasJWTTokens() bool {
|
296 |
+
m.mutex.RLock()
|
297 |
+
defer m.mutex.RUnlock()
|
298 |
+
|
299 |
+
return len(m.config.JetbrainsTokens) > 0
|
300 |
+
}
|
301 |
+
|
302 |
+
// SaveConfig 保存配置到文件
|
303 |
+
func (m *Manager) SaveConfig() error {
|
304 |
+
m.mutex.RLock()
|
305 |
+
defer m.mutex.RUnlock()
|
306 |
+
|
307 |
+
if m.configPath == "" {
|
308 |
+
m.configPath = "config.json"
|
309 |
+
}
|
310 |
+
|
311 |
+
// 确保目录存在
|
312 |
+
if err := os.MkdirAll(filepath.Dir(m.configPath), 0755); err != nil {
|
313 |
+
return fmt.Errorf("failed to create config directory: %v", err)
|
314 |
+
}
|
315 |
+
|
316 |
+
data, err := json.MarshalIndent(m.config, "", " ")
|
317 |
+
if err != nil {
|
318 |
+
return fmt.Errorf("failed to marshal config: %v", err)
|
319 |
+
}
|
320 |
+
|
321 |
+
if err := ioutil.WriteFile(m.configPath, data, 0644); err != nil {
|
322 |
+
return fmt.Errorf("failed to write config file: %v", err)
|
323 |
+
}
|
324 |
+
|
325 |
+
log.Printf("Config saved to: %s", m.configPath)
|
326 |
+
return nil
|
327 |
+
}
|
328 |
+
|
329 |
+
// GenerateExampleConfig 生成示例配置文件
|
330 |
+
func (m *Manager) GenerateExampleConfig(path string) error {
|
331 |
+
exampleConfig := &Config{
|
332 |
+
JetbrainsTokens: []JWTTokenConfig{
|
333 |
+
{
|
334 |
+
Token: "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
335 |
+
Name: "Primary_JWT",
|
336 |
+
Description: "Primary JWT token for JetBrains AI",
|
337 |
+
Priority: 1,
|
338 |
+
Metadata: map[string]string{
|
339 |
+
"environment": "production",
|
340 |
+
"region": "us-east-1",
|
341 |
+
},
|
342 |
+
},
|
343 |
+
{
|
344 |
+
Token: "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...",
|
345 |
+
Name: "Secondary_JWT",
|
346 |
+
Description: "Secondary JWT token for load balancing",
|
347 |
+
Priority: 2,
|
348 |
+
Metadata: map[string]string{
|
349 |
+
"environment": "production",
|
350 |
+
"region": "us-west-2",
|
351 |
+
},
|
352 |
+
},
|
353 |
+
},
|
354 |
+
BearerToken: "your_bearer_token_here",
|
355 |
+
LoadBalanceStrategy: RoundRobin,
|
356 |
+
HealthCheckInterval: 30 * time.Second,
|
357 |
+
ServerPort: 8080,
|
358 |
+
ServerHost: "0.0.0.0",
|
359 |
+
}
|
360 |
+
|
361 |
+
// 确保目录存在
|
362 |
+
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
363 |
+
return fmt.Errorf("failed to create directory: %v", err)
|
364 |
+
}
|
365 |
+
|
366 |
+
data, err := json.MarshalIndent(exampleConfig, "", " ")
|
367 |
+
if err != nil {
|
368 |
+
return fmt.Errorf("failed to marshal example config: %v", err)
|
369 |
+
}
|
370 |
+
|
371 |
+
if err := ioutil.WriteFile(path, data, 0644); err != nil {
|
372 |
+
return fmt.Errorf("failed to write example config: %v", err)
|
373 |
+
}
|
374 |
+
|
375 |
+
log.Printf("Example config generated: %s", path)
|
376 |
+
return nil
|
377 |
+
}
|
378 |
+
|
379 |
+
// PrintConfig 打印当前配置信息
|
380 |
+
func (m *Manager) PrintConfig() {
|
381 |
+
m.mutex.RLock()
|
382 |
+
defer m.mutex.RUnlock()
|
383 |
+
|
384 |
+
fmt.Println("=== Current Configuration ===")
|
385 |
+
fmt.Printf("JWT Tokens: %d configured\n", len(m.config.JetbrainsTokens))
|
386 |
+
for i, token := range m.config.JetbrainsTokens {
|
387 |
+
fmt.Printf(" %d. %s (%s...)\n", i+1, token.Name, token.Token[:min(len(token.Token), 20)])
|
388 |
+
}
|
389 |
+
fmt.Printf("Bearer Token: %s...\n", m.config.BearerToken[:min(len(m.config.BearerToken), 20)])
|
390 |
+
fmt.Printf("Load Balance Strategy: %s\n", m.config.LoadBalanceStrategy)
|
391 |
+
fmt.Printf("Health Check Interval: %v\n", m.config.HealthCheckInterval)
|
392 |
+
fmt.Printf("Server: %s:%d\n", m.config.ServerHost, m.config.ServerPort)
|
393 |
+
if m.configPath != "" {
|
394 |
+
fmt.Printf("Config File: %s\n", m.configPath)
|
395 |
+
}
|
396 |
+
fmt.Println("=============================")
|
397 |
+
}
|
398 |
+
|
399 |
+
// 辅助函数
|
400 |
+
func parsePort(portStr string) (int, error) {
|
401 |
+
var port int
|
402 |
+
if _, err := fmt.Sscanf(portStr, "%d", &port); err != nil {
|
403 |
+
return 0, err
|
404 |
+
}
|
405 |
+
if port <= 0 || port > 65535 {
|
406 |
+
return 0, fmt.Errorf("invalid port: %d", port)
|
407 |
+
}
|
408 |
+
return port, nil
|
409 |
+
}
|
410 |
+
|
411 |
+
func min(a, b int) int {
|
412 |
+
if a < b {
|
413 |
+
return a
|
414 |
+
}
|
415 |
+
return b
|
416 |
+
}
|
417 |
+
|
418 |
+
// 向后兼容的全局变量和函数
|
419 |
+
var JetbrainsAiConfig *Config
|
420 |
+
|
421 |
+
// LoadConfig 向后兼容的配置加载函数
|
422 |
+
func LoadConfig() *Config {
|
423 |
+
manager := GetGlobalConfig()
|
424 |
+
if err := manager.LoadConfig(); err != nil {
|
425 |
+
log.Printf("Warning: %v", err)
|
426 |
+
}
|
427 |
+
|
428 |
+
JetbrainsAiConfig = manager.GetConfig()
|
429 |
+
return JetbrainsAiConfig
|
430 |
+
}
|
internal/config/discovery.go
ADDED
@@ -0,0 +1,293 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package config
|
2 |
+
|
3 |
+
import (
|
4 |
+
"encoding/json"
|
5 |
+
"fmt"
|
6 |
+
"io/ioutil"
|
7 |
+
"log"
|
8 |
+
"os"
|
9 |
+
"path/filepath"
|
10 |
+
"time"
|
11 |
+
)
|
12 |
+
|
13 |
+
// ConfigDiscovery 配置发现器
|
14 |
+
type ConfigDiscovery struct {
|
15 |
+
searchPaths []string
|
16 |
+
manager *Manager
|
17 |
+
}
|
18 |
+
|
19 |
+
// NewConfigDiscovery 创建配置发现器
|
20 |
+
func NewConfigDiscovery(manager *Manager) *ConfigDiscovery {
|
21 |
+
return &ConfigDiscovery{
|
22 |
+
manager: manager,
|
23 |
+
searchPaths: []string{
|
24 |
+
// 当前目录
|
25 |
+
"config.json",
|
26 |
+
"jetbrains-ai-proxy.json",
|
27 |
+
".jetbrains-ai-proxy.json",
|
28 |
+
|
29 |
+
// config 目录
|
30 |
+
"config/config.json",
|
31 |
+
"config/jetbrains-ai-proxy.json",
|
32 |
+
"configs/config.json",
|
33 |
+
"configs/jetbrains-ai-proxy.json",
|
34 |
+
|
35 |
+
// 隐藏配置目录
|
36 |
+
".config/config.json",
|
37 |
+
".config/jetbrains-ai-proxy.json",
|
38 |
+
|
39 |
+
// 用户主目录
|
40 |
+
os.ExpandEnv("$HOME/.config/jetbrains-ai-proxy/config.json"),
|
41 |
+
os.ExpandEnv("$HOME/.jetbrains-ai-proxy/config.json"),
|
42 |
+
os.ExpandEnv("$HOME/.jetbrains-ai-proxy.json"),
|
43 |
+
|
44 |
+
// 系统配置目录 (Linux/macOS)
|
45 |
+
"/etc/jetbrains-ai-proxy/config.json",
|
46 |
+
"/usr/local/etc/jetbrains-ai-proxy/config.json",
|
47 |
+
},
|
48 |
+
}
|
49 |
+
}
|
50 |
+
|
51 |
+
// DiscoverAndLoad 发现并加载配置文件
|
52 |
+
func (cd *ConfigDiscovery) DiscoverAndLoad() error {
|
53 |
+
log.Println("Starting configuration discovery...")
|
54 |
+
|
55 |
+
// 1. 尝试从环境变量指定的配置文件加载
|
56 |
+
if configPath := os.Getenv("CONFIG_FILE"); configPath != "" {
|
57 |
+
if err := cd.loadConfigFile(configPath); err != nil {
|
58 |
+
log.Printf("Failed to load config from CONFIG_FILE=%s: %v", configPath, err)
|
59 |
+
} else {
|
60 |
+
log.Printf("Successfully loaded config from CONFIG_FILE: %s", configPath)
|
61 |
+
return nil
|
62 |
+
}
|
63 |
+
}
|
64 |
+
|
65 |
+
// 2. 搜索预定义路径
|
66 |
+
for _, path := range cd.searchPaths {
|
67 |
+
if cd.fileExists(path) {
|
68 |
+
if err := cd.loadConfigFile(path); err != nil {
|
69 |
+
log.Printf("Failed to load config from %s: %v", path, err)
|
70 |
+
continue
|
71 |
+
}
|
72 |
+
log.Printf("Successfully loaded config from: %s", path)
|
73 |
+
return nil
|
74 |
+
}
|
75 |
+
}
|
76 |
+
|
77 |
+
// 3. 尝试从当前目录的 .env 文件加载
|
78 |
+
if cd.fileExists(".env") {
|
79 |
+
log.Println("Found .env file, loading environment variables...")
|
80 |
+
return nil // .env 文件会在 LoadConfig 中自动加载
|
81 |
+
}
|
82 |
+
|
83 |
+
// 4. 如果没有找到配置文件,生成示例配置
|
84 |
+
log.Println("No configuration file found, generating example config...")
|
85 |
+
return cd.generateDefaultConfig()
|
86 |
+
}
|
87 |
+
|
88 |
+
// loadConfigFile 加载指定的配置文件
|
89 |
+
func (cd *ConfigDiscovery) loadConfigFile(path string) error {
|
90 |
+
data, err := ioutil.ReadFile(path)
|
91 |
+
if err != nil {
|
92 |
+
return fmt.Errorf("failed to read config file: %v", err)
|
93 |
+
}
|
94 |
+
|
95 |
+
var config Config
|
96 |
+
if err := json.Unmarshal(data, &config); err != nil {
|
97 |
+
return fmt.Errorf("failed to parse config file: %v", err)
|
98 |
+
}
|
99 |
+
|
100 |
+
// 验证配置
|
101 |
+
if err := cd.validateLoadedConfig(&config); err != nil {
|
102 |
+
return fmt.Errorf("invalid config: %v", err)
|
103 |
+
}
|
104 |
+
|
105 |
+
// 合并到管理器
|
106 |
+
cd.manager.mutex.Lock()
|
107 |
+
cd.manager.mergeConfig(&config)
|
108 |
+
cd.manager.configPath = path
|
109 |
+
cd.manager.mutex.Unlock()
|
110 |
+
|
111 |
+
return nil
|
112 |
+
}
|
113 |
+
|
114 |
+
// validateLoadedConfig 验证加载的配置
|
115 |
+
func (cd *ConfigDiscovery) validateLoadedConfig(config *Config) error {
|
116 |
+
if len(config.JetbrainsTokens) == 0 {
|
117 |
+
return fmt.Errorf("no JWT tokens found in config")
|
118 |
+
}
|
119 |
+
|
120 |
+
// 验证每个JWT token
|
121 |
+
for i, tokenConfig := range config.JetbrainsTokens {
|
122 |
+
if tokenConfig.Token == "" {
|
123 |
+
return fmt.Errorf("JWT token %d is empty", i+1)
|
124 |
+
}
|
125 |
+
if len(tokenConfig.Token) < 10 {
|
126 |
+
return fmt.Errorf("JWT token %d appears to be invalid (too short)", i+1)
|
127 |
+
}
|
128 |
+
}
|
129 |
+
|
130 |
+
if config.BearerToken == "" {
|
131 |
+
log.Println("Warning: No bearer token found in config file")
|
132 |
+
}
|
133 |
+
|
134 |
+
return nil
|
135 |
+
}
|
136 |
+
|
137 |
+
// generateDefaultConfig 生成默认配置
|
138 |
+
func (cd *ConfigDiscovery) generateDefaultConfig() error {
|
139 |
+
configDir := "config"
|
140 |
+
configPath := filepath.Join(configDir, "config.json")
|
141 |
+
|
142 |
+
// 创建配置目录
|
143 |
+
if err := os.MkdirAll(configDir, 0755); err != nil {
|
144 |
+
return fmt.Errorf("failed to create config directory: %v", err)
|
145 |
+
}
|
146 |
+
|
147 |
+
// 生成示例配置
|
148 |
+
if err := cd.manager.GenerateExampleConfig(configPath); err != nil {
|
149 |
+
return fmt.Errorf("failed to generate example config: %v", err)
|
150 |
+
}
|
151 |
+
|
152 |
+
// 同时生成 .env 示例文件
|
153 |
+
envPath := ".env.example"
|
154 |
+
if err := cd.generateEnvExample(envPath); err != nil {
|
155 |
+
log.Printf("Warning: Failed to generate .env example: %v", err)
|
156 |
+
}
|
157 |
+
|
158 |
+
log.Printf("Generated example configuration files:")
|
159 |
+
log.Printf(" - %s (JSON format)", configPath)
|
160 |
+
log.Printf(" - %s (Environment variables)", envPath)
|
161 |
+
log.Printf("Please edit these files with your actual JWT tokens and restart the application.")
|
162 |
+
|
163 |
+
return fmt.Errorf("no valid configuration found, example files generated")
|
164 |
+
}
|
165 |
+
|
166 |
+
// generateEnvExample 生成 .env 示例文件
|
167 |
+
func (cd *ConfigDiscovery) generateEnvExample(path string) error {
|
168 |
+
envContent := `# JetBrains AI Proxy Configuration
|
169 |
+
# Copy this file to .env and fill in your actual values
|
170 |
+
|
171 |
+
# Multiple JWT tokens (comma-separated)
|
172 |
+
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
173 |
+
|
174 |
+
# Or single JWT token (for backward compatibility)
|
175 |
+
# JWT_TOKEN=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
176 |
+
|
177 |
+
# Bearer token for API authentication
|
178 |
+
BEARER_TOKEN=your_bearer_token_here
|
179 |
+
|
180 |
+
# Load balancing strategy: round_robin or random
|
181 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
182 |
+
|
183 |
+
# Server configuration
|
184 |
+
SERVER_HOST=0.0.0.0
|
185 |
+
SERVER_PORT=8080
|
186 |
+
|
187 |
+
# Alternative: specify config file path
|
188 |
+
# CONFIG_FILE=config/config.json
|
189 |
+
`
|
190 |
+
|
191 |
+
return ioutil.WriteFile(path, []byte(envContent), 0644)
|
192 |
+
}
|
193 |
+
|
194 |
+
// fileExists 检查文件是否存在
|
195 |
+
func (cd *ConfigDiscovery) fileExists(path string) bool {
|
196 |
+
_, err := os.Stat(path)
|
197 |
+
return err == nil
|
198 |
+
}
|
199 |
+
|
200 |
+
// WatchConfig 监控配置文件变化(简单实现)
|
201 |
+
func (cd *ConfigDiscovery) WatchConfig() {
|
202 |
+
if cd.manager.configPath == "" {
|
203 |
+
return
|
204 |
+
}
|
205 |
+
|
206 |
+
go func() {
|
207 |
+
var lastModTime time.Time
|
208 |
+
|
209 |
+
// 获取初始修改时间
|
210 |
+
if stat, err := os.Stat(cd.manager.configPath); err == nil {
|
211 |
+
lastModTime = stat.ModTime()
|
212 |
+
}
|
213 |
+
|
214 |
+
ticker := time.NewTicker(5 * time.Second)
|
215 |
+
defer ticker.Stop()
|
216 |
+
|
217 |
+
for range ticker.C {
|
218 |
+
stat, err := os.Stat(cd.manager.configPath)
|
219 |
+
if err != nil {
|
220 |
+
continue
|
221 |
+
}
|
222 |
+
|
223 |
+
if stat.ModTime().After(lastModTime) {
|
224 |
+
log.Printf("Config file changed, reloading: %s", cd.manager.configPath)
|
225 |
+
if err := cd.loadConfigFile(cd.manager.configPath); err != nil {
|
226 |
+
log.Printf("Failed to reload config: %v", err)
|
227 |
+
} else {
|
228 |
+
log.Println("Config reloaded successfully")
|
229 |
+
}
|
230 |
+
lastModTime = stat.ModTime()
|
231 |
+
}
|
232 |
+
}
|
233 |
+
}()
|
234 |
+
}
|
235 |
+
|
236 |
+
// ListAvailableConfigs 列出可用的配置文件
|
237 |
+
func (cd *ConfigDiscovery) ListAvailableConfigs() []string {
|
238 |
+
var available []string
|
239 |
+
|
240 |
+
for _, path := range cd.searchPaths {
|
241 |
+
if cd.fileExists(path) {
|
242 |
+
available = append(available, path)
|
243 |
+
}
|
244 |
+
}
|
245 |
+
|
246 |
+
return available
|
247 |
+
}
|
248 |
+
|
249 |
+
// ValidateConfigFile 验证配置文件格式
|
250 |
+
func (cd *ConfigDiscovery) ValidateConfigFile(path string) error {
|
251 |
+
if !cd.fileExists(path) {
|
252 |
+
return fmt.Errorf("config file does not exist: %s", path)
|
253 |
+
}
|
254 |
+
|
255 |
+
data, err := ioutil.ReadFile(path)
|
256 |
+
if err != nil {
|
257 |
+
return fmt.Errorf("failed to read config file: %v", err)
|
258 |
+
}
|
259 |
+
|
260 |
+
var config Config
|
261 |
+
if err := json.Unmarshal(data, &config); err != nil {
|
262 |
+
return fmt.Errorf("invalid JSON format: %v", err)
|
263 |
+
}
|
264 |
+
|
265 |
+
return cd.validateLoadedConfig(&config)
|
266 |
+
}
|
267 |
+
|
268 |
+
// GetConfigSummary 获取配置摘要信息
|
269 |
+
func (cd *ConfigDiscovery) GetConfigSummary() map[string]interface{} {
|
270 |
+
config := cd.manager.GetConfig()
|
271 |
+
|
272 |
+
// 隐藏敏感信息
|
273 |
+
tokenSummary := make([]map[string]interface{}, len(config.JetbrainsTokens))
|
274 |
+
for i, token := range config.JetbrainsTokens {
|
275 |
+
tokenSummary[i] = map[string]interface{}{
|
276 |
+
"name": token.Name,
|
277 |
+
"description": token.Description,
|
278 |
+
"priority": token.Priority,
|
279 |
+
"token_preview": token.Token[:min(len(token.Token), 20)] + "...",
|
280 |
+
}
|
281 |
+
}
|
282 |
+
|
283 |
+
return map[string]interface{}{
|
284 |
+
"jwt_tokens_count": len(config.JetbrainsTokens),
|
285 |
+
"jwt_tokens": tokenSummary,
|
286 |
+
"bearer_token_set": config.BearerToken != "",
|
287 |
+
"load_balance_strategy": config.LoadBalanceStrategy,
|
288 |
+
"health_check_interval": config.HealthCheckInterval.String(),
|
289 |
+
"server_host": config.ServerHost,
|
290 |
+
"server_port": config.ServerPort,
|
291 |
+
"config_file": cd.manager.configPath,
|
292 |
+
}
|
293 |
+
}
|
internal/jetbrains/client.go
ADDED
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package jetbrains
|
2 |
+
|
3 |
+
import (
|
4 |
+
"context"
|
5 |
+
"fmt"
|
6 |
+
"github.com/go-resty/resty/v2"
|
7 |
+
"jetbrains-ai-proxy/internal/balancer"
|
8 |
+
"jetbrains-ai-proxy/internal/config"
|
9 |
+
"jetbrains-ai-proxy/internal/types"
|
10 |
+
"jetbrains-ai-proxy/internal/utils"
|
11 |
+
"log"
|
12 |
+
"sync"
|
13 |
+
)
|
14 |
+
|
15 |
+
var (
|
16 |
+
jwtBalancer balancer.JWTBalancer
|
17 |
+
healthChecker *balancer.HealthChecker
|
18 |
+
initOnce sync.Once
|
19 |
+
configManager *config.Manager
|
20 |
+
)
|
21 |
+
|
22 |
+
// InitializeFromConfig 从配置管理器初始化JWT负载均衡器
|
23 |
+
func InitializeFromConfig() error {
|
24 |
+
var initErr error
|
25 |
+
|
26 |
+
initOnce.Do(func() {
|
27 |
+
configManager = config.GetGlobalConfig()
|
28 |
+
|
29 |
+
// 加载配置
|
30 |
+
if err := configManager.LoadConfig(); err != nil {
|
31 |
+
initErr = fmt.Errorf("failed to load config: %v", err)
|
32 |
+
return
|
33 |
+
}
|
34 |
+
|
35 |
+
// 获取配置
|
36 |
+
cfg := configManager.GetConfig()
|
37 |
+
tokens := configManager.GetJWTTokens()
|
38 |
+
|
39 |
+
if len(tokens) == 0 {
|
40 |
+
initErr = fmt.Errorf("no JWT tokens configured")
|
41 |
+
return
|
42 |
+
}
|
43 |
+
|
44 |
+
// 创建负载均衡器
|
45 |
+
jwtBalancer = balancer.NewJWTBalancer(tokens, cfg.LoadBalanceStrategy)
|
46 |
+
|
47 |
+
// 创建并启动健康检查器
|
48 |
+
healthChecker = balancer.NewHealthChecker(jwtBalancer)
|
49 |
+
if cfg.HealthCheckInterval > 0 {
|
50 |
+
healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
|
51 |
+
}
|
52 |
+
healthChecker.Start()
|
53 |
+
|
54 |
+
log.Printf("JWT balancer initialized from config:")
|
55 |
+
log.Printf(" - Tokens: %d", len(tokens))
|
56 |
+
log.Printf(" - Strategy: %s", cfg.LoadBalanceStrategy)
|
57 |
+
log.Printf(" - Health check interval: %v", cfg.HealthCheckInterval)
|
58 |
+
})
|
59 |
+
|
60 |
+
return initErr
|
61 |
+
}
|
62 |
+
|
63 |
+
// InitializeBalancer 初始化JWT负载均衡器(向后兼容)
|
64 |
+
func InitializeBalancer(tokens []string, strategy string) error {
|
65 |
+
if len(tokens) == 0 {
|
66 |
+
return fmt.Errorf("no JWT tokens provided")
|
67 |
+
}
|
68 |
+
|
69 |
+
var balanceStrategy config.LoadBalanceStrategy
|
70 |
+
switch strategy {
|
71 |
+
case "random":
|
72 |
+
balanceStrategy = config.Random
|
73 |
+
case "round_robin", "":
|
74 |
+
balanceStrategy = config.RoundRobin
|
75 |
+
default:
|
76 |
+
balanceStrategy = config.RoundRobin
|
77 |
+
}
|
78 |
+
|
79 |
+
// 创建负载均衡器
|
80 |
+
jwtBalancer = balancer.NewJWTBalancer(tokens, balanceStrategy)
|
81 |
+
|
82 |
+
// 创建并启动健康检查器
|
83 |
+
healthChecker = balancer.NewHealthChecker(jwtBalancer)
|
84 |
+
healthChecker.Start()
|
85 |
+
|
86 |
+
log.Printf("JWT balancer initialized with %d tokens, strategy: %s", len(tokens), string(balanceStrategy))
|
87 |
+
return nil
|
88 |
+
}
|
89 |
+
|
90 |
+
// ReloadConfig 重新加载配置
|
91 |
+
func ReloadConfig() error {
|
92 |
+
if configManager == nil {
|
93 |
+
return fmt.Errorf("config manager not initialized")
|
94 |
+
}
|
95 |
+
|
96 |
+
// 重新加载配置
|
97 |
+
if err := configManager.LoadConfig(); err != nil {
|
98 |
+
return fmt.Errorf("failed to reload config: %v", err)
|
99 |
+
}
|
100 |
+
|
101 |
+
// 获取新配置
|
102 |
+
cfg := configManager.GetConfig()
|
103 |
+
tokens := configManager.GetJWTTokens()
|
104 |
+
|
105 |
+
if len(tokens) == 0 {
|
106 |
+
return fmt.Errorf("no JWT tokens in reloaded config")
|
107 |
+
}
|
108 |
+
|
109 |
+
// 更新负载均衡器
|
110 |
+
if jwtBalancer != nil {
|
111 |
+
jwtBalancer.RefreshTokens(tokens)
|
112 |
+
}
|
113 |
+
|
114 |
+
// 更新健康检查间隔
|
115 |
+
if healthChecker != nil && cfg.HealthCheckInterval > 0 {
|
116 |
+
healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
|
117 |
+
}
|
118 |
+
|
119 |
+
log.Printf("Config reloaded successfully:")
|
120 |
+
log.Printf(" - Tokens: %d", len(tokens))
|
121 |
+
log.Printf(" - Strategy: %s", cfg.LoadBalanceStrategy)
|
122 |
+
|
123 |
+
return nil
|
124 |
+
}
|
125 |
+
|
126 |
+
// StopBalancer 停止负载均衡器
|
127 |
+
func StopBalancer() {
|
128 |
+
if healthChecker != nil {
|
129 |
+
healthChecker.Stop()
|
130 |
+
}
|
131 |
+
}
|
132 |
+
|
133 |
+
// GetConfigManager 获取配置管理器
|
134 |
+
func GetConfigManager() *config.Manager {
|
135 |
+
return configManager
|
136 |
+
}
|
137 |
+
|
138 |
+
func SendJetbrainsRequest(ctx context.Context, req *types.JetbrainsRequest) (*resty.Response, error) {
|
139 |
+
// 获取一个可用的JWT token
|
140 |
+
token, err := jwtBalancer.GetToken()
|
141 |
+
if err != nil {
|
142 |
+
log.Printf("failed to get JWT token: %v", err)
|
143 |
+
return nil, fmt.Errorf("no available JWT tokens: %v", err)
|
144 |
+
}
|
145 |
+
|
146 |
+
resp, err := utils.RestySSEClient.R().
|
147 |
+
SetContext(ctx).
|
148 |
+
SetHeader(types.JwtTokenKey, token).
|
149 |
+
SetDoNotParseResponse(true).
|
150 |
+
SetBody(req).
|
151 |
+
Post(types.ChatStreamV7)
|
152 |
+
|
153 |
+
if err != nil {
|
154 |
+
log.Printf("jetbrains ai req error: %v", err)
|
155 |
+
// 标记token为不健康
|
156 |
+
jwtBalancer.MarkTokenUnhealthy(token)
|
157 |
+
return nil, err
|
158 |
+
}
|
159 |
+
|
160 |
+
// 检查响应状态码
|
161 |
+
if resp.StatusCode() == 401 {
|
162 |
+
// 401表示token无效,标记为不健康
|
163 |
+
jwtBalancer.MarkTokenUnhealthy(token)
|
164 |
+
log.Printf("JWT token invalid (401): %s...", token[:min(len(token), 10)])
|
165 |
+
return nil, fmt.Errorf("JWT token invalid")
|
166 |
+
} else if resp.StatusCode() == 200 {
|
167 |
+
// 成功响应,确保token标记为健康
|
168 |
+
jwtBalancer.MarkTokenHealthy(token)
|
169 |
+
}
|
170 |
+
|
171 |
+
return resp, nil
|
172 |
+
}
|
173 |
+
|
174 |
+
// GetBalancerStats 获取负载均衡器统计信息
|
175 |
+
func GetBalancerStats() (int, int) {
|
176 |
+
if jwtBalancer == nil {
|
177 |
+
return 0, 0
|
178 |
+
}
|
179 |
+
return jwtBalancer.GetHealthyTokenCount(), jwtBalancer.GetTotalTokenCount()
|
180 |
+
}
|
181 |
+
|
182 |
+
// min 辅助函数
|
183 |
+
func min(a, b int) int {
|
184 |
+
if a < b {
|
185 |
+
return a
|
186 |
+
}
|
187 |
+
return b
|
188 |
+
}
|
internal/jetbrains/sse.go
ADDED
@@ -0,0 +1,331 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package jetbrains
|
2 |
+
|
3 |
+
import (
|
4 |
+
"bufio"
|
5 |
+
"context"
|
6 |
+
"fmt"
|
7 |
+
"github.com/bytedance/sonic"
|
8 |
+
"github.com/sashabaranov/go-openai"
|
9 |
+
"io"
|
10 |
+
"jetbrains-ai-proxy/internal/utils"
|
11 |
+
"log"
|
12 |
+
"math"
|
13 |
+
"net/http"
|
14 |
+
"strconv"
|
15 |
+
"strings"
|
16 |
+
"time"
|
17 |
+
)
|
18 |
+
|
19 |
+
const (
|
20 |
+
sseObject = "chat.completion.chunk"
|
21 |
+
completionsObject = "chat.completions"
|
22 |
+
sseFinish = "[DONE]"
|
23 |
+
initialBufferSize = 4096
|
24 |
+
maxBufferSize = 1024 * 1024 // 1MB
|
25 |
+
flushThreshold = 10
|
26 |
+
heartbeatInterval = 30 * time.Second
|
27 |
+
)
|
28 |
+
|
29 |
+
type SSEData struct {
|
30 |
+
Type string `json:"type"`
|
31 |
+
EventType string `json:"event_type"`
|
32 |
+
Content string `json:"content,omitempty"`
|
33 |
+
Reason string `json:"reason,omitempty"`
|
34 |
+
Updated *UpdatedData `json:"updated,omitempty"`
|
35 |
+
Spent *SpentData `json:"spent,omitempty"`
|
36 |
+
}
|
37 |
+
|
38 |
+
type UpdatedData struct {
|
39 |
+
License string `json:"license"`
|
40 |
+
Current AmountData `json:"current"`
|
41 |
+
Maximum AmountData `json:"maximum"`
|
42 |
+
Until int64 `json:"until"`
|
43 |
+
QuotaID QuotaInfo `json:"quotaID"`
|
44 |
+
}
|
45 |
+
|
46 |
+
type AmountData struct {
|
47 |
+
Amount string `json:"amount"`
|
48 |
+
}
|
49 |
+
|
50 |
+
type QuotaInfo struct {
|
51 |
+
QuotaId string `json:"quotaId"`
|
52 |
+
}
|
53 |
+
|
54 |
+
type SpentData struct {
|
55 |
+
Amount string `json:"amount"`
|
56 |
+
}
|
57 |
+
|
58 |
+
// ResponseJetbrainsAIToClient 处理非流式响应
|
59 |
+
func ResponseJetbrainsAIToClient(ctx context.Context, req openai.ChatCompletionRequest, r io.Reader, fp string) (openai.ChatCompletionResponse, error) {
|
60 |
+
reader := bufio.NewReader(r)
|
61 |
+
var fullContent strings.Builder
|
62 |
+
|
63 |
+
now := time.Now().Unix()
|
64 |
+
chatId := strconv.Itoa(int(now))
|
65 |
+
|
66 |
+
for {
|
67 |
+
select {
|
68 |
+
case <-ctx.Done():
|
69 |
+
return openai.ChatCompletionResponse{}, ctx.Err()
|
70 |
+
default:
|
71 |
+
}
|
72 |
+
|
73 |
+
line, err := reader.ReadString('\n')
|
74 |
+
if err != nil {
|
75 |
+
if err == io.EOF {
|
76 |
+
log.Printf("Reached EOF for non-streaming response")
|
77 |
+
break
|
78 |
+
}
|
79 |
+
return openai.ChatCompletionResponse{}, fmt.Errorf("读取错误: %w", err)
|
80 |
+
}
|
81 |
+
|
82 |
+
if !strings.HasPrefix(line, "data: ") {
|
83 |
+
continue
|
84 |
+
}
|
85 |
+
|
86 |
+
jsonStr := strings.TrimSpace(strings.TrimPrefix(line, "data: "))
|
87 |
+
if jsonStr == "" || jsonStr == sseFinish || jsonStr == "end" {
|
88 |
+
continue
|
89 |
+
}
|
90 |
+
|
91 |
+
var sseData SSEData
|
92 |
+
if err := sonic.UnmarshalString(jsonStr, &sseData); err != nil {
|
93 |
+
log.Printf("解析SSE数据错误: %v", err)
|
94 |
+
continue
|
95 |
+
}
|
96 |
+
|
97 |
+
if sseData.Type == "Content" {
|
98 |
+
fullContent.WriteString(sseData.Content)
|
99 |
+
}
|
100 |
+
|
101 |
+
if sseData.Type == "QuotaMetadata" {
|
102 |
+
var spentAmount float64
|
103 |
+
if sseData.Spent != nil {
|
104 |
+
if amount, err := strconv.ParseFloat(sseData.Spent.Amount, 64); err == nil {
|
105 |
+
spentAmount = amount
|
106 |
+
} else {
|
107 |
+
log.Printf("Warning: failed to parse spent amount '%s': %v", sseData.Spent.Amount, err)
|
108 |
+
}
|
109 |
+
}
|
110 |
+
usage := utils.CalculateJetbrainsUsage(fullContent.String(), int(math.Round(spentAmount)))
|
111 |
+
return createMessage(chatId, now, req, usage, fullContent.String(), fp), nil
|
112 |
+
}
|
113 |
+
}
|
114 |
+
|
115 |
+
// 如果没有收到 QuotaMetadata,返回默认响应
|
116 |
+
usage := utils.CalculateJetbrainsUsage(fullContent.String(), 0)
|
117 |
+
return createMessage(chatId, now, req, usage, fullContent.String(), fp), nil
|
118 |
+
}
|
119 |
+
|
120 |
+
// StreamJetbrainsAISSEToClient 处理流式响应
|
121 |
+
func StreamJetbrainsAISSEToClient(ctx context.Context, req openai.ChatCompletionRequest, w io.Writer, r io.Reader, fp string) error {
|
122 |
+
log.Printf("=== Starting SSE Stream Processing for model: %s ===", req.Model)
|
123 |
+
|
124 |
+
reader := bufio.NewReaderSize(r, initialBufferSize)
|
125 |
+
writer := bufio.NewWriterSize(w, initialBufferSize)
|
126 |
+
|
127 |
+
now := time.Now().Unix()
|
128 |
+
chatId := strconv.Itoa(int(now))
|
129 |
+
fingerprint := fp
|
130 |
+
|
131 |
+
log.Printf("Session initialized - ChatID: %s, Fingerprint: %s", chatId, fingerprint)
|
132 |
+
|
133 |
+
var completionBuilder strings.Builder
|
134 |
+
messageCount := 0
|
135 |
+
totalBufferSize := 0
|
136 |
+
|
137 |
+
// 创建心跳检测器
|
138 |
+
heartbeat := time.NewTicker(heartbeatInterval)
|
139 |
+
defer heartbeat.Stop()
|
140 |
+
|
141 |
+
for {
|
142 |
+
select {
|
143 |
+
case <-ctx.Done():
|
144 |
+
return ctx.Err()
|
145 |
+
case <-heartbeat.C:
|
146 |
+
if err := sendHeartbeat(writer, w); err != nil {
|
147 |
+
log.Printf("Heartbeat error: %v", err)
|
148 |
+
}
|
149 |
+
continue
|
150 |
+
default:
|
151 |
+
}
|
152 |
+
|
153 |
+
line, err := reader.ReadString('\n')
|
154 |
+
if err != nil {
|
155 |
+
if err == io.EOF {
|
156 |
+
log.Printf("Reached EOF after %d messages", messageCount)
|
157 |
+
return nil
|
158 |
+
}
|
159 |
+
return fmt.Errorf("read error: %w", err)
|
160 |
+
}
|
161 |
+
|
162 |
+
log.Printf("Received line: %s", strings.TrimSpace(line))
|
163 |
+
|
164 |
+
// 检查缓冲区大小
|
165 |
+
totalBufferSize += len(line)
|
166 |
+
if totalBufferSize > maxBufferSize {
|
167 |
+
log.Printf("Buffer overflow: current size %d exceeds max size %d", totalBufferSize, maxBufferSize)
|
168 |
+
return fmt.Errorf("buffer overflow: exceeded maximum buffer size of %d bytes", maxBufferSize)
|
169 |
+
}
|
170 |
+
|
171 |
+
if !strings.HasPrefix(line, "data: ") {
|
172 |
+
continue
|
173 |
+
}
|
174 |
+
|
175 |
+
jsonStr := strings.TrimSpace(strings.TrimPrefix(line, "data: "))
|
176 |
+
if jsonStr == "" || jsonStr == "end" {
|
177 |
+
continue
|
178 |
+
}
|
179 |
+
|
180 |
+
var sseData SSEData
|
181 |
+
if err := sonic.UnmarshalString(jsonStr, &sseData); err != nil {
|
182 |
+
log.Printf("Error unmarshaling SSE data: %v", err)
|
183 |
+
continue
|
184 |
+
}
|
185 |
+
|
186 |
+
log.Printf("Received SSE data: %+v", sseData)
|
187 |
+
|
188 |
+
messageCount++
|
189 |
+
|
190 |
+
if err := processMessage(writer, w, sseData, chatId, fingerprint, now, &completionBuilder, req); err != nil {
|
191 |
+
log.Printf("Failed to process message: %v", err)
|
192 |
+
return err
|
193 |
+
}
|
194 |
+
|
195 |
+
// 定期刷新缓冲区
|
196 |
+
if messageCount >= flushThreshold {
|
197 |
+
if err := flushWriter(writer, w); err != nil {
|
198 |
+
return fmt.Errorf("flush error: %w", err)
|
199 |
+
}
|
200 |
+
messageCount = 0
|
201 |
+
}
|
202 |
+
|
203 |
+
// 检查是否结束
|
204 |
+
if sseData.Type == "QuotaMetadata" {
|
205 |
+
if err := sendFinishSignal(writer, w); err != nil {
|
206 |
+
return fmt.Errorf("finish signal error: %w", err)
|
207 |
+
}
|
208 |
+
log.Printf("Stream completed successfully")
|
209 |
+
return nil
|
210 |
+
}
|
211 |
+
}
|
212 |
+
}
|
213 |
+
|
214 |
+
// processMessage 处理单个消息
|
215 |
+
func processMessage(writer *bufio.Writer, w io.Writer, sseData SSEData, chatId, fingerprint string, now int64, completionBuilder *strings.Builder, req openai.ChatCompletionRequest) error {
|
216 |
+
switch sseData.Type {
|
217 |
+
case "Content":
|
218 |
+
completionBuilder.WriteString(sseData.Content)
|
219 |
+
sseMsg := createStreamMessage(chatId, now, req, fingerprint, sseData.Content, "")
|
220 |
+
return sendMessage(writer, w, sseMsg)
|
221 |
+
|
222 |
+
case "QuotaMetadata":
|
223 |
+
var spentAmount float64
|
224 |
+
if sseData.Spent != nil {
|
225 |
+
if amount, err := strconv.ParseFloat(sseData.Spent.Amount, 64); err == nil {
|
226 |
+
spentAmount = amount
|
227 |
+
} else {
|
228 |
+
log.Printf("Warning: failed to parse spent amount '%s': %v", sseData.Spent.Amount, err)
|
229 |
+
}
|
230 |
+
}
|
231 |
+
|
232 |
+
usage := utils.CalculateJetbrainsUsage(completionBuilder.String(), int(math.Round(spentAmount)))
|
233 |
+
sseMsg := createStreamMessage(chatId, now, req, fingerprint, "", "")
|
234 |
+
sseMsg.Choices[0].FinishReason = openai.FinishReasonStop
|
235 |
+
sseMsg.Usage = &usage
|
236 |
+
return sendMessage(writer, w, sseMsg)
|
237 |
+
|
238 |
+
default:
|
239 |
+
// 忽略其他类型的消息
|
240 |
+
log.Printf("Ignoring message type: %s", sseData.Type)
|
241 |
+
return nil
|
242 |
+
}
|
243 |
+
}
|
244 |
+
|
245 |
+
// createStreamMessage 创建流式消息
|
246 |
+
func createStreamMessage(chatId string, now int64, req openai.ChatCompletionRequest, fingerPrint string, content string, reasoningContent string) openai.ChatCompletionStreamResponse {
|
247 |
+
choice := openai.ChatCompletionStreamChoice{
|
248 |
+
Index: 0,
|
249 |
+
Delta: openai.ChatCompletionStreamChoiceDelta{
|
250 |
+
Role: openai.ChatMessageRoleAssistant,
|
251 |
+
Content: content,
|
252 |
+
ReasoningContent: reasoningContent,
|
253 |
+
},
|
254 |
+
ContentFilterResults: openai.ContentFilterResults{},
|
255 |
+
FinishReason: openai.FinishReasonNull,
|
256 |
+
}
|
257 |
+
|
258 |
+
return openai.ChatCompletionStreamResponse{
|
259 |
+
ID: "chatcmpl-" + chatId,
|
260 |
+
Object: sseObject,
|
261 |
+
Created: now,
|
262 |
+
Model: req.Model,
|
263 |
+
Choices: []openai.ChatCompletionStreamChoice{choice},
|
264 |
+
SystemFingerprint: fingerPrint,
|
265 |
+
}
|
266 |
+
}
|
267 |
+
|
268 |
+
// createMessage 创建非流式消息响应
|
269 |
+
func createMessage(chatId string, now int64, req openai.ChatCompletionRequest, usage openai.Usage, content string, fp string) openai.ChatCompletionResponse {
|
270 |
+
choice := openai.ChatCompletionChoice{
|
271 |
+
Index: 0,
|
272 |
+
Message: openai.ChatCompletionMessage{
|
273 |
+
Role: openai.ChatMessageRoleAssistant,
|
274 |
+
Content: content,
|
275 |
+
},
|
276 |
+
FinishReason: openai.FinishReasonStop,
|
277 |
+
}
|
278 |
+
|
279 |
+
return openai.ChatCompletionResponse{
|
280 |
+
ID: "chatcmpl-" + chatId,
|
281 |
+
Object: completionsObject,
|
282 |
+
Created: now,
|
283 |
+
Model: req.Model,
|
284 |
+
Choices: []openai.ChatCompletionChoice{choice},
|
285 |
+
SystemFingerprint: fp,
|
286 |
+
Usage: usage,
|
287 |
+
}
|
288 |
+
}
|
289 |
+
|
290 |
+
// sendMessage 发送消息到客户端
|
291 |
+
func sendMessage(writer *bufio.Writer, w io.Writer, sseMsg openai.ChatCompletionStreamResponse) error {
|
292 |
+
sendLine, err := sonic.MarshalString(sseMsg)
|
293 |
+
if err != nil {
|
294 |
+
return fmt.Errorf("marshal error: %w", err)
|
295 |
+
}
|
296 |
+
|
297 |
+
outputMsg := fmt.Sprintf("data: %s\n\n", sendLine)
|
298 |
+
if _, err := writer.WriteString(outputMsg); err != nil {
|
299 |
+
return fmt.Errorf("write error: %w", err)
|
300 |
+
}
|
301 |
+
|
302 |
+
return flushWriter(writer, w)
|
303 |
+
}
|
304 |
+
|
305 |
+
// sendHeartbeat 发送心跳包
|
306 |
+
func sendHeartbeat(writer *bufio.Writer, w io.Writer) error {
|
307 |
+
if _, err := writer.WriteString(": keepalive\n\n"); err != nil {
|
308 |
+
return fmt.Errorf("heartbeat write error: %w", err)
|
309 |
+
}
|
310 |
+
return flushWriter(writer, w)
|
311 |
+
}
|
312 |
+
|
313 |
+
// sendFinishSignal 发送结束信号
|
314 |
+
func sendFinishSignal(writer *bufio.Writer, w io.Writer) error {
|
315 |
+
finishMsg := fmt.Sprintf("data: %s\n\n", sseFinish)
|
316 |
+
if _, err := writer.WriteString(finishMsg); err != nil {
|
317 |
+
return fmt.Errorf("write finish signal error: %w", err)
|
318 |
+
}
|
319 |
+
return flushWriter(writer, w)
|
320 |
+
}
|
321 |
+
|
322 |
+
// flushWriter 刷新写入器
|
323 |
+
func flushWriter(writer *bufio.Writer, w io.Writer) error {
|
324 |
+
if err := writer.Flush(); err != nil {
|
325 |
+
return fmt.Errorf("flush error: %w", err)
|
326 |
+
}
|
327 |
+
if f, ok := w.(http.Flusher); ok {
|
328 |
+
f.Flush()
|
329 |
+
}
|
330 |
+
return nil
|
331 |
+
}
|
internal/middleware/auth.go
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package middleware
|
2 |
+
|
3 |
+
import (
|
4 |
+
"github.com/labstack/echo"
|
5 |
+
"jetbrains-ai-proxy/internal/config"
|
6 |
+
"log"
|
7 |
+
"net/http"
|
8 |
+
"strings"
|
9 |
+
)
|
10 |
+
|
11 |
+
func BearerAuth() echo.MiddlewareFunc {
|
12 |
+
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
13 |
+
return func(c echo.Context) error {
|
14 |
+
// 获取Authorization header
|
15 |
+
auth := c.Request().Header.Get("Authorization")
|
16 |
+
|
17 |
+
if auth == "" || !strings.HasPrefix(auth, "Bearer ") {
|
18 |
+
return echo.NewHTTPError(http.StatusUnauthorized, "invalid authorization header")
|
19 |
+
}
|
20 |
+
|
21 |
+
token := strings.TrimPrefix(auth, "Bearer ")
|
22 |
+
|
23 |
+
if token != config.JetbrainsAiConfig.BearerToken || token == "" {
|
24 |
+
log.Printf("invalid token: %s", token)
|
25 |
+
return echo.NewHTTPError(http.StatusUnauthorized, "invalid token")
|
26 |
+
}
|
27 |
+
|
28 |
+
return next(c)
|
29 |
+
}
|
30 |
+
}
|
31 |
+
}
|
internal/types/jetbrains.go
ADDED
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package types
|
2 |
+
|
3 |
+
import (
|
4 |
+
"encoding/json"
|
5 |
+
"fmt"
|
6 |
+
"github.com/sashabaranov/go-openai"
|
7 |
+
)
|
8 |
+
|
9 |
+
const (
|
10 |
+
ChatStreamV7 = "https://api.jetbrains.ai/user/v5/llm/chat/stream/v7"
|
11 |
+
PROMPT = "ij.chat.request.new-chat"
|
12 |
+
JwtTokenKey = "grazie-authenticate-jwt"
|
13 |
+
)
|
14 |
+
|
15 |
+
var modelMap = map[string]OpenAIModel{
|
16 |
+
"gpt-4o": {Object: "model", OwnedBy: "openai", Profile: "openai-gpt-4o"},
|
17 |
+
"o1": {Object: "model", OwnedBy: "openai", Profile: "openai-o1"},
|
18 |
+
"o3": {Object: "model", OwnedBy: "openai", Profile: "openai-o3"},
|
19 |
+
"o3-mini": {Object: "model", OwnedBy: "openai", Profile: "openai-o3-mini"},
|
20 |
+
"o4-mini": {Object: "model", OwnedBy: "openai", Profile: "openai-o4-mini"},
|
21 |
+
"gpt4.1": {Object: "model", OwnedBy: "openai", Profile: "openai-gpt4.1"},
|
22 |
+
"gpt4.1-mini": {Object: "model", OwnedBy: "openai", Profile: "openai-gpt4.1-mini"},
|
23 |
+
"gpt4.1-nano": {Object: "model", OwnedBy: "openai", Profile: "openai-gpt4.1-nano"},
|
24 |
+
|
25 |
+
"gemini-pro-2.5": {Object: "model", OwnedBy: "google", Profile: "google-chat-gemini-pro-2.5"},
|
26 |
+
"gemini-flash-2.0": {Object: "model", OwnedBy: "google", Profile: "google-chat-gemini-flash-2.0"},
|
27 |
+
"gemini-flash-2.5": {Object: "model", OwnedBy: "google", Profile: "google-chat-gemini-flash-2.5"},
|
28 |
+
|
29 |
+
"claude-3.5-haiku": {Object: "model", OwnedBy: "anthropic", Profile: "anthropic-claude-3.5-haiku"},
|
30 |
+
"claude-3.5-sonnet": {Object: "model", OwnedBy: "anthropic", Profile: "anthropic-claude-3.5-sonnet"},
|
31 |
+
"claude-3.7-sonnet": {Object: "model", OwnedBy: "anthropic", Profile: "anthropic-claude-3.7-sonnet"},
|
32 |
+
"claude-4-sonnet": {Object: "model", OwnedBy: "anthropic", Profile: "anthropic-claude-4-sonnet"},
|
33 |
+
}
|
34 |
+
|
35 |
+
type OpenAIModel struct {
|
36 |
+
ID string `json:"id"`
|
37 |
+
Object string `json:"object"`
|
38 |
+
OwnedBy string `json:"owned_by"`
|
39 |
+
Profile string `json:"profile"`
|
40 |
+
}
|
41 |
+
|
42 |
+
type OpenAIModelList struct {
|
43 |
+
Object string `json:"object"`
|
44 |
+
Data []OpenAIModel `json:"data"`
|
45 |
+
}
|
46 |
+
|
47 |
+
type MessageField struct {
|
48 |
+
Type string `json:"type"`
|
49 |
+
Content string `json:"content,omitempty"`
|
50 |
+
}
|
51 |
+
|
52 |
+
type JetbrainsRequest struct {
|
53 |
+
Prompt string `json:"prompt"`
|
54 |
+
Profile string `json:"profile"`
|
55 |
+
Chat ChatField `json:"chat"`
|
56 |
+
}
|
57 |
+
|
58 |
+
type ChatField struct {
|
59 |
+
MessageField []MessageField `json:"messages"`
|
60 |
+
}
|
61 |
+
|
62 |
+
func ChatGPTToJetbrainsAI(chatReq openai.ChatCompletionRequest) (*JetbrainsRequest, error) {
|
63 |
+
messageFields, err := convertOpenAIMessagesToJetbrains(chatReq.Messages)
|
64 |
+
if err != nil {
|
65 |
+
return nil, fmt.Errorf("failed to convert messages: %w", err)
|
66 |
+
}
|
67 |
+
|
68 |
+
openaiModel, err := GetModelByName(chatReq.Model)
|
69 |
+
if err != nil {
|
70 |
+
return nil, fmt.Errorf("failed to get model: %w", err)
|
71 |
+
}
|
72 |
+
|
73 |
+
mReq := &JetbrainsRequest{
|
74 |
+
Prompt: PROMPT,
|
75 |
+
Profile: openaiModel.Profile,
|
76 |
+
Chat: ChatField{
|
77 |
+
MessageField: messageFields,
|
78 |
+
},
|
79 |
+
}
|
80 |
+
if jsonData, err := json.MarshalIndent(mReq, "", " "); err == nil {
|
81 |
+
fmt.Printf("mReq JSON: %s\n", string(jsonData))
|
82 |
+
}
|
83 |
+
|
84 |
+
return mReq, nil
|
85 |
+
}
|
86 |
+
|
87 |
+
func convertOpenAIMessagesToJetbrains(openaiMessages []openai.ChatCompletionMessage) ([]MessageField, error) {
|
88 |
+
var messageField []MessageField
|
89 |
+
|
90 |
+
for _, msg := range openaiMessages {
|
91 |
+
if msg.Role == "system" {
|
92 |
+
messageField = append(messageField, MessageField{
|
93 |
+
Type: "system_message",
|
94 |
+
Content: msg.Content,
|
95 |
+
})
|
96 |
+
} else if msg.Role == "user" {
|
97 |
+
messageField = append(messageField, MessageField{
|
98 |
+
Type: "user_message",
|
99 |
+
Content: msg.Content,
|
100 |
+
})
|
101 |
+
} else if msg.Role == "assistant" {
|
102 |
+
messageField = append(messageField, MessageField{
|
103 |
+
Type: "assistant_message",
|
104 |
+
Content: msg.Content,
|
105 |
+
})
|
106 |
+
}
|
107 |
+
}
|
108 |
+
return messageField, nil
|
109 |
+
}
|
110 |
+
|
111 |
+
func GetModelByName(modelName string) (OpenAIModel, error) {
|
112 |
+
model, exists := modelMap[modelName]
|
113 |
+
if !exists {
|
114 |
+
return OpenAIModel{}, fmt.Errorf("model '%s' not found", modelName)
|
115 |
+
}
|
116 |
+
return model, nil
|
117 |
+
}
|
118 |
+
|
119 |
+
func GetSupportedModels() OpenAIModelList {
|
120 |
+
var modelSlice []OpenAIModel
|
121 |
+
for id, model := range modelMap {
|
122 |
+
modelWithID := model
|
123 |
+
modelWithID.ID = id
|
124 |
+
modelSlice = append(modelSlice, modelWithID)
|
125 |
+
}
|
126 |
+
|
127 |
+
return OpenAIModelList{
|
128 |
+
Object: "list",
|
129 |
+
Data: modelSlice,
|
130 |
+
}
|
131 |
+
}
|
internal/utils/req_client.go
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package utils
|
2 |
+
|
3 |
+
import (
|
4 |
+
"crypto/tls"
|
5 |
+
"fmt"
|
6 |
+
"github.com/go-resty/resty/v2"
|
7 |
+
"time"
|
8 |
+
)
|
9 |
+
|
10 |
+
var (
|
11 |
+
RestySSEClient = resty.New().
|
12 |
+
SetTimeout(1 * time.Minute).
|
13 |
+
SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true}).
|
14 |
+
SetDoNotParseResponse(true).
|
15 |
+
SetHeaders(map[string]string{
|
16 |
+
"Content-Type": "application/json",
|
17 |
+
}).
|
18 |
+
OnAfterResponse(func(c *resty.Client, resp *resty.Response) error {
|
19 |
+
if resp.StatusCode() != 200 {
|
20 |
+
return fmt.Errorf("Jetbrains API error: status %d, body: %s",
|
21 |
+
resp.StatusCode(), resp.String())
|
22 |
+
}
|
23 |
+
return nil
|
24 |
+
})
|
25 |
+
)
|
internal/utils/string.go
ADDED
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package utils
|
2 |
+
|
3 |
+
import (
|
4 |
+
"math/rand"
|
5 |
+
"time"
|
6 |
+
)
|
7 |
+
|
8 |
+
var randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
9 |
+
|
10 |
+
func RandStringUsingMathRand(n int) string {
|
11 |
+
var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
|
12 |
+
|
13 |
+
result := make([]rune, n)
|
14 |
+
for i := 0; i < n; i++ {
|
15 |
+
result[i] = letters[randSource.Intn(len(letters))]
|
16 |
+
}
|
17 |
+
return string(result)
|
18 |
+
}
|
internal/utils/token.go
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package utils
|
2 |
+
|
3 |
+
import (
|
4 |
+
"fmt"
|
5 |
+
"github.com/pkoukk/tiktoken-go"
|
6 |
+
"github.com/sashabaranov/go-openai"
|
7 |
+
)
|
8 |
+
|
9 |
+
func CalculateTokens(text string) int {
|
10 |
+
encoding := "cl100k_base"
|
11 |
+
tke, err := tiktoken.GetEncoding(encoding)
|
12 |
+
if err != nil {
|
13 |
+
err = fmt.Errorf("getEncoding: %v", err)
|
14 |
+
return 0
|
15 |
+
}
|
16 |
+
token := tke.Encode(text, nil, nil)
|
17 |
+
return len(token)
|
18 |
+
}
|
19 |
+
|
20 |
+
func CalculateJetbrainsUsage(completionText string, spent int) openai.Usage {
|
21 |
+
completionTokens := CalculateTokens(completionText)
|
22 |
+
return openai.Usage{
|
23 |
+
PromptTokens: spent - completionTokens,
|
24 |
+
CompletionTokens: spent - completionTokens,
|
25 |
+
TotalTokens: spent,
|
26 |
+
}
|
27 |
+
}
|
main.go
ADDED
@@ -0,0 +1,278 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
package main
|
2 |
+
|
3 |
+
import (
|
4 |
+
"errors"
|
5 |
+
"flag"
|
6 |
+
"fmt"
|
7 |
+
"github.com/labstack/echo"
|
8 |
+
"github.com/labstack/echo/middleware"
|
9 |
+
"jetbrains-ai-proxy/internal/apiserver"
|
10 |
+
"jetbrains-ai-proxy/internal/config"
|
11 |
+
"jetbrains-ai-proxy/internal/jetbrains"
|
12 |
+
"log"
|
13 |
+
"net/http"
|
14 |
+
"os"
|
15 |
+
"os/signal"
|
16 |
+
"syscall"
|
17 |
+
)
|
18 |
+
|
19 |
+
func main() {
|
20 |
+
// 定义命令行参数
|
21 |
+
configFile := flag.String("config", "", "配置文件路径")
|
22 |
+
port := flag.Int("p", 0, "服务器监听端口 (覆盖配置文件)")
|
23 |
+
host := flag.String("h", "", "服务器监听地址 (覆盖配置文件)")
|
24 |
+
jwtTokens := flag.String("c", "", "JWT Tokens值,多个token用逗号分隔 (覆盖配置文件)")
|
25 |
+
bearerToken := flag.String("k", "", "Bearer Token值 (覆盖配置文件)")
|
26 |
+
loadBalanceStrategy := flag.String("s", "", "负载均衡策略: round_robin 或 random (覆盖配置文件)")
|
27 |
+
generateConfig := flag.Bool("generate-config", false, "生成示例配置文件")
|
28 |
+
printConfig := flag.Bool("print-config", false, "打印当前配置信息")
|
29 |
+
|
30 |
+
flag.Usage = func() {
|
31 |
+
fmt.Printf("用法: %s [选项]\n\n", flag.CommandLine.Name())
|
32 |
+
fmt.Println("选项:")
|
33 |
+
flag.PrintDefaults()
|
34 |
+
fmt.Println("\n配置优先级 (从高到低):")
|
35 |
+
fmt.Println(" 1. 命令行参数")
|
36 |
+
fmt.Println(" 2. 环境变量")
|
37 |
+
fmt.Println(" 3. 配置文件")
|
38 |
+
fmt.Println(" 4. 默认值")
|
39 |
+
fmt.Println("\n配置方式:")
|
40 |
+
fmt.Println(" 方式1 - 使用配置文件:")
|
41 |
+
fmt.Println(" ./jetbrains-ai-proxy --generate-config # 生成示例配置")
|
42 |
+
fmt.Println(" # 编辑 config/config.json")
|
43 |
+
fmt.Println(" ./jetbrains-ai-proxy")
|
44 |
+
fmt.Println("")
|
45 |
+
fmt.Println(" 方式2 - 使用环境变量:")
|
46 |
+
fmt.Println(" export JWT_TOKENS=\"jwt1,jwt2,jwt3\"")
|
47 |
+
fmt.Println(" export BEARER_TOKEN=\"your_token\"")
|
48 |
+
fmt.Println(" ./jetbrains-ai-proxy")
|
49 |
+
fmt.Println("")
|
50 |
+
fmt.Println(" 方式3 - 使用命令行参数:")
|
51 |
+
fmt.Println(" ./jetbrains-ai-proxy -c \"jwt1,jwt2,jwt3\" -k \"bearer_token\"")
|
52 |
+
fmt.Println("")
|
53 |
+
fmt.Println("负载均衡策略:")
|
54 |
+
fmt.Println(" round_robin: 轮询策略(默认)")
|
55 |
+
fmt.Println(" random: 随机策略")
|
56 |
+
}
|
57 |
+
|
58 |
+
flag.Parse()
|
59 |
+
|
60 |
+
// 处理特殊命令
|
61 |
+
if *generateConfig {
|
62 |
+
if err := generateExampleConfig(); err != nil {
|
63 |
+
log.Fatalf("Failed to generate config: %v", err)
|
64 |
+
}
|
65 |
+
return
|
66 |
+
}
|
67 |
+
|
68 |
+
// 获取配置管理器
|
69 |
+
configManager := config.GetGlobalConfig()
|
70 |
+
|
71 |
+
// 如果指定了配置文件,设置环境变量
|
72 |
+
if *configFile != "" {
|
73 |
+
os.Setenv("CONFIG_FILE", *configFile)
|
74 |
+
}
|
75 |
+
|
76 |
+
// 加载配置
|
77 |
+
if err := configManager.LoadConfig(); err != nil {
|
78 |
+
log.Printf("Warning: %v", err)
|
79 |
+
log.Println("Continuing with command line arguments and environment variables...")
|
80 |
+
}
|
81 |
+
|
82 |
+
// 应用命令行参数覆盖
|
83 |
+
applyCommandLineOverrides(configManager, port, host, jwtTokens, bearerToken, loadBalanceStrategy)
|
84 |
+
|
85 |
+
// 打印配置信息
|
86 |
+
if *printConfig {
|
87 |
+
configManager.PrintConfig()
|
88 |
+
return
|
89 |
+
}
|
90 |
+
|
91 |
+
// 验证配置
|
92 |
+
if !configManager.HasJWTTokens() {
|
93 |
+
log.Fatal("No JWT tokens configured. Use --generate-config to create example configuration.")
|
94 |
+
}
|
95 |
+
|
96 |
+
cfg := configManager.GetConfig()
|
97 |
+
if cfg.BearerToken == "" {
|
98 |
+
log.Fatal("Bearer token is required. Please configure it in config file, environment variable, or command line.")
|
99 |
+
}
|
100 |
+
|
101 |
+
// 初始化JWT负载均衡器
|
102 |
+
if err := jetbrains.InitializeFromConfig(); err != nil {
|
103 |
+
log.Fatalf("Failed to initialize JWT balancer: %v", err)
|
104 |
+
}
|
105 |
+
|
106 |
+
// 设置优雅关闭
|
107 |
+
setupGracefulShutdown()
|
108 |
+
|
109 |
+
// 启动配置文件监控
|
110 |
+
discovery := config.NewConfigDiscovery(configManager)
|
111 |
+
discovery.WatchConfig()
|
112 |
+
|
113 |
+
// 创建Echo实例
|
114 |
+
e := echo.New()
|
115 |
+
e.Use(middleware.Logger())
|
116 |
+
e.Use(middleware.Recover())
|
117 |
+
|
118 |
+
// 添加管理端点
|
119 |
+
setupManagementEndpoints(e, configManager)
|
120 |
+
|
121 |
+
// 注册API路由
|
122 |
+
apiserver.RegisterRoutes(e)
|
123 |
+
|
124 |
+
// 启动服务器
|
125 |
+
addr := fmt.Sprintf("%s:%d", cfg.ServerHost, cfg.ServerPort)
|
126 |
+
log.Printf("Server starting on %s", addr)
|
127 |
+
configManager.PrintConfig()
|
128 |
+
|
129 |
+
if err := e.Start(addr); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
130 |
+
log.Fatalf("start server error: %v", err)
|
131 |
+
}
|
132 |
+
}
|
133 |
+
|
134 |
+
// generateExampleConfig 生成示例配置
|
135 |
+
func generateExampleConfig() error {
|
136 |
+
manager := config.NewManager()
|
137 |
+
|
138 |
+
// 生成JSON配置文件
|
139 |
+
if err := manager.GenerateExampleConfig("config/config.json"); err != nil {
|
140 |
+
return fmt.Errorf("failed to generate JSON config: %v", err)
|
141 |
+
}
|
142 |
+
|
143 |
+
// 生成.env示例文件
|
144 |
+
config.NewConfigDiscovery(manager)
|
145 |
+
envContent := `# JetBrains AI Proxy Configuration
|
146 |
+
# Copy this file to .env and fill in your actual values
|
147 |
+
|
148 |
+
# Multiple JWT tokens (comma-separated)
|
149 |
+
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...
|
150 |
+
|
151 |
+
# Bearer token for API authentication
|
152 |
+
BEARER_TOKEN=your_bearer_token_here
|
153 |
+
|
154 |
+
# Load balancing strategy: round_robin or random
|
155 |
+
LOAD_BALANCE_STRATEGY=round_robin
|
156 |
+
|
157 |
+
# Server configuration
|
158 |
+
SERVER_HOST=0.0.0.0
|
159 |
+
SERVER_PORT=8080
|
160 |
+
`
|
161 |
+
|
162 |
+
if err := os.WriteFile(".env.example", []byte(envContent), 0644); err != nil {
|
163 |
+
return fmt.Errorf("failed to generate .env example: %v", err)
|
164 |
+
}
|
165 |
+
|
166 |
+
fmt.Println("✅ Example configuration files generated:")
|
167 |
+
fmt.Println(" 📄 config/config.json - JSON configuration file")
|
168 |
+
fmt.Println(" 📄 .env.example - Environment variables example")
|
169 |
+
fmt.Println("")
|
170 |
+
fmt.Println("📝 Next steps:")
|
171 |
+
fmt.Println(" 1. Edit config/config.json with your JWT tokens")
|
172 |
+
fmt.Println(" 2. Or copy .env.example to .env and edit it")
|
173 |
+
fmt.Println(" 3. Run: ./jetbrains-ai-proxy")
|
174 |
+
|
175 |
+
return nil
|
176 |
+
}
|
177 |
+
|
178 |
+
// applyCommandLineOverrides 应用命令行参数覆盖
|
179 |
+
func applyCommandLineOverrides(manager *config.Manager, port *int, host, jwtTokens, bearerToken, strategy *string) {
|
180 |
+
if *jwtTokens != "" {
|
181 |
+
manager.SetJWTTokens(*jwtTokens)
|
182 |
+
log.Printf("JWT tokens overridden by command line")
|
183 |
+
}
|
184 |
+
|
185 |
+
if *bearerToken != "" {
|
186 |
+
manager.SetBearerToken(*bearerToken)
|
187 |
+
log.Printf("Bearer token overridden by command line")
|
188 |
+
}
|
189 |
+
|
190 |
+
if *strategy != "" {
|
191 |
+
manager.SetLoadBalanceStrategy(*strategy)
|
192 |
+
log.Printf("Load balance strategy overridden by command line: %s", *strategy)
|
193 |
+
}
|
194 |
+
|
195 |
+
// 覆盖服务器配置
|
196 |
+
cfg := manager.GetConfig()
|
197 |
+
if *port > 0 {
|
198 |
+
cfg.ServerPort = *port
|
199 |
+
log.Printf("Server port overridden by command line: %d", *port)
|
200 |
+
}
|
201 |
+
|
202 |
+
if *host != "" {
|
203 |
+
cfg.ServerHost = *host
|
204 |
+
log.Printf("Server host overridden by command line: %s", *host)
|
205 |
+
}
|
206 |
+
}
|
207 |
+
|
208 |
+
// setupManagementEndpoints 设置管理端点
|
209 |
+
func setupManagementEndpoints(e *echo.Echo, manager *config.Manager) {
|
210 |
+
// 健康检查端点
|
211 |
+
e.GET("/health", func(c echo.Context) error {
|
212 |
+
healthy, total := jetbrains.GetBalancerStats()
|
213 |
+
cfg := manager.GetConfig()
|
214 |
+
|
215 |
+
return c.JSON(http.StatusOK, map[string]interface{}{
|
216 |
+
"status": "ok",
|
217 |
+
"healthy_tokens": healthy,
|
218 |
+
"total_tokens": total,
|
219 |
+
"strategy": cfg.LoadBalanceStrategy,
|
220 |
+
"server_info": map[string]interface{}{
|
221 |
+
"host": cfg.ServerHost,
|
222 |
+
"port": cfg.ServerPort,
|
223 |
+
},
|
224 |
+
})
|
225 |
+
})
|
226 |
+
|
227 |
+
// 配置信息端点
|
228 |
+
e.GET("/config", func(c echo.Context) error {
|
229 |
+
discovery := config.NewConfigDiscovery(manager)
|
230 |
+
summary := discovery.GetConfigSummary()
|
231 |
+
return c.JSON(http.StatusOK, summary)
|
232 |
+
})
|
233 |
+
|
234 |
+
// 重载配置端点
|
235 |
+
e.POST("/reload", func(c echo.Context) error {
|
236 |
+
if err := jetbrains.ReloadConfig(); err != nil {
|
237 |
+
return c.JSON(http.StatusInternalServerError, map[string]interface{}{
|
238 |
+
"error": err.Error(),
|
239 |
+
})
|
240 |
+
}
|
241 |
+
|
242 |
+
return c.JSON(http.StatusOK, map[string]interface{}{
|
243 |
+
"message": "Configuration reloaded successfully",
|
244 |
+
})
|
245 |
+
})
|
246 |
+
|
247 |
+
// 负载均衡器统计端点
|
248 |
+
e.GET("/stats", func(c echo.Context) error {
|
249 |
+
healthy, total := jetbrains.GetBalancerStats()
|
250 |
+
cfg := manager.GetConfig()
|
251 |
+
|
252 |
+
return c.JSON(http.StatusOK, map[string]interface{}{
|
253 |
+
"balancer": map[string]interface{}{
|
254 |
+
"healthy_tokens": healthy,
|
255 |
+
"total_tokens": total,
|
256 |
+
"strategy": cfg.LoadBalanceStrategy,
|
257 |
+
},
|
258 |
+
"config": map[string]interface{}{
|
259 |
+
"health_check_interval": cfg.HealthCheckInterval.String(),
|
260 |
+
"server_host": cfg.ServerHost,
|
261 |
+
"server_port": cfg.ServerPort,
|
262 |
+
},
|
263 |
+
})
|
264 |
+
})
|
265 |
+
}
|
266 |
+
|
267 |
+
// setupGracefulShutdown 设置优雅关闭
|
268 |
+
func setupGracefulShutdown() {
|
269 |
+
c := make(chan os.Signal, 1)
|
270 |
+
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
|
271 |
+
|
272 |
+
go func() {
|
273 |
+
<-c
|
274 |
+
log.Println("Shutting down gracefully...")
|
275 |
+
jetbrains.StopBalancer()
|
276 |
+
os.Exit(0)
|
277 |
+
}()
|
278 |
+
}
|
start.sh
ADDED
@@ -0,0 +1,200 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/bin/bash
|
2 |
+
|
3 |
+
# JetBrains AI Proxy 启动脚本
|
4 |
+
# 支持自动配置发现和多种配置方式
|
5 |
+
|
6 |
+
set -e
|
7 |
+
|
8 |
+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
9 |
+
cd "$SCRIPT_DIR"
|
10 |
+
|
11 |
+
# 颜色定义
|
12 |
+
RED='\033[0;31m'
|
13 |
+
GREEN='\033[0;32m'
|
14 |
+
YELLOW='\033[1;33m'
|
15 |
+
BLUE='\033[0;34m'
|
16 |
+
NC='\033[0m' # No Color
|
17 |
+
|
18 |
+
# 打印带颜色的消息
|
19 |
+
print_info() {
|
20 |
+
echo -e "${BLUE}ℹ️ $1${NC}"
|
21 |
+
}
|
22 |
+
|
23 |
+
print_success() {
|
24 |
+
echo -e "${GREEN}✅ $1${NC}"
|
25 |
+
}
|
26 |
+
|
27 |
+
print_warning() {
|
28 |
+
echo -e "${YELLOW}⚠️ $1${NC}"
|
29 |
+
}
|
30 |
+
|
31 |
+
print_error() {
|
32 |
+
echo -e "${RED}❌ $1${NC}"
|
33 |
+
}
|
34 |
+
|
35 |
+
# 检查可执行文件
|
36 |
+
check_executable() {
|
37 |
+
if [ ! -f "./jetbrains-ai-proxy" ]; then
|
38 |
+
print_error "可执行文件 './jetbrains-ai-proxy' 不存在"
|
39 |
+
print_info "请先编译项目: go build -o jetbrains-ai-proxy"
|
40 |
+
exit 1
|
41 |
+
fi
|
42 |
+
|
43 |
+
if [ ! -x "./jetbrains-ai-proxy" ]; then
|
44 |
+
print_info "设置可执行权限..."
|
45 |
+
chmod +x ./jetbrains-ai-proxy
|
46 |
+
fi
|
47 |
+
}
|
48 |
+
|
49 |
+
# 检查配置
|
50 |
+
check_configuration() {
|
51 |
+
print_info "检查配置..."
|
52 |
+
|
53 |
+
# 检查是否存在配置文件
|
54 |
+
config_files=(
|
55 |
+
"config.json"
|
56 |
+
"config/config.json"
|
57 |
+
"configs/config.json"
|
58 |
+
".config/jetbrains-ai-proxy.json"
|
59 |
+
)
|
60 |
+
|
61 |
+
config_found=false
|
62 |
+
for config_file in "${config_files[@]}"; do
|
63 |
+
if [ -f "$config_file" ]; then
|
64 |
+
print_success "找到配置文件: $config_file"
|
65 |
+
config_found=true
|
66 |
+
break
|
67 |
+
fi
|
68 |
+
done
|
69 |
+
|
70 |
+
# 检查环境变量
|
71 |
+
env_configured=false
|
72 |
+
if [ -n "$JWT_TOKENS" ] || [ -n "$JWT_TOKEN" ]; then
|
73 |
+
if [ -n "$BEARER_TOKEN" ]; then
|
74 |
+
print_success "检测到环境变量配置"
|
75 |
+
env_configured=true
|
76 |
+
else
|
77 |
+
print_warning "检测到JWT tokens但缺少BEARER_TOKEN环境变量"
|
78 |
+
fi
|
79 |
+
fi
|
80 |
+
|
81 |
+
# 检查.env文件
|
82 |
+
if [ -f ".env" ]; then
|
83 |
+
print_success "找到 .env 文件"
|
84 |
+
env_configured=true
|
85 |
+
fi
|
86 |
+
|
87 |
+
# 如果没有找到任何配置,生成示例配置
|
88 |
+
if [ "$config_found" = false ] && [ "$env_configured" = false ]; then
|
89 |
+
print_warning "未找到配置文件或环境变量配置"
|
90 |
+
print_info "生成示例配置文件..."
|
91 |
+
|
92 |
+
if ./jetbrains-ai-proxy --generate-config; then
|
93 |
+
print_success "示例配置文件已生成"
|
94 |
+
print_info "请编辑 config/config.json 或 .env.example 文件"
|
95 |
+
print_info "然后重新运行此脚本"
|
96 |
+
exit 0
|
97 |
+
else
|
98 |
+
print_error "生成示例配置失败"
|
99 |
+
exit 1
|
100 |
+
fi
|
101 |
+
fi
|
102 |
+
}
|
103 |
+
|
104 |
+
# 显示配置信息
|
105 |
+
show_config() {
|
106 |
+
print_info "当前配置信息:"
|
107 |
+
./jetbrains-ai-proxy --print-config
|
108 |
+
}
|
109 |
+
|
110 |
+
# 启动服务
|
111 |
+
start_service() {
|
112 |
+
print_info "启动 JetBrains AI Proxy..."
|
113 |
+
|
114 |
+
# 如果有命令行参数,直接传递
|
115 |
+
if [ $# -gt 0 ]; then
|
116 |
+
print_info "使用命令行参数: $*"
|
117 |
+
exec ./jetbrains-ai-proxy "$@"
|
118 |
+
else
|
119 |
+
# 使用配置文件启动
|
120 |
+
exec ./jetbrains-ai-proxy
|
121 |
+
fi
|
122 |
+
}
|
123 |
+
|
124 |
+
# 显示帮助信息
|
125 |
+
show_help() {
|
126 |
+
echo "JetBrains AI Proxy 启动脚本"
|
127 |
+
echo ""
|
128 |
+
echo "用法:"
|
129 |
+
echo " $0 # 使用配置文件启动"
|
130 |
+
echo " $0 [options] # 使用命令行参数启动"
|
131 |
+
echo " $0 --help # 显示帮助信息"
|
132 |
+
echo " $0 --config # 显示当前配置"
|
133 |
+
echo " $0 --generate # 生成示例配置文件"
|
134 |
+
echo ""
|
135 |
+
echo "配置方式 (优先级从高到低):"
|
136 |
+
echo " 1. 命令行参数"
|
137 |
+
echo " 2. 环境变量"
|
138 |
+
echo " 3. 配置文件 (config.json, config/config.json 等)"
|
139 |
+
echo " 4. 默认值"
|
140 |
+
echo ""
|
141 |
+
echo "示例:"
|
142 |
+
echo " # 生成配置文件"
|
143 |
+
echo " $0 --generate"
|
144 |
+
echo ""
|
145 |
+
echo " # 使用配置文件启动"
|
146 |
+
echo " $0"
|
147 |
+
echo ""
|
148 |
+
echo " # 使用命令行参数启动"
|
149 |
+
echo " $0 -c \"jwt1,jwt2,jwt3\" -k \"bearer_token\" -s random"
|
150 |
+
echo ""
|
151 |
+
echo " # 使用环境变量启动"
|
152 |
+
echo " export JWT_TOKENS=\"jwt1,jwt2,jwt3\""
|
153 |
+
echo " export BEARER_TOKEN=\"your_token\""
|
154 |
+
echo " $0"
|
155 |
+
echo ""
|
156 |
+
echo "管理端点:"
|
157 |
+
echo " GET /health - 健康检查"
|
158 |
+
echo " GET /config - 配置信息"
|
159 |
+
echo " GET /stats - 统计信息"
|
160 |
+
echo " POST /reload - 重载配置"
|
161 |
+
}
|
162 |
+
|
163 |
+
# 主函数
|
164 |
+
main() {
|
165 |
+
echo "🚀 JetBrains AI Proxy 启动脚本"
|
166 |
+
echo "================================"
|
167 |
+
|
168 |
+
# 处理特殊参数
|
169 |
+
case "${1:-}" in
|
170 |
+
--help|-h)
|
171 |
+
show_help
|
172 |
+
exit 0
|
173 |
+
;;
|
174 |
+
--config)
|
175 |
+
check_executable
|
176 |
+
show_config
|
177 |
+
exit 0
|
178 |
+
;;
|
179 |
+
--generate)
|
180 |
+
check_executable
|
181 |
+
./jetbrains-ai-proxy --generate-config
|
182 |
+
exit 0
|
183 |
+
;;
|
184 |
+
esac
|
185 |
+
|
186 |
+
# 检查可执行文件
|
187 |
+
check_executable
|
188 |
+
|
189 |
+
# 检查配置
|
190 |
+
check_configuration
|
191 |
+
|
192 |
+
# 启动服务
|
193 |
+
start_service "$@"
|
194 |
+
}
|
195 |
+
|
196 |
+
# 捕获中断信号
|
197 |
+
trap 'print_info "正在停止服务..."; exit 0' INT TERM
|
198 |
+
|
199 |
+
# 运行主函数
|
200 |
+
main "$@"
|