github-actions[bot] commited on
Commit
6fefda3
·
1 Parent(s): daebe81

Update from GitHub Actions

Browse files
.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 "$@"