Skip to content

feat(gateway): P3-B streaming SSE cross-protocol conversion (v0.5.0) #85

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
argszero merged 1 commit into from
Aug 18, 2026
Merged

feat(gateway): P3-B streaming SSE cross-protocol conversion (v0.5.0) #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 25 additions & 1 deletion Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 4 additions & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "aitokenpool"
version = "0.4.1"
version = "0.5.0"
edition = "2021"
description = "AI Token 共享池 — 企业 key 池 + 公共共享市场"
license = "MIT"
Expand All @@ -20,6 +20,9 @@ reqwest = { version = "0.12", features = ["rustls-tls", "json", "stream"] }

# SSE 流式（bytes_stream + Body::from_stream）
futures-util = "0.3"
# SSE 流式协议转换（P3-B：跨协议流事件互转）
async-stream = "0.3"
bytes = "1"

# 上游 key 加密（AES-256-GCM，RustCrypto）
aes-gcm = "0.10"
Expand Down
7 changes: 4 additions & 3 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ AITokenPool 是一个开源的 **AI Token 共享平台**：企业版（内部 ke

## 状态

后端 + 前端已全部实现，当前版本 **v0.4.1**（94/94 测试全绿）：
后端 + 前端已全部实现，当前版本 **v0.5.0**（102/102 测试全绿）：

- ✅ **P0-A（v0.2.0）**：后端骨架 + TOML 配置（`Config::validate`：points_per_unit>0 / plan→provider 存在 / endpoints≥1 / protocol 枚举）+ SQLite 数据层（幂等迁移 v4 + demo/admin/ops 种子）+ 认证（argon2 + Bearer API Key）+ API Key 端点
- ✅ **P0-B（v0.2.1）**：网关转发（OpenAI Chat Completions + Anthropic Messages）+ 路由故障转移（粘性 / 健康冷却 / 3 次切换上限）+ 计量账本（点数计算、90/10 分成、事务性 settle）
Expand All @@ -28,6 +28,7 @@ AITokenPool 是一个开源的 **AI Token 共享平台**：企业版（内部 ke
- ✅ **P2-C（v0.3.3）**：部门/成员管理 + 加额审批 + 用量报表（users/models/departments 三组聚合）+ 运营者视图（runtime/credits/users）
- ✅ **P3-A（v0.4.0）**：网关三协议互转——OpenAI Chat / OpenAI Responses / Anthropic 任一端点可调用只暴露其他协议的 plan（自动转换，同协议透传零损耗）+ 新增 `/v1/responses` 端点
- ✅ **P3-A 补充（v0.4.1）**：`GET /v1/models` OpenAI 兼容模型列表（认证可选，带 Bearer 附加 available_keys；`/models` 别名）
- ✅ **P3-B（v0.5.0）**：流式 SSE 跨协议转换——openai/anthropic/responses 任意协议的 `stream:true` 请求可转发到任意协议上游，响应流事件实时互转（openai delta ↔ anthropic content_block_delta ↔ responses output_text.delta，含工具调用/thinking 增量与流内 usage 计量；responses→anthropic 流式暂延后）

`ui/` 已由纯静态原型升级为**对接真实 API**（登录、钱包、市场、共享、交易、设置、管理、运营全部真实数据），由后端 `ServeDir` 静态托管，无需单独部署前端。

Expand Down Expand Up @@ -70,10 +71,10 @@ open http://localhost:8080/ # 浏览器访问

## API 端点（Bearer 认证）

- `GET /healthz` → `{"status":"ok","version":"0.4.1"}`
- `GET /healthz` → `{"status":"ok","version":"0.5.0"}`
- `POST /api/auth/login` → `{api_key}`；`GET /api/me` → `{id,email,name,role}`
- `POST|GET /api/api-keys`（key 脱敏 `atk_live_****xxxx`）；`DELETE /api/api-keys/:id`（撤销）
- `POST /v1/chat/completions` / `POST /anthropic/v1/messages` / `POST /v1/responses`（网关，三协议互转；SSE 流式同协议透传）；`GET /v1/models`（OpenAI 兼容模型列表，认证可选）
- `POST /v1/chat/completions` / `POST /anthropic/v1/messages` / `POST /v1/responses`（网关，三协议互转，非流式 + 流式 SSE 跨协议转换）；`GET /v1/models`（OpenAI 兼容模型列表，认证可选）
- `GET /api/models`（模型市场）
- `POST|GET /api/sharings` + `PATCH /api/sharings/:id`（上架 / 列表 / 暂停 / 恢复 / 删除）
- `GET /api/wallet` / `GET /api/transactions?type=` / `GET /api/dashboard`（钱包 / 交易 / 仪表盘）
Expand Down
209 changes: 153 additions & 56 deletions src/gateway.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,12 +319,16 @@ async fn forward(
))
}

