添加qiming-rcoder模块

This commit is contained in:
Codex
2026-06-01 13:54:52 +08:00
parent 8092c4b1f8
commit 4b1a580132
539 changed files with 151650 additions and 0 deletions

View 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 连接池 TTL5分钟
///
/// 连接超过此时间未被使用则自动清理,防止内存泄漏。
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);
}
}

View 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)
}

View 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(),
}
}

View 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));
}
}

View 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");
}
}

View 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::*;

View 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",
}
}