systemprompt_agent/services/mcp/task_helper/
mod.rs

1//! Helpers used by the MCP bridge to ensure-or-create tasks, broadcast
2//! completion, and persist tool-execution messages.
3
4mod completion;
5mod messages;
6
7pub use completion::complete_task;
8pub use messages::{SaveMessagesForToolExecutionParams, save_messages_for_tool_execution};
9
10use crate::models::a2a::{Task, TaskState, TaskStatus};
11use crate::repository::context::ContextRepository;
12use crate::repository::task::TaskRepository;
13use rmcp::ErrorData as McpError;
14use systemprompt_database::DbPool;
15use systemprompt_identifiers::TaskId;
16use systemprompt_models::TaskMetadata;
17
18#[derive(Debug)]
19pub struct TaskResult {
20    pub task_id: TaskId,
21    pub is_owner: bool,
22}
23
24pub async fn ensure_task_exists(
25    db_pool: &DbPool,
26    request_context: &mut systemprompt_models::execution::context::RequestContext,
27    tool_name: &str,
28    mcp_server_name: &str,
29) -> Result<TaskResult, McpError> {
30    if let Some(task_id) = request_context.task_id() {
31        tracing::info!(task_id = %task_id.as_str(), "Task reused from parent");
32        return Ok(TaskResult {
33            task_id: task_id.clone(),
34            is_owner: false,
35        });
36    }
37
38    let initial_context_id = request_context.context_id().as_str().to_owned();
39    let context_repo = ContextRepository::new(db_pool).map_err(|e| {
40        McpError::internal_error(format!("Failed to create context repository: {e}"), None)
41    })?;
42
43    let context_id = if initial_context_id.is_empty() {
44        if let Ok(Some(existing)) = context_repo
45            .find_by_session_id(request_context.session_id())
46            .await
47        {
48            tracing::debug!(
49                context_id = %existing.context_id,
50                session_id = %request_context.session_id(),
51                "Reusing existing context for MCP session"
52            );
53            request_context.execution.context_id = existing.context_id.clone();
54            existing.context_id
55        } else {
56            let new_context_id = context_repo
57                .create_context(
58                    request_context.user_id(),
59                    Some(request_context.session_id()),
60                    &format!("MCP Session: {}", request_context.session_id()),
61                )
62                .await
63                .map_err(|e| {
64                    tracing::error!(error = %e, "Failed to auto-create context for MCP session");
65                    McpError::internal_error(format!("Failed to create context: {e}"), None)
66                })?;
67
68            request_context.execution.context_id = new_context_id.clone();
69            tracing::info!(
70                context_id = %new_context_id,
71                session_id = %request_context.session_id(),
72                "Auto-created context for MCP session"
73            );
74            new_context_id
75        }
76    } else {
77        let old_context_id = request_context.context_id().clone();
78        match context_repo
79            .validate_context_ownership(&old_context_id, request_context.user_id())
80            .await
81        {
82            Ok(()) => old_context_id,
83            Err(e) => {
84                tracing::warn!(
85                    context_id = %old_context_id,
86                    user_id = %request_context.user_id(),
87                    error = %e,
88                    "Context validation failed, auto-creating new context"
89                );
90                let new_context_id = context_repo
91                    .create_context(
92                        request_context.user_id(),
93                        Some(request_context.session_id()),
94                        &format!("MCP Session: {}", request_context.session_id()),
95                    )
96                    .await
97                    .map_err(|e| {
98                        tracing::error!(error = %e, "Failed to auto-create replacement context");
99                        McpError::internal_error(format!("Failed to create context: {e}"), None)
100                    })?;
101
102                request_context.execution.context_id = new_context_id.clone();
103                tracing::info!(
104                    old_context_id = %old_context_id,
105                    new_context_id = %new_context_id,
106                    session_id = %request_context.session_id(),
107                    "Auto-created replacement context for invalid context_id"
108                );
109                new_context_id
110            },
111        }
112    };
113
114    let task_repo = TaskRepository::new(db_pool).map_err(|e| {
115        McpError::internal_error(format!("Failed to create task repository: {e}"), None)
116    })?;
117
118    let task_id = TaskId::generate();
119
120    let agent_name = request_context.agent_name().to_string();
121
122    let metadata = TaskMetadata::new_mcp_execution(
123        agent_name.clone(),
124        tool_name.to_owned(),
125        mcp_server_name.to_owned(),
126    );
127
128    let task = Task {
129        id: task_id.clone(),
130        context_id: context_id.clone(),
131        status: TaskStatus {
132            state: TaskState::Submitted,
133            message: None,
134            timestamp: Some(chrono::Utc::now()),
135        },
136        history: None,
137        artifacts: None,
138        metadata: Some(metadata),
139        created_at: Some(chrono::Utc::now()),
140        last_modified: Some(chrono::Utc::now()),
141    };
142
143    task_repo
144        .create_task(crate::repository::task::RepoCreateTaskParams {
145            task: &task,
146            user_id: request_context.user_id(),
147            session_id: request_context.session_id(),
148            trace_id: request_context.trace_id(),
149            agent_name: &agent_name,
150        })
151        .await
152        .map_err(|e| McpError::internal_error(format!("Failed to create task: {e}"), None))?;
153
154    request_context.execution.task_id = Some(task_id.clone());
155
156    tracing::info!(
157        task_id = %task_id.as_str(),
158        tool = %tool_name,
159        agent = %agent_name,
160        "Task created"
161    );
162
163    Ok(TaskResult {
164        task_id,
165        is_owner: true,
166    })
167}
systemprompt_agent/services/mcp/task_helper/mod.rs

systemprompt_agent/services/mcp/task_helper/
mod.rs
