velesdb_memory/mcp/
migration_tools.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use rmcp::handler::server::wrapper::{Json, Parameters};
5use rmcp::model::ErrorCode;
6use rmcp::{tool, tool_router, ErrorData};
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10use super::{to_error, McpServer};
11use crate::mutation::catchup::CatchUpConfig;
12use crate::mutation::controller::ControllerConfig;
13use crate::service::{JobPhase, MigrationStartConfig, MigrationStatus, OnlineMigrationManager};
14use crate::DynEmbedder;
15
16const DEFAULT_JOURNAL_BYTES: u64 = 64 * 1024 * 1024;
17const DEFAULT_FACT_BATCH: usize = 256;
18const DEFAULT_REPLAY_BATCH: usize = 256;
19const DEFAULT_EDGE_CAP: usize = 4_096;
20const DEFAULT_OBSERVATION_WINDOW: usize = 3;
21const DEFAULT_VERIFICATION_RESERVE_MS: u64 = 100;
22
23#[derive(Debug, Deserialize, JsonSchema)]
24#[schemars(transform = crate::schema::strip_int_formats)]
25struct MigrationStartParams {
26    target_backend: String,
27    pause_budget_ms: u64,
28    journal_max_bytes: Option<u64>,
29    fact_batch: Option<usize>,
30    replay_batch: Option<usize>,
31    edge_cap: Option<usize>,
32    observation_window: Option<usize>,
33    verification_reserve_ms: Option<u64>,
34}
35
36#[derive(Debug, Serialize, JsonSchema)]
37#[schemars(transform = crate::schema::strip_int_formats)]
38struct MigrationStatusResult {
39    configured: bool,
40    job: Option<MigrationJobResult>,
41}
42
43#[derive(Debug, Serialize, JsonSchema)]
44#[schemars(transform = crate::schema::strip_int_formats)]
45struct MigrationJobResult {
46    epoch_id: String,
47    phase: String,
48    running: bool,
49    target_backend: String,
50    target_model: String,
51    target_dimension: usize,
52    destination: String,
53    cancellation_requested: bool,
54    recovery_action: Option<String>,
55    last_error: Option<String>,
56    progress: MigrationProgressResult,
57}
58
59#[derive(Debug, Serialize, JsonSchema)]
60#[schemars(transform = crate::schema::strip_int_formats)]
61struct MigrationProgressResult {
62    base_facts: u64,
63    base_edge_sets: u64,
64    base_batches: u64,
65    input_watermark: u64,
66    output_watermark: u64,
67    distinct_dirty_facts: u64,
68    distinct_edge_sources: u64,
69    pending_journal_bytes: u64,
70    estimated_pause_ms: Option<u64>,
71    measured_cutover_ms: Option<u64>,
72}
73
74#[tool_router(router = migration_tool_router, vis = "pub(super)")]
75impl McpServer {
76    #[tool(
77        name = "migration_start",
78        output_schema = crate::schema::wire_safe_output_schema::<MigrationStatusResult>(),
79        description = "Start one daemon-owned online embedding migration in the background. The existing MCP transport remains the only control boundary; no listener or credential is persisted. `target_backend` selects a backend configured in this daemon's environment. `pause_budget_ms` is the maximum measured request pause. Returns immediately with durable job status; poll migration_status. Refuses a pre-existing destination or another non-terminal epoch."
80    )]
81    async fn migration_start(
82        &self,
83        Parameters(params): Parameters<MigrationStartParams>,
84    ) -> Result<Json<MigrationStatusResult>, ErrorData> {
85        let manager = self.require_migration_manager()?;
86        let status = tokio::task::spawn_blocking(move || {
87            manager.start(&params.target_backend, params.config())?;
88            current_status(&manager)
89        })
90        .await
91        .map_err(super::join_error)?
92        .map_err(to_error)?;
93        Ok(Json(status))
94    }
95
96    #[tool(
97        name = "migration_status",
98        output_schema = crate::schema::wire_safe_output_schema::<MigrationStatusResult>(),
99        description = "Read the durable online embedding migration phase, progress watermarks, convergence estimate, measured cutover, cancellation flag, last error and required recovery action. Returns immediately and performs no migration work."
100    )]
101    async fn migration_status(&self) -> Result<Json<MigrationStatusResult>, ErrorData> {
102        let Some(manager) = self.online_migration.clone() else {
103            return Ok(Json(MigrationStatusResult {
104                configured: false,
105                job: None,
106            }));
107        };
108        let status = tokio::task::spawn_blocking(move || current_status(&manager))
109            .await
110            .map_err(super::join_error)?
111            .map_err(to_error)?;
112        Ok(Json(status))
113    }
114
115    #[tool(
116        name = "migration_cancel",
117        output_schema = crate::schema::wire_safe_output_schema::<MigrationStatusResult>(),
118        description = "Durably request cancellation while the source remains authoritative. A running worker observes the request at its next bounded batch; a paused job is cancelled immediately. From quiescing onward cancellation refuses and reports the required recovery action instead of guessing a rollback."
119    )]
120    async fn migration_cancel(&self) -> Result<Json<MigrationStatusResult>, ErrorData> {
121        let manager = self.require_migration_manager()?;
122        let status = tokio::task::spawn_blocking(move || {
123            manager.cancel()?;
124            current_status(&manager)
125        })
126        .await
127        .map_err(super::join_error)?
128        .map_err(to_error)?;
129        Ok(Json(status))
130    }
131
132    #[tool(
133        name = "migration_recover",
134        output_schema = crate::schema::wire_safe_output_schema::<MigrationStatusResult>(),
135        description = "Resume a durable prepared, capturing, base-copied, catching-up, non-converging or cutover-ready job after verifying the environment-backed target model, dimension and vector witness. Quiescing or activated jobs refuse here until crash-safe cutover recovery has restored a single authoritative generation."
136    )]
137    async fn migration_recover(&self) -> Result<Json<MigrationStatusResult>, ErrorData> {
138        let manager = self.require_migration_manager()?;
139        let status = tokio::task::spawn_blocking(move || {
140            manager.recover()?;
141            current_status(&manager)
142        })
143        .await
144        .map_err(super::join_error)?
145        .map_err(to_error)?;
146        Ok(Json(status))
147    }
148
149    fn require_migration_manager(
150        &self,
151    ) -> Result<Arc<OnlineMigrationManager<DynEmbedder>>, ErrorData> {
152        self.online_migration.clone().ok_or_else(|| {
153            ErrorData::new(
154                ErrorCode::INTERNAL_ERROR,
155                "online migration is not configured for this server".to_owned(),
156                None,
157            )
158        })
159    }
160}
161
162impl MigrationStartParams {
163    fn config(&self) -> MigrationStartConfig {
164        MigrationStartConfig {
165            journal_max_bytes: self.journal_max_bytes.unwrap_or(DEFAULT_JOURNAL_BYTES),
166            catch_up: CatchUpConfig {
167                fact_batch: self.fact_batch.unwrap_or(DEFAULT_FACT_BATCH),
168                replay_batch: self.replay_batch.unwrap_or(DEFAULT_REPLAY_BATCH),
169                edge_cap: self.edge_cap.unwrap_or(DEFAULT_EDGE_CAP),
170            },
171            controller: ControllerConfig {
172                observation_window: self
173                    .observation_window
174                    .unwrap_or(DEFAULT_OBSERVATION_WINDOW),
175                pause_budget: Duration::from_millis(self.pause_budget_ms),
176                verification_reserve: Duration::from_millis(
177                    self.verification_reserve_ms
178                        .unwrap_or(DEFAULT_VERIFICATION_RESERVE_MS),
179                ),
180            },
181        }
182    }
183}
184
185fn current_status(
186    manager: &OnlineMigrationManager<DynEmbedder>,
187) -> Result<MigrationStatusResult, crate::MemoryError> {
188    Ok(MigrationStatusResult {
189        configured: true,
190        job: manager.status()?.map(MigrationJobResult::from),
191    })
192}
193
194impl From<MigrationStatus> for MigrationJobResult {
195    fn from(status: MigrationStatus) -> Self {
196        let record = status.record;
197        let progress = MigrationProgressResult {
198            base_facts: record.progress.base_facts,
199            base_edge_sets: record.progress.base_edge_sets,
200            base_batches: record.progress.base_batches,
201            input_watermark: record.progress.input_watermark,
202            output_watermark: record.progress.output_watermark,
203            distinct_dirty_facts: record.progress.distinct_dirty_facts,
204            distinct_edge_sources: record.progress.distinct_edge_sources,
205            pending_journal_bytes: record.progress.pending_journal_bytes,
206            estimated_pause_ms: record.progress.estimated_pause.map(duration_millis),
207            measured_cutover_ms: record.progress.measured_cutover.map(duration_millis),
208        };
209        let identity = &record.spec.identity;
210        Self {
211            epoch_id: identity.epoch_id().to_owned(),
212            phase: phase_name(record.phase).to_owned(),
213            running: status.running,
214            target_backend: record.spec.target_backend.clone(),
215            target_model: identity.target_model().to_owned(),
216            target_dimension: identity.target_dimension(),
217            destination: identity.destination_path().display().to_string(),
218            cancellation_requested: record.cancellation_requested,
219            recovery_action: record.recovery_action,
220            last_error: record.last_error,
221            progress,
222        }
223    }
224}
225
226fn phase_name(phase: JobPhase) -> &'static str {
227    match phase {
228        JobPhase::Prepared => "prepared",
229        JobPhase::Capturing => "capturing",
230        JobPhase::BaseCopied => "base_copied",
231        JobPhase::CatchingUp => "catching_up",
232        JobPhase::NonConverging => "non_converging",
233        JobPhase::CutoverReady => "cutover_ready",
234        JobPhase::Quiescing => "quiescing",
235        JobPhase::Activated => "activated",
236        JobPhase::Committed => "committed",
237        JobPhase::Cancelled => "cancelled",
238    }
239}
240
241fn duration_millis(duration: Duration) -> u64 {
242    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
243}
244
245#[cfg(test)]
246#[path = "migration_tools_tests.rs"]
247mod tests;
velesdb_memory/mcp/migration_tools.rs

velesdb_memory/mcp/
migration_tools.rs
