Skip to main content

headless_lms_mcp/staging/
mod.rs

1//! The write pipeline every write tool ends in: find what the call changed, validate written
2//! blocks with the CMS, and stage the call, or on a dry run tell what staging it would do.
3
4use std::collections::{HashMap, HashSet};
5
6use headless_lms_models::{
7    agent_changeset_entities::ChangesetEntitySummary,
8    exercise_slide_submissions,
9    library::{
10        agent_changesets::{
11            RiskBaseline, StageCall, StageCallOutcome, StagedEntity, StagedStateRead,
12            course_state::{CourseState, CourseView, PageRecord},
13            stage_call, stage_call_dry_run,
14        },
15        entity_snapshots::{ContentEntityType, EntityKey, EntitySnapshot, PageSnapshot},
16    },
17};
18use serde_json::{Value, json};
19use sqlx::PgConnection;
20use uuid::Uuid;
21
22use crate::access::connection;
23use crate::diff;
24use crate::engine::blocks::ExerciseCopy;
25use crate::explanations::{ExplainedEntity, explain_blockers};
26use crate::ids::RandomIds;
27use crate::state::McpState;
28use crate::tools::{RequestContext, ToolError, ToolErrorCode};
29
30mod page_writes;
31mod report;
32mod validation;
33
34use page_writes::page_writes;
35pub use report::WriteReport;
36use report::{
37    BriefDryRunSummary, DryRunEntity, DryRunSummary, NothingStaged, StagedSummary, StagingSummary,
38    page_hints, write_hints,
39};
40use validation::{ValidatedBlocks, changeset_block_names, validate_changed_blocks};
41
42/// Where a write call's entity writes come from.
43pub enum WriteSource {
44    /// Found by comparing the staged state with the state after the call.
45    Derived {
46        exercise_copies: Vec<ExerciseCopy>,
47        /// The client ids of blocks the server wrote rather than the agent, besides the exercise
48        /// copies; see [`validate_changed_blocks`].
49        generated_blocks: HashSet<Uuid>,
50    },
51    /// A revert of the call `call_id`. Skips validation: it restores content that was valid when
52    /// it was written.
53    Revert {
54        call_id: Uuid,
55        writes: Vec<StagedEntity>,
56        /// The state each written entity had when the caller read it: staged where the changeset
57        /// has it, live otherwise.
58        read_states: HashMap<EntityKey, Option<EntitySnapshot>>,
59    },
60}
61
62/// Whether a write call stages its changes or only tells what staging them would do.
63#[derive(Clone, Copy, PartialEq, Eq)]
64pub enum WriteMode {
65    Stage,
66    DryRun,
67}
68
69impl WriteMode {
70    /// The mode a tool's `dry_run` argument asks for.
71    pub fn from_dry_run(dry_run: bool) -> Self {
72        if dry_run { Self::DryRun } else { Self::Stage }
73    }
74}
75
76/// One write call for [`commit`].
77pub struct WriteRequest {
78    /// The operations as the agent sent them, for the audit log.
79    pub ops: Value,
80    pub summary: String,
81    pub mode: WriteMode,
82    /// On a dry run, report counts and risky entities instead of every entity and block change.
83    pub is_dry_run_brief: bool,
84    pub source: WriteSource,
85}
86
87impl WriteRequest {
88    /// A write whose changes are found by comparing states.
89    pub fn new(ops: Value, summary: String, mode: WriteMode) -> Self {
90        Self {
91            ops,
92            summary,
93            mode,
94            is_dry_run_brief: false,
95            source: WriteSource::Derived {
96                exercise_copies: Vec::new(),
97                generated_blocks: HashSet::new(),
98            },
99        }
100    }
101}
102
103/// Runs the pipeline from the user's staged state to `after`: validate, find the written
104/// entities and stage them as one call, or on a dry run tell what staging them would do. Holds no
105/// pooled connection while the validator runs.
106pub async fn commit(
107    state: &McpState,
108    ctx: &RequestContext,
109    view: &CourseView,
110    mut after: CourseState,
111    request: WriteRequest,
112) -> Result<WriteReport, ToolError> {
113    let WriteRequest {
114        ops,
115        summary,
116        mode,
117        is_dry_run_brief,
118        source,
119    } = request;
120    let validated = match &source {
121        WriteSource::Derived {
122            exercise_copies,
123            generated_blocks,
124        } => {
125            // A copied exercise block carries the attributes the editor stored, which agents
126            // may not write.
127            let server_written: HashSet<Uuid> = generated_blocks
128                .iter()
129                .copied()
130                .chain(exercise_copies.iter().map(|copy| copy.new_block_id))
131                .collect();
132            validate_changed_blocks(
133                state.validator.as_ref(),
134                &view.live,
135                &view.staged,
136                &mut after,
137                &server_written,
138                &ctx.progress,
139            )
140            .await?
141        }
142        WriteSource::Revert { .. } => ValidatedBlocks::default(),
143    };
144    let mut conn = connection(state).await?;
145    let DerivedWrites {
146        writes,
147        live_snapshots,
148        revert_of_call_id,
149    } = derive_writes(&mut conn, view, &after, source).await?;
150    let unchanged_block_warnings = validated.unchanged_block_warnings;
151    if writes.is_empty() {
152        return Ok(WriteReport {
153            summary: StagingSummary::Nothing(NothingStaged {
154                staged: [],
155                message: "The call changed nothing, so nothing was staged.",
156                unchanged_block_warnings,
157            }),
158            hints: Vec::new(),
159            after,
160        });
161    }
162
163    let written_pages: Vec<Uuid> = writes
164        .iter()
165        .map(StagedEntity::key)
166        .filter(|key| key.entity_type == ContentEntityType::Page)
167        .map(|key| key.entity_id)
168        .collect();
169    let block_changes = diff::diff_pages(&view.staged, &after, &written_pages);
170    let read = staged_state_read(view, &writes, live_snapshots);
171    let call = StageCall {
172        actor: &ctx.principal.actor,
173        course_id: view.course.id,
174        tool_name: ctx.tool_name,
175        ops: &ops,
176        summary: Some(&summary),
177        entities: writes,
178        revert_of_call_id,
179        read: &read,
180        risk: &RiskBaseline::View(view),
181    };
182    let outcome = if mode == WriteMode::DryRun {
183        stage_call_dry_run(&mut conn, call).await?
184    } else {
185        stage_call(&mut conn, call).await?
186    };
187    let staged = match outcome {
188        StageCallOutcome::Staged(staged) => staged,
189        StageCallOutcome::ChangesetClosed => {
190            return Err(ToolError::StagedStateChanged {
191                message: "Nothing was staged: your changeset was published, discarded or expired while this call ran. Read the course again and repeat the call on the current state.".into(),
192                entities: Vec::new(),
193            });
194        }
195        StageCallOutcome::StaleRead { entities } => return Err(stale_read(&entities)),
196        StageCallOutcome::WouldAddBlockers(blockers) => {
197            return Err(ToolError::refused(
198                ToolErrorCode::RevertConflict,
199                "Nothing was staged: reverting the call would clash with later changes to the course structure. Change the entities named here first, or revert the later calls too.",
200                json!({"blockers": explain_blockers(blockers)}),
201            ));
202        }
203    };
204    let entities = staged.entities.iter().map(|row| {
205        let mut entity = ExplainedEntity::from(ChangesetEntitySummary::from(row));
206        let key = entity.entity.key();
207        entity.is_chapter_front_page = key.entity_type == ContentEntityType::Page
208            && (after.is_chapter_front_page(key.entity_id)
209                || view.staged.is_chapter_front_page(key.entity_id));
210        entity
211    });
212    if mode == WriteMode::DryRun && is_dry_run_brief {
213        let entities: Vec<ExplainedEntity> = entities.collect();
214        return Ok(WriteReport {
215            summary: StagingSummary::BriefDryRun(BriefDryRunSummary {
216                dry_run: true,
217                would_stage_count: entities.len(),
218                block_change_count: block_changes.len(),
219                risky_entities: entities
220                    .into_iter()
221                    .filter(|entity| !entity.risk_reasons.is_empty())
222                    .collect(),
223                unchanged_block_warnings,
224            }),
225            hints: Vec::new(),
226            after,
227        });
228    }
229    if mode == WriteMode::DryRun {
230        let would_stage = entities
231            .map(|entity| DryRunEntity {
232                field_changes: after.field_changes_since(&view.staged, entity.entity.key()),
233                entity,
234            })
235            .collect();
236        return Ok(WriteReport {
237            summary: StagingSummary::DryRun(DryRunSummary {
238                dry_run: true,
239                would_stage,
240                block_changes,
241                unchanged_block_warnings,
242            }),
243            hints: Vec::new(),
244            after,
245        });
246    }
247
248    let removed_exercises: Vec<Uuid> = block_changes
249        .iter()
250        .filter(|change| change.kind == diff::ChangeKind::Removed)
251        .filter_map(|change| change.exercise_id)
252        .collect();
253    let removed_exercises_with_submissions = if removed_exercises.is_empty() {
254        0
255    } else {
256        exercise_slide_submissions::get_exercise_ids_with_submissions_by_course_id(
257            &mut conn,
258            view.course.id,
259            &removed_exercises,
260        )
261        .await?
262        .len()
263    };
264    // Each block type's format tips come once per changeset.
265    let tipped_block_names = changeset_block_names(&view.live, &view.staged);
266    let new_block_names: Vec<&str> = validated
267        .block_names
268        .iter()
269        .map(String::as_str)
270        .filter(|name| !tipped_block_names.contains(*name))
271        .collect();
272    let mut hints = write_hints(
273        state,
274        &staged,
275        view.course.is_draft,
276        removed_exercises_with_submissions,
277        &new_block_names,
278    )
279    .await;
280    hints.extend(page_hints(
281        ctx.client.client_kind(),
282        &block_changes,
283        &view.live,
284        &view.staged,
285        &after,
286        &written_pages,
287    ));
288    let summary = StagedSummary {
289        changeset: staged.changeset.id,
290        review_url: state.review_url(staged.changeset.id),
291        changeset_revision: staged.changeset.revision,
292        call: staged.call_id,
293        staged: entities.collect(),
294        changeset_expires_at: staged.changeset.expires_at,
295        unchanged_block_warnings,
296    };
297    Ok(WriteReport {
298        summary: StagingSummary::Staged(summary),
299        hints,
300        after,
301    })
302}
303
304/// What a write call stages, found from its [`WriteSource`].
305struct DerivedWrites {
306    writes: Vec<StagedEntity>,
307    /// The live state of the written entities the call read live.
308    live_snapshots: HashMap<EntityKey, Option<EntitySnapshot>>,
309    revert_of_call_id: Option<Uuid>,
310}
311
312async fn derive_writes(
313    conn: &mut PgConnection,
314    view: &CourseView,
315    after: &CourseState,
316    source: WriteSource,
317) -> Result<DerivedWrites, ToolError> {
318    match source {
319        WriteSource::Revert {
320            call_id,
321            writes,
322            read_states,
323        } => {
324            let live_snapshots = read_states
325                .into_iter()
326                .filter(|(key, _)| view.staged_row(*key).is_none())
327                .collect();
328            Ok(DerivedWrites {
329                writes,
330                live_snapshots,
331                revert_of_call_id: Some(call_id),
332            })
333        }
334        WriteSource::Derived {
335            exercise_copies, ..
336        } => {
337            let changes = after.changes_since(&view.staged);
338            let mut writes = changes.entities;
339            let pages = page_writes(
340                conn,
341                view,
342                after,
343                &changes.pages,
344                &exercise_copies,
345                &RandomIds,
346            )
347            .await?;
348            writes.extend(pages.writes);
349            let live_snapshots = live_snapshots(view, &writes, pages.live_pages)?;
350            Ok(DerivedWrites {
351                writes,
352                live_snapshots,
353                revert_of_call_id: None,
354            })
355        }
356    }
357}
358
359fn stale_read(entities: &[EntityKey]) -> ToolError {
360    ToolError::StagedStateChanged {
361        message: format!(
362            "Nothing was staged: another call or a live edit changed {} of the entities this call writes after this call read them. Read them again and repeat the call.",
363            entities.len()
364        ),
365        entities: entities
366            .iter()
367            .map(|key| json!({"entity_type": key.entity_type, "entity_id": key.entity_id}))
368            .collect(),
369    }
370}
371
372/// The live state of each written entity the user has no staged changes to, as the call read it.
373/// Pages were read twice, as blocks with the view and in full when their snapshots were built;
374/// a page that changed in between is a stale read.
375fn live_snapshots(
376    view: &CourseView,
377    writes: &[StagedEntity],
378    mut live_pages: HashMap<Uuid, PageSnapshot>,
379) -> Result<HashMap<EntityKey, Option<EntitySnapshot>>, ToolError> {
380    let mut snapshots = HashMap::new();
381    let mut stale = Vec::new();
382    for write in writes {
383        let key = write.key();
384        if view.staged_row(key).is_some() {
385            continue;
386        }
387        let EntityKey {
388            entity_type,
389            entity_id,
390        } = key;
391        let snapshot = if entity_type == ContentEntityType::Page {
392            match (view.live.page(entity_id), live_pages.remove(&entity_id)) {
393                (None, None) => None,
394                (Some(read), Some(loaded)) if PageRecord::from(&loaded) == *read => {
395                    Some(EntitySnapshot::Page(loaded))
396                }
397                _ => {
398                    stale.push(key);
399                    continue;
400                }
401            }
402        } else {
403            view.live.entity_snapshot(key)
404        };
405        snapshots.insert(key, snapshot);
406    }
407    if stale.is_empty() {
408        Ok(snapshots)
409    } else {
410        Err(stale_read(&stale))
411    }
412}
413
414/// The staged state `writes` were built on: the changeset rows of the view and the live state of
415/// the other entities.
416fn staged_state_read(
417    view: &CourseView,
418    writes: &[StagedEntity],
419    live_snapshots: HashMap<EntityKey, Option<EntitySnapshot>>,
420) -> StagedStateRead {
421    StagedStateRead {
422        changeset_id: view.changeset.as_ref().map(|changeset| changeset.id),
423        entity_revisions: writes
424            .iter()
425            .filter_map(|write| Some((write.key(), view.staged_row(write.key())?.revision)))
426            .collect(),
427        live_snapshots,
428    }
429}
430
431#[cfg(test)]
432mod tests;