enact_grpc/

lib.rs

pub mod proto {
    pub mod runtime {
        pub mod v1 {
            tonic::include_proto!("enact.runtime.v1");
        }
    }
}

use enact_cluster::{
    load_cluster_config, DispatchStrategy, DispatchTask, Dispatcher, RedisNodeRegistry,
    RedisTaskQueue, TaskId,
};
use enact_core::callable::{CallableRegistry, LlmCallable};
use enact_core::graph::{Checkpoint, CheckpointStore, InMemoryCheckpointStore};
use enact_core::inbox::{
    ControlAction, ControlMessage, GuidanceMessage, GuidancePriority, InMemoryInboxStore,
    InboxMessage, InboxStore,
};
use enact_core::kernel::ExecutionId;
use enact_core::kernel::SpawnMode;
use enact_core::providers::ModelProvider;
use enact_providers::factory::create_default_provider;
use proto::runtime::v1::*;
use std::pin::Pin;
use std::sync::Arc;
use tokio_stream::{wrappers::ReceiverStream, Stream};
use tonic::{Request, Response, Status};
use tracing::{error, info};

pub use proto::runtime::v1::runtime_service_server::{RuntimeService, RuntimeServiceServer};

#[derive(Clone)]
pub struct GrpcState {
    pub registry: Arc<CallableRegistry>,
    pub provider: Arc<dyn ModelProvider>,
    pub inbox_store: Arc<dyn InboxStore>,
    pub checkpoint_store: Arc<dyn CheckpointStore>,
}

impl Default for GrpcState {
    fn default() -> Self {
        Self::new()
    }
}

impl GrpcState {
    fn create_provider() -> Arc<dyn ModelProvider> {
        create_default_provider().unwrap_or_else(|e| {
            panic!(
                "Failed to create provider from config: {}. Set default_model_id in ~/.enact/config.yaml, or add providers.yaml with models and required env vars.",
                e
            )
        })
    }

    /// Create state with provider from config (default_model_id in config.yaml).
    pub fn new() -> Self {
        let provider = Self::create_provider();

        let registry = Arc::new(CallableRegistry::new());
        register_default_callables(&registry, provider.clone());

        let inbox_store = InMemoryInboxStore::shared();
        let checkpoint_store = Arc::new(InMemoryCheckpointStore::new()) as Arc<dyn CheckpointStore>;

        Self {
            registry,
            provider,
            inbox_store,
            checkpoint_store,
        }
    }

    /// Create GrpcState with a custom inbox store (provider from config).
    pub fn with_inbox_store(inbox_store: Arc<dyn InboxStore>) -> Self {
        let provider = Self::create_provider();

        let registry = Arc::new(CallableRegistry::new());
        register_default_callables(&registry, provider.clone());

        let checkpoint_store = Arc::new(InMemoryCheckpointStore::new()) as Arc<dyn CheckpointStore>;

        Self {
            registry,
            provider,
            inbox_store,
            checkpoint_store,
        }
    }

    /// Create GrpcState with custom inbox and checkpoint stores (provider from config).
    pub fn with_stores(
        inbox_store: Arc<dyn InboxStore>,
        checkpoint_store: Arc<dyn CheckpointStore>,
    ) -> Self {
        let provider = Self::create_provider();

        let registry = Arc::new(CallableRegistry::new());
        register_default_callables(&registry, provider.clone());

        Self {
            registry,
            provider,
            inbox_store,
            checkpoint_store,
        }
    }

    /// Create GrpcState with an explicit provider (e.g. for tests without config).
    pub fn with_provider(
        provider: Arc<dyn ModelProvider>,
        inbox_store: Arc<dyn InboxStore>,
    ) -> Self {
        let registry = Arc::new(CallableRegistry::new());
        register_default_callables(&registry, provider.clone());
        let checkpoint_store = Arc::new(InMemoryCheckpointStore::new()) as Arc<dyn CheckpointStore>;
        Self {
            registry,
            provider,
            inbox_store,
            checkpoint_store,
        }
    }
}

fn register_default_callables(registry: &CallableRegistry, provider: Arc<dyn ModelProvider>) {
    let assistant = LlmCallable::with_provider(
        "assistant".to_string(),
        "You are a helpful assistant.".to_string(),
        provider.clone(),
    );
    registry.register("assistant".to_string(), Arc::new(assistant));

    let coder = LlmCallable::with_provider(
        "coder".to_string(),
        "You are an expert programmer. Help with coding tasks, debugging, and code review."
            .to_string(),
        provider.clone(),
    );
    registry.register("coder".to_string(), Arc::new(coder));

    info!("Registered {} callables for gRPC", registry.len());
}

