Skip to main content

tower_mcp/
async_task.rs

1//! Async task management for long-running MCP operations
2//!
3//! This module provides task lifecycle management for operations that may take
4//! longer than a typical request/response cycle. Legacy clients request task
5//! augmentation explicitly; final-protocol servers elect tasks after extension
6//! negotiation. Tasks can be tracked, polled, updated with input, and cancelled.
7//!
8//! Task state lives behind the pluggable [`TaskStore`] trait, mirroring the
9//! shape of [`crate::session_store`] and [`crate::event_store`]: a trait, an
10//! error enum, and an in-memory default. By default routers use
11//! [`MemoryTaskStore`], which keeps tasks in an in-process map (behavior
12//! identical to earlier versions). External stores (Redis, Postgres, etc.) can
13//! be plugged in so `tasks/get` works on any instance behind a load balancer
14//! in the sessionless 2026-07-28 flows (SEP-2663).
15//!
16//! # Example
17//!
18//! ```rust,no_run
19//! use std::sync::Arc;
20//! use tower_mcp::async_task::{MemoryTaskStore, TaskStore};
21//! use tower_mcp::McpRouter;
22//!
23//! let store: Arc<dyn TaskStore> = Arc::new(MemoryTaskStore::new());
24//! let router = McpRouter::new().task_store(store);
25//! ```
26//!
27//! See `examples/tasks.rs` for a runnable server.
28//!
29//! # Authorization
30//!
31//! SEP-2663 requires servers to authorize every task request, and warns that a
32//! task ID can act as a bearer token: whoever holds it can poll, update, or
33//! cancel the task. This module answers that in two layers.
34//!
35//! [`generate_task_id`] draws 128 bits from the system CSPRNG, so IDs cannot
36//! be enumerated or guessed. That only protects IDs nobody has seen, so each
37//! task also records the principal that created it (see [`TaskOwner`]), and
38//! every later operation must match under [`owner_matches`].
39//!
40//! Matching is equality, not "protect owned tasks and leave unowned ones
41//! open":
42//!
43//! | Task owner | Caller  | Result                                |
44//! |------------|---------|---------------------------------------|
45//! | none       | none    | allowed, no authentication configured |
46//! | `alice`    | `alice` | allowed                               |
47//! | `alice`    | `bob`   | denied                                |
48//! | `alice`    | none    | denied                                |
49//! | none       | `alice` | denied                                |
50//!
51//! The last row is deliberate. An unowned task can only exist if it was
52//! created with no authenticated context, so a request that now carries a
53//! principal is a different security context rather than an upgrade of the
54//! same one. Servers mixing public and authenticated paths (see
55//! [`AuthConfig::public_path`](crate::auth::AuthConfig::public_path)) should
56//! expect a task created anonymously to be unreachable once a token is
57//! presented.
58//!
59//! The principal comes from the OAuth `sub` claim that the HTTP and WebSocket
60//! transports bridge into request extensions. Without the `oauth` feature
61//! there is no principal, so every task is unowned and servers with no
62//! authentication behave as they did before ownership existed.
63//!
64//! ## Why a denial looks like a missing task
65//!
66//! A refused operation returns exactly what an unknown task returns: `-32602`
67//! with "Task not found".
68//!
69//! SEP-2663 mandates `-32602` for an invalid or nonexistent task ID, but
70//! leaves the authorization failure to the server: tasks should be bound to
71//! "some sort of authorization context, the implementation of which is left to
72//! individual servers according to their existing bespoke permission models".
73//! Reusing `-32602` is therefore tower-mcp policy, not a spec requirement.
74//!
75//! The reasoning is that answering "forbidden" would confirm the ID is real,
76//! which is what unguessable IDs exist to prevent. The same SEP notes that
77//! where binding is impossible "the task ID becomes the only line of defense
78//! against contamination". A server that prefers a distinguishable error can
79//! wrap the router and translate.
80//!
81//! Expiry follows the same rule: [`Task::is_expired`] runs from creation, and
82//! an expired task reads as absent rather than as expired, so a retention
83//! window cannot be probed either.
84//!
85//! # Status notifications
86//!
87//! A client may watch a task instead of polling it, by naming its ID in the
88//! `taskIds` filter of a `subscriptions/listen` stream. Each
89//! `notifications/tasks` carries the complete task, identical to the
90//! `tasks/get` response at that moment, so a client that hears about a
91//! completion already holds the result.
92//!
93//! The router announces the transitions it drives. A server that drives one
94//! itself, most commonly [`TaskStore::require_input`], announces it with
95//! [`McpRouter::notify_task_status_changed`](crate::McpRouter::notify_task_status_changed).
96//!
97//! Notifications are best effort and `tasks/get` stays authoritative: a task
98//! outlives the request that created it, so there may be no subscriber at the
99//! moment a transition happens, and a client that missed one loses nothing but
100//! time.
101
102use std::collections::{BTreeSet, HashMap};
103use std::fmt::Write as _;
104use std::sync::atomic::{AtomicBool, Ordering};
105use std::sync::{Arc, RwLock};
106use std::time::{Duration, Instant};
107
108use async_trait::async_trait;
109
110use crate::error::JsonRpcError;
111use crate::protocol::{CallToolResult, InputRequests, InputResponses, TaskObject, TaskStatus};
112
113/// Default time-to-live for a task (5 minutes, in milliseconds).
114///
115/// Per SEP-2663 the TTL runs from task creation, not from the moment the task
116/// reaches a terminal state.
117const DEFAULT_TTL_MS: u64 = 300_000;
118
119/// Default poll interval suggestion (2 seconds, in milliseconds)
120const DEFAULT_POLL_INTERVAL_MS: u64 = 2_000;
121
122/// Internal task representation with full state
123#[derive(Debug)]
124pub struct Task {
125    /// Unique task identifier
126    pub id: String,
127    /// Name of the tool being executed
128    pub tool_name: String,
129    /// Arguments passed to the tool
130    pub arguments: serde_json::Value,
131    /// Current task status
132    pub status: TaskStatus,
133    /// When the task was created
134    pub created_at: Instant,
135    /// ISO 8601 timestamp string
136    pub created_at_str: String,
137    /// ISO 8601 timestamp of last state change
138    pub last_updated_at_str: String,
139    /// Time-to-live in milliseconds (for cleanup after completion)
140    pub ttl: u64,
141    /// Suggested polling interval in milliseconds
142    pub poll_interval: u64,
143    /// Human-readable status message
144    pub status_message: Option<String>,
145    /// The result of the tool call (when completed)
146    pub result: Option<CallToolResult>,
147    /// Structured execution error (when failed).
148    ///
149    /// SEP-2663 requires `tasks/get` to surface a JSON-RPC error object, not a
150    /// message string. A tool that returns `CallToolResult { isError: true }`
151    /// is a *completed* task carrying an error result, so it never sets this.
152    pub error: Option<JsonRpcError>,
153    /// Principal that created the task, or `None` when it was created
154    /// without an authenticated context.
155    ///
156    /// Never serialized: ownership is an authorization fact, not wire state.
157    pub owner: TaskOwner,
158    /// Input requests currently awaiting a client response, keyed as sent.
159    pub input_requests: InputRequests,
160    /// Keys answered by a previous `tasks/update`.
161    pub answered_input_keys: BTreeSet<String>,
162    /// Keys displaced by a later [`TaskStore::require_input`] before being
163    /// answered.
164    pub superseded_input_keys: BTreeSet<String>,
165    /// Cancellation token for aborting the task
166    pub cancellation_token: CancellationToken,
167    /// When the task reached terminal status (for TTL tracking)
168    pub completed_at: Option<Instant>,
169    /// Notified when task reaches a terminal state
170    pub completion_notify: Arc<tokio::sync::Notify>,
171}
172
173impl Task {
174    /// Create a new task
175    fn new(
176        id: String,
177        tool_name: String,
178        arguments: serde_json::Value,
179        ttl: Option<u64>,
180        owner: TaskOwner,
181    ) -> Self {
182        let cancelled = Arc::new(AtomicBool::new(false));
183        let now_str = chrono_now_iso8601();
184        Self {
185            id,
186            tool_name,
187            arguments,
188            status: TaskStatus::Working,
189            created_at: Instant::now(),
190            created_at_str: now_str.clone(),
191            last_updated_at_str: now_str,
192            ttl: ttl.unwrap_or(DEFAULT_TTL_MS),
193            poll_interval: DEFAULT_POLL_INTERVAL_MS,
194            status_message: Some("Task started".to_string()),
195            result: None,
196            error: None,
197            owner,
198            input_requests: InputRequests::new(),
199            answered_input_keys: BTreeSet::new(),
200            superseded_input_keys: BTreeSet::new(),
201            cancellation_token: CancellationToken { cancelled },
202            completed_at: None,
203            completion_notify: Arc::new(tokio::sync::Notify::new()),
204        }
205    }
206
207    /// Convert to TaskObject for API responses
208    pub fn to_task_object(&self) -> TaskObject {
209        TaskObject {
210            task_id: self.id.clone(),
211            status: self.status,
212            status_message: self.status_message.clone(),
213            created_at: self.created_at_str.clone(),
214            last_updated_at: self.last_updated_at_str.clone(),
215            ttl: Some(self.ttl),
216            poll_interval: Some(self.poll_interval),
217            result: None,
218            error: None,
219            meta: None,
220        }
221    }
222
223    /// Check if this task should be cleaned up (TTL expired).
224    ///
225    /// The clock runs from creation, per SEP-2663. A long-running task can
226    /// therefore expire while still working, which is the intended behavior:
227    /// `ttlMs` bounds how long the server retains the task, not how long it
228    /// lingers after finishing.
229    pub fn is_expired(&self) -> bool {
230        self.created_at.elapsed() > Duration::from_millis(self.ttl)
231    }
232
233    /// Outstanding input requests, if the task is waiting on the client.
234    pub fn outstanding_input_requests(&self) -> &InputRequests {
235        &self.input_requests
236    }
237
238    /// Check if the task has been cancelled
239    pub fn is_cancelled(&self) -> bool {
240        self.cancellation_token.is_cancelled()
241    }
242}
243
244/// Generate an unguessable task identifier.
245///
246/// SEP-2663 notes that a task ID can function as a bearer token: anything that
247/// knows the ID can poll, update, or cancel the task. Identifiers are therefore
248/// 128 random bits from the system CSPRNG, rendered as hex, rather than a
249/// sequential counter.
250///
251/// # Panics
252///
253/// Panics if the operating system entropy source is unavailable. A server that
254/// cannot generate unguessable identifiers must not fall back to guessable
255/// ones.
256pub fn generate_task_id() -> String {
257    let mut bytes = [0u8; 16];
258    getrandom::fill(&mut bytes).expect("system entropy source unavailable for task ID generation");
259    let mut id = String::with_capacity(2 * bytes.len());
260    for byte in bytes {
261        let _ = write!(id, "{byte:02x}");
262    }
263    id
264}
265
266/// The principal a task belongs to.
267///
268/// `None` means the task was created without an authenticated context, which
269/// is the normal case for a server with no authentication configured.
270///
271/// SEP-2663 notes that a task ID can behave as a bearer token. Recording the
272/// owner is what stops the ID from being sufficient authority on its own once
273/// a second principal learns it.
274pub type TaskOwner = Option<String>;
275
276/// Whether `principal` may act on a task owned by `owner`.
277///
278/// Matching is equality, not "protect owned tasks and leave unowned ones
279/// open". An unowned task can only exist if it was created with no
280/// authenticated context, so a request that now carries a principal is a
281/// different security context and is refused.
282pub fn owner_matches(owner: &TaskOwner, principal: Option<&str>) -> bool {
283    owner.as_deref() == principal
284}
285
286/// Outcome of applying `tasks/update.inputResponses` to a task.
287///
288/// SEP-2663 requires partial responses to be honored: keys that match an
289/// outstanding request are consumed, everything else is ignored rather than
290/// rejected, and any request left unanswered stays outstanding.
291#[derive(Debug, Clone, Default, PartialEq, Eq)]
292pub struct AppliedInputResponses {
293    /// Keys matched to an outstanding request and consumed.
294    pub accepted: BTreeSet<String>,
295    /// Keys ignored because they were never issued, were already answered, or
296    /// were superseded by a later request.
297    pub ignored: BTreeSet<String>,
298    /// Requests still awaiting a response after this update.
299    pub still_outstanding: BTreeSet<String>,
300}
301
302impl AppliedInputResponses {
303    /// Whether every outstanding request has now been answered.
304    pub fn is_complete(&self) -> bool {
305        self.still_outstanding.is_empty()
306    }
307}
308
309/// A shareable cancellation token for task management
310#[derive(Debug, Clone)]
311pub struct CancellationToken {
312    cancelled: Arc<AtomicBool>,
313}
314
315impl CancellationToken {
316    /// Check if cancellation has been requested
317    pub fn is_cancelled(&self) -> bool {
318        self.cancelled.load(Ordering::Relaxed)
319    }
320
321    /// Request cancellation
322    pub fn cancel(&self) {
323        self.cancelled.store(true, Ordering::Relaxed);
324    }
325}
326
327/// Errors returned by [`TaskStore`] implementations.
328///
329/// Mirrors the three-variant shape of
330/// [`SessionStoreError`](crate::session_store::SessionStoreError): encode and
331/// decode errors from (de)serializing task state, and catch-all backend errors
332/// from the storage layer. [`MemoryTaskStore`] never returns errors; the
333/// variants exist for external implementations.
334#[derive(Debug, thiserror::Error)]
335#[non_exhaustive]
336pub enum TaskStoreError {
337    /// Failed to encode task state (e.g. serde serialization error).
338    #[error("encode error: {0}")]
339    Encode(String),
340    /// Failed to decode task state (e.g. corrupt data in the backend).
341    #[error("decode error: {0}")]
342    Decode(String),
343    /// Backend error (e.g. connection failure, transient storage error).
344    #[error("backend error: {0}")]
345    Backend(String),
346}
347
348/// Result alias for task store operations.
349pub type Result<T> = std::result::Result<T, TaskStoreError>;
350
351/// A task's current snapshot: the task object plus any result or error
352/// captured so far.
353///
354/// The error is a structured [`JsonRpcError`] because SEP-2663 requires
355/// `tasks/get` on a failed task to return a JSON-RPC error object.
356pub type TaskSnapshot = (TaskObject, Option<CallToolResult>, Option<JsonRpcError>);
357
358/// Storage backend for async task state.
359///
360/// Implementations persist task lifecycle state keyed by task ID. The default
361/// implementation is [`MemoryTaskStore`]; external stores (Redis, Postgres,
362/// etc.) typically live in separate crates.
363///
364/// # Semantics
365///
366/// - Terminal states ([`TaskStatus::is_terminal`]) are immutable: once a task
367///   is completed, failed, or cancelled, further transitions must be rejected
368///   (`Ok(false)` from the transition methods).
369/// - An expired task is indistinguishable from an unknown one. Reads return
370///   `None` once `ttlMs` has elapsed since creation, whether or not the entry
371///   has actually been reclaimed, so callers cannot probe for the existence of
372///   a task whose retention window has closed.
373/// - [`cancel_task`](Self::cancel_task) must signal the task's
374///   [`CancellationToken`] even if the task is already terminal.
375/// - [`wait_for_completion`](Self::wait_for_completion) blocks until the task
376///   reaches a terminal state; how an implementation waits (notification,
377///   polling, pub/sub) is an implementation detail and must not leak into the
378///   trait.
379#[async_trait]
380pub trait TaskStore: Send + Sync + 'static {
381    /// Create and store a new task owned by `owner`.
382    ///
383    /// Returns the task ID and a cancellation token for the spawned work.
384    /// `owner` is the authenticated principal responsible for the task, or
385    /// `None` when the request carried no authenticated context.
386    async fn create_task(
387        &self,
388        tool_name: &str,
389        arguments: serde_json::Value,
390        ttl: Option<u64>,
391        owner: TaskOwner,
392    ) -> Result<(String, CancellationToken)>;
393
394    /// Read a task's owner.
395    ///
396    /// The outer `Option` distinguishes a known task from an unknown or
397    /// expired one; the inner [`TaskOwner`] distinguishes an owned task from
398    /// one created without an authenticated principal.
399    async fn task_owner(&self, task_id: &str) -> Result<Option<TaskOwner>>;
400
401    /// Get task object by ID. Returns `None` if unknown.
402    async fn get_task(&self, task_id: &str) -> Result<Option<TaskObject>>;
403
404    /// Get a task's full snapshot (task object, result, error) by ID.
405    async fn get_task_result(&self, task_id: &str) -> Result<Option<TaskSnapshot>>;
406
407    /// Wait for a task to reach a terminal state, then return its snapshot.
408    ///
409    /// If the task is already terminal, returns immediately. Otherwise blocks
410    /// until the task completes, fails, or is cancelled. Returns `None` if
411    /// the task is unknown.
412    async fn wait_for_completion(&self, task_id: &str) -> Result<Option<TaskSnapshot>>;
413
414    /// List all tasks, optionally filtered by status.
415    async fn list_tasks(&self, status_filter: Option<TaskStatus>) -> Result<Vec<TaskObject>>;
416
417    /// Mark a task as requiring input, recording the requests to be answered.
418    ///
419    /// `requests` replaces the outstanding set. Any key that was outstanding
420    /// and is not re-issued becomes superseded; a re-issued key is a fresh
421    /// question and becomes outstanding again even if previously answered.
422    ///
423    /// Returns `Ok(false)` if the task is unknown, expired, or already
424    /// terminal.
425    async fn require_input(
426        &self,
427        task_id: &str,
428        requests: InputRequests,
429        message: Option<&str>,
430    ) -> Result<bool>;
431
432    /// Read the requests a task is currently waiting on.
433    ///
434    /// Returns an empty map when the task is not `input_required`, and `None`
435    /// when the task is unknown or expired.
436    async fn outstanding_input_requests(&self, task_id: &str) -> Result<Option<InputRequests>>;
437
438    /// Apply `tasks/update.inputResponses` to a task.
439    ///
440    /// Consumes the keys that match an outstanding request and ignores the
441    /// rest. When the last outstanding request is answered the task returns to
442    /// [`TaskStatus::Working`].
443    ///
444    /// Returns `None` if the task is unknown, expired, or already terminal.
445    async fn apply_input_responses(
446        &self,
447        task_id: &str,
448        responses: InputResponses,
449    ) -> Result<Option<AppliedInputResponses>>;
450
451    /// Update a task's time-to-live, measured from creation.
452    ///
453    /// SEP-2663 allows `ttlMs` to change over a task's lifetime. Returns
454    /// `Ok(false)` if the task is unknown or already expired.
455    async fn set_ttl(&self, task_id: &str, ttl_ms: u64) -> Result<bool>;
456
457    /// Mark a task as completed with a result.
458    ///
459    /// A result carrying `isError: true` still completes the task: the tool
460    /// ran and produced a domain error, which SEP-2663 distinguishes from an
461    /// execution failure.
462    ///
463    /// Returns `Ok(false)` if the task is unknown, expired, or already
464    /// terminal.
465    async fn complete_task(&self, task_id: &str, result: CallToolResult) -> Result<bool>;
466
467    /// Mark a task as failed with a structured execution error.
468    ///
469    /// Returns `Ok(false)` if the task is unknown, expired, or already
470    /// terminal.
471    async fn fail_task(&self, task_id: &str, error: JsonRpcError) -> Result<bool>;
472
473    /// Cancel a task.
474    ///
475    /// Signals the task's [`CancellationToken`] and, if the task is not
476    /// already terminal, marks it cancelled. Returns the updated task object,
477    /// or `None` if the task is unknown.
478    async fn cancel_task(&self, task_id: &str, reason: Option<&str>) -> Result<Option<TaskObject>>;
479}
480
481/// In-memory [`TaskStore`] backed by a `HashMap`.
482///
483/// This is the default store. Suitable for single-instance deployments. For
484/// horizontal scaling, use an external store that shares state across
485/// instances. Completion wakeups for
486/// [`wait_for_completion`](TaskStore::wait_for_completion) use a per-task
487/// [`tokio::sync::Notify`], which is an implementation detail of this store.
488#[derive(Debug, Clone)]
489pub struct MemoryTaskStore {
490    tasks: Arc<RwLock<HashMap<String, Task>>>,
491}
492
493impl Default for MemoryTaskStore {
494    fn default() -> Self {
495        Self::new()
496    }
497}
498
499impl MemoryTaskStore {
500    /// Create a new task store
501    pub fn new() -> Self {
502        Self {
503            tasks: Arc::new(RwLock::new(HashMap::new())),
504        }
505    }
506
507    /// Remove expired tasks (call periodically for cleanup).
508    ///
509    /// Returns the number removed. Not part of the [`TaskStore`] trait;
510    /// external backends typically expire entries natively (e.g. Redis TTL).
511    ///
512    /// Calling this is an optimization, not a correctness requirement: reads
513    /// already treat an expired task as absent.
514    pub fn cleanup_expired(&self) -> usize {
515        if let Ok(mut tasks) = self.tasks.write() {
516            let before = tasks.len();
517            tasks.retain(|_, t| !t.is_expired());
518            before - tasks.len()
519        } else {
520            0
521        }
522    }
523
524    /// Get the number of tasks in the store
525    #[cfg(test)]
526    pub fn len(&self) -> usize {
527        if let Ok(tasks) = self.tasks.read() {
528            tasks.len()
529        } else {
530            0
531        }
532    }
533
534    /// Check if the store is empty
535    #[cfg(test)]
536    pub fn is_empty(&self) -> bool {
537        self.len() == 0
538    }
539}
540
541#[async_trait]
542impl TaskStore for MemoryTaskStore {
543    async fn create_task(
544        &self,
545        tool_name: &str,
546        arguments: serde_json::Value,
547        ttl: Option<u64>,
548        owner: TaskOwner,
549    ) -> Result<(String, CancellationToken)> {
550        let id = generate_task_id();
551        let task = Task::new(id.clone(), tool_name.to_string(), arguments, ttl, owner);
552        let token = task.cancellation_token.clone();
553
554        if let Ok(mut tasks) = self.tasks.write() {
555            tasks.insert(id.clone(), task);
556        }
557
558        Ok((id, token))
559    }
560
561    async fn get_task(&self, task_id: &str) -> Result<Option<TaskObject>> {
562        Ok(if let Ok(tasks) = self.tasks.read() {
563            tasks
564                .get(task_id)
565                .filter(|t| !t.is_expired())
566                .map(|t| t.to_task_object())
567        } else {
568            None
569        })
570    }
571
572    async fn task_owner(&self, task_id: &str) -> Result<Option<TaskOwner>> {
573        Ok(if let Ok(tasks) = self.tasks.read() {
574            tasks
575                .get(task_id)
576                .filter(|t| !t.is_expired())
577                .map(|t| t.owner.clone())
578        } else {
579            None
580        })
581    }
582
583    async fn get_task_result(&self, task_id: &str) -> Result<Option<TaskSnapshot>> {
584        Ok(if let Ok(tasks) = self.tasks.read() {
585            tasks
586                .get(task_id)
587                .filter(|t| !t.is_expired())
588                .map(|t| (t.to_task_object(), t.result.clone(), t.error.clone()))
589        } else {
590            None
591        })
592    }
593
594    async fn wait_for_completion(&self, task_id: &str) -> Result<Option<TaskSnapshot>> {
595        // First check if already terminal and get the notify handle
596        let notify = {
597            let Ok(tasks) = self.tasks.read() else {
598                return Ok(None);
599            };
600            let Some(task) = tasks.get(task_id).filter(|t| !t.is_expired()) else {
601                return Ok(None);
602            };
603            if task.status.is_terminal() {
604                return Ok(Some((
605                    task.to_task_object(),
606                    task.result.clone(),
607                    task.error.clone(),
608                )));
609            }
610            task.completion_notify.clone()
611        };
612
613        // Wait for completion notification
614        notify.notified().await;
615
616        // Read the result
617        self.get_task_result(task_id).await
618    }
619
620    async fn list_tasks(&self, status_filter: Option<TaskStatus>) -> Result<Vec<TaskObject>> {
621        Ok(if let Ok(tasks) = self.tasks.read() {
622            tasks
623                .values()
624                .filter(|t| !t.is_expired())
625                .filter(|t| status_filter.is_none() || status_filter == Some(t.status))
626                .map(|t| t.to_task_object())
627                .collect()
628        } else {
629            vec![]
630        })
631    }
632
633    async fn require_input(
634        &self,
635        task_id: &str,
636        requests: InputRequests,
637        message: Option<&str>,
638    ) -> Result<bool> {
639        let Ok(mut tasks) = self.tasks.write() else {
640            return Ok(false);
641        };
642        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
643            return Ok(false);
644        };
645        if task.status.is_terminal() {
646            return Ok(false);
647        }
648
649        // Outstanding requests the server did not re-issue are superseded.
650        for key in std::mem::take(&mut task.input_requests).into_keys() {
651            if !requests.contains_key(&key) {
652                task.superseded_input_keys.insert(key);
653            }
654        }
655        // A re-issued key is a fresh question, whatever its prior fate.
656        for key in requests.keys() {
657            task.answered_input_keys.remove(key);
658            task.superseded_input_keys.remove(key);
659        }
660
661        task.input_requests = requests;
662        task.status = TaskStatus::InputRequired;
663        task.status_message = Some(
664            message
665                .map(str::to_string)
666                .unwrap_or_else(|| "Awaiting client input".to_string()),
667        );
668        task.last_updated_at_str = chrono_now_iso8601();
669        Ok(true)
670    }
671
672    async fn outstanding_input_requests(&self, task_id: &str) -> Result<Option<InputRequests>> {
673        Ok(if let Ok(tasks) = self.tasks.read() {
674            tasks
675                .get(task_id)
676                .filter(|t| !t.is_expired())
677                .map(|t| t.input_requests.clone())
678        } else {
679            None
680        })
681    }
682
683    async fn apply_input_responses(
684        &self,
685        task_id: &str,
686        responses: InputResponses,
687    ) -> Result<Option<AppliedInputResponses>> {
688        let Ok(mut tasks) = self.tasks.write() else {
689            return Ok(None);
690        };
691        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
692            return Ok(None);
693        };
694        if task.status.is_terminal() {
695            return Ok(None);
696        }
697
698        let mut applied = AppliedInputResponses::default();
699        for key in responses.into_keys() {
700            if task.input_requests.remove(&key).is_some() {
701                task.answered_input_keys.insert(key.clone());
702                applied.accepted.insert(key);
703            } else {
704                // Never issued, already answered, or superseded. All three are
705                // ignored rather than rejected, so a client replaying a stale
706                // update does not fail the task.
707                applied.ignored.insert(key);
708            }
709        }
710        applied.still_outstanding = task.input_requests.keys().cloned().collect();
711
712        if !applied.accepted.is_empty() {
713            task.last_updated_at_str = chrono_now_iso8601();
714        }
715        if applied.is_complete() && task.status == TaskStatus::InputRequired {
716            task.status = TaskStatus::Working;
717            task.status_message = Some("Task resumed".to_string());
718        }
719        Ok(Some(applied))
720    }
721
722    async fn set_ttl(&self, task_id: &str, ttl_ms: u64) -> Result<bool> {
723        let Ok(mut tasks) = self.tasks.write() else {
724            return Ok(false);
725        };
726        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
727            return Ok(false);
728        };
729        task.ttl = ttl_ms;
730        task.last_updated_at_str = chrono_now_iso8601();
731        Ok(true)
732    }
733
734    async fn complete_task(&self, task_id: &str, result: CallToolResult) -> Result<bool> {
735        let Ok(mut tasks) = self.tasks.write() else {
736            return Ok(false);
737        };
738        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
739            return Ok(false);
740        };
741        if task.status.is_terminal() {
742            return Ok(false);
743        }
744        task.status = TaskStatus::Completed;
745        task.status_message = Some("Task completed".to_string());
746        task.result = Some(result);
747        task.input_requests.clear();
748        task.completed_at = Some(Instant::now());
749        task.last_updated_at_str = chrono_now_iso8601();
750        task.completion_notify.notify_waiters();
751        Ok(true)
752    }
753
754    async fn fail_task(&self, task_id: &str, error: JsonRpcError) -> Result<bool> {
755        let Ok(mut tasks) = self.tasks.write() else {
756            return Ok(false);
757        };
758        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
759            return Ok(false);
760        };
761        if task.status.is_terminal() {
762            return Ok(false);
763        }
764        task.status = TaskStatus::Failed;
765        task.status_message = Some(format!("Task failed: {}", error.message));
766        task.error = Some(error);
767        task.input_requests.clear();
768        task.completed_at = Some(Instant::now());
769        task.last_updated_at_str = chrono_now_iso8601();
770        task.completion_notify.notify_waiters();
771        Ok(true)
772    }
773
774    async fn cancel_task(&self, task_id: &str, reason: Option<&str>) -> Result<Option<TaskObject>> {
775        let Ok(mut tasks) = self.tasks.write() else {
776            return Ok(None);
777        };
778        let Some(task) = tasks.get_mut(task_id).filter(|t| !t.is_expired()) else {
779            return Ok(None);
780        };
781
782        // Signal cancellation
783        task.cancellation_token.cancel();
784
785        // If not already terminal, mark as cancelled
786        if !task.status.is_terminal() {
787            task.input_requests.clear();
788            task.status = TaskStatus::Cancelled;
789            task.status_message = Some(
790                reason
791                    .map(|r| format!("Cancelled: {}", r))
792                    .unwrap_or_else(|| "Task cancelled".to_string()),
793            );
794            task.completed_at = Some(Instant::now());
795            task.last_updated_at_str = chrono_now_iso8601();
796            task.completion_notify.notify_waiters();
797        }
798        Ok(Some(task.to_task_object()))
799    }
800}
801
802/// Build the validated extension declaration for the final Tasks extension.
803///
804/// The SEP-2663 capability shape is an empty object: support is declared by
805/// the identifier's presence, with no settings to negotiate.
806pub fn tasks_extension() -> crate::ExtensionDeclaration {
807    crate::ExtensionDeclaration::empty(crate::protocol::TASKS_EXTENSION_ID)
808        .expect("the built-in Tasks extension declaration is valid")
809}
810
811impl crate::McpRouter {
812    /// Advertise final Tasks support (SEP-2663) from this server.
813    ///
814    /// Compiling the task APIs does not advertise them. A server opts in here,
815    /// and only then does the final protocol path advertise
816    /// `io.modelcontextprotocol/tasks`, elect to return tasks from ordinary
817    /// `tools/call` requests, or serve the final task methods. Legacy
818    /// 2025-11-25 task behavior is unaffected either way.
819    pub fn with_tasks(self) -> Self {
820        self.with_protocol_extension(tasks_extension())
821    }
822}
823
824impl crate::McpClientBuilder {
825    /// Declare final Tasks support (SEP-2663) from this client.
826    pub fn with_tasks(self) -> Self {
827        self.with_protocol_extension(tasks_extension())
828    }
829}
830
831impl crate::RequestContext {
832    /// Whether both peers negotiated the final Tasks extension.
833    ///
834    /// Task dispatch keys off this rather than off the protocol version: a
835    /// 2026-07-28 request from a client that did not declare the extension
836    /// must not receive a task.
837    pub fn supports_tasks(&self) -> bool {
838        self.negotiated_extensions()
839            .is_some_and(|extensions| extensions.contains(crate::protocol::TASKS_EXTENSION_ID))
840    }
841}
842
843/// Generate ISO 8601 timestamp for current time
844fn chrono_now_iso8601() -> String {
845    use std::time::SystemTime;
846
847    let now = SystemTime::now();
848    let duration = now
849        .duration_since(SystemTime::UNIX_EPOCH)
850        .unwrap_or_default();
851
852    let secs = duration.as_secs();
853    let millis = duration.subsec_millis();
854
855    // Simple ISO 8601 format (UTC)
856    // Calculate date/time components
857    let days = secs / 86400;
858    let remaining = secs % 86400;
859    let hours = remaining / 3600;
860    let remaining = remaining % 3600;
861    let minutes = remaining / 60;
862    let seconds = remaining % 60;
863
864    // Calculate year/month/day from days since epoch (1970-01-01)
865    // This is a simplified calculation that handles leap years
866    let mut year = 1970i32;
867    let mut remaining_days = days as i32;
868
869    loop {
870        let days_in_year = if is_leap_year(year) { 366 } else { 365 };
871        if remaining_days < days_in_year {
872            break;
873        }
874        remaining_days -= days_in_year;
875        year += 1;
876    }
877
878    let days_in_months: [i32; 12] = if is_leap_year(year) {
879        [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
880    } else {
881        [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
882    };
883
884    let mut month = 1;
885    for days_in_month in days_in_months.iter() {
886        if remaining_days < *days_in_month {
887            break;
888        }
889        remaining_days -= days_in_month;
890        month += 1;
891    }
892
893    let day = remaining_days + 1;
894
895    format!(
896        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:03}Z",
897        year, month, day, hours, minutes, seconds, millis
898    )
899}
900
901fn is_leap_year(year: i32) -> bool {
902    (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0)
903}
904
905#[cfg(test)]
906mod tests {
907    use super::*;
908    use crate::protocol::{
909        ElicitAction, ElicitResult, InputRequest, InputResponse, ListRootsParams,
910    };
911
912    #[tokio::test]
913    async fn test_create_task() {
914        let store = MemoryTaskStore::new();
915        let (id, token) = store
916            .create_task("test-tool", serde_json::json!({"a": 1}), None, None)
917            .await
918            .unwrap();
919
920        assert!(!id.is_empty());
921        assert!(!token.is_cancelled());
922
923        let info = store
924            .get_task(&id)
925            .await
926            .unwrap()
927            .expect("task should exist");
928        assert_eq!(info.task_id, id);
929        assert_eq!(info.status, TaskStatus::Working);
930    }
931
932    #[tokio::test]
933    async fn test_task_lifecycle() {
934        let store = MemoryTaskStore::new();
935        let (id, _) = store
936            .create_task("test-tool", serde_json::json!({}), None, None)
937            .await
938            .unwrap();
939
940        // Complete task
941        assert!(
942            store
943                .complete_task(&id, CallToolResult::text("Done"))
944                .await
945                .unwrap()
946        );
947
948        let info = store.get_task(&id).await.unwrap().unwrap();
949        assert_eq!(info.status, TaskStatus::Completed);
950    }
951
952    #[tokio::test]
953    async fn test_task_cancellation() {
954        let store = MemoryTaskStore::new();
955        let (id, token) = store
956            .create_task("test-tool", serde_json::json!({}), None, None)
957            .await
958            .unwrap();
959
960        assert!(!token.is_cancelled());
961
962        let task_obj = store
963            .cancel_task(&id, Some("User requested"))
964            .await
965            .unwrap();
966        assert!(task_obj.is_some());
967        assert_eq!(task_obj.unwrap().status, TaskStatus::Cancelled);
968        assert!(token.is_cancelled());
969
970        let info = store.get_task(&id).await.unwrap().unwrap();
971        assert_eq!(info.status, TaskStatus::Cancelled);
972    }
973
974    #[tokio::test]
975    async fn test_task_failure() {
976        let store = MemoryTaskStore::new();
977        let (id, _) = store
978            .create_task("test-tool", serde_json::json!({}), None, None)
979            .await
980            .unwrap();
981
982        assert!(
983            store
984                .fail_task(&id, JsonRpcError::internal_error("Something went wrong"))
985                .await
986                .unwrap()
987        );
988
989        let info = store.get_task(&id).await.unwrap().unwrap();
990        assert_eq!(info.status, TaskStatus::Failed);
991        assert!(info.status_message.as_ref().unwrap().contains("failed"));
992    }
993
994    #[tokio::test]
995    async fn test_list_tasks() {
996        let store = MemoryTaskStore::new();
997        store
998            .create_task("tool1", serde_json::json!({}), None, None)
999            .await
1000            .unwrap();
1001        store
1002            .create_task("tool2", serde_json::json!({}), None, None)
1003            .await
1004            .unwrap();
1005        let (id3, _) = store
1006            .create_task("tool3", serde_json::json!({}), None, None)
1007            .await
1008            .unwrap();
1009
1010        // Complete one task
1011        store
1012            .complete_task(&id3, CallToolResult::text("Done"))
1013            .await
1014            .unwrap();
1015
1016        // List all tasks
1017        let all = store.list_tasks(None).await.unwrap();
1018        assert_eq!(all.len(), 3);
1019
1020        // List only working tasks
1021        let working = store.list_tasks(Some(TaskStatus::Working)).await.unwrap();
1022        assert_eq!(working.len(), 2);
1023
1024        // List only completed tasks
1025        let completed = store.list_tasks(Some(TaskStatus::Completed)).await.unwrap();
1026        assert_eq!(completed.len(), 1);
1027    }
1028
1029    #[tokio::test]
1030    async fn test_terminal_state_immutable() {
1031        let store = MemoryTaskStore::new();
1032        let (id, _) = store
1033            .create_task("test-tool", serde_json::json!({}), None, None)
1034            .await
1035            .unwrap();
1036
1037        // Complete the task
1038        store
1039            .complete_task(&id, CallToolResult::text("Done"))
1040            .await
1041            .unwrap();
1042
1043        // Try to fail - should fail
1044        assert!(
1045            !store
1046                .fail_task(&id, JsonRpcError::internal_error("Error"))
1047                .await
1048                .unwrap()
1049        );
1050
1051        // Status should still be completed
1052        let info = store.get_task(&id).await.unwrap().unwrap();
1053        assert_eq!(info.status, TaskStatus::Completed);
1054    }
1055
1056    #[tokio::test]
1057    async fn test_task_ids_unique() {
1058        let store = MemoryTaskStore::new();
1059        let (id1, _) = store
1060            .create_task("tool", serde_json::json!({}), None, None)
1061            .await
1062            .unwrap();
1063        let (id2, _) = store
1064            .create_task("tool", serde_json::json!({}), None, None)
1065            .await
1066            .unwrap();
1067        let (id3, _) = store
1068            .create_task("tool", serde_json::json!({}), None, None)
1069            .await
1070            .unwrap();
1071
1072        assert_ne!(id1, id2);
1073        assert_ne!(id2, id3);
1074        assert_ne!(id1, id3);
1075    }
1076
1077    #[tokio::test]
1078    async fn test_get_task_result() {
1079        let store = MemoryTaskStore::new();
1080        let (id, _) = store
1081            .create_task("test-tool", serde_json::json!({}), None, None)
1082            .await
1083            .unwrap();
1084
1085        // Complete with result
1086        let result = CallToolResult::text("The result");
1087        store.complete_task(&id, result).await.unwrap();
1088
1089        let (task_obj, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1090        assert_eq!(task_obj.status, TaskStatus::Completed);
1091        assert!(result.is_some());
1092        assert!(error.is_none());
1093    }
1094
1095    #[tokio::test]
1096    async fn test_wait_for_completion_returns_terminal_snapshot() {
1097        let store = MemoryTaskStore::new();
1098        let (id, _) = store
1099            .create_task("test-tool", serde_json::json!({}), None, None)
1100            .await
1101            .unwrap();
1102
1103        // Complete the task from another task while a waiter is blocked.
1104        let waiter_store = store.clone();
1105        let waiter_id = id.clone();
1106        let waiter =
1107            tokio::spawn(async move { waiter_store.wait_for_completion(&waiter_id).await });
1108
1109        tokio::time::sleep(Duration::from_millis(10)).await;
1110        store
1111            .complete_task(&id, CallToolResult::text("Done"))
1112            .await
1113            .unwrap();
1114
1115        let (task_obj, result, error) = waiter.await.unwrap().unwrap().unwrap();
1116        assert_eq!(task_obj.status, TaskStatus::Completed);
1117        assert!(result.is_some());
1118        assert!(error.is_none());
1119    }
1120
1121    #[tokio::test]
1122    async fn dyn_task_store_object_safe() {
1123        // Compile-time check that TaskStore is object-safe.
1124        let store: Arc<dyn TaskStore> = Arc::new(MemoryTaskStore::new());
1125        let (id, _) = store
1126            .create_task("tool", serde_json::json!({}), None, None)
1127            .await
1128            .unwrap();
1129        assert!(store.get_task(&id).await.unwrap().is_some());
1130    }
1131
1132    #[test]
1133    fn test_iso8601_timestamp() {
1134        let ts = chrono_now_iso8601();
1135        // Basic format check
1136        assert!(ts.ends_with('Z'));
1137        assert!(ts.contains('T'));
1138        assert_eq!(ts.len(), 24); // YYYY-MM-DDTHH:MM:SS.mmmZ
1139    }
1140
1141    #[test]
1142    fn test_task_status_display() {
1143        assert_eq!(TaskStatus::Working.to_string(), "working");
1144        assert_eq!(TaskStatus::InputRequired.to_string(), "input_required");
1145        assert_eq!(TaskStatus::Completed.to_string(), "completed");
1146        assert_eq!(TaskStatus::Failed.to_string(), "failed");
1147        assert_eq!(TaskStatus::Cancelled.to_string(), "cancelled");
1148    }
1149
1150    #[test]
1151    fn test_task_status_is_terminal() {
1152        assert!(!TaskStatus::Working.is_terminal());
1153        assert!(!TaskStatus::InputRequired.is_terminal());
1154        assert!(TaskStatus::Completed.is_terminal());
1155        assert!(TaskStatus::Failed.is_terminal());
1156        assert!(TaskStatus::Cancelled.is_terminal());
1157    }
1158
1159    fn requests(keys: &[&str]) -> InputRequests {
1160        keys.iter()
1161            .map(|k| {
1162                (
1163                    k.to_string(),
1164                    InputRequest::ListRoots(ListRootsParams { meta: None }),
1165                )
1166            })
1167            .collect()
1168    }
1169
1170    fn accept(key: &str) -> (String, InputResponse) {
1171        (
1172            key.to_string(),
1173            InputResponse::Elicit(ElicitResult {
1174                action: ElicitAction::Accept,
1175                content: None,
1176                meta: None,
1177            }),
1178        )
1179    }
1180
1181    async fn working_task(store: &MemoryTaskStore, ttl: Option<u64>) -> String {
1182        store
1183            .create_task("tool", serde_json::json!({}), ttl, None)
1184            .await
1185            .unwrap()
1186            .0
1187    }
1188
1189    #[tokio::test]
1190    async fn task_ids_are_unguessable_not_sequential() {
1191        let store = MemoryTaskStore::new();
1192        let mut ids = BTreeSet::new();
1193        for _ in 0..64 {
1194            ids.insert(working_task(&store, None).await);
1195        }
1196        assert_eq!(ids.len(), 64, "task IDs collided");
1197
1198        for id in &ids {
1199            assert_eq!(id.len(), 32, "expected 128 bits of hex: {id}");
1200            assert!(id.chars().all(|c| c.is_ascii_hexdigit()), "{id}");
1201            assert!(!id.starts_with("task-"), "sequential-looking ID: {id}");
1202        }
1203
1204        // A counter would make every ID a near-neighbor of the last. Require
1205        // the set to span a wide range of leading bytes instead.
1206        let leading: BTreeSet<&str> = ids.iter().map(|id| &id[..2]).collect();
1207        assert!(
1208            leading.len() > 32,
1209            "only {} distinct leading bytes across 64 IDs",
1210            leading.len()
1211        );
1212    }
1213
1214    #[tokio::test]
1215    async fn ttl_runs_from_creation_and_expired_tasks_read_as_absent() {
1216        let store = MemoryTaskStore::new();
1217        let id = working_task(&store, Some(0)).await;
1218
1219        // TTL of 0 expires immediately, while the task is still working, so
1220        // the clock plainly is not waiting for a terminal state.
1221        tokio::time::sleep(Duration::from_millis(5)).await;
1222
1223        assert!(store.get_task(&id).await.unwrap().is_none());
1224        assert!(store.get_task_result(&id).await.unwrap().is_none());
1225        assert!(store.list_tasks(None).await.unwrap().is_empty());
1226        assert!(
1227            store
1228                .outstanding_input_requests(&id)
1229                .await
1230                .unwrap()
1231                .is_none()
1232        );
1233        assert!(store.cancel_task(&id, None).await.unwrap().is_none());
1234        assert!(!store.set_ttl(&id, 60_000).await.unwrap());
1235        assert!(
1236            !store
1237                .complete_task(&id, CallToolResult::text("late"))
1238                .await
1239                .unwrap()
1240        );
1241    }
1242
1243    #[tokio::test]
1244    async fn ttl_is_mutable_over_the_task_lifetime() {
1245        let store = MemoryTaskStore::new();
1246        let id = working_task(&store, Some(60_000)).await;
1247
1248        assert!(store.set_ttl(&id, 120_000).await.unwrap());
1249        let task = store.get_task(&id).await.unwrap().unwrap();
1250        assert_eq!(task.ttl, Some(120_000));
1251
1252        // Shortening the window to zero retires the task immediately.
1253        assert!(store.set_ttl(&id, 0).await.unwrap());
1254        tokio::time::sleep(Duration::from_millis(5)).await;
1255        assert!(store.get_task(&id).await.unwrap().is_none());
1256    }
1257
1258    #[tokio::test]
1259    async fn require_input_records_requests_and_exposes_them() {
1260        let store = MemoryTaskStore::new();
1261        let id = working_task(&store, None).await;
1262
1263        assert!(
1264            store
1265                .require_input(&id, requests(&["approval", "region"]), Some("need input"))
1266                .await
1267                .unwrap()
1268        );
1269
1270        let task = store.get_task(&id).await.unwrap().unwrap();
1271        assert_eq!(task.status, TaskStatus::InputRequired);
1272        assert_eq!(task.status_message.as_deref(), Some("need input"));
1273
1274        let outstanding = store
1275            .outstanding_input_requests(&id)
1276            .await
1277            .unwrap()
1278            .unwrap();
1279        assert_eq!(
1280            outstanding.keys().collect::<Vec<_>>(),
1281            vec!["approval", "region"],
1282            "every outstanding request must be exposed, not just the newest"
1283        );
1284    }
1285
1286    #[tokio::test]
1287    async fn partial_input_responses_leave_the_rest_outstanding() {
1288        let store = MemoryTaskStore::new();
1289        let id = working_task(&store, None).await;
1290        store
1291            .require_input(&id, requests(&["approval", "region"]), None)
1292            .await
1293            .unwrap();
1294
1295        let applied = store
1296            .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1297            .await
1298            .unwrap()
1299            .unwrap();
1300
1301        assert_eq!(applied.accepted, ["approval".to_string()].into());
1302        assert!(applied.ignored.is_empty());
1303        assert_eq!(applied.still_outstanding, ["region".to_string()].into());
1304        assert!(!applied.is_complete());
1305
1306        // The task stays blocked while anything is unanswered.
1307        let task = store.get_task(&id).await.unwrap().unwrap();
1308        assert_eq!(task.status, TaskStatus::InputRequired);
1309        assert_eq!(
1310            store
1311                .outstanding_input_requests(&id)
1312                .await
1313                .unwrap()
1314                .unwrap()
1315                .keys()
1316                .collect::<Vec<_>>(),
1317            vec!["region"]
1318        );
1319
1320        // Answering the last one resumes the task.
1321        let applied = store
1322            .apply_input_responses(&id, [accept("region")].into_iter().collect())
1323            .await
1324            .unwrap()
1325            .unwrap();
1326        assert!(applied.is_complete());
1327        assert_eq!(
1328            store.get_task(&id).await.unwrap().unwrap().status,
1329            TaskStatus::Working
1330        );
1331    }
1332
1333    #[tokio::test]
1334    async fn unknown_answered_and_superseded_response_keys_are_ignored() {
1335        let store = MemoryTaskStore::new();
1336        let id = working_task(&store, None).await;
1337        store
1338            .require_input(&id, requests(&["approval", "stale"]), None)
1339            .await
1340            .unwrap();
1341
1342        // Answer one, then re-issue a set that drops `stale`, superseding it.
1343        store
1344            .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1345            .await
1346            .unwrap()
1347            .unwrap();
1348        store
1349            .require_input(&id, requests(&["region"]), None)
1350            .await
1351            .unwrap();
1352
1353        let applied = store
1354            .apply_input_responses(
1355                &id,
1356                [accept("never-issued"), accept("approval"), accept("stale")]
1357                    .into_iter()
1358                    .collect(),
1359            )
1360            .await
1361            .unwrap()
1362            .unwrap();
1363
1364        assert!(
1365            applied.accepted.is_empty(),
1366            "none of these keys are outstanding"
1367        );
1368        assert_eq!(
1369            applied.ignored,
1370            [
1371                "never-issued".to_string(),
1372                "approval".to_string(),
1373                "stale".to_string()
1374            ]
1375            .into(),
1376            "unknown, already-answered, and superseded keys are all ignored"
1377        );
1378        assert_eq!(applied.still_outstanding, ["region".to_string()].into());
1379        assert_eq!(
1380            store.get_task(&id).await.unwrap().unwrap().status,
1381            TaskStatus::InputRequired,
1382            "ignoring a stale update must not resume or fail the task"
1383        );
1384    }
1385
1386    #[tokio::test]
1387    async fn reissued_key_becomes_a_fresh_question() {
1388        let store = MemoryTaskStore::new();
1389        let id = working_task(&store, None).await;
1390        store
1391            .require_input(&id, requests(&["approval"]), None)
1392            .await
1393            .unwrap();
1394        store
1395            .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1396            .await
1397            .unwrap()
1398            .unwrap();
1399
1400        // The server asks the same key again: the earlier answer must not
1401        // satisfy it.
1402        store
1403            .require_input(&id, requests(&["approval"]), None)
1404            .await
1405            .unwrap();
1406        let applied = store
1407            .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1408            .await
1409            .unwrap()
1410            .unwrap();
1411        assert_eq!(applied.accepted, ["approval".to_string()].into());
1412        assert!(applied.is_complete());
1413    }
1414
1415    #[tokio::test]
1416    async fn failed_tasks_preserve_the_structured_error() {
1417        let store = MemoryTaskStore::new();
1418        let id = working_task(&store, None).await;
1419
1420        let mut error = JsonRpcError::invalid_params("bad region");
1421        error.data = Some(serde_json::json!({"field": "region"}));
1422        assert!(store.fail_task(&id, error).await.unwrap());
1423
1424        let (_, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1425        assert!(result.is_none());
1426        let error = error.expect("structured error must survive the store");
1427        assert_eq!(
1428            error.code, -32602,
1429            "the original code must not be flattened"
1430        );
1431        assert_eq!(error.message, "bad region");
1432        assert_eq!(error.data.unwrap()["field"], "region");
1433    }
1434
1435    #[tokio::test]
1436    async fn tool_error_results_complete_the_task() {
1437        let store = MemoryTaskStore::new();
1438        let id = working_task(&store, None).await;
1439
1440        let mut result = CallToolResult::text("domain failure");
1441        result.is_error = true;
1442        assert!(store.complete_task(&id, result).await.unwrap());
1443
1444        let (task, result, error) = store.get_task_result(&id).await.unwrap().unwrap();
1445        assert_eq!(
1446            task.status,
1447            TaskStatus::Completed,
1448            "isError is a domain error, not an execution failure"
1449        );
1450        assert!(result.unwrap().is_error);
1451        assert!(error.is_none(), "no JSON-RPC error accompanies isError");
1452    }
1453
1454    #[tokio::test]
1455    async fn tasks_record_their_creating_principal() {
1456        let store = MemoryTaskStore::new();
1457        let (owned, _) = store
1458            .create_task("tool", serde_json::json!({}), None, Some("alice".into()))
1459            .await
1460            .unwrap();
1461        let (unowned, _) = store
1462            .create_task("tool", serde_json::json!({}), None, None)
1463            .await
1464            .unwrap();
1465
1466        assert_eq!(
1467            store.task_owner(&owned).await.unwrap(),
1468            Some(Some("alice".to_string()))
1469        );
1470        assert_eq!(store.task_owner(&unowned).await.unwrap(), Some(None));
1471        assert_eq!(
1472            store.task_owner("does-not-exist").await.unwrap(),
1473            None,
1474            "an unknown task has no owner record at all"
1475        );
1476
1477        // Ownership is an authorization fact and must not reach the wire.
1478        let wire = serde_json::to_value(store.get_task(&owned).await.unwrap().unwrap()).unwrap();
1479        assert!(
1480            wire.get("owner").is_none(),
1481            "owner leaked to the wire: {wire}"
1482        );
1483        assert!(!wire.to_string().contains("alice"));
1484    }
1485
1486    #[test]
1487    fn owner_matching_is_equality_not_leniency() {
1488        assert!(owner_matches(&None, None), "no auth configured");
1489        assert!(owner_matches(&Some("alice".into()), Some("alice")));
1490
1491        assert!(
1492            !owner_matches(&Some("alice".into()), Some("bob")),
1493            "a different principal must not inherit the task"
1494        );
1495        assert!(
1496            !owner_matches(&Some("alice".into()), None),
1497            "dropping the token must not grant access"
1498        );
1499        assert!(
1500            !owner_matches(&None, Some("alice")),
1501            "an unowned task belongs to a different security context"
1502        );
1503    }
1504
1505    #[tokio::test]
1506    async fn terminal_states_clear_outstanding_requests() {
1507        for (label, terminate) in [("completed", true), ("cancelled", false)] {
1508            let store = MemoryTaskStore::new();
1509            let id = working_task(&store, None).await;
1510            store
1511                .require_input(&id, requests(&["approval"]), None)
1512                .await
1513                .unwrap();
1514
1515            if terminate {
1516                store
1517                    .complete_task(&id, CallToolResult::text("done"))
1518                    .await
1519                    .unwrap();
1520            } else {
1521                store.cancel_task(&id, None).await.unwrap();
1522            }
1523
1524            assert!(
1525                store
1526                    .outstanding_input_requests(&id)
1527                    .await
1528                    .unwrap()
1529                    .unwrap()
1530                    .is_empty(),
1531                "{label} task still advertises outstanding input requests"
1532            );
1533            assert!(
1534                store
1535                    .apply_input_responses(&id, [accept("approval")].into_iter().collect())
1536                    .await
1537                    .unwrap()
1538                    .is_none(),
1539                "{label} task accepted a late input response"
1540            );
1541        }
1542    }
1543}