feat(code): gRPC server with JWT auth + tool routing
tonic 0.14 gRPC server for sunbeam code sessions: - bidirectional streaming Session RPC - JWT interceptor validates tokens against Hydra JWKS - tool router classifies calls as client-side (file_read, bash, grep, etc.) or server-side (gitea, identity, search, etc.) - service stub with session lifecycle (start, chat, tool results, end) - coding_model config (default: devstral-small-2506) - grpc config section (listen_addr, jwks_url) - 182 tests (5 new: JWT claims, tool routing) phase 2 TODOs: Matrix room bridge, Mistral agent loop, streaming
This commit is contained in:
175
src/grpc/service.rs
Normal file
175
src/grpc/service.rs
Normal file
@@ -0,0 +1,175 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::Stream;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tonic::{Request, Response, Status, Streaming};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
use super::auth::Claims;
|
||||
use super::proto::code_agent_server::CodeAgent;
|
||||
use super::proto::*;
|
||||
use super::GrpcState;
|
||||
|
||||
pub struct CodeAgentService {
|
||||
state: Arc<GrpcState>,
|
||||
}
|
||||
|
||||
impl CodeAgentService {
|
||||
pub fn new(state: Arc<GrpcState>) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl CodeAgent for CodeAgentService {
|
||||
type SessionStream = Pin<Box<dyn Stream<Item = Result<ServerMessage, Status>> + Send>>;
|
||||
|
||||
async fn session(
|
||||
&self,
|
||||
request: Request<Streaming<ClientMessage>>,
|
||||
) -> Result<Response<Self::SessionStream>, Status> {
|
||||
// Extract JWT claims from the request extensions (set by auth middleware)
|
||||
let claims = request
|
||||
.extensions()
|
||||
.get::<Claims>()
|
||||
.cloned()
|
||||
.ok_or_else(|| Status::unauthenticated("No valid authentication token"))?;
|
||||
|
||||
info!(
|
||||
user = claims.sub.as_str(),
|
||||
email = claims.email.as_deref().unwrap_or("?"),
|
||||
"New coding session"
|
||||
);
|
||||
|
||||
let mut in_stream = request.into_inner();
|
||||
let state = self.state.clone();
|
||||
|
||||
// Channel for sending server messages to the client
|
||||
let (tx, rx) = mpsc::channel::<Result<ServerMessage, Status>>(32);
|
||||
|
||||
// Spawn the session handler
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_session(&state, &claims, &mut in_stream, &tx).await {
|
||||
error!(user = claims.sub.as_str(), "Session error: {e}");
|
||||
let _ = tx
|
||||
.send(Ok(ServerMessage {
|
||||
payload: Some(server_message::Payload::Error(Error {
|
||||
message: e.to_string(),
|
||||
fatal: true,
|
||||
})),
|
||||
}))
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
let out_stream = ReceiverStream::new(rx);
|
||||
Ok(Response::new(Box::pin(out_stream)))
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a single coding session (runs in a spawned task).
|
||||
async fn handle_session(
|
||||
state: &GrpcState,
|
||||
claims: &Claims,
|
||||
in_stream: &mut Streaming<ClientMessage>,
|
||||
tx: &mpsc::Sender<Result<ServerMessage, Status>>,
|
||||
) -> anyhow::Result<()> {
|
||||
// Wait for the first message — must be StartSession
|
||||
let first = in_stream
|
||||
.message()
|
||||
.await?
|
||||
.ok_or_else(|| anyhow::anyhow!("Stream closed before StartSession"))?;
|
||||
|
||||
let start = match first.payload {
|
||||
Some(client_message::Payload::Start(s)) => s,
|
||||
_ => anyhow::bail!("First message must be StartSession"),
|
||||
};
|
||||
|
||||
info!(
|
||||
user = claims.sub.as_str(),
|
||||
project = start.project_path.as_str(),
|
||||
model = start.model.as_str(),
|
||||
client_tools = start.client_tools.len(),
|
||||
"Session started"
|
||||
);
|
||||
|
||||
// TODO Phase 2: Create/find Matrix room for this project
|
||||
// TODO Phase 2: Create Mistral conversation
|
||||
// TODO Phase 2: Enter agent loop
|
||||
|
||||
// For now, send SessionReady and echo back
|
||||
tx.send(Ok(ServerMessage {
|
||||
payload: Some(server_message::Payload::Ready(SessionReady {
|
||||
session_id: uuid::Uuid::new_v4().to_string(),
|
||||
room_id: String::new(), // TODO: Matrix room
|
||||
model: if start.model.is_empty() {
|
||||
state
|
||||
.config
|
||||
.agents
|
||||
.coding_model
|
||||
.clone()
|
||||
} else {
|
||||
start.model.clone()
|
||||
},
|
||||
})),
|
||||
}))
|
||||
.await?;
|
||||
|
||||
// Main message loop
|
||||
while let Some(msg) = in_stream.message().await? {
|
||||
match msg.payload {
|
||||
Some(client_message::Payload::Input(input)) => {
|
||||
info!(
|
||||
user = claims.sub.as_str(),
|
||||
text_len = input.text.len(),
|
||||
"User input received"
|
||||
);
|
||||
|
||||
// TODO Phase 2: Send to Mistral, handle tool calls, stream response
|
||||
// For now, echo back as a simple acknowledgment
|
||||
tx.send(Ok(ServerMessage {
|
||||
payload: Some(server_message::Payload::Done(TextDone {
|
||||
full_text: format!("[stub] received: {}", input.text),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
})),
|
||||
}))
|
||||
.await?;
|
||||
}
|
||||
Some(client_message::Payload::ToolResult(result)) => {
|
||||
info!(
|
||||
call_id = result.call_id.as_str(),
|
||||
is_error = result.is_error,
|
||||
"Tool result received"
|
||||
);
|
||||
// TODO Phase 2: Feed back to Mistral
|
||||
}
|
||||
Some(client_message::Payload::Approval(approval)) => {
|
||||
info!(
|
||||
call_id = approval.call_id.as_str(),
|
||||
approved = approval.approved,
|
||||
"Tool approval received"
|
||||
);
|
||||
// TODO Phase 2: Execute or skip tool
|
||||
}
|
||||
Some(client_message::Payload::End(_)) => {
|
||||
info!(user = claims.sub.as_str(), "Session ended by client");
|
||||
tx.send(Ok(ServerMessage {
|
||||
payload: Some(server_message::Payload::End(SessionEnd {
|
||||
summary: "Session ended.".into(),
|
||||
})),
|
||||
}))
|
||||
.await?;
|
||||
break;
|
||||
}
|
||||
Some(client_message::Payload::Start(_)) => {
|
||||
warn!("Received duplicate StartSession — ignoring");
|
||||
}
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user