async fn try_dispatch_background_task(req: &RunAgentRequest) -> anyhow::Result<Option<TaskId>> {
    if !req.background {
        return Ok(None);
    }

    let cluster = load_cluster_config()?;
    let node_registry = Arc::new(RedisNodeRegistry::new(
        &cluster.redis_url,
        (cluster.worker.heartbeat_interval_secs.max(1)) * 3,
    )?);
    let task_queue = Arc::new(RedisTaskQueue::new(
        &cluster.redis_url,
        cluster.dispatcher.task_stream.clone(),
    )?);
    let dispatcher = Dispatcher::new(
        match cluster.dispatcher.strategy {
            DispatchStrategy::LeastLoaded => DispatchStrategy::LeastLoaded,
            DispatchStrategy::RoundRobin => DispatchStrategy::RoundRobin,
            DispatchStrategy::CapabilityMatch => DispatchStrategy::CapabilityMatch,
        },
        node_registry,
        task_queue,
    );

    let context_json = serde_json::to_value(&req.context).unwrap_or_else(|_| serde_json::json!({}));
    let spawn_mode = SpawnMode::Child {
        background: req.background,
        inherit_inbox: req.inherit_inbox,
        policies: None,
    };

    let task = DispatchTask {
        task_id: TaskId::new(),
        agent_name: req.agent_name.clone(),
        input: serde_json::json!({
            "input": req.input,
            "context": context_json,
            "parent_execution_id": req.parent_execution_id,
            "checkpoint_id": req.checkpoint_id,
        }),
        spawn_mode,
        priority: 5,
    };

    let task_id = dispatcher.dispatch(task, None).await?;
    Ok(Some(task_id))
}

pub struct RuntimeServiceImpl {
    state: Arc<GrpcState>,
}

impl RuntimeServiceImpl {
    pub fn new(state: Arc<GrpcState>) -> Self {
        Self { state }
    }
}

#[tonic::async_trait]
impl RuntimeService for RuntimeServiceImpl {
    async fn run_agent(
        &self,
        request: Request<RunAgentRequest>,
    ) -> Result<Response<RunAgentResponse>, Status> {
        let req = request.into_inner();
        let run_id = uuid::Uuid::new_v4().to_string();

        info!("RunAgent: {} (run_id: {})", req.agent_name, run_id);

        let callable =
            self.state.registry.get(&req.agent_name).ok_or_else(|| {
                Status::not_found(format!("Agent '{}' not found", req.agent_name))
            })?;

        if let Ok(Some(task_id)) = try_dispatch_background_task(&req).await {
            return Ok(Response::new(RunAgentResponse {
                run_id,
                output: format!("Queued background task {}", task_id),
                iterations: 0,
                tool_calls: vec![],
            }));
        }

        let input = if !req.context.is_empty() {
            let context_str = req
                .context
                .iter()
                .map(|(k, v)| format!("{}: {}", k, v))
                .collect::<Vec<_>>()
                .join("\n");
            format!("Context:\n{}\n\nTask: {}", context_str, req.input)
        } else {
            req.input
        };

        match callable.run(&input).await {
            Ok(output) => {
                let response = RunAgentResponse {
                    run_id,
                    output,
                    iterations: 1,
                    tool_calls: vec![],
                };
                Ok(Response::new(response))
            }
            Err(e) => {
                error!("Agent execution failed: {}", e);
                Err(Status::internal(format!("Execution failed: {}", e)))
            }
        }
    }

    type RunAgentStreamStream = Pin<Box<dyn Stream<Item = Result<StreamEvent, Status>> + Send>>;

    type ResumeFromCheckpointStream =
        Pin<Box<dyn Stream<Item = Result<StreamEvent, Status>> + Send>>;

