Skip to content
Clone

feat(sagent): add sandbox agent crate skeleton

ZhenYi committed
parent 17c9a63 commit 156e960 Browse files

14 files changed

+664 −0

File tree 14
Added lib/sagent/src/gateway/mod.rs View lib/sagent/src/gateway/mod.rs
@@ -0,0 +1,204 @@
1use std::future::Future;
2use std::time::Duration;
3
4use agent_sdk::{AdminClient, Client, EventStream, SdkError};
5use agent_types::api::{Page, RunDto};
6use agent_types::ids::{RunId, SessionId, TenantId};
7use agent_types::session::{MessageAccepted, MessageDto, SessionDto, SessionTokenDto};
8use config::sagent::SagentConfig;
9use smol_str::SmolStr;
10use time::OffsetDateTime;
11use uuid::Uuid;
12
13use crate::errors::SagentError;
14
15pub const INVOKE_SCOPES: [&str; 6] = [
16 "sessions:read",
17 "sessions:write",
18 "sessions:events:subscribe",
19 "runs:read",
20 "runs:write",
21 "runs:cancel",
22];
23
24pub struct AgentGateway {
25 admin: AdminClient,
26 base: String,
27 tenant: TenantId,
28 token_ttl_secs: u64,
29 timeout: Duration,
30}
31
32impl AgentGateway {
33 pub fn new(cfg: &SagentConfig) -> Result<Self, SagentError> {
34 if !cfg.is_on() {
35 return Err(SagentError::Unavailable("sagent is not enabled".into()));
36 }
37 let base = cfg.base();
38 let admin = AdminClient::new(&base, cfg.admin_key.clone())?;
39 let tenant = Uuid::parse_str(cfg.tenant_id.trim())
40 .map_err(|_| SagentError::BadRequest("sagent tenant id is not a uuid".into()))?;
41 Ok(Self {
42 admin,
43 base,
44 tenant: TenantId::from_uuid(tenant),
45 token_ttl_secs: cfg.session_ttl_secs,
46 timeout: Duration::from_secs(cfg.request_timeout_secs.max(1)),
47 })
48 }
49
50 pub fn base(&self) -> &str {
51 &self.base
52 }
53
54 pub fn tenant(&self) -> TenantId {
55 self.tenant
56 }
57
58 pub fn token_ttl_secs(&self) -> u64 {
59 self.token_ttl_secs
60 }
61
62 async fn timed<T>(
63 &self,
64 future: impl Future<Output = Result<T, SdkError>>,
65 ) -> Result<T, SagentError> {
66 match tokio::time::timeout(self.timeout, future).await {
67 Ok(result) => result.map_err(SagentError::from),
68 Err(_) => Err(SagentError::Unavailable("agent gateway request timed out".into())),
69 }
70 }
71
72 pub async fn ensure_session(
73 &self,
74 existing: Option<Uuid>,
75 title: &str,
76 metadata: Option<serde_json::Value>,
77 ) -> Result<SessionDto, SagentError> {
78 if let Some(id) = existing {
79 match self.timed(self.admin.get_session(SessionId::from_uuid(id))).await {
80 Ok(session) => return Ok(session),
81 Err(SagentError::NotFound) => {}
82 Err(error) => return Err(error),
83 }
84 }
85 self.timed(self.admin.create_session(None, Some(title), metadata)).await
86 }
87
88 pub async fn list_session_messages(
89 &self,
90 session_id: Uuid,
91 limit: Option<u32>,
92 cursor: Option<u64>,
93 ) -> Result<Page<MessageDto>, SagentError> {
94 self.timed(self.admin.list_messages(SessionId::from_uuid(session_id), limit, cursor)).await
95 }
96
97 pub async fn open_session(
98 &self,
99 session_id: Uuid,
100 subject: Uuid,
101 scopes: &[&str],
102 ) -> Result<SessionHandle, SagentError> {
103 let session = SessionId::from_uuid(session_id);
104 let token = self
105 .timed(self.admin.issue_session_token_with_scopes(
106 session,
107 self.tenant,
108 &subject.to_string(),
109 Some(self.token_ttl_secs),
110 Some(scopes.to_vec()),
111 ))
112 .await?;
113 let expires_in = token.expires_in_secs.min(i64::MAX as u64) as i64;
114 let expires_at = OffsetDateTime::now_utc() + time::Duration::seconds(expires_in);
115 let client = Client::new(&self.base, token.token.to_string(), session)?;
116 Ok(SessionHandle { client, session, token, expires_at })
117 }
118}
119
120pub struct SessionHandle {
121 client: Client,
122 session: SessionId,
123 token: SessionTokenDto,
124 expires_at: OffsetDateTime,
125}
126
127impl SessionHandle {
128 pub fn session_id(&self) -> Uuid {
129 self.session.as_uuid()
130 }
131
132 pub fn token(&self) -> &str {
133 self.token.token.as_str()
134 }
135
136 pub fn grant_id(&self) -> Uuid {
137 self.token.grant_id.as_uuid()
138 }
139
140 pub fn expires_in_secs(&self) -> u64 {
141 self.token.expires_in_secs
142 }
143
144 pub fn expires_at(&self) -> OffsetDateTime {
145 self.expires_at
146 }
147
148 pub fn is_expired(&self) -> bool {
149 OffsetDateTime::now_utc() >= self.expires_at
150 }
151
152 pub async fn get_session(&self) -> Result<SessionDto, SagentError> {
153 self.client.get_session().await.map_err(SagentError::from)
154 }
155
156 pub async fn list_messages(
157 &self,
158 limit: Option<u32>,
159 cursor: Option<u64>,
160 ) -> Result<Page<MessageDto>, SagentError> {
161 self.client.list_messages(limit, cursor).await.map_err(SagentError::from)
162 }
163
164 pub async fn send_message(
165 &self,
166 text: impl Into<SmolStr>,
167 idempotency_key: Option<SmolStr>,
168 ) -> Result<MessageAccepted, SagentError> {
169 self.client.send_message(text, idempotency_key).await.map_err(SagentError::from)
170 }
171
172 pub async fn get_run(&self, run_id: Uuid) -> Result<RunDto, SagentError> {
173 self.client.get_run(RunId::from_uuid(run_id)).await.map_err(SagentError::from)
174 }
175
176 pub async fn cancel_run(&self, run_id: Uuid) -> Result<RunDto, SagentError> {
177 self.client.cancel_run(RunId::from_uuid(run_id)).await.map_err(SagentError::from)
178 }
179
180 pub async fn stream_events(&self, last_event_id: Option<u64>) -> Result<EventStream, SagentError> {
181 self.client.stream_events(last_event_id).await.map_err(SagentError::from)
182 }
183}
184
185impl From<SdkError> for SagentError {
186 fn from(error: SdkError) -> Self {
187 match error {
188 SdkError::Transport(message) | SdkError::Decode(message) => Self::Unavailable(message),
189 SdkError::RateLimited(message) => Self::RateLimited(message),
190 SdkError::NotFound(_) | SdkError::Gone(_) => Self::NotFound,
191 SdkError::Conflict(message) => Self::Conflict(message),
192 SdkError::BadRequest(message) => Self::BadRequest(message),
193 SdkError::Unauthorized(message) | SdkError::Forbidden(message) => Self::Upstream(message),
194 SdkError::Invalid(message) => Self::Internal(message),
195 SdkError::Server { status, message } => {
196 if status >= 500 {
197 Self::Unavailable(message)
198 } else {
199 Self::Upstream(message)
200 }
201 }
202 }
203 }
204}
No newline at end of file
Added lib/sagent/src/handlers/mod.rs View lib/sagent/src/handlers/mod.rs

