添加qiming-rcoder模块
This commit is contained in:
235
qiming-rcoder/crates/rcoder/src/grpc/channel_pool.rs
Normal file
235
qiming-rcoder/crates/rcoder/src/grpc/channel_pool.rs
Normal file
@@ -0,0 +1,235 @@
|
||||
//! gRPC Channel 连接池
|
||||
//!
|
||||
//! 管理到各个 agent_runner 容器的 gRPC 连接,支持 TTL 自动清理失效连接。
|
||||
|
||||
use anyhow::Result;
|
||||
use dashmap::DashMap;
|
||||
use shared_types::grpc::agent_service_client::AgentServiceClient;
|
||||
use std::time::{Duration, Instant};
|
||||
use tonic::transport::Channel;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// gRPC 连接池 TTL(5分钟)
|
||||
///
|
||||
/// 连接超过此时间未被使用则自动清理,防止内存泄漏。
|
||||
const CHANNEL_TTL_SECS: u64 = 300;
|
||||
|
||||
/// gRPC 连接池最大容量
|
||||
const MAX_CAPACITY: usize = 10000;
|
||||
|
||||
/// 创建配置好的 gRPC 客户端(设置消息大小限制)
|
||||
///
|
||||
/// tonic 的消息大小限制是在 AgentServiceClient 级别配置的,
|
||||
/// 无法在 Channel 或 Endpoint 级别统一配置,所以需要这个辅助函数。
|
||||
fn create_configured_client(channel: Channel) -> AgentServiceClient<Channel> {
|
||||
AgentServiceClient::new(channel)
|
||||
.max_decoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE)
|
||||
.max_encoding_message_size(shared_types::GRPC_MAX_MESSAGE_SIZE)
|
||||
}
|
||||
|
||||
/// Channel 元数据(包含创建时间)
|
||||
#[derive(Clone)]
|
||||
struct ChannelEntry {
|
||||
channel: Channel,
|
||||
created_at: Instant,
|
||||
}
|
||||
|
||||
impl ChannelEntry {
|
||||
fn is_expired(&self) -> bool {
|
||||
self.created_at.elapsed() > Duration::from_secs(CHANNEL_TTL_SECS)
|
||||
}
|
||||
}
|
||||
|
||||
/// gRPC 连接池
|
||||
///
|
||||
/// 为每个容器维护独立的 gRPC 连接,支持:
|
||||
/// - 连接复用:相同地址的请求复用同一连接
|
||||
/// - TTL 自动清理:5分钟未使用的连接自动移除,防止内存泄漏
|
||||
/// - 并发安全:支持高并发下的安全连接创建
|
||||
pub struct GrpcChannelPool {
|
||||
/// 容器地址到 Channel 的映射
|
||||
channels: DashMap<String, ChannelEntry>,
|
||||
}
|
||||
|
||||
impl GrpcChannelPool {
|
||||
/// 创建新的连接池
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
channels: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取指定地址的 gRPC 客户端
|
||||
///
|
||||
/// 如果连接不存在则创建新连接。过期连接会被自动清理。
|
||||
pub async fn get_client(&self, addr: &str) -> Result<AgentServiceClient<Channel>> {
|
||||
// 先检查缓存,同时清理过期条目
|
||||
// 使用 remove_if_available 模式避免 TOCTOU 竞态
|
||||
let should_remove = {
|
||||
self.channels.get(addr).map(|entry| entry.is_expired()).unwrap_or(false)
|
||||
};
|
||||
|
||||
if should_remove {
|
||||
// 过期则移除
|
||||
self.channels.remove(addr);
|
||||
}
|
||||
|
||||
// 再次检查(可能被其他线程修改)
|
||||
if let Some(entry) = self.channels.get(addr) {
|
||||
debug!("📡 [gRPC] reuse connection: {}", addr);
|
||||
return Ok(create_configured_client(entry.channel.clone()));
|
||||
}
|
||||
|
||||
// 缓存未命中或已过期,创建新连接
|
||||
info!("🔌 [gRPC] creating connection: {}", addr);
|
||||
let endpoint = format!("http://{}", addr);
|
||||
let channel = Channel::from_shared(endpoint)
|
||||
.map_err(|e| anyhow::anyhow!("Invalid URI: {}", e))?
|
||||
.connect_timeout(Duration::from_secs(
|
||||
shared_types::GRPC_CONNECT_TIMEOUT_SECS,
|
||||
))
|
||||
.timeout(Duration::from_secs(
|
||||
shared_types::GRPC_REQUEST_TIMEOUT_SECS,
|
||||
))
|
||||
// HTTP/2 Keepalive 配置
|
||||
.http2_keep_alive_interval(Duration::from_secs(30))
|
||||
.keep_alive_timeout(Duration::from_secs(10))
|
||||
.keep_alive_while_idle(true)
|
||||
// TCP Keepalive 配置
|
||||
.tcp_keepalive(Some(Duration::from_secs(60)))
|
||||
.tcp_nodelay(true)
|
||||
.connect()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Connection failed: {}", e))?;
|
||||
|
||||
// 清理过期连接
|
||||
self.cleanup_expired();
|
||||
|
||||
// 插入缓存
|
||||
self.channels.insert(
|
||||
addr.to_string(),
|
||||
ChannelEntry {
|
||||
channel: channel.clone(),
|
||||
created_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
|
||||
debug!("📡 [gRPC] connection ready: {}", addr);
|
||||
Ok(create_configured_client(channel))
|
||||
}
|
||||
|
||||
/// 清理过期的连接
|
||||
///
|
||||
/// 每次调用时检查并清理过期连接,保持缓存高效。
|
||||
fn cleanup_expired(&self) {
|
||||
let len = self.channels.len();
|
||||
|
||||
// 收集所有过期连接的键
|
||||
let expired: Vec<String> = self
|
||||
.channels
|
||||
.iter()
|
||||
.filter(|e| e.is_expired())
|
||||
.map(|e| e.key().clone())
|
||||
.collect();
|
||||
|
||||
// 移除过期连接
|
||||
for key in &expired {
|
||||
self.channels.remove(key);
|
||||
}
|
||||
|
||||
if !expired.is_empty() {
|
||||
debug!(
|
||||
"🔌 [gRPC] cleaned up {} expired connections (cache size: {})",
|
||||
expired.len(),
|
||||
self.channels.len()
|
||||
);
|
||||
}
|
||||
|
||||
// 如果清理后仍然满(>= 80%),再清理一半的非过期连接
|
||||
if self.channels.len() >= MAX_CAPACITY * 8 / 10 {
|
||||
let to_remove: Vec<_> = self
|
||||
.channels
|
||||
.iter()
|
||||
.take(MAX_CAPACITY / 2)
|
||||
.map(|e| e.key().clone())
|
||||
.collect();
|
||||
|
||||
for key in &to_remove {
|
||||
self.channels.remove(key);
|
||||
}
|
||||
warn!(
|
||||
"🔌 [gRPC] cache still full after cleanup, evicted {} entries",
|
||||
to_remove.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取指定容器端口的 gRPC 客户端
|
||||
///
|
||||
/// 假设容器 IP 为 localhost,端口为 gRPC 端口(默认 50051)
|
||||
pub async fn get_client_for_container(
|
||||
&self,
|
||||
container_ip: &str,
|
||||
grpc_port: u16,
|
||||
) -> Result<AgentServiceClient<Channel>> {
|
||||
let addr = format!("{}:{}", container_ip, grpc_port);
|
||||
self.get_client(&addr).await
|
||||
}
|
||||
|
||||
/// 移除指定地址的连接
|
||||
pub fn remove(&self, addr: &str) {
|
||||
if self.channels.remove(addr).is_some() {
|
||||
info!("🔌 [gRPC] removed connection: {}", addr);
|
||||
}
|
||||
}
|
||||
|
||||
/// 清空所有连接
|
||||
pub fn clear(&self) {
|
||||
self.channels.clear();
|
||||
info!("🔌 [gRPC] cleared all connections");
|
||||
}
|
||||
|
||||
/// 获取当前连接数
|
||||
pub fn connection_count(&self) -> usize {
|
||||
self.channels.len()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for GrpcChannelPool {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for GrpcChannelPool {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("GrpcChannelPool")
|
||||
.field("connection_count", &self.connection_count())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_new_pool() {
|
||||
let pool = GrpcChannelPool::new();
|
||||
assert_eq!(pool.connection_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_non_existent() {
|
||||
let pool = GrpcChannelPool::new();
|
||||
pool.remove("non_existent");
|
||||
assert_eq!(pool.connection_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let pool = GrpcChannelPool::new();
|
||||
pool.clear();
|
||||
assert_eq!(pool.connection_count(), 0);
|
||||
}
|
||||
}
|
||||
198
qiming-rcoder/crates/rcoder/src/grpc/chat_client.rs
Normal file
198
qiming-rcoder/crates/rcoder/src/grpc/chat_client.rs
Normal file
@@ -0,0 +1,198 @@
|
||||
//! gRPC Chat 客户端
|
||||
//!
|
||||
//! 通过 gRPC 调用 agent_runner 的 Chat RPC
|
||||
|
||||
use crate::grpc::GrpcChannelPool;
|
||||
use shared_types::ChatAgentConfig;
|
||||
use shared_types::grpc::{
|
||||
CancelRequest, CancelResponse, ChatRequest as GrpcChatRequest, ChatResponse as GrpcChatResponse,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
/// 通过 gRPC 发送 Chat 请求到 agent_runner (使用连接池)
|
||||
pub async fn grpc_chat_with_pool(
|
||||
pool: &Arc<GrpcChannelPool>,
|
||||
grpc_addr: &str,
|
||||
project_id: String,
|
||||
session_id: Option<String>,
|
||||
prompt: String,
|
||||
attachments: Vec<shared_types::Attachment>,
|
||||
data_source_attachments: Vec<String>,
|
||||
model_config: Option<shared_types::ModelProviderConfig>,
|
||||
request_id: Option<String>,
|
||||
request_timeout: Option<std::time::Duration>,
|
||||
// 新增参数 (v2)
|
||||
system_prompt: Option<String>,
|
||||
user_prompt: Option<String>,
|
||||
agent_config: Option<ChatAgentConfig>,
|
||||
service_type: Option<shared_types::ServiceType>,
|
||||
user_id: Option<String>, // 新增:用于 ComputerAgentRunner 模式
|
||||
) -> anyhow::Result<GrpcChatResponse> {
|
||||
info!(
|
||||
"🚀 [gRPC_CHAT] Sending Chat request (connection pool): addr={}, project_id={}",
|
||||
grpc_addr, project_id
|
||||
);
|
||||
|
||||
// 使用连接池获取客户端
|
||||
let mut client = pool.get_client(grpc_addr).await?;
|
||||
|
||||
// 构建 gRPC 请求
|
||||
let grpc_request = GrpcChatRequest {
|
||||
project_id,
|
||||
session_id: session_id.unwrap_or_default(),
|
||||
prompt,
|
||||
model_config: model_config.map(super::converters::to_grpc_model_config),
|
||||
attachments: attachments
|
||||
.into_iter()
|
||||
.map(super::converters::to_grpc_attachment)
|
||||
.collect(),
|
||||
request_id,
|
||||
data_source_attachments,
|
||||
// 新增字段 (v2)
|
||||
system_prompt,
|
||||
user_prompt,
|
||||
agent_config: agent_config.map(super::converters::to_grpc_chat_agent_config),
|
||||
service_type: service_type.map(|st| format!("{:?}", st)),
|
||||
user_id, // 传递 user_id
|
||||
};
|
||||
|
||||
debug!("[gRPC_CHAT] sendrequest: {:?}", grpc_request);
|
||||
|
||||
// 构建 tonic Request 并设置请求级别超时
|
||||
let locale = super::current_grpc_locale();
|
||||
let mut request = super::new_request_with_locale(grpc_request, locale);
|
||||
|
||||
// ✅ 使用 Tonic 原生 API 设置请求超时
|
||||
if let Some(timeout) = request_timeout {
|
||||
request.set_timeout(timeout);
|
||||
debug!("⏱️ [gRPC_CHAT] request timeout: {:?}", timeout);
|
||||
}
|
||||
|
||||
// 发送请求
|
||||
let response = client.chat(request).await.map_err(|e| {
|
||||
error!("[gRPC_CHAT] Chat RPC call failed: {}", e);
|
||||
anyhow::anyhow!("gRPC Chat call failed: {}", e)
|
||||
})?;
|
||||
|
||||
let chat_response = response.into_inner();
|
||||
|
||||
info!(
|
||||
"✅ [gRPC_CHAT] Received response: project_id={}, session_id={}, success={}",
|
||||
chat_response.project_id, chat_response.session_id, chat_response.success
|
||||
);
|
||||
|
||||
Ok(chat_response)
|
||||
}
|
||||
|
||||
/// 将 gRPC ChatResponse 转换为内部 ChatResponse
|
||||
pub fn grpc_response_to_chat_response(grpc_resp: GrpcChatResponse) -> shared_types::ChatResponse {
|
||||
shared_types::ChatResponse {
|
||||
project_id: grpc_resp.project_id,
|
||||
session_id: grpc_resp.session_id,
|
||||
error: grpc_resp.error,
|
||||
request_id: grpc_resp.request_id,
|
||||
need_fallback: None,
|
||||
fallback_reason: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 通过 gRPC 取消会话(使用连接池)
|
||||
pub async fn grpc_cancel_session_with_pool(
|
||||
pool: &Arc<GrpcChannelPool>,
|
||||
grpc_addr: &str,
|
||||
session_id: String,
|
||||
reason: String,
|
||||
project_id: String,
|
||||
) -> anyhow::Result<CancelResponse> {
|
||||
info!(
|
||||
"🛑 [gRPC_CANCEL] Sending cancel session request (connection pool): addr={}, session_id={}, project_id={}",
|
||||
grpc_addr, session_id, project_id
|
||||
);
|
||||
|
||||
// 使用连接池获取客户端
|
||||
let mut client = pool.get_client(grpc_addr).await?;
|
||||
|
||||
// 构建 gRPC 请求
|
||||
let grpc_request = CancelRequest {
|
||||
session_id,
|
||||
reason,
|
||||
project_id,
|
||||
};
|
||||
|
||||
debug!("[gRPC_CANCEL] sendrequest: {:?}", grpc_request);
|
||||
|
||||
// 发送请求
|
||||
let locale = super::current_grpc_locale();
|
||||
let request = super::new_request_with_locale(grpc_request, locale);
|
||||
|
||||
let response = client.cancel_session(request).await.map_err(|e| {
|
||||
error!("[gRPC_CANCEL] CancelSession RPC call failed: {}", e);
|
||||
anyhow::anyhow!("gRPC CancelSession call failed: {}", e)
|
||||
})?;
|
||||
|
||||
let cancel_response = response.into_inner();
|
||||
|
||||
info!(
|
||||
"✅ [gRPC_CANCEL] Received response: success={}, message={:?}",
|
||||
cancel_response.success, cancel_response.message
|
||||
);
|
||||
|
||||
Ok(cancel_response)
|
||||
}
|
||||
|
||||
/// 通过 gRPC 取消会话(不使用连接池,兼容旧接口)
|
||||
pub async fn grpc_cancel_session(
|
||||
grpc_addr: &str,
|
||||
session_id: String,
|
||||
reason: String,
|
||||
project_id: String,
|
||||
) -> anyhow::Result<CancelResponse> {
|
||||
// 创建临时连接池(单次使用)
|
||||
let pool = Arc::new(GrpcChannelPool::new());
|
||||
grpc_cancel_session_with_pool(&pool, grpc_addr, session_id, reason, project_id).await
|
||||
}
|
||||
|
||||
/// 通过 gRPC 停止 Agent(使用连接池)
|
||||
pub async fn grpc_stop_agent_with_pool(
|
||||
pool: &Arc<GrpcChannelPool>,
|
||||
grpc_addr: &str,
|
||||
project_id: String,
|
||||
reason: Option<String>,
|
||||
force: bool,
|
||||
) -> anyhow::Result<shared_types::grpc::StopAgentResponse> {
|
||||
info!(
|
||||
"🔄 [gRPC_STOP_AGENT] Sending stop Agent request (connection pool): addr={}, project_id={}, force={}",
|
||||
grpc_addr, project_id, force
|
||||
);
|
||||
|
||||
// 使用连接池获取客户端
|
||||
let mut client = pool.get_client(grpc_addr).await?;
|
||||
|
||||
// 构建 gRPC 请求
|
||||
let grpc_request = shared_types::grpc::StopAgentRequest {
|
||||
project_id: project_id.clone(),
|
||||
reason,
|
||||
force,
|
||||
};
|
||||
|
||||
debug!("[gRPC_STOP_AGENT] sendrequest: {:?}", grpc_request);
|
||||
|
||||
// 发送请求
|
||||
let locale = super::current_grpc_locale();
|
||||
let request = super::new_request_with_locale(grpc_request, locale);
|
||||
|
||||
let response = client.stop_agent(request).await.map_err(|e| {
|
||||
error!("[gRPC_STOP_AGENT] StopAgent RPC call failed: {}", e);
|
||||
anyhow::anyhow!("gRPC StopAgent call failed: {}", e)
|
||||
})?;
|
||||
|
||||
let stop_response = response.into_inner();
|
||||
|
||||
info!(
|
||||
"✅ [gRPC_STOP_AGENT] Received response: result={}, success={}, message={:?}",
|
||||
stop_response.result, stop_response.success, stop_response.message
|
||||
);
|
||||
|
||||
Ok(stop_response)
|
||||
}
|
||||
231
qiming-rcoder/crates/rcoder/src/grpc/converters.rs
Normal file
231
qiming-rcoder/crates/rcoder/src/grpc/converters.rs
Normal file
@@ -0,0 +1,231 @@
|
||||
//! 类型转换器
|
||||
//!
|
||||
//! 在 Rust 内部类型和 gRPC Protobuf 类型之间进行转换
|
||||
|
||||
use shared_types::grpc::{
|
||||
Attachment as GrpcAttachment, AttachmentSource as GrpcAttachmentSource,
|
||||
AudioAttachment as GrpcAudioAttachment, Base64Data, ChatAgentConfig as GrpcChatAgentConfig,
|
||||
ChatAgentServerConfig as GrpcChatAgentServerConfig,
|
||||
ChatContextServerConfig as GrpcChatContextServerConfig, ChatRequest as GrpcChatRequest,
|
||||
ChatResponse as GrpcChatResponse, DocumentAttachment as GrpcDocumentAttachment,
|
||||
ImageAttachment as GrpcImageAttachment, ImageDimensions as GrpcImageDimensions,
|
||||
ModelProviderConfig as GrpcModelProviderConfig, ProgressEvent,
|
||||
TextAttachment as GrpcTextAttachment, attachment, attachment_source,
|
||||
};
|
||||
use shared_types::{Attachment, AttachmentSource, ModelProviderConfig, UnifiedSessionMessage};
|
||||
use shared_types::{ChatAgentConfig, ChatAgentServerConfig, ChatContextServerConfig};
|
||||
|
||||
/// 将内部 ChatRequest 转换为 gRPC ChatRequest
|
||||
///
|
||||
/// 注意:此函数目前未被使用,chat_client.rs 直接构建 GrpcChatRequest。
|
||||
/// 保留以备将来使用或重构。
|
||||
#[allow(dead_code)]
|
||||
pub fn to_grpc_chat_request(
|
||||
project_id: String,
|
||||
session_id: String,
|
||||
prompt: String,
|
||||
attachments: Vec<Attachment>,
|
||||
data_source_attachments: Vec<String>,
|
||||
model_config: Option<ModelProviderConfig>,
|
||||
request_id: Option<String>,
|
||||
// 新增参数 (v2)
|
||||
system_prompt: Option<String>,
|
||||
user_prompt: Option<String>,
|
||||
agent_config: Option<ChatAgentConfig>,
|
||||
service_type: Option<shared_types::ServiceType>,
|
||||
user_id: Option<String>, // 新增:用于 ComputerAgentRunner 模式
|
||||
) -> GrpcChatRequest {
|
||||
GrpcChatRequest {
|
||||
project_id,
|
||||
session_id,
|
||||
prompt,
|
||||
model_config: model_config.map(to_grpc_model_config),
|
||||
attachments: attachments.into_iter().map(to_grpc_attachment).collect(),
|
||||
request_id,
|
||||
data_source_attachments,
|
||||
// 新增字段 (v2)
|
||||
system_prompt,
|
||||
user_prompt,
|
||||
agent_config: agent_config.map(to_grpc_chat_agent_config),
|
||||
service_type: service_type.map(|st| format!("{:?}", st)),
|
||||
user_id, // 传递 user_id
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 ModelProviderConfig 转换为 gRPC 格式
|
||||
pub fn to_grpc_model_config(config: ModelProviderConfig) -> GrpcModelProviderConfig {
|
||||
GrpcModelProviderConfig {
|
||||
id: config.id, // 保留原始 ID,用于会话复用判断
|
||||
provider: config.name,
|
||||
model: config.default_model,
|
||||
api_key: Some(config.api_key),
|
||||
api_base: Some(config.base_url),
|
||||
requires_openai_auth: Some(config.requires_openai_auth),
|
||||
api_protocol: config.api_protocol,
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 AttachmentSource 转换为 gRPC 格式
|
||||
///
|
||||
/// 辅助函数,将 Rust 侧的 AttachmentSource 枚举转换为 gRPC 的 AttachmentSource
|
||||
fn to_grpc_attachment_source(source: AttachmentSource) -> Option<GrpcAttachmentSource> {
|
||||
let grpc_source = match source {
|
||||
AttachmentSource::FilePath { path } => attachment_source::Source::FilePath(path),
|
||||
AttachmentSource::Base64 { data, mime_type } => {
|
||||
attachment_source::Source::Base64(Base64Data { data, mime_type })
|
||||
}
|
||||
AttachmentSource::Url { url } => attachment_source::Source::Url(url),
|
||||
};
|
||||
|
||||
Some(GrpcAttachmentSource {
|
||||
source: Some(grpc_source),
|
||||
})
|
||||
}
|
||||
|
||||
/// 将 Attachment 转换为 gRPC 格式
|
||||
///
|
||||
/// 完整实现:根据 Attachment 类型正确填充 gRPC 的 oneof 字段
|
||||
pub fn to_grpc_attachment(attachment: Attachment) -> GrpcAttachment {
|
||||
let attachment_type = match attachment {
|
||||
Attachment::Text(text) => attachment::AttachmentType::Text(GrpcTextAttachment {
|
||||
id: text.id,
|
||||
source: to_grpc_attachment_source(text.source),
|
||||
filename: text.filename,
|
||||
description: text.description,
|
||||
}),
|
||||
Attachment::Image(image) => attachment::AttachmentType::Image(GrpcImageAttachment {
|
||||
id: image.id,
|
||||
source: to_grpc_attachment_source(image.source),
|
||||
mime_type: image.mime_type,
|
||||
filename: image.filename,
|
||||
description: image.description,
|
||||
dimensions: image.dimensions.map(|d| GrpcImageDimensions {
|
||||
width: d.width,
|
||||
height: d.height,
|
||||
}),
|
||||
}),
|
||||
Attachment::Audio(audio) => attachment::AttachmentType::Audio(GrpcAudioAttachment {
|
||||
id: audio.id,
|
||||
source: to_grpc_attachment_source(audio.source),
|
||||
mime_type: audio.mime_type,
|
||||
filename: audio.filename,
|
||||
description: audio.description,
|
||||
duration: audio.duration,
|
||||
}),
|
||||
Attachment::Document(doc) => attachment::AttachmentType::Document(GrpcDocumentAttachment {
|
||||
id: doc.id,
|
||||
source: to_grpc_attachment_source(doc.source),
|
||||
mime_type: doc.mime_type,
|
||||
filename: doc.filename,
|
||||
description: doc.description,
|
||||
size: doc.size,
|
||||
}),
|
||||
};
|
||||
|
||||
GrpcAttachment {
|
||||
attachment_type: Some(attachment_type),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 gRPC ChatResponse 提取结果
|
||||
pub fn from_grpc_chat_response(resp: GrpcChatResponse) -> (String, String, bool, Option<String>) {
|
||||
(resp.project_id, resp.session_id, resp.success, resp.error)
|
||||
}
|
||||
|
||||
/// 从 gRPC ProgressEvent 转换为 UnifiedSessionMessage
|
||||
///
|
||||
/// 简化版:直接使用透传的字段
|
||||
pub fn from_grpc_progress_event(
|
||||
event: ProgressEvent,
|
||||
session_id: &str,
|
||||
) -> Option<UnifiedSessionMessage> {
|
||||
use chrono::Utc;
|
||||
use shared_types::SessionMessageType;
|
||||
use tracing::warn;
|
||||
|
||||
let timestamp = match chrono::DateTime::from_timestamp_millis(event.timestamp) {
|
||||
Some(ts) => ts,
|
||||
None => {
|
||||
warn!(
|
||||
"⚠️ [CONVERTER] Invalid timestamp: session_id={}, timestamp={}, using current time",
|
||||
session_id, event.timestamp
|
||||
);
|
||||
Utc::now()
|
||||
}
|
||||
};
|
||||
|
||||
// 从 message_type 字符串解析枚举
|
||||
let message_type = match event.message_type.as_str() {
|
||||
"SessionPromptStart" => SessionMessageType::SessionPromptStart,
|
||||
"SessionPromptEnd" => SessionMessageType::SessionPromptEnd,
|
||||
"AgentSessionUpdate" => SessionMessageType::AgentSessionUpdate,
|
||||
"Heartbeat" => SessionMessageType::Heartbeat,
|
||||
_ => SessionMessageType::AgentSessionUpdate, // 默认为 AgentSessionUpdate
|
||||
};
|
||||
|
||||
// 解析 payload JSON
|
||||
let data = match serde_json::from_str(&event.payload) {
|
||||
Ok(data) => data,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [CONVERTER] Failed to parse gRPC payload: session_id={}, payload_preview={}, error={}",
|
||||
session_id,
|
||||
event.payload.chars().take(100).collect::<String>(),
|
||||
e
|
||||
);
|
||||
// 返回包含原始 payload 的错误对象
|
||||
serde_json::json!({
|
||||
"_parse_error": e.to_string(),
|
||||
"_original_payload": event.payload
|
||||
})
|
||||
}
|
||||
};
|
||||
|
||||
Some(UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type,
|
||||
sub_type: event.sub_type,
|
||||
data,
|
||||
timestamp,
|
||||
})
|
||||
}
|
||||
|
||||
// === ChatAgentConfig 类型转换 (v2) ===
|
||||
|
||||
/// 将 ChatAgentConfig 转换为 gRPC 格式
|
||||
pub fn to_grpc_chat_agent_config(config: ChatAgentConfig) -> GrpcChatAgentConfig {
|
||||
GrpcChatAgentConfig {
|
||||
agent_server: config.agent_server.map(to_grpc_chat_agent_server_config),
|
||||
context_servers: config
|
||||
.context_servers
|
||||
.into_iter()
|
||||
.map(|(k, v)| (k, to_grpc_chat_context_server_config(v)))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 ChatAgentServerConfig 转换为 gRPC 格式
|
||||
pub fn to_grpc_chat_agent_server_config(
|
||||
config: ChatAgentServerConfig,
|
||||
) -> GrpcChatAgentServerConfig {
|
||||
GrpcChatAgentServerConfig {
|
||||
agent_id: config.agent_id,
|
||||
command: config.command,
|
||||
args: config.args.unwrap_or_default(),
|
||||
env: config.env.unwrap_or_default(),
|
||||
metadata: config.metadata.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 ChatContextServerConfig 转换为 gRPC 格式
|
||||
pub fn to_grpc_chat_context_server_config(
|
||||
config: ChatContextServerConfig,
|
||||
) -> GrpcChatContextServerConfig {
|
||||
GrpcChatContextServerConfig {
|
||||
source: config.source,
|
||||
enabled: config.enabled,
|
||||
command: config.command,
|
||||
args: config.args.unwrap_or_default(),
|
||||
env: config.env.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
209
qiming-rcoder/crates/rcoder/src/grpc/error.rs
Normal file
209
qiming-rcoder/crates/rcoder/src/grpc/error.rs
Normal file
@@ -0,0 +1,209 @@
|
||||
//! gRPC 错误分类和处理
|
||||
//!
|
||||
//! 基于 Tonic 的 Status Code 进行智能错误分类,优化重试策略
|
||||
|
||||
use tonic::{Code, Status};
|
||||
|
||||
/// gRPC 错误分类
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum GrpcErrorCategory {
|
||||
/// 可重试错误(网络问题、资源不足等临时性错误)
|
||||
Retryable,
|
||||
/// 不可重试错误(参数错误、权限问题等客户端错误)
|
||||
NonRetryable,
|
||||
/// 永久性错误(未找到、未实现等服务端永久性问题)
|
||||
Permanent,
|
||||
}
|
||||
|
||||
/// 基于 Tonic Status Code 分类 gRPC 错误
|
||||
///
|
||||
/// 根据 gRPC 标准错误码判断错误是否应该重试
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `status` - gRPC Status 对象
|
||||
///
|
||||
/// # Returns
|
||||
/// 错误分类结果
|
||||
///
|
||||
/// # Examples
|
||||
/// ```
|
||||
/// use tonic::{Code, Status};
|
||||
/// use rcoder::grpc::GrpcErrorCategory;
|
||||
///
|
||||
/// let status = Status::unavailable("服务不可用");
|
||||
/// let category = rcoder::grpc::categorize_grpc_error(&status);
|
||||
/// assert_eq!(category, GrpcErrorCategory::Retryable);
|
||||
/// ```
|
||||
pub fn categorize_grpc_error(status: &Status) -> GrpcErrorCategory {
|
||||
match status.code() {
|
||||
// ✅ 可重试错误:网络问题、资源不足、瞬时故障
|
||||
Code::Unavailable | // 服务不可用(最常见的网络问题)
|
||||
Code::DeadlineExceeded | // 超时(可能是临时性网络延迟)
|
||||
Code::ResourceExhausted | // 资源耗尽(服务器过载,可能恢复)
|
||||
Code::Aborted | // 操作被中止(可能是并发冲突,重试可能成功)
|
||||
Code::Internal | // 内部错误(可能是临时性服务器问题)
|
||||
Code::Unknown => // 未知错误(保守策略:允许重试)
|
||||
GrpcErrorCategory::Retryable,
|
||||
|
||||
// ❌ 永久性错误:服务端不支持或资源不存在
|
||||
Code::NotFound | // 未找到资源
|
||||
Code::Unimplemented | // 方法未实现
|
||||
Code::OutOfRange => // 超出范围(通常是客户端逻辑错误)
|
||||
GrpcErrorCategory::Permanent,
|
||||
|
||||
// ❌ 不可重试错误:客户端问题,重试也不会成功
|
||||
Code::InvalidArgument | // 参数错误
|
||||
Code::Unauthenticated | // 未认证
|
||||
Code::PermissionDenied | // 权限不足
|
||||
Code::FailedPrecondition | // 前置条件失败
|
||||
Code::AlreadyExists | // 资源已存在
|
||||
Code::Cancelled => // 用户取消(不应重试)
|
||||
GrpcErrorCategory::NonRetryable,
|
||||
|
||||
// ✅ OK - 理论上不应该走到这里
|
||||
Code::Ok => GrpcErrorCategory::NonRetryable,
|
||||
|
||||
// ⚠️ DataLoss - 严重错误,但可能是临时性的
|
||||
Code::DataLoss => GrpcErrorCategory::Retryable,
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断 gRPC 错误是否应该重试
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `status` - gRPC Status 对象
|
||||
///
|
||||
/// # Returns
|
||||
/// `true` 如果错误可以重试,`false` 否则
|
||||
pub fn should_retry_grpc_error(status: &Status) -> bool {
|
||||
matches!(categorize_grpc_error(status), GrpcErrorCategory::Retryable)
|
||||
}
|
||||
|
||||
/// 从 anyhow::Error 中提取 Tonic Status(如果存在)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `error` - anyhow Error 对象
|
||||
///
|
||||
/// # Returns
|
||||
/// `Some(Status)` 如果错误包含 Tonic Status,`None` 否则
|
||||
pub fn extract_grpc_status(error: &anyhow::Error) -> Option<&Status> {
|
||||
error.downcast_ref::<Status>()
|
||||
}
|
||||
|
||||
/// 判断 anyhow::Error 是否应该重试(自动提取 Tonic Status)
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `error` - anyhow Error 对象
|
||||
///
|
||||
/// # Returns
|
||||
/// `true` 如果错误包含可重试的 gRPC 错误,`false` 否则
|
||||
pub fn should_retry_error(error: &anyhow::Error) -> bool {
|
||||
if let Some(status) = extract_grpc_status(error) {
|
||||
should_retry_grpc_error(status)
|
||||
} else {
|
||||
// 非 gRPC 错误,保守策略:不重试
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取错误的友好描述
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `status` - gRPC Status 对象
|
||||
///
|
||||
/// # Returns
|
||||
/// 错误的中文描述
|
||||
pub fn get_error_description(status: &Status) -> &'static str {
|
||||
match status.code() {
|
||||
Code::Ok => "Success",
|
||||
Code::Cancelled => "Operation cancelled",
|
||||
Code::Unknown => "Unknown error",
|
||||
Code::InvalidArgument => "Invalid argument",
|
||||
Code::DeadlineExceeded => "Request timeout",
|
||||
Code::NotFound => "Resource not found",
|
||||
Code::AlreadyExists => "Resource already exists",
|
||||
Code::PermissionDenied => "Permission denied",
|
||||
Code::ResourceExhausted => "Resource exhausted",
|
||||
Code::FailedPrecondition => "Precondition failed",
|
||||
Code::Aborted => "Operation aborted",
|
||||
Code::OutOfRange => "Out of range",
|
||||
Code::Unimplemented => "Method not implemented",
|
||||
Code::Internal => "Internal server error",
|
||||
Code::Unavailable => "Service unavailable",
|
||||
Code::DataLoss => "Data loss",
|
||||
Code::Unauthenticated => "Unauthenticated",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_categorize_retryable_errors() {
|
||||
let retryable_codes = vec![
|
||||
Code::Unavailable,
|
||||
Code::DeadlineExceeded,
|
||||
Code::ResourceExhausted,
|
||||
Code::Aborted,
|
||||
Code::Internal,
|
||||
Code::Unknown,
|
||||
];
|
||||
|
||||
for code in retryable_codes {
|
||||
let status = Status::new(code, "test error");
|
||||
assert_eq!(
|
||||
categorize_grpc_error(&status),
|
||||
GrpcErrorCategory::Retryable,
|
||||
"Code {:?} should be retryable",
|
||||
code
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_categorize_non_retryable_errors() {
|
||||
let non_retryable_codes = vec![
|
||||
Code::InvalidArgument,
|
||||
Code::Unauthenticated,
|
||||
Code::PermissionDenied,
|
||||
Code::FailedPrecondition,
|
||||
Code::AlreadyExists,
|
||||
Code::Cancelled,
|
||||
];
|
||||
|
||||
for code in non_retryable_codes {
|
||||
let status = Status::new(code, "test error");
|
||||
assert_eq!(
|
||||
categorize_grpc_error(&status),
|
||||
GrpcErrorCategory::NonRetryable,
|
||||
"Code {:?} should not be retryable",
|
||||
code
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_categorize_permanent_errors() {
|
||||
let permanent_codes = vec![Code::NotFound, Code::Unimplemented, Code::OutOfRange];
|
||||
|
||||
for code in permanent_codes {
|
||||
let status = Status::new(code, "test error");
|
||||
assert_eq!(
|
||||
categorize_grpc_error(&status),
|
||||
GrpcErrorCategory::Permanent,
|
||||
"Code {:?} should be permanent",
|
||||
code
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_should_retry_grpc_error() {
|
||||
let retryable = Status::unavailable("service unavailable");
|
||||
assert!(should_retry_grpc_error(&retryable));
|
||||
|
||||
let non_retryable = Status::invalid_argument("bad request");
|
||||
assert!(!should_retry_grpc_error(&non_retryable));
|
||||
}
|
||||
}
|
||||
54
qiming-rcoder/crates/rcoder/src/grpc/locale_metadata.rs
Normal file
54
qiming-rcoder/crates/rcoder/src/grpc/locale_metadata.rs
Normal file
@@ -0,0 +1,54 @@
|
||||
use tonic::metadata::MetadataValue;
|
||||
|
||||
const ACCEPT_LANGUAGE_METADATA_KEY: &str = "accept-language";
|
||||
|
||||
/// Get locale from current HTTP request context (task-local).
|
||||
pub fn current_grpc_locale() -> &'static str {
|
||||
shared_types::current_request_locale()
|
||||
}
|
||||
|
||||
/// Inject `accept-language` metadata into a gRPC request.
|
||||
pub fn inject_accept_language_metadata<T>(request: &mut tonic::Request<T>, locale: &'static str) {
|
||||
match MetadataValue::try_from(locale) {
|
||||
Ok(value) => {
|
||||
request
|
||||
.metadata_mut()
|
||||
.insert(ACCEPT_LANGUAGE_METADATA_KEY, value);
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
"[GRPC_LOCALE] Invalid accept-language metadata value: locale={}, error={}",
|
||||
locale,
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a gRPC request with `accept-language` metadata.
|
||||
pub fn new_request_with_locale<T>(message: T, locale: &'static str) -> tonic::Request<T> {
|
||||
let mut request = tonic::Request::new(message);
|
||||
inject_accept_language_metadata(&mut request, locale);
|
||||
request
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{current_grpc_locale, inject_accept_language_metadata};
|
||||
|
||||
#[test]
|
||||
fn test_inject_accept_language_metadata() {
|
||||
let mut request = tonic::Request::new(());
|
||||
inject_accept_language_metadata(&mut request, "zh-CN");
|
||||
let value = request
|
||||
.metadata()
|
||||
.get("accept-language")
|
||||
.and_then(|v| v.to_str().ok());
|
||||
assert_eq!(value, Some("zh-CN"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_current_grpc_locale_fallback_default() {
|
||||
assert_eq!(current_grpc_locale(), "en-US");
|
||||
}
|
||||
}
|
||||
17
qiming-rcoder/crates/rcoder/src/grpc/mod.rs
Normal file
17
qiming-rcoder/crates/rcoder/src/grpc/mod.rs
Normal file
@@ -0,0 +1,17 @@
|
||||
//! gRPC 客户端模块
|
||||
//!
|
||||
//! 提供 rcoder 与 agent_runner 之间的 gRPC 通信客户端实现
|
||||
|
||||
pub mod channel_pool;
|
||||
pub mod chat_client;
|
||||
pub mod converters;
|
||||
pub mod error;
|
||||
pub mod locale_metadata;
|
||||
pub mod sse_stream;
|
||||
|
||||
pub use channel_pool::GrpcChannelPool;
|
||||
pub use chat_client::*;
|
||||
pub use converters::*;
|
||||
pub use error::*;
|
||||
pub use locale_metadata::*;
|
||||
pub use sse_stream::*;
|
||||
441
qiming-rcoder/crates/rcoder/src/grpc/sse_stream.rs
Normal file
441
qiming-rcoder/crates/rcoder/src/grpc/sse_stream.rs
Normal file
@@ -0,0 +1,441 @@
|
||||
//! gRPC SSE 流处理器
|
||||
//!
|
||||
//! 通过 gRPC SubscribeProgress 接收 agent_runner 的进度事件,
|
||||
//! 并转换为 SSE 事件返回给客户端
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use shared_types::grpc::{GetStatusRequest, ProgressRequest};
|
||||
use shared_types::{SessionMessageType, UnifiedSessionMessage};
|
||||
use tonic::Code;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
/// 创建基于 gRPC 的 SSE 代理流
|
||||
///
|
||||
/// 通过 gRPC `SubscribeProgress` 方法订阅 agent_runner 的进度事件,
|
||||
/// 并将事件转换为 SSE 格式返回
|
||||
///
|
||||
/// 🚀 优化:使用连接池 + 智能重试机制
|
||||
/// 🆕 新增:在建立流之前检查 Agent 状态,如果 Agent 闲置则直接发送 SessionPromptEnd 并关闭
|
||||
pub async fn create_grpc_sse_stream(
|
||||
grpc_addr: String,
|
||||
session_id: String,
|
||||
project_id: String,
|
||||
pool: std::sync::Arc<crate::grpc::GrpcChannelPool>,
|
||||
locale: &'static str,
|
||||
) -> impl futures_util::Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>
|
||||
{
|
||||
let (tx, rx) = tokio::sync::mpsc::channel(100);
|
||||
let session_id_clone = session_id.clone();
|
||||
|
||||
// 在后台任务中处理 gRPC 流
|
||||
tokio::spawn(async move {
|
||||
info!(
|
||||
"🔗 [gRPC_SSE] Starting connection to agent_runner gRPC: addr={}, session_id={}, project_id={}",
|
||||
grpc_addr, session_id_clone, project_id
|
||||
);
|
||||
|
||||
let max_retries = 2;
|
||||
let mut last_error_msg = String::new();
|
||||
|
||||
for attempt in 1..=max_retries {
|
||||
// 1. 从连接池获取客户端
|
||||
let mut client = match pool.get_client(&grpc_addr).await {
|
||||
Ok(client) => client,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to get client (attempt {}/{}): {}, cleaning connection pool and retrying...",
|
||||
attempt, max_retries, e
|
||||
);
|
||||
pool.remove(&grpc_addr);
|
||||
last_error_msg = format!("failed to get client: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 🆕 2. 先检查 Agent 状态(使用 session_id 查询)
|
||||
let status_request = crate::grpc::new_request_with_locale(
|
||||
GetStatusRequest {
|
||||
project_id: String::new(), // 不使用 project_id
|
||||
session_id: session_id_clone.clone(), // 使用 session_id 查询
|
||||
},
|
||||
locale,
|
||||
);
|
||||
|
||||
match client.get_status(status_request).await {
|
||||
Ok(response) => {
|
||||
let status = response.into_inner().status;
|
||||
if status == "idle" {
|
||||
// Agent 闲置,发送 SessionPromptEnd 并关闭连接
|
||||
info!(
|
||||
"💤 [gRPC_SSE] Agent is idle, sending SessionPromptEnd and closing: session_id={}",
|
||||
session_id_clone
|
||||
);
|
||||
let end_event = create_session_prompt_end_event(&session_id_clone);
|
||||
if let Err(e) = tx.send(Ok(end_event)).await {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to send SessionPromptEnd event: session_id={}, error={}",
|
||||
session_id_clone, e
|
||||
);
|
||||
}
|
||||
return; // 直接结束,不建立流
|
||||
}
|
||||
info!(
|
||||
"🔄 [gRPC_SSE] Agent status is {}, continuing to establish stream: session_id={}",
|
||||
status, session_id_clone
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
// 状态检查失败,记录警告但继续尝试建立流
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Agent status check failed: {}, continuing to try establishing stream: session_id={}",
|
||||
e, session_id_clone
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 发送 SubscribeProgress 请求
|
||||
let request = crate::grpc::new_request_with_locale(
|
||||
ProgressRequest {
|
||||
session_id: session_id_clone.clone(),
|
||||
},
|
||||
locale,
|
||||
);
|
||||
|
||||
match client.subscribe_progress(request).await {
|
||||
Ok(response) => {
|
||||
info!(
|
||||
"✅ [gRPC_SSE] Successfully established SubscribeProgress stream: session_id={}",
|
||||
session_id_clone
|
||||
);
|
||||
|
||||
let mut stream = response.into_inner();
|
||||
|
||||
// 持续接收 gRPC 流中的事件
|
||||
loop {
|
||||
match stream.message().await {
|
||||
Ok(Some(progress_event)) => {
|
||||
debug!(
|
||||
"📨 [gRPC_SSE] Received progress event: session_id={}, message_type={}, sub_type={}",
|
||||
session_id_clone,
|
||||
progress_event.message_type,
|
||||
progress_event.sub_type
|
||||
);
|
||||
|
||||
// 将 ProgressEvent 转换为 SSE Event(传入 session_id 以重建完整消息结构)
|
||||
let sse_event =
|
||||
progress_event_to_sse(&progress_event, &session_id_clone);
|
||||
|
||||
if tx.send(Ok(sse_event)).await.is_err() {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Client disconnected: session_id={}",
|
||||
session_id_clone
|
||||
);
|
||||
// 客户端断开,直接退出任务
|
||||
return;
|
||||
}
|
||||
}
|
||||
Ok(None) => {
|
||||
// 流正常结束(agent_runner 主动关闭)
|
||||
info!(
|
||||
"✅ [gRPC_SSE] gRPC stream ended normally: session_id={}",
|
||||
session_id_clone
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(e) => {
|
||||
// 流异常结束(连接中断、超时等)
|
||||
error!(
|
||||
"❌ [gRPC_SSE] gRPC stream error: session_id={}, code={}, message={}",
|
||||
session_id_clone,
|
||||
e.code(),
|
||||
e.message()
|
||||
);
|
||||
|
||||
// 发送标准格式的错误消息
|
||||
let error_event = create_grpc_stream_error_event(
|
||||
&session_id_clone,
|
||||
e.code(),
|
||||
e.message(),
|
||||
);
|
||||
if let Err(e) = tx.send(Ok(error_event)).await {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to send error event: session_id={}, error={}",
|
||||
session_id_clone, e
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] SubscribeProgress call failed (attempt {}/{}): {}",
|
||||
attempt, max_retries, e
|
||||
);
|
||||
|
||||
// 如果不是最后一次尝试,清理连接池并重试
|
||||
if attempt < max_retries {
|
||||
info!(
|
||||
"🔌 [gRPC_SSE] Possibly connection broken, removing {} from connection pool and retrying...",
|
||||
grpc_addr
|
||||
);
|
||||
pool.remove(&grpc_addr);
|
||||
last_error_msg = format!("stream subscription failed: {}", e);
|
||||
continue;
|
||||
}
|
||||
|
||||
last_error_msg = format!("stream subscription ultimately failed: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 如果循环结束还没有 return,说明所有重试都失败了
|
||||
error!(
|
||||
"❌ [gRPC_SSE] Retried {} times ultimately failed: session_id={}, error={}",
|
||||
max_retries, session_id_clone, last_error_msg
|
||||
);
|
||||
|
||||
let error_event = create_connection_error_event(&session_id_clone, &last_error_msg);
|
||||
if let Err(e) = tx.send(Ok(error_event)).await {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to send error event: session_id={}, error={}",
|
||||
session_id_clone, e
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
tokio_stream::wrappers::ReceiverStream::new(rx)
|
||||
}
|
||||
|
||||
/// 创建 Agent 闲置时的 SessionPromptEnd SSE 事件
|
||||
///
|
||||
/// 当 Agent 处于闲置状态时,发送此事件通知前端没有正在执行的任务
|
||||
fn create_session_prompt_end_event(session_id: &str) -> axum::response::sse::Event {
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type: SessionMessageType::SessionPromptEnd,
|
||||
sub_type: "end_turn".to_string(),
|
||||
data: serde_json::json!({
|
||||
"reason": "EndTurn",
|
||||
"description": "Agent has no task in execution"
|
||||
}),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to serialize SessionPromptEnd message: {}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
// 返回包含 session_id 的最小可用结构
|
||||
format!(
|
||||
r#"{{"session_id":"{}","message_type":"SessionPromptEnd","sub_type":"end_turn","data":null}}"#,
|
||||
session_id
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
axum::response::sse::Event::default()
|
||||
.event("prompt_end")
|
||||
.data(json_data)
|
||||
}
|
||||
|
||||
/// 将 gRPC ProgressEvent 转换为 SSE Event
|
||||
///
|
||||
/// 使用 UnifiedSessionMessage 结构体重建完整消息,包含 sessionId、messageType、subType、data、timestamp
|
||||
/// 使用 sub_type 作为 SSE 事件名,前端通过 eventSource.addEventListener(sub_type, ...) 监听
|
||||
fn progress_event_to_sse(
|
||||
event: &shared_types::grpc::ProgressEvent,
|
||||
session_id: &str,
|
||||
) -> axum::response::sse::Event {
|
||||
// 解析 payload 为 data 字段
|
||||
let data: serde_json::Value =
|
||||
serde_json::from_str(&event.payload).unwrap_or(serde_json::Value::Null);
|
||||
|
||||
// 将 gRPC 时间戳(毫秒)转换为 DateTime<Utc>
|
||||
let timestamp = match DateTime::<Utc>::from_timestamp_millis(event.timestamp) {
|
||||
Some(ts) => ts,
|
||||
None => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Invalid timestamp: session_id={}, timestamp={}, using current time",
|
||||
session_id, event.timestamp
|
||||
);
|
||||
Utc::now()
|
||||
}
|
||||
};
|
||||
|
||||
// 将 message_type 字符串转换为 SessionMessageType 枚举
|
||||
let message_type = parse_message_type(&event.message_type);
|
||||
|
||||
// 使用 UnifiedSessionMessage 结构体构建完整消息
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type,
|
||||
sub_type: event.sub_type.clone(),
|
||||
data,
|
||||
timestamp,
|
||||
};
|
||||
|
||||
// 序列化为 JSON
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to serialize ProgressEvent message: session_id={}, message_type={}, error={}",
|
||||
session_id, event.message_type, e
|
||||
);
|
||||
// 返回包含 session_id 的最小可用结构
|
||||
format!(
|
||||
r#"{{"session_id":"{}","message_type":"Unknown","sub_type":"{}","data":null}}"#,
|
||||
session_id, event.sub_type
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
// 使用 sub_type 作为 SSE 事件名
|
||||
// 前端通过 eventSource.addEventListener('agent_message_chunk', ...) 等方式监听
|
||||
axum::response::sse::Event::default()
|
||||
.event(&event.sub_type)
|
||||
.data(json_data)
|
||||
}
|
||||
|
||||
/// 将 message_type 字符串解析为 SessionMessageType 枚举
|
||||
///
|
||||
/// 支持的格式:
|
||||
/// - "SessionPromptStart" -> SessionMessageType::SessionPromptStart
|
||||
/// - "SessionPromptEnd" -> SessionMessageType::SessionPromptEnd
|
||||
/// - "AgentSessionUpdate" -> SessionMessageType::AgentSessionUpdate
|
||||
/// - "Heartbeat" -> SessionMessageType::Heartbeat
|
||||
fn parse_message_type(message_type: &str) -> SessionMessageType {
|
||||
match message_type {
|
||||
"SessionPromptStart" => SessionMessageType::SessionPromptStart,
|
||||
"SessionPromptEnd" => SessionMessageType::SessionPromptEnd,
|
||||
"AgentSessionUpdate" => SessionMessageType::AgentSessionUpdate,
|
||||
"Heartbeat" => SessionMessageType::Heartbeat,
|
||||
// 默认作为 AgentSessionUpdate 处理
|
||||
_ => {
|
||||
debug!(
|
||||
"⚠️ [gRPC_SSE] Unknown message_type: {}, using AgentSessionUpdate as default",
|
||||
message_type
|
||||
);
|
||||
SessionMessageType::AgentSessionUpdate
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取容器的 gRPC 地址
|
||||
///
|
||||
/// 返回格式: `{container_ip}:{grpc_port}`
|
||||
/// 默认 gRPC 端口为 50051
|
||||
pub async fn get_container_grpc_addr(project_id: &str, grpc_port: u16) -> anyhow::Result<String> {
|
||||
info!(
|
||||
"🔍 [CONTAINER] Getting container gRPC address: project_id={}",
|
||||
project_id
|
||||
);
|
||||
|
||||
let runtime = docker_manager::runtime::RuntimeManager::get()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to get runtime: {}", e))?;
|
||||
|
||||
let agent_info = runtime
|
||||
.get_container_info_by_identifier(project_id, &shared_types::ServiceType::RCoder)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("Failed to get container info: {}", e))?
|
||||
.ok_or_else(|| anyhow::anyhow!("Container info not found: project_id={}", project_id))?;
|
||||
|
||||
let grpc_addr = format!("{}:{}", agent_info.container_ip, grpc_port);
|
||||
|
||||
info!("[CONTAINER] get container gRPC addr: {}", grpc_addr);
|
||||
Ok(grpc_addr)
|
||||
}
|
||||
|
||||
/// 创建 gRPC 流异常错误事件
|
||||
///
|
||||
/// 当 gRPC 流在传输过程中异常结束时发送此事件
|
||||
fn create_grpc_stream_error_event(
|
||||
session_id: &str,
|
||||
code: Code,
|
||||
message: &str,
|
||||
) -> axum::response::sse::Event {
|
||||
// 使用项目标准的错误码映射
|
||||
let error_code = map_tonic_code_to_error_code(code);
|
||||
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type: SessionMessageType::SessionPromptEnd,
|
||||
sub_type: "error".to_string(),
|
||||
data: serde_json::json!({
|
||||
"code": error_code,
|
||||
"message": "Agent computer execution error, please retry (tasks consuming too much memory may cause the agent computer process to terminate).",
|
||||
}),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to serialize gRPC stream error event: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
// 返回包含基本信息的最小结构
|
||||
format!(
|
||||
r#"{{"session_id":"{}","message_type":"SessionPromptEnd","sub_type":"error","data":{{"code":"{}","message":"Agent computer execution error, please retry (tasks consuming too much memory may cause the agent computer process to terminate)."}}}}"#,
|
||||
session_id, error_code
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
axum::response::sse::Event::default()
|
||||
.event("error")
|
||||
.data(json_data)
|
||||
}
|
||||
|
||||
/// 创建连接失败错误事件
|
||||
///
|
||||
/// 当 gRPC 连接建立失败(重试后)时发送此事件
|
||||
fn create_connection_error_event(session_id: &str, message: &str) -> axum::response::sse::Event {
|
||||
let unified_message = UnifiedSessionMessage {
|
||||
session_id: session_id.to_string(),
|
||||
message_type: SessionMessageType::SessionPromptEnd,
|
||||
sub_type: "error".to_string(),
|
||||
data: serde_json::json!({
|
||||
"code": "GRPC_CONNECTION_FAILED",
|
||||
"message": message,
|
||||
}),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
|
||||
let json_data = match serde_json::to_string(&unified_message) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"⚠️ [gRPC_SSE] Failed to serialize connection error event: session_id={}, error={}",
|
||||
session_id, e
|
||||
);
|
||||
// 返回包含基本信息的最小结构
|
||||
format!(
|
||||
r#"{{"session_id":"{}","message_type":"SessionPromptEnd","sub_type":"error","data":{{"code":"GRPC_CONNECTION_FAILED","message":"Connection failed"}}}}"#,
|
||||
session_id
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
axum::response::sse::Event::default()
|
||||
.event("error")
|
||||
.data(json_data)
|
||||
}
|
||||
|
||||
/// 将 tonic::Code 映射为业务错误码
|
||||
fn map_tonic_code_to_error_code(code: Code) -> &'static str {
|
||||
match code {
|
||||
Code::Unavailable => "GRPC_SERVICE_UNAVAILABLE",
|
||||
Code::Cancelled => "GRPC_CANCELLED",
|
||||
Code::DeadlineExceeded => "GRPC_TIMEOUT",
|
||||
Code::Unknown => "GRPC_UNKNOWN_ERROR",
|
||||
_ => "GRPC_ERROR",
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user