Implement SEP-2663 Tasks Extension (#1579)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Pranav Senthilnathan
GitHub
parent
ed07fdb6f4
commit
dbb7a20176
1 changed files with 8190 additions and 11405 deletions
|
|
@ -1,543 +1,222 @@
|
|||
using ModelContextProtocol.Protocol;
|
||||
using System.Collections.Concurrent;
|
||||
using System.Collections.Immutable;
|
||||
using System.Diagnostics.CodeAnalysis;
|
||||
using System.Text.Json;
|
||||
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
namespace ModelContextProtocol.Tests.Internal;
|
||||
#else
|
||||
namespace ModelContextProtocol;
|
||||
#endif
|
||||
namespace ModelContextProtocol.Server;
|
||||
|
||||
/// <summary>
|
||||
/// Provides an in-memory implementation of <see cref="IMcpTaskStore"/> for development and testing.
|
||||
/// Provides an in-memory implementation of <see cref="IMcpTaskStore"/> for development and testing scenarios.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// This implementation uses thread-safe concurrent collections and is suitable for single-server
|
||||
/// scenarios and testing. It is not recommended for production multi-server deployments as tasks
|
||||
/// are stored only in memory and are lost on server restart.
|
||||
/// This implementation stores all task state in memory using immutable snapshots and
|
||||
/// compare-and-swap updates for thread safety without locks.
|
||||
/// Tasks are not persisted across process restarts.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// Features:
|
||||
/// <list type="bullet">
|
||||
/// <item><description>Thread-safe operations using <see cref="ConcurrentDictionary{TKey, TValue}"/></description></item>
|
||||
/// <item><description>Automatic TTL-based cleanup via background task</description></item>
|
||||
/// <item><description>Session-based isolation when sessionId is provided</description></item>
|
||||
/// <item><description>Configurable default TTL and maximum TTL limits</description></item>
|
||||
/// </list>
|
||||
/// For production scenarios requiring durability, session isolation, or TTL-based cleanup,
|
||||
/// implement a custom <see cref="IMcpTaskStore"/>.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
[Experimental(Experimentals.Tasks_DiagnosticId, UrlFormat = Experimentals.Tasks_Url)]
|
||||
public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
|
||||
[Experimental(Experimentals.Extensions_DiagnosticId, UrlFormat = Experimentals.Extensions_Url)]
|
||||
public class InMemoryMcpTaskStore : IMcpTaskStore
|
||||
{
|
||||
private readonly ConcurrentDictionary<string, TaskEntry> _tasks = new();
|
||||
private readonly TimeSpan? _defaultTtl;
|
||||
private readonly TimeSpan? _maxTtl;
|
||||
private readonly TimeSpan _pollInterval;
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
private readonly ITimer? _cleanupTimer;
|
||||
#else
|
||||
private readonly Timer? _cleanupTimer;
|
||||
#endif
|
||||
private readonly int _pageSize;
|
||||
private readonly int? _maxTasks;
|
||||
private readonly int? _maxTasksPerSession;
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
private readonly TimeProvider _timeProvider;
|
||||
#endif
|
||||
private readonly ConcurrentDictionary<string, McpTaskInfo> _tasks = new();
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="InMemoryMcpTaskStore"/> class.
|
||||
/// Gets or sets the default poll interval in milliseconds for new tasks.
|
||||
/// </summary>
|
||||
/// <param name="defaultTtl">
|
||||
/// Default TTL to use when task creation does not specify a TTL. Null means unlimited.
|
||||
/// </param>
|
||||
/// <param name="maxTtl">
|
||||
/// Maximum TTL allowed. If a task requests a longer TTL, it will be capped to this value.
|
||||
/// Null means no maximum limit.
|
||||
/// </param>
|
||||
/// <param name="pollInterval">
|
||||
/// Advertised polling interval for tasks. Default is 1 second.
|
||||
/// This value is used when creating new tasks to indicate how frequently clients should poll for updates.
|
||||
/// </param>
|
||||
/// <param name="cleanupInterval">
|
||||
/// Interval for running background cleanup of expired tasks. Default is 1 minute.
|
||||
/// Pass <see cref="Timeout.InfiniteTimeSpan"/> to disable automatic cleanup.
|
||||
/// </param>
|
||||
/// <param name="pageSize">
|
||||
/// Maximum number of tasks to return per page in <see cref="ListTasksAsync"/>. Default is 100.
|
||||
/// </param>
|
||||
/// <param name="maxTasks">
|
||||
/// Maximum number of tasks allowed in the store globally. Null means unlimited.
|
||||
/// When the limit is reached, <see cref="CreateTaskAsync"/> will throw <see cref="InvalidOperationException"/>.
|
||||
/// </param>
|
||||
/// <param name="maxTasksPerSession">
|
||||
/// Maximum number of tasks allowed per session. Null means unlimited.
|
||||
/// When the limit is reached for a session, <see cref="CreateTaskAsync"/> will throw <see cref="InvalidOperationException"/>.
|
||||
/// </param>
|
||||
public InMemoryMcpTaskStore(
|
||||
TimeSpan? defaultTtl = null,
|
||||
TimeSpan? maxTtl = null,
|
||||
TimeSpan? pollInterval = null,
|
||||
TimeSpan? cleanupInterval = null,
|
||||
int pageSize = 100,
|
||||
int? maxTasks = null,
|
||||
int? maxTasksPerSession = null)
|
||||
{
|
||||
if (defaultTtl.HasValue && maxTtl.HasValue && defaultTtl.Value > maxTtl.Value)
|
||||
{
|
||||
throw new ArgumentException(
|
||||
$"Default TTL ({defaultTtl.Value}) cannot exceed maximum TTL ({maxTtl.Value}).",
|
||||
nameof(defaultTtl));
|
||||
}
|
||||
/// <value>The default is 1000 milliseconds.</value>
|
||||
public long DefaultPollIntervalMs { get; set; } = 1000;
|
||||
|
||||
pollInterval ??= TimeSpan.FromSeconds(1);
|
||||
if (pollInterval <= TimeSpan.Zero)
|
||||
{
|
||||
throw new ArgumentOutOfRangeException(
|
||||
nameof(pollInterval),
|
||||
pollInterval,
|
||||
"Poll interval must be positive.");
|
||||
}
|
||||
|
||||
if (pageSize <= 0)
|
||||
{
|
||||
throw new ArgumentOutOfRangeException(
|
||||
nameof(pageSize),
|
||||
pageSize,
|
||||
"Page size must be positive.");
|
||||
}
|
||||
|
||||
if (maxTasks is <= 0)
|
||||
{
|
||||
throw new ArgumentOutOfRangeException(
|
||||
nameof(maxTasks),
|
||||
maxTasks,
|
||||
"Max tasks must be positive.");
|
||||
}
|
||||
|
||||
if (maxTasksPerSession is <= 0)
|
||||
{
|
||||
throw new ArgumentOutOfRangeException(
|
||||
nameof(maxTasksPerSession),
|
||||
maxTasksPerSession,
|
||||
"Max tasks per session must be positive.");
|
||||
}
|
||||
|
||||
_defaultTtl = defaultTtl;
|
||||
_maxTtl = maxTtl;
|
||||
_pollInterval = pollInterval.Value;
|
||||
_pageSize = pageSize;
|
||||
_maxTasks = maxTasks;
|
||||
_maxTasksPerSession = maxTasksPerSession;
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
_timeProvider = TimeProvider.System;
|
||||
#endif
|
||||
|
||||
cleanupInterval ??= TimeSpan.FromMinutes(1);
|
||||
if (cleanupInterval.Value != Timeout.InfiniteTimeSpan)
|
||||
{
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
_cleanupTimer = _timeProvider.CreateTimer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
|
||||
#else
|
||||
_cleanupTimer = new Timer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="InMemoryMcpTaskStore"/> class with a custom time provider.
|
||||
/// This constructor is only available for testing purposes.
|
||||
/// Gets or sets the default time-to-live in milliseconds for new tasks, or <see langword="null"/> for unlimited.
|
||||
/// </summary>
|
||||
internal InMemoryMcpTaskStore(
|
||||
TimeSpan? defaultTtl,
|
||||
TimeSpan? maxTtl,
|
||||
TimeSpan? pollInterval,
|
||||
TimeSpan? cleanupInterval,
|
||||
int pageSize,
|
||||
int? maxTasks,
|
||||
int? maxTasksPerSession,
|
||||
TimeProvider timeProvider)
|
||||
: this(defaultTtl, maxTtl, pollInterval, cleanupInterval, pageSize, maxTasks, maxTasksPerSession)
|
||||
{
|
||||
_timeProvider = timeProvider ?? TimeProvider.System;
|
||||
}
|
||||
#endif
|
||||
public long? DefaultTtlMs { get; set; }
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<McpTask> CreateTaskAsync(
|
||||
McpTaskMetadata taskParams,
|
||||
RequestId requestId,
|
||||
JsonRpcRequest request,
|
||||
string? sessionId = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
public Task<McpTaskInfo> CreateTaskAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Check global task limit
|
||||
if (_maxTasks is { } maxTasks && _tasks.Count >= maxTasks)
|
||||
{
|
||||
throw new InvalidOperationException(
|
||||
$"Maximum number of tasks ({maxTasks}) has been reached. Cannot create new task.");
|
||||
}
|
||||
var taskId = Guid.NewGuid().ToString("N");
|
||||
var now = DateTimeOffset.UtcNow;
|
||||
|
||||
// Check per-session task limit
|
||||
if (_maxTasksPerSession is { } maxPerSession && sessionId is not null)
|
||||
{
|
||||
var sessionTaskCount = _tasks.Values.Count(e => e.SessionId == sessionId && !IsExpired(e));
|
||||
if (sessionTaskCount >= maxPerSession)
|
||||
{
|
||||
throw new InvalidOperationException(
|
||||
$"Maximum number of tasks per session ({maxPerSession}) has been reached for session '{sessionId}'. Cannot create new task.");
|
||||
}
|
||||
}
|
||||
var info = new McpTaskInfo(taskId, McpTaskStatus.Working, now, now, DefaultTtlMs, DefaultPollIntervalMs);
|
||||
_tasks[taskId] = info;
|
||||
|
||||
var taskId = GenerateTaskId();
|
||||
var now = GetUtcNow();
|
||||
|
||||
// Determine TTL: use requested, fall back to default, respect max limit
|
||||
var ttl = taskParams.TimeToLive ?? _defaultTtl;
|
||||
if (ttl is { } ttlValue && _maxTtl is { } maxTtlValue && ttlValue > maxTtlValue)
|
||||
{
|
||||
ttl = maxTtlValue;
|
||||
}
|
||||
|
||||
TaskEntry entry = new()
|
||||
{
|
||||
TaskId = taskId,
|
||||
Status = McpTaskStatus.Working,
|
||||
CreatedAt = now,
|
||||
LastUpdatedAt = now,
|
||||
TimeToLive = ttl,
|
||||
PollInterval = _pollInterval,
|
||||
RequestId = requestId,
|
||||
Request = request,
|
||||
SessionId = sessionId
|
||||
};
|
||||
|
||||
if (!_tasks.TryAdd(taskId, entry))
|
||||
{
|
||||
// This should be extremely rare with GUID-based IDs
|
||||
throw new InvalidOperationException($"Task ID collision: {taskId}");
|
||||
}
|
||||
|
||||
return Task.FromResult(entry.ToMcpTask());
|
||||
return Task.FromResult(info);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<McpTask?> GetTaskAsync(string taskId, string? sessionId = null, CancellationToken cancellationToken = default)
|
||||
public Task<McpTaskInfo?> GetTaskAsync(string taskId, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (!_tasks.TryGetValue(taskId, out var entry))
|
||||
{
|
||||
return Task.FromResult<McpTask?>(null);
|
||||
}
|
||||
|
||||
// Enforce session isolation if sessionId is provided
|
||||
if (sessionId != null && entry.SessionId != sessionId)
|
||||
{
|
||||
return Task.FromResult<McpTask?>(null);
|
||||
}
|
||||
|
||||
return Task.FromResult<McpTask?>(entry.ToMcpTask());
|
||||
_tasks.TryGetValue(taskId, out var info);
|
||||
return Task.FromResult<McpTaskInfo?>(info);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<McpTask> StoreTaskResultAsync(
|
||||
string taskId,
|
||||
McpTaskStatus status,
|
||||
JsonElement result,
|
||||
string? sessionId = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
public Task SetCompletedAsync(string taskId, JsonElement result, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (status is not (McpTaskStatus.Completed or McpTaskStatus.Failed))
|
||||
{
|
||||
throw new ArgumentException(
|
||||
$"Status must be {nameof(McpTaskStatus.Completed)} or {nameof(McpTaskStatus.Failed)}.",
|
||||
nameof(status));
|
||||
}
|
||||
|
||||
// Retry loop for optimistic concurrency
|
||||
while (true)
|
||||
{
|
||||
if (!_tasks.TryGetValue(taskId, out var entry))
|
||||
Update(taskId, entry => IsTerminal(entry.Status)
|
||||
? entry
|
||||
: entry with
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
}
|
||||
Status = McpTaskStatus.Completed,
|
||||
Result = result,
|
||||
LastUpdatedAt = DateTimeOffset.UtcNow,
|
||||
});
|
||||
|
||||
// Enforce session isolation
|
||||
if (sessionId != null && entry.SessionId != sessionId)
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
}
|
||||
|
||||
// Prevent overwriting terminal state
|
||||
if (IsTerminalStatus(entry.Status))
|
||||
{
|
||||
throw new InvalidOperationException(
|
||||
$"Cannot store result for task in terminal state: {entry.Status}");
|
||||
}
|
||||
|
||||
var updatedEntry = new TaskEntry(entry)
|
||||
{
|
||||
Status = status,
|
||||
LastUpdatedAt = GetUtcNow(),
|
||||
StoredResult = result
|
||||
};
|
||||
|
||||
if (_tasks.TryUpdate(taskId, updatedEntry, entry))
|
||||
{
|
||||
return Task.FromResult(updatedEntry.ToMcpTask());
|
||||
}
|
||||
|
||||
// Entry was modified by another thread, retry
|
||||
}
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<JsonElement> GetTaskResultAsync(string taskId, string? sessionId = null, CancellationToken cancellationToken = default)
|
||||
public Task SetFailedAsync(string taskId, JsonElement error, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (!_tasks.TryGetValue(taskId, out var entry))
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
}
|
||||
Update(taskId, entry => IsTerminal(entry.Status)
|
||||
? entry
|
||||
: entry with
|
||||
{
|
||||
Status = McpTaskStatus.Failed,
|
||||
Error = error,
|
||||
LastUpdatedAt = DateTimeOffset.UtcNow,
|
||||
});
|
||||
|
||||
// Enforce session isolation
|
||||
if (sessionId != entry.SessionId)
|
||||
{
|
||||
throw new InvalidOperationException($"Invalid sessionId: {sessionId} provided for {taskId}");
|
||||
}
|
||||
|
||||
if (entry.StoredResult is not { } storedResult)
|
||||
{
|
||||
throw new InvalidOperationException($"No result stored for task: {taskId}");
|
||||
}
|
||||
|
||||
return Task.FromResult(storedResult);
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<McpTask> UpdateTaskStatusAsync(
|
||||
string taskId,
|
||||
McpTaskStatus status,
|
||||
string? statusMessage,
|
||||
string? sessionId = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
public Task<bool> SetCancelledAsync(string taskId, CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Retry loop for optimistic concurrency
|
||||
while (true)
|
||||
if (!_tasks.TryGetValue(taskId, out var entry) || IsTerminal(entry.Status))
|
||||
{
|
||||
if (!_tasks.TryGetValue(taskId, out var entry))
|
||||
return Task.FromResult(false);
|
||||
}
|
||||
|
||||
bool transitioned = false;
|
||||
Update(taskId, e =>
|
||||
{
|
||||
if (IsTerminal(e.Status))
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
return e;
|
||||
}
|
||||
|
||||
// Enforce session isolation
|
||||
if (sessionId != null && entry.SessionId != sessionId)
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
}
|
||||
|
||||
var updatedEntry = new TaskEntry(entry)
|
||||
{
|
||||
Status = status,
|
||||
StatusMessage = statusMessage,
|
||||
LastUpdatedAt = GetUtcNow(),
|
||||
};
|
||||
|
||||
if (_tasks.TryUpdate(taskId, updatedEntry, entry))
|
||||
{
|
||||
return Task.FromResult(updatedEntry.ToMcpTask());
|
||||
}
|
||||
|
||||
// Entry was modified by another thread, retry
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<ListTasksResult> ListTasksAsync(
|
||||
string? cursor = null,
|
||||
string? sessionId = null,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Stream enumeration - filter by session, exclude expired, apply keyset pagination
|
||||
var query = _tasks.Values
|
||||
.Where(e => sessionId == null || e.SessionId == sessionId)
|
||||
.Where(e => !IsExpired(e));
|
||||
|
||||
// Apply keyset filter if cursor provided: TaskId > cursor
|
||||
// UUID v7 task IDs are monotonically increasing and inherently time-ordered
|
||||
if (cursor != null)
|
||||
{
|
||||
query = query.Where(e => string.CompareOrdinal(e.TaskId, cursor) > 0);
|
||||
}
|
||||
|
||||
// Order by TaskId for stable, deterministic pagination
|
||||
// UUID v7 task IDs sort chronologically due to embedded timestamp
|
||||
var page = query
|
||||
.OrderBy(e => e.TaskId, StringComparer.Ordinal)
|
||||
.Take(_pageSize + 1) // Take one extra to check if there's a next page
|
||||
.Select(e => e.ToMcpTask())
|
||||
.ToList();
|
||||
|
||||
// Set nextCursor if we have more results
|
||||
string? nextCursor;
|
||||
if (page.Count > _pageSize)
|
||||
{
|
||||
var lastItemInPage = page[_pageSize - 1]; // Last item we'll actually return
|
||||
nextCursor = lastItemInPage.TaskId;
|
||||
page.RemoveAt(_pageSize); // Remove the extra item
|
||||
}
|
||||
else
|
||||
{
|
||||
nextCursor = null;
|
||||
}
|
||||
|
||||
return Task.FromResult(new ListTasksResult
|
||||
{
|
||||
Tasks = page.ToArray(),
|
||||
NextCursor = nextCursor
|
||||
transitioned = true;
|
||||
return e with { Status = McpTaskStatus.Cancelled, LastUpdatedAt = DateTimeOffset.UtcNow };
|
||||
});
|
||||
|
||||
return Task.FromResult(transitioned);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task<McpTask> CancelTaskAsync(string taskId, string? sessionId = null, CancellationToken cancellationToken = default)
|
||||
public event Action<InputResponseReceivedEventArgs>? InputResponseReceived;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task ResolveInputRequestsAsync(
|
||||
string taskId,
|
||||
IDictionary<string, InputResponse> inputResponses,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Retry loop for optimistic concurrency
|
||||
while (true)
|
||||
bool wasTerminal = false;
|
||||
Update(taskId, entry =>
|
||||
{
|
||||
if (!_tasks.TryGetValue(taskId, out var entry))
|
||||
if (IsTerminal(entry.Status))
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
wasTerminal = true;
|
||||
return entry;
|
||||
}
|
||||
|
||||
// Enforce session isolation
|
||||
if (sessionId != null && entry.SessionId != sessionId)
|
||||
var requests = entry.InputRequests as ImmutableDictionary<string, InputRequest>
|
||||
?? entry.InputRequests?.ToImmutableDictionary()
|
||||
?? ImmutableDictionary<string, InputRequest>.Empty;
|
||||
|
||||
foreach (var key in inputResponses.Keys)
|
||||
{
|
||||
throw new InvalidOperationException($"Task not found: {taskId}");
|
||||
requests = requests.Remove(key);
|
||||
}
|
||||
|
||||
// If already in terminal state, return unchanged
|
||||
if (IsTerminalStatus(entry.Status))
|
||||
{
|
||||
return Task.FromResult(entry.ToMcpTask());
|
||||
}
|
||||
var status = requests.IsEmpty ? McpTaskStatus.Working : entry.Status;
|
||||
|
||||
var updatedEntry = new TaskEntry(entry)
|
||||
return entry with
|
||||
{
|
||||
Status = McpTaskStatus.Cancelled,
|
||||
LastUpdatedAt = GetUtcNow(),
|
||||
InputRequests = requests,
|
||||
Status = status,
|
||||
LastUpdatedAt = DateTimeOffset.UtcNow,
|
||||
};
|
||||
});
|
||||
|
||||
if (_tasks.TryUpdate(taskId, updatedEntry, entry))
|
||||
if (wasTerminal)
|
||||
{
|
||||
// Drop responses targeting a terminal task — there are no listeners that can act on them.
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
foreach (var kvp in inputResponses)
|
||||
{
|
||||
InputResponseReceived?.Invoke(new InputResponseReceivedEventArgs
|
||||
{
|
||||
return Task.FromResult(updatedEntry.ToMcpTask());
|
||||
TaskId = taskId,
|
||||
RequestId = kvp.Key,
|
||||
Response = kvp.Value,
|
||||
});
|
||||
}
|
||||
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task SetInputRequestsAsync(
|
||||
string taskId,
|
||||
IDictionary<string, InputRequest> inputRequests,
|
||||
CancellationToken cancellationToken = default)
|
||||
{
|
||||
Update(taskId, entry =>
|
||||
{
|
||||
if (IsTerminal(entry.Status))
|
||||
{
|
||||
return entry;
|
||||
}
|
||||
|
||||
// Entry was modified by another thread, retry
|
||||
}
|
||||
var requests = entry.InputRequests as ImmutableDictionary<string, InputRequest>
|
||||
?? entry.InputRequests?.ToImmutableDictionary()
|
||||
?? ImmutableDictionary<string, InputRequest>.Empty;
|
||||
|
||||
foreach (var kvp in inputRequests)
|
||||
{
|
||||
requests = requests.SetItem(kvp.Key, kvp.Value);
|
||||
}
|
||||
|
||||
return entry with
|
||||
{
|
||||
InputRequests = requests,
|
||||
Status = McpTaskStatus.InputRequired,
|
||||
LastUpdatedAt = DateTimeOffset.UtcNow,
|
||||
};
|
||||
});
|
||||
|
||||
return Task.CompletedTask;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Disposes the task store and stops background cleanup.
|
||||
/// </summary>
|
||||
public void Dispose()
|
||||
{
|
||||
_cleanupTimer?.Dispose();
|
||||
}
|
||||
|
||||
private string GenerateTaskId() =>
|
||||
IdHelpers.CreateMonotonicId(GetUtcNow());
|
||||
|
||||
private static bool IsTerminalStatus(McpTaskStatus status) =>
|
||||
private static bool IsTerminal(McpTaskStatus status) =>
|
||||
status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled;
|
||||
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
private DateTimeOffset GetUtcNow() => _timeProvider.GetUtcNow();
|
||||
#else
|
||||
private static DateTimeOffset GetUtcNow() => DateTimeOffset.UtcNow;
|
||||
#endif
|
||||
|
||||
#if MCP_TEST_TIME_PROVIDER
|
||||
private bool IsExpired(TaskEntry entry)
|
||||
#else
|
||||
private static bool IsExpired(TaskEntry entry)
|
||||
#endif
|
||||
private void Update(string taskId, Func<McpTaskInfo, McpTaskInfo> transform)
|
||||
{
|
||||
if (entry.TimeToLive == null)
|
||||
SpinWait spin = default;
|
||||
while (true)
|
||||
{
|
||||
return false; // Unlimited lifetime
|
||||
if (!_tasks.TryGetValue(taskId, out var current))
|
||||
{
|
||||
throw new InvalidOperationException($"Task '{taskId}' not found.");
|
||||
}
|
||||
|
||||
var updated = transform(current);
|
||||
if (ReferenceEquals(updated, current) || _tasks.TryUpdate(taskId, updated, current))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
spin.SpinOnce();
|
||||
}
|
||||
|
||||
var expirationTime = entry.CreatedAt + entry.TimeToLive.Value;
|
||||
return GetUtcNow() >= expirationTime;
|
||||
}
|
||||
|
||||
private void CleanupExpiredTasks(object? state)
|
||||
{
|
||||
var expiredTaskIds = _tasks
|
||||
.Where(kvp => IsExpired(kvp.Value))
|
||||
.Select(kvp => kvp.Key)
|
||||
.ToList();
|
||||
|
||||
foreach (var taskId in expiredTaskIds)
|
||||
{
|
||||
_tasks.TryRemove(taskId, out _);
|
||||
}
|
||||
}
|
||||
|
||||
private sealed class TaskEntry
|
||||
{
|
||||
// Flattened McpTask properties
|
||||
public required string TaskId { get; init; }
|
||||
public required McpTaskStatus Status { get; init; }
|
||||
public string? StatusMessage { get; init; }
|
||||
public required DateTimeOffset CreatedAt { get; init; }
|
||||
public required DateTimeOffset LastUpdatedAt { get; init; }
|
||||
public TimeSpan? TimeToLive { get; init; }
|
||||
public TimeSpan? PollInterval { get; init; }
|
||||
|
||||
// Request metadata
|
||||
public required RequestId RequestId { get; init; }
|
||||
public required JsonRpcRequest Request { get; init; }
|
||||
public required string? SessionId { get; init; }
|
||||
public JsonElement? StoredResult { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Copy constructor for creating modified copies.
|
||||
/// </summary>
|
||||
[SetsRequiredMembers]
|
||||
public TaskEntry(TaskEntry source)
|
||||
{
|
||||
TaskId = source.TaskId;
|
||||
Status = source.Status;
|
||||
StatusMessage = source.StatusMessage;
|
||||
CreatedAt = source.CreatedAt;
|
||||
LastUpdatedAt = source.LastUpdatedAt;
|
||||
TimeToLive = source.TimeToLive;
|
||||
PollInterval = source.PollInterval;
|
||||
RequestId = source.RequestId;
|
||||
Request = source.Request;
|
||||
SessionId = source.SessionId;
|
||||
StoredResult = source.StoredResult;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Default constructor for initial creation.
|
||||
/// </summary>
|
||||
public TaskEntry() { }
|
||||
|
||||
/// <summary>
|
||||
/// Converts this entry back to an McpTask for external consumption.
|
||||
/// </summary>
|
||||
public McpTask ToMcpTask() => new()
|
||||
{
|
||||
TaskId = TaskId,
|
||||
Status = Status,
|
||||
StatusMessage = StatusMessage,
|
||||
CreatedAt = CreatedAt,
|
||||
LastUpdatedAt = LastUpdatedAt,
|
||||
TimeToLive = TimeToLive,
|
||||
PollInterval = PollInterval
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue