systemprompt-mcp 0.36.0

Native Model Context Protocol (MCP) implementation for systemprompt.io. Orchestration, per-server OAuth2, RBAC middleware, and tool-call governance — the core of the AI governance pipeline.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

//! Client-side task polling for the `io.modelcontextprotocol/tasks`
//! extension (SEP-2663).
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use std::time::Duration;

use rmcp::model::{CallToolResult, CreateTaskResult, GetTaskParams, TaskPayload};
use rmcp::service::{RoleClient, RunningService};

use crate::error::{McpDomainError, McpDomainResult};

const DEFAULT_POLL_INTERVAL: Duration = Duration::from_millis(500);

// Why: a server that omits both ttlMs and pollIntervalMs could otherwise pin
// this client in an unbounded poll loop; the deadline is ours, not the spec's.
const MAX_POLL_BUDGET: Duration = Duration::from_secs(600);

pub(super) async fn poll_task_to_completion<S>(
    client: &RunningService<RoleClient, S>,
    server: &str,
    created: CreateTaskResult,
) -> McpDomainResult<CallToolResult>
where
    S: rmcp::Service<RoleClient>,
{
    let task_id = created.task.task_id.clone();
    let poll_interval = created
        .task
        .poll_interval_ms
        .map_or(DEFAULT_POLL_INTERVAL, Duration::from_millis);
    let budget = created.task.ttl_ms.map_or(MAX_POLL_BUDGET, |ttl| {
        Duration::from_millis(ttl).min(MAX_POLL_BUDGET)
    });
    let deadline = tokio::time::Instant::now() + budget;

    loop {
        if tokio::time::Instant::now() >= deadline {
            return Err(McpDomainError::Timeout {
                server: server.to_owned(),
                after_ms: u64::try_from(budget.as_millis()).unwrap_or(u64::MAX),
            });
        }
        tokio::time::sleep(poll_interval).await;

        let polled = client
            .get_task(GetTaskParams::new(task_id.clone()))
            .await
            .map_err(|e| McpDomainError::ToolExecutionFailed(format!("tasks/get failed: {e}")))?;

        match polled.task.payload {
            TaskPayload::Working => {},
            TaskPayload::InputRequired { .. } => {
                tracing::debug!(task_id = %task_id, "Task awaiting input; continuing to poll");
            },
            TaskPayload::Completed { result } => {
                return serde_json::from_value(serde_json::Value::Object(result)).map_err(|e| {
                    McpDomainError::ToolExecutionFailed(format!(
                        "task completed with a non-CallToolResult payload: {e}"
                    ))
                });
            },
            TaskPayload::Failed { error } => {
                return Err(McpDomainError::ToolExecutionFailed(format!(
                    "task failed: {}",
                    serde_json::Value::Object(error)
                )));
            },
            TaskPayload::Cancelled => {
                return Err(McpDomainError::ToolExecutionFailed(
                    "task was cancelled by the server".to_owned(),
                ));
            },
            other => {
                tracing::warn!(status = ?other.status(), "Unrecognized task payload; continuing to poll");
            },
        }
    }
}