/// SSE 流式转发（P0-C）：请求体带 stream:true 时走此分支。
/// SSE 响应流（P3-B：透传或协议转换后的字节流）
type SseStream = std::pin::Pin<
Box<dyn futures_util::Stream<Item = Result<axum::body::Bytes, std::io::Error>> + Send>,
>;

/// SSE 流式转发（P0-C 透传 + P3-B 跨协议转换）：请求体带 stream:true 时走此分支。
///
/// 流程：余额预检 → 路由选 key（初始连接失败可故障转移，最高 3 次）→ 拿到 200 后
/// 逐块透传上游响应体（data: 行原样转发，保持 event/comment 原始格式）→
/// 流尾解析 usage（openai 最后 chunk 的 usage / anthropic message_delta 的
/// output_tokens + message_start 的 input_tokens）→ 复用 settle_usage 入账 →
/// 逐块转发上游响应体；跨协议时（出站协议 ≠ 入站协议）把上游 SSE 流经 src/sse.rs
/// 转换器转为客户端协议 SSE 流（P3-B）→ 流内提取 usage → 复用 settle_usage 入账 →
/// 记粘性。客户端提前断开 → 响应体被 drop → 上游连接自动中止，不入账。
async fn forward_stream(
st: &AppState,
Expand All @@ -351,19 +355,6 @@ async fn forward_stream(
"该模型暂无可用 key",
));
}
// P3-A：流式只支持同协议透传（SSE 协议转换留 P3-B）。
// 滤掉需要跨协议转换的 key（其 plan 无入站协议端点）——避免随机选 key 抖动。
let keys: Vec<dao::KeyRow> = keys
.into_iter()
.filter(|k| resolve_outbound(&st.cfg, &k.plan, protocol) == Some(protocol))
.collect();
if keys.is_empty() {
// 模型有 key 但均需协议转换 → 明确报错（P3-B 后支持）
return Err(err_json(
StatusCode::BAD_REQUEST,
"流式请求需要协议转换，暂不支持（P3-B 实现）",
));
}

