dk_protocol/
server.rs

1use std::sync::Arc;
2
3use dk_engine::conflict::{ClaimTracker, LocalClaimTracker};
4use dk_engine::repo::Engine;
5use tonic::{Request, Response, Status};
6
7use crate::auth::AuthConfig;
8use crate::events::EventBus;
9use crate::session::{AgentSession, SessionManager};
10
11/// The gRPC server that implements the `AgentService` trait generated by
12/// tonic from the `agent.proto` definition.
13///
14/// Holds a shared reference to the [`Engine`] (internally concurrent via
15/// fine-grained locks) and a [`SessionManager`] for stateful agent sessions.
16pub struct ProtocolServer {
17    pub(crate) engine: Arc<Engine>,
18    pub(crate) session_mgr: Arc<SessionManager>,
19    pub(crate) auth_config: AuthConfig,
20    pub(crate) event_bus: Arc<EventBus>,
21    pub(crate) claim_tracker: Arc<dyn ClaimTracker>,
22}
23
24impl ProtocolServer {
25    /// Create a new `ProtocolServer`.
26    ///
27    /// `auth_config` controls how agents authenticate on every
28    /// `ConnectRequest`.  The session timeout is fixed at 30 minutes.
29    pub fn new(engine: Arc<Engine>, auth_config: AuthConfig) -> Self {
30        Self {
31            engine,
32            session_mgr: Arc::new(SessionManager::new(std::time::Duration::from_secs(
33                30 * 60,
34            ))),
35            auth_config,
36            event_bus: Arc::new(EventBus::new()),
37            claim_tracker: Arc::new(LocalClaimTracker::new()),
38        }
39    }
40
41    /// Borrow the engine.
42    pub fn engine(&self) -> &Engine {
43        &self.engine
44    }
45
46    /// Borrow the session manager.
47    pub fn session_mgr(&self) -> &SessionManager {
48        &self.session_mgr
49    }
50
51    /// Borrow the shared event bus.
52    pub fn event_bus(&self) -> &EventBus {
53        &self.event_bus
54    }
55
56    /// Return a cloned `Arc` handle to the shared event bus.
57    ///
58    /// Callers that need to publish to the engine event bus from a spawned
59    /// task (where a borrow of `&ProtocolServer` is not available) should
60    /// use this to obtain a cheaply-clonable, ownership-safe handle.
61    pub fn event_bus_arc(&self) -> Arc<EventBus> {
62        Arc::clone(&self.event_bus)
63    }
64
65    /// Borrow the shared symbol claim tracker.
66    pub fn claim_tracker(&self) -> &dyn ClaimTracker {
67        &*self.claim_tracker
68    }
69
70    /// Create a `ProtocolServer` with a custom claim tracker implementation.
71    pub fn with_claim_tracker(
72        engine: Arc<Engine>,
73        auth_config: AuthConfig,
74        claim_tracker: Arc<dyn ClaimTracker>,
75    ) -> Self {
76        Self {
77            engine,
78            session_mgr: Arc::new(SessionManager::new(std::time::Duration::from_secs(
79                30 * 60,
80            ))),
81            auth_config,
82            event_bus: Arc::new(EventBus::new()),
83            claim_tracker,
84        }
85    }
86
87    /// Validate an auth token against the configured secret.
88    pub(crate) fn validate_auth(&self, token: &str) -> Result<String, Status> {
89        self.auth_config.validate(token)
90    }
91
92    /// Check whether the bearer token in `metadata` carries admin scope.
93    ///
94    /// Extracts the `authorization: Bearer <token>` header, performs a full
95    /// signature + expiry validation using the server's configured secret, and
96    /// returns `true` when the `scope` claim equals `"admin"` or contains the
97    /// word `"admin"`.
98    ///
99    /// Returns `false` (not an error) if there is no auth header, the token
100    /// is not a JWT, or signature validation fails — callers can then fall
101    /// through to the owner-check.
102    pub(crate) fn has_admin_scope(&self, metadata: &tonic::metadata::MetadataMap) -> bool {
103        let Some(val) = metadata.get("authorization") else {
104            return false;
105        };
106        let Ok(header) = val.to_str() else {
107            return false;
108        };
109        let token = header.strip_prefix("Bearer ").unwrap_or(header);
110        // Perform full signature validation via the configured AuthConfig; only
111        // tokens that pass cryptographic verification can claim admin scope.
112        self.auth_config
113            .validate_claims(token)
114            .as_ref()
115            .map(crate::auth::claims_have_admin_scope)
116            .unwrap_or(false)
117    }
118
119    /// Look up a session by its string-encoded UUID.  Returns an error if the
120    /// ID is malformed or the session has expired / does not exist.
121    pub(crate) fn validate_session(&self, session_id_str: &str) -> Result<AgentSession, Status> {
122        let sid = session_id_str
123            .parse::<uuid::Uuid>()
124            .map_err(|_| Status::invalid_argument("Invalid session ID format"))?;
125        self.session_mgr
126            .get_session(&sid)
127            .ok_or_else(|| Status::not_found("Session not found or expired"))
128    }
129}
130
131// ── AgentService tonic trait implementation ──
132
133#[tonic::async_trait]
134impl crate::agent_service_server::AgentService for ProtocolServer {
135    async fn connect(
136        &self,
137        request: Request<crate::ConnectRequest>,
138    ) -> Result<Response<crate::ConnectResponse>, Status> {
139        crate::connect::handle_connect(self, request.into_inner()).await
140    }
141
142    async fn context(
143        &self,
144        request: Request<crate::ContextRequest>,
145    ) -> Result<Response<crate::ContextResponse>, Status> {
146        crate::context::handle_context(self, request.into_inner()).await
147    }
148
149    async fn submit(
150        &self,
151        request: Request<crate::SubmitRequest>,
152    ) -> Result<Response<crate::SubmitResponse>, Status> {
153        crate::submit::handle_submit(self, request.into_inner()).await
154    }
155
156    type VerifyStream = tokio_stream::wrappers::ReceiverStream<Result<crate::VerifyStepResult, Status>>;
157
158    async fn verify(
159        &self,
160        request: Request<crate::VerifyRequest>,
161    ) -> Result<Response<Self::VerifyStream>, Status> {
162        let req = request.into_inner();
163        let (tx, rx) = tokio::sync::mpsc::channel(32);
164
165        let server_clone = ProtocolServer {
166            engine: self.engine.clone(),
167            session_mgr: self.session_mgr.clone(),
168            auth_config: self.auth_config.clone(),
169            event_bus: self.event_bus.clone(),
170            claim_tracker: self.claim_tracker.clone(),
171        };
172
173        tokio::spawn(async move {
174            crate::verify::handle_verify(&server_clone, req, tx).await;
175        });
176
177        Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(rx)))
178    }
179
180    async fn merge(
181        &self,
182        request: Request<crate::MergeRequest>,
183    ) -> Result<Response<crate::MergeResponse>, Status> {
184        let resp = crate::merge::handle_merge(self, request.into_inner()).await?;
185        Ok(Response::new(resp))
186    }
187
188    type WatchStream = tokio_stream::wrappers::ReceiverStream<Result<crate::WatchEvent, Status>>;
189
190    async fn watch(
191        &self,
192        request: Request<crate::WatchRequest>,
193    ) -> Result<Response<Self::WatchStream>, Status> {
194        let req = request.into_inner();
195        let (tx, rx) = tokio::sync::mpsc::channel(64);
196        let server_clone = ProtocolServer {
197            engine: self.engine.clone(),
198            session_mgr: self.session_mgr.clone(),
199            auth_config: self.auth_config.clone(),
200            event_bus: self.event_bus.clone(),
201            claim_tracker: self.claim_tracker.clone(),
202        };
203        tokio::spawn(async move {
204            crate::watch::handle_watch(&server_clone, req, tx).await;
205        });
206        Ok(Response::new(tokio_stream::wrappers::ReceiverStream::new(rx)))
207    }
208
209    async fn file_read(
210        &self,
211        request: Request<crate::FileReadRequest>,
212    ) -> Result<Response<crate::FileReadResponse>, Status> {
213        crate::file_read::handle_file_read(self, request.into_inner()).await
214    }
215
216    async fn file_write(
217        &self,
218        request: Request<crate::FileWriteRequest>,
219    ) -> Result<Response<crate::FileWriteResponse>, Status> {
220        crate::file_write::handle_file_write(self, request.into_inner()).await
221    }
222
223    async fn file_list(
224        &self,
225        request: Request<crate::FileListRequest>,
226    ) -> Result<Response<crate::FileListResponse>, Status> {
227        crate::file_list::handle_file_list(self, request.into_inner()).await
228    }
229
230    async fn pre_submit_check(
231        &self,
232        request: Request<crate::PreSubmitCheckRequest>,
233    ) -> Result<Response<crate::PreSubmitCheckResponse>, Status> {
234        crate::pre_submit::handle_pre_submit_check(self, request.into_inner()).await
235    }
236
237    async fn get_session_status(
238        &self,
239        request: Request<crate::SessionStatusRequest>,
240    ) -> Result<Response<crate::SessionStatusResponse>, Status> {
241        crate::session_status::handle_get_session_status(self, request.into_inner()).await
242    }
243
244    async fn push(
245        &self,
246        request: Request<crate::PushRequest>,
247    ) -> Result<Response<crate::PushResponse>, Status> {
248        let resp = crate::push::handle_push(self, request.into_inner()).await?;
249        Ok(Response::new(resp))
250    }
251
252    async fn approve(
253        &self,
254        _request: Request<crate::ApproveRequest>,
255    ) -> Result<Response<crate::ApproveResponse>, Status> {
256        Err(Status::unimplemented(
257            "approve is a platform-level operation; use the managed server",
258        ))
259    }
260
261    async fn resolve(
262        &self,
263        _request: Request<crate::ResolveRequest>,
264    ) -> Result<Response<crate::ResolveResponse>, Status> {
265        Err(Status::unimplemented(
266            "resolve is a platform-level operation; use the managed server",
267        ))
268    }
269
270    async fn close(
271        &self,
272        _request: Request<crate::CloseRequest>,
273    ) -> Result<Response<crate::CloseResponse>, Status> {
274        Err(Status::unimplemented(
275            "close is a platform-level operation; use the managed server",
276        ))
277    }
278
279    async fn abandon(
280        &self,
281        request: Request<crate::AbandonRequest>,
282    ) -> Result<Response<crate::AbandonResponse>, Status> {
283        crate::abandon::handle_abandon(self, request).await
284    }
285
286    async fn review(
287        &self,
288        _request: Request<crate::ReviewRequest>,
289    ) -> Result<Response<crate::ReviewResponse>, Status> {
290        Err(Status::unimplemented(
291            "review is a platform-level operation; use the managed server",
292        ))
293    }
294
295    async fn record_review(
296        &self,
297        _request: Request<crate::RecordReviewRequest>,
298    ) -> Result<Response<crate::RecordReviewResponse>, Status> {
299        Err(Status::unimplemented(
300            "record_review is a platform-level operation; use the managed server",
301        ))
302    }
303}
dk_protocol/server.rs

dk_protocol/
server.rs