    async fn run_agent_stream(
        &self,
        request: Request<RunAgentRequest>,
    ) -> Result<Response<Self::RunAgentStreamStream>, Status> {
        let req = request.into_inner();
        let run_id = uuid::Uuid::new_v4().to_string();

        info!("RunAgentStream: {} (run_id: {})", req.agent_name, run_id);

        let callable =
            self.state.registry.get(&req.agent_name).ok_or_else(|| {
                Status::not_found(format!("Agent '{}' not found", req.agent_name))
            })?;

        if let Ok(Some(task_id)) = try_dispatch_background_task(&req).await {
            let (tx, rx) = tokio::sync::mpsc::channel(8);
            let run_id_clone = run_id.clone();
            let _ = tx
                .send(Ok(StreamEvent {
                    run_id: run_id_clone.clone(),
                    event: Some(stream_event::Event::RunStarted(RunStarted {})),
                }))
                .await;
            let _ = tx
                .send(Ok(StreamEvent {
                    run_id: run_id_clone,
                    event: Some(stream_event::Event::RunCompleted(RunCompleted {
                        final_output: format!("Queued background task {}", task_id),
                    })),
                }))
                .await;
            let stream = ReceiverStream::new(rx);
            return Ok(Response::new(Box::pin(stream)));
        }

        let input = if !req.context.is_empty() {
            let context_str = req
                .context
                .iter()
                .map(|(k, v)| format!("{}: {}", k, v))
                .collect::<Vec<_>>()
                .join("\n");
            format!("Context:\n{}\n\nTask: {}", context_str, req.input)
        } else {
            req.input
        };

        let (tx, rx) = tokio::sync::mpsc::channel(100);
        let run_id_clone = run_id.clone();
        let agent_name = req.agent_name.clone();
        let checkpoint_store = self.state.checkpoint_store.clone();
        let execution_id = ExecutionId::from_string(&run_id);

        tokio::spawn(async move {
            let _ = tx
                .send(Ok(StreamEvent {
                    run_id: run_id_clone.clone(),
                    event: Some(stream_event::Event::RunStarted(RunStarted {})),
                }))
                .await;

            match callable.run(&input).await {
                Ok(output) => {
                    let checkpoint =
                        Checkpoint::new(execution_id.clone()).with_agent_name(&agent_name);
                    if let Err(e) = checkpoint_store.save(checkpoint).await {
                        error!("Failed to save checkpoint: {}", e);
                    }
                    let _ = tx
                        .send(Ok(StreamEvent {
                            run_id: run_id_clone,
                            event: Some(stream_event::Event::RunCompleted(RunCompleted {
                                final_output: output,
                            })),
                        }))
                        .await;
                }
                Err(e) => {
                    let _ = tx
                        .send(Ok(StreamEvent {
                            run_id: run_id_clone,
                            event: Some(stream_event::Event::Error(ErrorEvent {
                                message: e.to_string(),
                                code: "EXECUTION_ERROR".to_string(),
                            })),
                        }))
                        .await;
                }
            }
        });

        let stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn resume_from_checkpoint(
        &self,
        request: Request<ResumeRequest>,
    ) -> Result<Response<Self::ResumeFromCheckpointStream>, Status> {
        let req = request.into_inner();
        let checkpoint_id = req.checkpoint_id;

        info!("ResumeFromCheckpoint: {}", checkpoint_id);

        // Load checkpoint from store
        let checkpoint = self
            .state
            .checkpoint_store
            .load(&checkpoint_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to load checkpoint: {}", e)))?
            .ok_or_else(|| {
                Status::not_found(format!("Checkpoint '{}' not found", checkpoint_id))
            })?;

        // Get the run_id from checkpoint to find the original agent
        let run_id = checkpoint.run_id.clone();

        // Build input from checkpoint state + optional new input
        let input = if let Some(new_input) = req.new_input {
            // If new input provided, append to checkpoint messages
            format!(
                "Previous context from checkpoint {}:\n{}\n\nNew input: {}",
                checkpoint_id,
                checkpoint
                    .messages
                    .iter()
                    .map(|m| format!("{}: {}", m.role, m.content))
                    .collect::<Vec<_>>()
                    .join("\n"),
                new_input
            )
        } else {
            // Resume with checkpoint state
            format!(
                "Resuming from checkpoint {}. Previous conversation:\n{}",
                checkpoint_id,
                checkpoint
                    .messages
                    .iter()
                    .map(|m| format!("{}: {}", m.role, m.content))
                    .collect::<Vec<_>>()
                    .join("\n")
            )
        };

        // Retrieve agent name from checkpoint metadata, defaulting to "assistant" for
        // checkpoints created before agent_name was stored in metadata
        let agent_name = checkpoint.agent_name().unwrap_or("assistant");

        let callable = self
            .state
            .registry
            .get(agent_name)
            .ok_or_else(|| Status::not_found(format!("Agent '{}' not found", agent_name)))?;

        let (tx, rx) = tokio::sync::mpsc::channel(32);
        let callable_clone = callable.clone();
        let run_id_clone = run_id.as_str().to_string();

        tokio::spawn(async move {
            // Send run started event
            let _ = tx
                .send(Ok(StreamEvent {
                    run_id: run_id_clone.clone(),
                    event: Some(stream_event::Event::RunStarted(RunStarted {})),
                }))
                .await;

            // Execute the resumed callable
            match callable_clone.run(&input).await {
                Ok(output) => {
                    // Stream the output as content delta
                    let _ = tx
                        .send(Ok(StreamEvent {
                            run_id: run_id_clone.clone(),
                            event: Some(stream_event::Event::ContentDelta(ContentDelta {
                                content: output.clone(),
                            })),
                        }))
                        .await;

                    // Send run completed
                    let _ = tx
                        .send(Ok(StreamEvent {
                            run_id: run_id_clone,
                            event: Some(stream_event::Event::RunCompleted(RunCompleted {
                                final_output: output,
                            })),
                        }))
                        .await;
                }
                Err(e) => {
                    let _ = tx
                        .send(Ok(StreamEvent {
                            run_id: run_id_clone,
                            event: Some(stream_event::Event::Error(ErrorEvent {
                                message: e.to_string(),
                                code: "EXECUTION_ERROR".to_string(),
                            })),
                        }))
                        .await;
                }
            }
        });

        let stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn get_agent_info(
        &self,
        request: Request<GetAgentInfoRequest>,
    ) -> Result<Response<AgentInfo>, Status> {
        let req = request.into_inner();

        let _callable =
            self.state.registry.get(&req.agent_name).ok_or_else(|| {
                Status::not_found(format!("Agent '{}' not found", req.agent_name))
            })?;

        let info = AgentInfo {
            name: req.agent_name,
            description: "Enact agent".to_string(),
            tools: vec![],
        };

        Ok(Response::new(info))
    }

    async fn list_agents(
        &self,
        _request: Request<ListAgentsRequest>,
    ) -> Result<Response<ListAgentsResponse>, Status> {
        let names = self.state.registry.list();
        let agents = names
            .into_iter()
            .map(|name| AgentInfo {
                name,
                description: "Enact agent".to_string(),
                tools: vec![],
            })
            .collect();

        Ok(Response::new(ListAgentsResponse { agents }))
    }

    async fn health_check(
        &self,
        _request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        Ok(Response::new(HealthCheckResponse {
            healthy: true,
            version: env!("CARGO_PKG_VERSION").to_string(),
        }))
    }

    async fn cancel(
        &self,
        request: Request<CancelRequest>,
    ) -> Result<Response<CancelResponse>, Status> {
        let req = request.into_inner();
        let execution_id = ExecutionId::from_string(&req.run_id);

        info!("Cancel request for execution: {}", execution_id);

        // Create control message with Cancel action
        let control_msg =
            ControlMessage::new(execution_id.clone(), ControlAction::Cancel, "grpc_api");

        // Add reason if provided
        let control_msg = if let Some(reason) = req.reason {
            control_msg.with_reason(reason)
        } else {
            control_msg
        };

        // Push to inbox store
        self.state
            .inbox_store
            .push(&execution_id, InboxMessage::Control(control_msg));

        info!(
            "Cancel message pushed to inbox for execution: {}",
            execution_id
        );

        Ok(Response::new(CancelResponse {
            success: true,
            run_id: req.run_id,
            message: Some("Cancel request sent".to_string()),
        }))
    }

    async fn pause(
        &self,
        request: Request<PauseRequest>,
    ) -> Result<Response<PauseResponse>, Status> {
        let req = request.into_inner();
        let execution_id = ExecutionId::from_string(&req.run_id);

        info!("Pause request for execution: {}", execution_id);

        // Create control message with Pause action
        let control_msg =
            ControlMessage::new(execution_id.clone(), ControlAction::Pause, "grpc_api");

        // Push to inbox store
        self.state
            .inbox_store
            .push(&execution_id, InboxMessage::Control(control_msg));

        info!(
            "Pause message pushed to inbox for execution: {}",
            execution_id
        );

        Ok(Response::new(PauseResponse {
            success: true,
            run_id: req.run_id,
            checkpoint_id: None, // Checkpoint will be created by the execution loop
        }))
    }

    type ResumeStream = Pin<Box<dyn Stream<Item = Result<StreamEvent, Status>> + Send>>;

    async fn resume(
        &self,
        request: Request<ResumeExecutionRequest>,
    ) -> Result<Response<Self::ResumeStream>, Status> {
        let req = request.into_inner();
        let execution_id = ExecutionId::from_string(&req.run_id);

        info!("Resume request for execution: {}", execution_id);

        // Create control message with Resume action
        let control_msg =
            ControlMessage::new(execution_id.clone(), ControlAction::Resume, "grpc_api");

        // Push to inbox store
        self.state
            .inbox_store
            .push(&execution_id, InboxMessage::Control(control_msg));

        info!(
            "Resume message pushed to inbox for execution: {}",
            execution_id
        );

        // Create a stream that will emit events as the execution resumes
        let (tx, rx) = tokio::sync::mpsc::channel(100);
        let run_id = req.run_id.clone();

        // Send initial event indicating resume was requested
        tokio::spawn(async move {
            let _ = tx
                .send(Ok(StreamEvent {
                    run_id: run_id.clone(),
                    event: Some(stream_event::Event::RunStarted(RunStarted {})),
                }))
                .await;

            // Note: The actual execution resumption will be handled by the execution loop
            // when it processes the Resume control message from the inbox.
            // This stream currently just acknowledges the resume request.
            // A full implementation would connect to the actual execution stream.
        });

        let stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(stream)))
    }

    async fn approve_plan(
        &self,
        request: Request<ApprovePlanRequest>,
    ) -> Result<Response<ApprovePlanResponse>, Status> {
        let req = request.into_inner();
        let execution_id = ExecutionId::from_string(&req.run_id);

        info!("Approve plan request for execution: {}", execution_id);

        // Create guidance message approving the plan
        let guidance_msg = GuidanceMessage::from_user(
            execution_id.clone(),
            "PLAN_APPROVED: User approved the proposed plan. Proceed with execution.",
        )
        .with_priority(GuidancePriority::High);

        // Push to inbox store
        self.state
            .inbox_store
            .push(&execution_id, InboxMessage::Guidance(guidance_msg));

        info!(
            "Plan approval message pushed to inbox for execution: {}",
            execution_id
        );

        Ok(Response::new(ApprovePlanResponse {
            success: true,
            run_id: req.run_id,
        }))
    }

    async fn reject_plan(
        &self,
        request: Request<RejectPlanRequest>,
    ) -> Result<Response<RejectPlanResponse>, Status> {
        let req = request.into_inner();
        let execution_id = ExecutionId::from_string(&req.run_id);

        info!("Reject plan request for execution: {}", execution_id);

        // Create guidance message rejecting the plan
        let reason = req
            .reason
            .unwrap_or_else(|| "No reason provided".to_string());
        let guidance_msg = GuidanceMessage::from_user(
            execution_id.clone(),
            format!(
                "PLAN_REJECTED: User rejected the proposed plan. Reason: {}",
                reason
            ),
        )
        .with_priority(GuidancePriority::High);

        // Push to inbox store
        self.state
            .inbox_store
            .push(&execution_id, InboxMessage::Guidance(guidance_msg));

        info!(
            "Plan rejection message pushed to inbox for execution: {}",
            execution_id
        );

        Ok(Response::new(RejectPlanResponse {
            success: true,
            run_id: req.run_id,
        }))
    }
}

pub async fn serve_grpc(state: Arc<GrpcState>, addr: &str) -> Result<(), anyhow::Error> {
    let service = RuntimeServiceImpl::new(state);
    let addr = addr.parse()?;

    info!("gRPC server listening on {}", addr);

    tonic::transport::Server::builder()
        .add_service(RuntimeServiceServer::new(service))
        .serve(addr)
        .await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use enact_core::inbox::InMemoryInboxStore;
    use enact_core::providers::{ChatRequest, ChatResponse};
    use tokio_stream::StreamExt;

    struct TestProvider;

    #[async_trait]
    impl ModelProvider for TestProvider {
        fn name(&self) -> &str {
            "test"
        }
        fn model(&self) -> &str {
            "test-model"
        }
        async fn chat(&self, _: ChatRequest) -> anyhow::Result<ChatResponse> {
            anyhow::bail!("test only")
        }
    }

    fn create_test_state() -> Arc<GrpcState> {
        let inbox_store = InMemoryInboxStore::shared();
        Arc::new(GrpcState::with_provider(
            Arc::new(TestProvider),
            inbox_store,
        ))
    }

    #[tokio::test]
    async fn test_health_check() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let request = Request::new(HealthCheckRequest {});

        let response = service.health_check(request).await.unwrap();
        let health = response.into_inner();

        assert!(health.healthy);
        assert!(!health.version.is_empty());
    }