No content changes.

Added lib/sagent/src/services/mod.rs View lib/sagent/src/services/mod.rs
@@ -0,0 +1 @@
1pub mod service;
No newline at end of file
Added lib/sagent/src/services/service.rs View lib/sagent/src/services/service.rs
@@ -0,0 +1,118 @@
1use std::sync::{Arc, OnceLock};
2
3use async_trait::async_trait;
4use config::AppConfig;
5use export::AppState;
6use export::database::AppDatabase;
7use models::sagent::agent_event_line::{self, AgentEventType};
8use models::users::user_notifications::{self, NotificationType};
9use sea_orm::{ActiveModelTrait, ConnectionTrait, Set};
10use session::SessionError;
11use session::share::guard::{SessionGuard, SessionWithUser};
12use time::OffsetDateTime;
13use uuid::Uuid;
14
15use crate::errors::SagentError;
16use crate::gateway::AgentGateway;
17
18#[derive(Clone)]
19pub struct SagentService {
20 pub state: AppState,
21 gateway: Arc<OnceLock<Result<AgentGateway, String>>>,
22}
23
24impl SagentService {
25 pub fn new(state: AppState) -> Self {
26 Self { state, gateway: Arc::new(OnceLock::new()) }
27 }
28
29 pub fn config(&self) -> &AppConfig {
30 &self.state.config
31 }
32
33 pub fn db(&self) -> &Arc<AppDatabase> {
34 &self.state.database
35 }
36
37 pub fn paginate(&self, limit: Option<u64>, offset: Option<u64>) -> (u64, u64) {
38 let max_pagination = self.config().sagent.max_pagination;
39 let default_pagination = self.config().sagent.default_pagination;
40 (
41 limit.unwrap_or(default_pagination).clamp(1, max_pagination),
42 offset.unwrap_or(0).clamp(0, 5000),
43 )
44 }
45
46 pub fn gateway(&self) -> Result<&AgentGateway, SagentError> {
47 let cached = self.gateway.get_or_init(|| {
48 AgentGateway::new(&self.config().sagent).map_err(|error| error.to_string())
49 });
50 match cached {
51 Ok(gateway) => Ok(gateway),
52 Err(message) => Err(SagentError::Unavailable(message.clone())),
53 }
54 }
55
56 pub(crate) async fn write_event(
57 &self,
58 txn: &impl ConnectionTrait,
59 agent_id: Uuid,
60 collection_id: Uuid,
61 user_id: Option<Uuid>,
62 event_type: AgentEventType,
63 run_id: Option<Uuid>,
64 content: Option<serde_json::Value>,
65 ) -> Result<(), SagentError> {
66 agent_event_line::ActiveModel {
67 id: Set(Uuid::now_v7()),
68 agent_id: Set(Some(agent_id)),
69 collection_id: Set(Some(collection_id)),
70 user_id: Set(user_id),
71 event_type: Set(event_type),
72 run_id: Set(run_id),
73 content: Set(content),
74 created_at: Set(OffsetDateTime::now_utc()),
75 }
76 .insert(txn)
77 .await?;
78 Ok(())
79 }
80
81 pub(crate) async fn notify_user(
82 &self,
83 txn: &impl ConnectionTrait,
84 pending: &export::notify::PendingNotifications,
85 recipient: Uuid,
86 ntype: NotificationType,
87 title: String,
88 body: Option<String>,
89 data: Option<serde_json::Value>,
90 actor_uid: Option<Uuid>,
91 ) -> Result<(), SagentError> {
92 let row = user_notifications::ActiveModel {
93 id: Set(Uuid::now_v7()),
94 user_id: Set(recipient),
95 notification_type: Set(ntype),
96 title: Set(title),
97 body: Set(body),
98 data_json: Set(data),
99 actor_uid: Set(actor_uid),
100 read_at: Set(None),
101 created_at: Set(OffsetDateTime::now_utc()),
102 }
103 .insert(txn)
104 .await?;
105 pending.push(recipient, row.to_event_value());
106 Ok(())
107 }
108}
109
110#[async_trait]
111impl SessionGuard for SagentService {
112 async fn lookup_session_with_user(
113 &self,
114 token_hash: &str,
115 ) -> Result<Option<SessionWithUser>, SessionError> {
116 session::lookup_session_with_user(self.db().read(), token_hash).await
117 }
118}
No newline at end of file
Added lib/sagent/src/share/authz.rs View lib/sagent/src/share/authz.rs
@@ -0,0 +1,32 @@
1use models::space::spaces;
2use space::share::authz::{MANAGE_CHANNELS, effective_permissions};
3use uuid::Uuid;
4
5use crate::errors::SagentError;
6use crate::services::service::SagentService;
7
8pub async fn require_collection_manage(
9 service: &SagentService,
10 space: &spaces::Model,
11 caller: Uuid,
12) -> Result<(), SagentError> {
13 if is_collection_manager(service, space, caller).await? {
14 Ok(())
15 } else {
16 Err(SagentError::Forbidden)
17 }
18}
19
20pub async fn is_collection_manager(
21 service: &SagentService,
22 space: &spaces::Model,
23 caller: Uuid,
24) -> Result<bool, SagentError> {
25 if space.owner_id == caller {
26 return Ok(true);
27 }
28 let perms = effective_permissions(service.db().read(), space, caller)
29 .await
30 .map_err(SagentError::from)?;
31 Ok(perms.has(MANAGE_CHANNELS))
32}
No newline at end of file
Added lib/sagent/src/share/guard.rs View lib/sagent/src/share/guard.rs
@@ -0,0 +1,16 @@
1use axum::extract::{Request, State};
2use axum::middleware::Next;
3use axum::response::Response;
4use session::Session;
5use session::share::guard::session_guard;
6
7use crate::services::service::SagentService;
8
9pub async fn sagent_guard(
10 State(service): State<SagentService>,
11 session: Session,
12 request: Request,
13 next: Next,
14) -> Response {
15 session_guard::<SagentService>(State(service), session, request, next).await
16}
No newline at end of file
Added lib/sagent/src/share/lookup.rs View lib/sagent/src/share/lookup.rs
@@ -0,0 +1,135 @@
1use models::sagent::{agent, agents};
2use models::space::spaces;
3use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
4use space::share::authz::{READ, VIEW_SPACE, effective_permissions};
5use uuid::Uuid;
6
7use crate::errors::SagentError;
8use crate::services::service::SagentService;
9
10pub async fn lookup_space_by_name(
11 service: &SagentService,
12 space_name: &str,
13) -> Result<spaces::Model, SagentError> {
14 let lower = space_name.trim().to_lowercase();
15 if lower.is_empty() {
16 return Err(SagentError::NotFound);
17 }
18 spaces::Entity::find()
19 .filter(spaces::Column::LowerName.eq(&lower))
20 .filter(spaces::Column::IsActive.eq(true))
21 .one(service.db().read())
22 .await?
23 .ok_or(SagentError::NotFound)
24}
25
26pub async fn require_member(
27 service: &SagentService,
28 space: &spaces::Model,
29 caller: Uuid,
30) -> Result<(), SagentError> {
31 let perms = effective_permissions(service.db().read(), space, caller)
32 .await
33 .map_err(SagentError::from)?;
34 if perms.has(VIEW_SPACE) || perms.has(READ) {
35 Ok(())
36 } else {
37 Err(SagentError::Forbidden)
38 }
39}
40
41pub async fn space_context(
42 service: &SagentService,
43 caller: Uuid,
44 space_name: &str,
45) -> Result<spaces::Model, SagentError> {
46 let space = lookup_space_by_name(service, space_name).await?;
47 require_member(service, &space, caller).await?;
48 Ok(space)
49}
50
51pub async fn lookup_collection(
52 service: &SagentService,
53 space: &spaces::Model,
54 collection_id: Uuid,
55 caller: Uuid,
56) -> Result<agents::Model, SagentError> {
57 let collection = agents::Entity::find_by_id(collection_id)
58 .one(service.db().read())
59 .await?
60 .ok_or(SagentError::NotFound)?;
61
62 let visible = models::dir::dirs::Entity::find_by_id(collection.dir_id)
63 .filter(models::dir::dirs::Column::SpaceId.eq(Some(space.id)))
64 .filter(models::dir::dirs::Column::Kind.eq(models::dir::dirs::DirKind::Object))
65 .one(service.db().read())
66 .await?
67 .is_some();
68 if !visible {
69 return Err(SagentError::NotFound);
70 }
71
72 if !space::share::dir_visibility::is_dir_visible(
73 service.db().read(),
74 space,
75 collection.dir_id,
76 caller,
77 )
78 .await?
79 {
80 return Err(SagentError::NotFound);
81 }
82
83 Ok(collection)
84}
85
86pub async fn lookup_task(
87 service: &SagentService,
88 collection: &agents::Model,
89 task_id: Uuid,
90) -> Result<agent::Model, SagentError> {
91 agent::Entity::find_by_id(task_id)
92 .filter(agent::Column::CollectionId.eq(collection.id))
93 .one(service.db().read())
94 .await?
95 .ok_or(SagentError::NotFound)
96}
97
98pub async fn lookup_task_by_ref(
99 service: &SagentService,
100 collection: &agents::Model,
101 task_ref: &str,
102) -> Result<agent::Model, SagentError> {
103 let base = agent::Entity::find().filter(agent::Column::CollectionId.eq(collection.id));
104 let found = if let Ok(id) = Uuid::parse_str(task_ref) {
105 base.filter(agent::Column::Id.eq(id)).one(service.db().read()).await?
106 } else if let Ok(number) = task_ref.parse::<i32>() {
107 base.filter(agent::Column::AutoincrementId.eq(number)).one(service.db().read()).await?
108 } else {
109 None
110 };
111 found.ok_or(SagentError::NotFound)
112}
113
114pub async fn collection_context(
115 service: &SagentService,
116 caller: Uuid,
117 space_name: &str,
118 collection_id: Uuid,
119) -> Result<(spaces::Model, agents::Model), SagentError> {
120 let space = space_context(service, caller, space_name).await?;
121 let collection = lookup_collection(service, &space, collection_id, caller).await?;
122 Ok((space, collection))
123}
124
125pub async fn task_context(
126 service: &SagentService,
127 caller: Uuid,
128 space_name: &str,
129 collection_id: Uuid,
130 task_ref: &str,
131) -> Result<(spaces::Model, agents::Model, agent::Model), SagentError> {
132 let (space, collection) = collection_context(service, caller, space_name, collection_id).await?;
133 let task = lookup_task_by_ref(service, &collection, task_ref).await?;
134 Ok((space, collection, task))
135}
No newline at end of file
Added lib/sagent/src/share/mod.rs View lib/sagent/src/share/mod.rs
@@ -0,0 +1,3 @@
1pub mod authz;
2pub mod guard;
3pub mod lookup;
No newline at end of file
Added lib/sagent/src/errors.rs View lib/sagent/src/errors.rs
@@ -0,0 +1,75 @@
1use std::fmt;
2
3#[derive(Debug)]
4pub enum SagentError {
5 BadRequest(String),
6 Unauthorized,
7 Forbidden,
8 NotFound,
9 Conflict(String),
10 RateLimited(String),
11 Upstream(String),
12 Unavailable(String),
13 Internal(String),
14}
15
16impl fmt::Display for SagentError {
17 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
18 match self {
19 Self::BadRequest(message) => write!(f, "{message}"),
20 Self::Unauthorized => write!(f, "unauthorized"),
21 Self::Forbidden => write!(f, "forbidden"),
22 Self::NotFound => write!(f, "not found"),
23 Self::Conflict(message) => write!(f, "{message}"),
24 Self::RateLimited(message) => write!(f, "{message}"),
25 Self::Upstream(message) => write!(f, "upstream rejected the request: {message}"),
26 Self::Unavailable(message) => write!(f, "agent gateway unavailable: {message}"),
27 Self::Internal(message) => write!(f, "{message}"),
28 }
29 }
30}
31
32impl std::error::Error for SagentError {}
33
34export::db_error_conversions!(SagentError);
35export::export_error_conversion!(SagentError);
36export::json_payload_conversion!(SagentError);
37export::domain_error_passthrough!(SagentError from space::errors::SpaceError);
38
39export::api_error! {
40 log: "sagent",
41 pub enum SagentApiError {
42 BadRequest(400, "request failed validation"),
43 Unauthorized(401, "authentication required"),
44 Forbidden(403, "operation not allowed"),
45 NotFound(404, "resource not found"),
46 Conflict(409, "resource conflict"),
47 TooManyRequests(429, "rate limit exceeded"),
48 Internal(500, "internal server error"),
49 BadGateway(502, "agent gateway rejected the request"),
50 ServiceUnavailable(503, "agent gateway unavailable"),
51 }
52}
53
54impl From<SagentError> for SagentApiError {
55 fn from(error: SagentError) -> Self {
56 match error {
57 SagentError::BadRequest(msg) => Self::bad_request(msg),
58 SagentError::Unauthorized => Self::unauthorized(),
59 SagentError::Forbidden => Self::forbidden(),
60 SagentError::NotFound => Self::not_found("sagent resource not found"),
61 SagentError::Conflict(msg) => Self::conflict(msg),
62 SagentError::RateLimited(msg) => Self::rate_limited(msg),
63 SagentError::Upstream(msg) => Self::build(
64 502,
65 export::rest::AppResponseData::err(502, msg).with_error_code("gateway_rejected"),
66 ),
67 SagentError::Unavailable(msg) => Self::build(
68 503,
69 export::rest::AppResponseData::err(503, msg)
70 .with_error_code("gateway_unavailable"),
71 ),
72 SagentError::Internal(message) => Self::internal(message),
73 }
74 }
75}
No newline at end of file
Added lib/sagent/src/lib.rs View lib/sagent/src/lib.rs
@@ -0,0 +1,7 @@
1pub mod errors;
2pub mod gateway;
3pub mod handlers;
4pub mod openapi;
5pub mod routes;
6pub mod services;
7pub mod share;
No newline at end of file
Added lib/sagent/src/openapi.rs View lib/sagent/src/openapi.rs
@@ -0,0 +1,12 @@
1use export::rest::AppResponseData;
2use utoipa::OpenApi;
3
4#[derive(OpenApi)]
5#[openapi(
6 paths(),
7 components(schemas(AppResponseData<export::rest::NoData>, export::rest::NoData)),
8 tags(
9 (name = "sagent", description = "Sandbox agent API")
10 )
11)]
12pub struct SagentsApiDoc;
No newline at end of file
Added lib/sagent/src/routes.rs View lib/sagent/src/routes.rs
@@ -0,0 +1,11 @@
1use axum::Router;
2
3use crate::services::service::SagentService;
4use crate::share::guard::sagent_guard;
5
6pub fn router(service: SagentService) -> Router<SagentService> {
7 Router::new().route_layer(axum::middleware::from_fn_with_state(
8 service.clone(),
9 sagent_guard,
10 ))
11}
No newline at end of file
Added lib/sagent/Cargo.toml View lib/sagent/Cargo.toml
@@ -0,0 +1,25 @@
1[package]
2name = "sagent"
3version.workspace = true
4edition.workspace = true
5
6[dependencies]
7agent-sdk = { workspace = true }
8agent-types = { workspace = true }
9async-trait = { workspace = true }
10axum = { workspace = true }
11config = { workspace = true }
12export = { workspace = true }
13futures = { workspace = true }
14models = { workspace = true }
15sea-orm = { workspace = true }
16serde = { workspace = true, features = ["derive"] }
17serde_json = { workspace = true }
18session = { workspace = true }
19smol_str = { workspace = true }
20space = { workspace = true }
21time = { workspace = true }
22tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
23tracing = { workspace = true }
24utoipa = { workspace = true }
25uuid = { workspace = true }
No newline at end of file
Modified Cargo.lock View Cargo.lock
@@ -8545,6 +8545,31 @@
85458545checksum = "ef703b7cb59335eae2eb93ceb664c0eb7ea6bf567079d843e09420219668e072"
85468546
85478547[[package]]
8548name = "sagent"
8549version = "1.0.0"
8550dependencies = [
8551 "agent-sdk",
8552 "agent-types",
8553 "async-trait",
8554 "axum",
8555 "config",
8556 "export",
8557 "futures",
8558 "models",
8559 "sea-orm",
8560 "serde",
8561 "serde_json",
8562 "session",
8563 "smol_str",
8564 "space",
8565 "time 0.3.55",
8566 "tokio",
8567 "tracing",
8568 "utoipa",
8569 "uuid",
8570]
8571
8572[[package]]
85488573name = "salsa20"
85498574version = "0.11.0"
85508575source = "registry+https://github.com/rust-lang/crates.io-index"