Completion complete

crates/contextforge-data-plane-lib/src/gateway/mcp_service/completion.rs

@@ -1,20 +1,69 @@
 use rmcp::{
     ErrorData, RoleServer,
-    model::{CompleteRequestParams, CompleteResult, ErrorCode},
+    model::{CompleteRequestParams, CompleteResult, ErrorCode, Reference},
     service::RequestContext,
 };
+use tracing::info;
+
+use crate::gateway::{
+    mcp_call_validator::AuthorizedCallValidator, mcp_service::initialization::connect_backend_for_request,
+    routing_error::backend_forward_error,
+};
 
 use super::McpService;
 
 #[allow(clippy::unused_async)]
 pub(super) async fn complete(
-    _: &McpService,
-    _: CompleteRequestParams,
-    _: RequestContext<RoleServer>,
+    mcp_service: &McpService,
+    request: CompleteRequestParams,
+    cx: RequestContext<RoleServer>,
 ) -> Result<CompleteResult, ErrorData> {
-    Err(ErrorData {
-        code: ErrorCode::INVALID_REQUEST,
-        message: "Fan out not supported at the moment. Go to control plane".into(),
+    let mcp_call_validator = AuthorizedCallValidator::new("complete", &cx);
+    let (virtual_host, _claims) = mcp_call_validator.validate_stateless()?;
+
+    let route = match &request.r#ref {
+        Reference::Prompt(prompt) => virtual_host.prompts.get(&prompt.name),
+        Reference::Resource(resource) => {
+            virtual_host.resource_templates.get(&resource.uri).or_else(|| virtual_host.resources.get(&resource.uri))
+        },
+        _ => None,
+    };
+    let Some(route) = route else {
+        return Err(ErrorData {
+            code: ErrorCode::INVALID_PARAMS,
+            message: "Routing problem... completion not found".into(),
+            data: None,
+        });
+    };
+
+    let backend_name = route.backend_name.clone();
+    let completion_name = route.upstream_name.clone();
+
+    let backend = virtual_host.backends.get(&backend_name).ok_or_else(|| ErrorData {
+        code: ErrorCode::INVALID_PARAMS,
+        message: "Routing problem... backend not found".into(),
         data: None,
-    })
+    })?;
+
+    let service_name = backend_name.clone();
+    let mut backend_service = connect_backend_for_request(mcp_service, &backend_name, backend, &cx).await?;
+
+    let mut routed_request = request;
+    match &mut routed_request.r#ref {
+        Reference::Prompt(prompt) => prompt.name = completion_name,
+        Reference::Resource(resource) => resource.uri = completion_name,
+        _ => {},
+    }
+
+    let response = backend_service.complete(routed_request).await;
+
+    if let Err(error) = backend_service.close().await {
+        tracing::warn!("complete: backend cleanup failed backend_name = {service_name} error = {error:?}");
+    }
+
+    let response = response.map_err(|error| backend_forward_error("complete", &service_name, &error))?;
+
+    info!("complete: backend {service_name} returned {} contents", response.completion.values.len());
+
+    Ok(response)
 }

crates/contextforge-data-plane-lib/tests/gateway/completions.rs

@@ -1,7 +1,11 @@
 use contextforge_data_plane_lib::Result;
+use rmcp::model::{ArgumentInfo, CompleteRequestParams, Reference};
 
 use crate::harness::{TEST_USER_ID, connect_modern_client, create_client, modern_client_info, start_counter_gateway};
 
+const EXAMPLE_PROMPT: &str = "00000000-0000-0000-0000-000000000001-example_prompt";
+const MEMO_RESOURCE: &str = "00000000-0000-0000-0000-000000000001-memo://insights";
+
 #[tokio::test]
 async fn plaintext_complete_for_unrouted_reference_errors() -> Result<()> {
     let fixture = start_counter_gateway(TEST_USER_ID).await?;
@@ -13,3 +17,27 @@ async fn plaintext_complete_for_unrouted_reference_errors() -> Result<()> {
         .expect_err("an unrouted completion reference must fail");
     Ok(())
 }
+
+#[tokio::test]
+async fn test_complete_for_promtp() -> Result<()> {
+    let fixture = start_counter_gateway(TEST_USER_ID).await?;
+    let service = connect_modern_client(&fixture.gateway_url, create_client(TEST_USER_ID), modern_client_info()).await;
+
+    let result = service
+        .complete(CompleteRequestParams::new(Reference::for_prompt(EXAMPLE_PROMPT), ArgumentInfo::new("message", "h")))
+        .await?;
+    assert_eq!(result.completion.values, ["hello", "hola"]);
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_complete_for_resource() -> Result<()> {
+    let fixture = start_counter_gateway(TEST_USER_ID).await?;
+    let service = connect_modern_client(&fixture.gateway_url, create_client(TEST_USER_ID), modern_client_info()).await;
+
+    let result = service
+        .complete(CompleteRequestParams::new(Reference::for_resource(MEMO_RESOURCE), ArgumentInfo::new("message", "h")))
+        .await?;
+    assert_eq!(result.completion.values, ["memo://insights"]);
+    Ok(())
+}