    #[tokio::test]
    async fn test_list_agents() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let request = Request::new(ListAgentsRequest {});

        let response = service.list_agents(request).await.unwrap();
        let list = response.into_inner();

        assert!(!list.agents.is_empty());
        // Should have at least assistant and coder
        assert!(list.agents.iter().any(|a| a.name == "assistant"));
        assert!(list.agents.iter().any(|a| a.name == "coder"));
    }

    #[tokio::test]
    async fn test_get_agent_info() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let request = Request::new(GetAgentInfoRequest {
            agent_name: "assistant".to_string(),
        });

        let response = service.get_agent_info(request).await.unwrap();
        let info = response.into_inner();

        assert_eq!(info.name, "assistant");
        assert!(!info.description.is_empty());
    }

    #[tokio::test]
    async fn test_get_agent_info_not_found() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let request = Request::new(GetAgentInfoRequest {
            agent_name: "nonexistent".to_string(),
        });

        let result = service.get_agent_info(request).await;
        assert!(result.is_err());

        let err = result.unwrap_err();
        assert_eq!(err.code(), tonic::Code::NotFound);
    }

    #[tokio::test]
    async fn test_cancel_execution() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let run_id = uuid::Uuid::new_v4().to_string();

        let request = Request::new(CancelRequest {
            run_id: run_id.clone(),
            reason: Some("Test cancellation".to_string()),
        });

        let response = service.cancel(request).await.unwrap();
        let cancel_resp = response.into_inner();

        assert!(cancel_resp.success);
        assert_eq!(cancel_resp.run_id, run_id);
        assert!(cancel_resp.message.is_some());
    }

    #[tokio::test]
    async fn test_pause_execution() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let run_id = uuid::Uuid::new_v4().to_string();

        let request = Request::new(PauseRequest {
            run_id: run_id.clone(),
        });

        let response = service.pause(request).await.unwrap();
        let pause_resp = response.into_inner();

        assert!(pause_resp.success);
        assert_eq!(pause_resp.run_id, run_id);
    }

    #[tokio::test]
    async fn test_resume_execution() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let run_id = uuid::Uuid::new_v4().to_string();

        let request = Request::new(ResumeExecutionRequest {
            run_id: run_id.clone(),
            checkpoint_id: None,
        });

        let response = service.resume(request).await.unwrap();
        let mut stream = response.into_inner();

        // Should receive at least RunStarted event
        let first_event = stream.next().await;
        assert!(first_event.is_some());

        let event = first_event.unwrap().unwrap();
        assert_eq!(event.run_id, run_id);
    }

    #[tokio::test]
    async fn test_approve_plan() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let run_id = uuid::Uuid::new_v4().to_string();

        let request = Request::new(ApprovePlanRequest {
            run_id: run_id.clone(),
        });

        let response = service.approve_plan(request).await.unwrap();
        let approve_resp = response.into_inner();

        assert!(approve_resp.success);
        assert_eq!(approve_resp.run_id, run_id);
    }

    #[tokio::test]
    async fn test_reject_plan() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);
        let run_id = uuid::Uuid::new_v4().to_string();

        let request = Request::new(RejectPlanRequest {
            run_id: run_id.clone(),
            reason: Some("Not aligned with goals".to_string()),
        });

        let response = service.reject_plan(request).await.unwrap();
        let reject_resp = response.into_inner();

        assert!(reject_resp.success);
        assert_eq!(reject_resp.run_id, run_id);
    }

    #[tokio::test]
    async fn test_run_agent_agent_not_found() {
        let state = create_test_state();
        let service = RuntimeServiceImpl::new(state);

        let request = Request::new(RunAgentRequest {
            agent_name: "nonexistent".to_string(),
            input: "Hello".to_string(),
            context: std::collections::HashMap::new(),
            checkpoint_id: None,
            background: false,
            inherit_inbox: false,
            parent_execution_id: None,
        });

        let result = service.run_agent(request).await;
        assert!(result.is_err());

        let err = result.unwrap_err();
        assert_eq!(err.code(), tonic::Code::NotFound);
    }
}