// 选 key 并建立上游连接（此阶段失败可切换；连接成功后不再切换）
for _ in 0..crate::router::MAX_SWITCHES {
Expand All @@ -378,7 +369,14 @@ async fn forward_stream(
.find(|k| k.id == key_id)
.expect("pick 返回的 key 必然在候选集");

let Some(url) = resolve_endpoint(&st.cfg, &key.plan, protocol) else {
// P3-B：定出站协议（同协议优先；跨协议按 anthropic → openai_chat → responses 降级）
let Some(outbound) = resolve_outbound(&st.cfg, &key.plan, protocol) else {
st.router.mark_unhealthy(key_id);
continue;
};
let needs_transform = outbound != protocol;

let Some(url) = resolve_endpoint(&st.cfg, &key.plan, outbound) else {
st.router.mark_unhealthy(key_id);
continue;
};
Expand All @@ -387,21 +385,37 @@ async fn forward_stream(
continue;
};

let resp = if protocol == "anthropic" {
// 跨协议 → 转换请求体（P3-A transform_request 已处理 stream:true）；同协议原样透传
let up_body = if needs_transform {
match serde_json::from_str::<serde_json::Value>(&body) {
Ok(v) => {
let transformed = crate::protocol::transform_request(&v, protocol, outbound);
match serde_json::to_string(&transformed) {
Ok(s) => s,
Err(_) => body.clone(),
}
}
Err(_) => body.clone(),
}
} else {
body.clone()
};

let resp = if outbound == "anthropic" {
st.http
.post(&url)
.header("x-api-key", &plain_key)
.header("content-type", "application/json")
.header("anthropic-version", "2023-06-01")
.body(body.clone())
.body(up_body)
.send()
.await
} else {
st.http
.post(&url)
.header("authorization", format!("Bearer {plain_key}"))
.header("content-type", "application/json")
.body(body.clone())
.body(up_body)
.send()
.await
};
Expand Down Expand Up @@ -429,32 +443,82 @@ async fn forward_stream(
return Ok(passthrough(status, bytes));
}

// 连接成功：构建 SSE 透传流（usage 捕获 + 流尾入账）
// 连接成功：构建 SSE 流（同协议透传 / 跨协议转换），流尾 usage 入账
let key = key.clone();
let st = st.clone();
let model = model.to_string();
let protocol = protocol.to_string();
let capture = std::sync::Arc::new(std::sync::Mutex::new(UsageCapture::new(&protocol)));
let cap_fwd = std::sync::Arc::clone(&capture);
let fwd = resp.bytes_stream().map(move |item| {
if let Ok(bytes) = &item {
if let Ok(mut cap) = cap_fwd.lock() {
cap.push(bytes);
}
let outbound = outbound.to_string();

let (fwd, finalize): (SseStream, SseStream) = if needs_transform {
// P3-B：responses 上游 → anthropic 客户端暂未实现，明确报错
if outbound == "responses" && protocol == "anthropic" {
return Err(err_json(
StatusCode::BAD_REQUEST,
"responses 上游流式转换到 anthropic 客户端暂未支持（P3-B 延后）",
));
}
item.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { e.into() })
});
// 流尾：解析 usage → 入账（客户端未完整接收时该 future 不会执行 → 不入账）
let finalize = futures_util::stream::once(async move {
let (input, output) = {
let mut cap = capture.lock().expect("usage capture lock");
cap.finish()
let usage_slot = crate::sse::usage_slot();
let raw = resp
.bytes_stream()
.map(|item| item.map_err(std::io::Error::other));
let converted: SseStream = match (outbound.as_str(), protocol.as_str()) {
("openai_chat", "anthropic") => {
Box::pin(crate::sse::openai_sse_to_anthropic(raw, usage_slot.clone()))
}
("openai_chat", "responses") => Box::pin(
crate::sse::openai_sse_to_openai_responses(raw, usage_slot.clone()),
),
("anthropic", "openai_chat") => {
Box::pin(crate::sse::anthropic_sse_to_openai(raw, usage_slot.clone()))
}
("anthropic", "responses") => Box::pin(crate::sse::anthropic_sse_to_responses(
raw,
usage_slot.clone(),
)),
("responses", "openai_chat") => Box::pin(crate::sse::responses_sse_to_openai_chat(
raw,
usage_slot.clone(),
)),
(u, c) => {
return Err(err_json(
StatusCode::BAD_REQUEST,
&format!("流式协议转换 {u} → {c} 暂未支持"),
));
}
};
settle_usage(&st, auth, &key, &model, input, output);
Ok::<axum::body::Bytes, Box<dyn std::error::Error + Send + Sync>>(
axum::body::Bytes::new(),
)
});
let slot_final = usage_slot.clone();
let finalize = futures_util::stream::once(async move {
let (input, output) = slot_final
.lock()
.map(|s| s.unwrap_or((0.0, 0.0)))
.unwrap_or((0.0, 0.0));
settle_usage(&st, auth, &key, &model, input, output);
Ok::<axum::body::Bytes, std::io::Error>(axum::body::Bytes::new())
});
(converted, Box::pin(finalize))
} else {
// 同协议透传：usage 捕获 + 流尾入账
let capture = std::sync::Arc::new(std::sync::Mutex::new(UsageCapture::new(&protocol)));
let cap_fwd = std::sync::Arc::clone(&capture);
let fwd = resp.bytes_stream().map(move |item| {
if let Ok(bytes) = &item {
if let Ok(mut cap) = cap_fwd.lock() {
cap.push(bytes);
}
}
item.map_err(|e| -> std::io::Error { std::io::Error::other(e) })
});
let finalize = futures_util::stream::once(async move {
let (input, output) = {
let mut cap = capture.lock().expect("usage capture lock");
cap.finish()
};
settle_usage(&st, auth, &key, &model, input, output);
Ok::<axum::body::Bytes, std::io::Error>(axum::body::Bytes::new())
});
(Box::pin(fwd), Box::pin(finalize))
};
let body = Body::from_stream(fwd.chain(finalize));
return Ok(Response::builder()
.status(StatusCode::OK)
Expand Down Expand Up @@ -1221,21 +1285,35 @@ mod tests {
up.abort();
}

/// P3-A：流式 + 跨协议转换 → 400 明确报错（SSE 转换留 P3-B）
/// P3-B（rant 2026-08-18T18:59:29）：流式跨协议转换 — openai 客户端 stream
/// → plan 只暴露 anthropic 端点 → 上游 anthropic SSE → 转回 openai SSE 事件。
#[tokio::test]
async fn streaming_with_conversion_rejected_400() {
async fn sse_cross_protocol_openai_to_anthropic_conversion() {
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
.await
.unwrap();
let port = listener.local_addr().unwrap().port();
let received = std::sync::Arc::new(std::sync::Mutex::new(None::<String>));
let recv2 = std::sync::Arc::clone(&received);
let up = tokio::spawn(fake_anthropic_upstream(listener, recv2));
let up = tokio::spawn(fake_anthropic_sse_upstream(listener));
let base = format!("http://127.0.0.1:{port}");

// plan 只暴露 anthropic 端点
// plan 只暴露 anthropic 端点 → 必须走协议转换
let st = test_state_eps("conv_sse", "test-anth-s", &["anthropic"], &base);
insert_key(&st, 404, 1, "test-model", "test-anth-s");
{
let conn = st.db.lock().unwrap();
models_row(&conn, "test", "test-model", 10.0, 20.0);
// 独立分享者（owner=2，与消费者 user=1 分离 → 净扣 1.4 点可断言）
conn.execute(
"INSERT OR IGNORE INTO users (id, email, password_hash, name, role) VALUES (2, 'owner2@t.local', 'x', '分享者', 'user')",
[],
)
.unwrap();
conn.execute(
"INSERT OR IGNORE INTO quotas (user_id, balance) VALUES (2, 0)",
[],
)
.unwrap();
}
insert_key(&st, 404, 2, "test-model", "test-anth-s");
let key = login_key(st.clone()).await;

let (s, body) = post_raw(
Expand All @@ -1247,16 +1325,33 @@ mod tests {
.await;
assert_eq!(
s,
StatusCode::BAD_REQUEST,
StatusCode::OK,
"body: {}",
String::from_utf8_lossy(&body)
);
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
let text = String::from_utf8_lossy(&body);
assert!(
v["error"]["message"].as_str().unwrap().contains("协议转换"),
"body: {}",
String::from_utf8_lossy(&body)
text.contains("\"role\":\"assistant\""),
"openai role chunk: {text}"
);
assert!(
text.contains("\"content\":\"hi\""),
"openai 文本增量: {text}"
);
assert!(text.contains("data: [DONE]"), "openai [DONE]: {text}");
assert!(
!text.contains("event: message_start"),
"不应透传 anthropic 事件"
);
// 入账：80×10/1e6 + 30×20/1e6 = 0.0014 USD × 1000 = 1.4 点
let conn = st.db.lock().unwrap();
let bal: f64 = conn
.query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| {
r.get(0)
})
.unwrap();
assert!((bal - (12471.0 - 1.4)).abs() < 1e-9, "consumer={bal}");
drop(conn);
up.abort();
}

Expand Down Expand Up @@ -1457,9 +1552,11 @@ mod tests {
let app = axum::Router::new().route(
"/v1/messages",
axum::routing::post(|_body: String| async {
let body = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":80,\"output_tokens\":1}}}\n\n\
event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"hi\"}}\n\n\
event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":30}}\n\n\
let body = "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"model\":\"test-model\",\"usage\":{\"input_tokens\":80,\"output_tokens\":1}}}\n\n\
event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n\
event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n\
event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n\
event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":30}}\n\n\
event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
(
[("content-type", "text/event-stream")],
Expand Down
1 change: 1 addition & 0 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ mod gift;
mod protocol;
mod router;
mod routes;
mod sse;

use std::sync::Arc;

Expand Down
Loading
Loading