Skip to content

Commit 224a0a0

Browse files
jackbatznerCopilot
jackbatzner
and
Copilot
committed
fix: address dotnet mcp review feedback
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 636fa65 commit 224a0a0

19 files changed

Lines changed: 336 additions & 51 deletions

‎packages/agent-governance-dotnet/src/AgentGovernance.ModelContextProtocol/McpSdkGovernanceExtensions.cs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,8 +137,7 @@ private static void AddCallToolGovernanceFilter(
137137
ex,
138138
"MCP governance threw during tool interception for {ToolName} ({AgentId}); denying",
139139
toolName, agentId);
140-
throw new McpException(
141-
$"Governance error: tool call denied (fail-closed). {ex.Message}");
140+
throw new McpException("Governance error: tool call denied (fail-closed).");
142141
}
143142

144143
if (!allowed)

‎packages/agent-governance-dotnet/src/AgentGovernance/Extensions/McpGovernanceExtensions.cs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,8 @@ public static class McpGovernanceExtensions
154154
/// Options for MCP-specific governance. When <c>null</c>, uses defaults.
155155
/// </param>
156156
/// <param name="agentId">
157-
/// The DID of the agent that will use the message handler.
157+
/// Optional DID of the agent that will use the message handler.
158+
/// When <c>null</c>, uses <see cref="McpGovernanceOptions.AgentId"/>.
158159
/// </param>
159160
/// <param name="timeProvider">Optional clock used for MCP timestamps and expiry checks.</param>
160161
/// <param name="sessionStore">Optional session store for session authentication state.</param>
@@ -167,7 +168,7 @@ public static class McpGovernanceExtensions
167168
public static McpGovernanceStack AddMcpGovernance(
168169
GovernanceOptions? kernelOptions = null,
169170
McpGovernanceOptions? mcpOptions = null,
170-
string agentId = "did:mesh:default",
171+
string? agentId = null,
171172
TimeProvider? timeProvider = null,
172173
IMcpSessionStore? sessionStore = null,
173174
IMcpNonceStore? nonceStore = null,
@@ -180,6 +181,7 @@ public static McpGovernanceStack AddMcpGovernance(
180181
var resolvedNonceStore = nonceStore ?? new InMemoryMcpNonceStore();
181182
var resolvedRateLimitStore = rateLimitStore ?? new InMemoryMcpRateLimitStore();
182183
var resolvedAuditSink = auditSink ?? new InMemoryMcpAuditSink();
184+
var resolvedAgentId = agentId ?? opts.AgentId;
183185

184186
var kernel = new GovernanceKernel(kernelOptions);
185187

@@ -213,7 +215,7 @@ public static McpGovernanceStack AddMcpGovernance(
213215

214216
var toolMapper = new McpToolMapper(opts.CustomToolMappings);
215217

216-
var handler = new McpMessageHandler(gateway, toolMapper, agentId);
218+
var handler = new McpMessageHandler(gateway, toolMapper, resolvedAgentId);
217219

218220
var responseScanner = opts.EnableResponseScanning ? new McpResponseScanner() : null;
219221

‎packages/agent-governance-dotnet/src/AgentGovernance/Extensions/McpGovernanceMiddleware.cs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,16 @@ public async Task InvokeAsync(HttpContext context, RequestDelegate next)
4646

4747
try
4848
{
49+
context.Request.EnableBuffering();
50+
4951
// Read the JSON-RPC request body
50-
using var reader = new StreamReader(context.Request.Body, encoding: System.Text.Encoding.UTF8);
52+
using var reader = new StreamReader(
53+
context.Request.Body,
54+
encoding: System.Text.Encoding.UTF8,
55+
detectEncodingFromByteOrderMarks: false,
56+
leaveOpen: true);
5157
var body = await reader.ReadToEndAsync();
58+
context.Request.Body.Position = 0;
5259
var message = JsonSerializer.Deserialize<Dictionary<string, object?>>(body,
5360
new JsonSerializerOptions { PropertyNameCaseInsensitive = true, MaxDepth = 32 });
5461

@@ -82,6 +89,7 @@ await context.Response.WriteAsync(
8289
catch (JsonException)
8390
{
8491
// Not valid JSON — pass through to next middleware
92+
context.Request.Body.Position = 0;
8593
await next(context);
8694
}
8795
}

‎packages/agent-governance-dotnet/src/AgentGovernance/Extensions/McpServiceCollectionExtensions.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ public static IServiceCollection AddMcpGovernance(
8080
services.AddSingleton(sp => new McpMessageHandler(
8181
sp.GetRequiredService<McpGateway>(),
8282
sp.GetRequiredService<McpToolMapper>(),
83-
"did:mesh:default"));
83+
options.AgentId));
8484

8585
if (options.EnableResponseScanning)
8686
services.AddSingleton<McpResponseScanner>();

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/CredentialRedactor.cs‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,9 @@ public static class CredentialRedactor
6767

6868
/// <summary>PEM-encoded private keys.</summary>
6969
public static readonly Regex PrivateKeyPattern =
70-
new(@"-----BEGIN\s+(RSA\s+|EC\s+|OPENSSH\s+)?PRIVATE\s+KEY-----", RegexOptions.Compiled, RegexTimeout);
70+
new(@"-----BEGIN(?:\s+[A-Z0-9]+)*\s+PRIVATE\s+KEY-----[\s\S]*?-----END(?:\s+[A-Z0-9]+)*\s+PRIVATE\s+KEY-----",
71+
RegexOptions.Compiled | RegexOptions.Singleline,
72+
RegexTimeout);
7173

7274
/// <summary>Azure/SQL connection strings with password.</summary>
7375
public static readonly Regex ConnectionStringPattern =
@@ -124,10 +126,10 @@ public static string Redact(string? input)
124126
if (!ReferenceEquals(before, result))
125127
count++;
126128
}
127-
catch (RegexMatchTimeoutException)
129+
catch (RegexMatchTimeoutException ex)
128130
{
129-
// If regex times out, redact entire value as precaution
130-
continue;
131+
Logger?.LogWarning(ex, "MCP credential redaction timed out; redacting entire value");
132+
return RedactedPlaceholder;
131133
}
132134
}
133135

@@ -218,9 +220,10 @@ public static bool ContainsCredentials(string? input)
218220
if (pattern.IsMatch(input))
219221
return true;
220222
}
221-
catch (RegexMatchTimeoutException)
223+
catch (RegexMatchTimeoutException ex)
222224
{
223-
continue;
225+
Logger?.LogWarning(ex, "MCP credential detection timed out; treating input as sensitive");
226+
return true;
224227
}
225228
}
226229

@@ -243,9 +246,10 @@ public static IReadOnlyList<string> DetectCredentialTypes(string? input)
243246
if (pattern.IsMatch(input))
244247
detected.Add(name);
245248
}
246-
catch (RegexMatchTimeoutException)
249+
catch (RegexMatchTimeoutException ex)
247250
{
248-
continue;
251+
Logger?.LogWarning(ex, "MCP credential type detection timed out; reporting unknown sensitive content");
252+
return ["Unknown sensitive content"];
249253
}
250254
}
251255

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/McpGateway.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ public McpGateway(
183183
Logger?.LogError(ex, "MCP gateway error for {ToolName} - failing closed", toolName);
184184

185185
// Fail-closed: any exception → deny.
186-
var failReason = $"Gateway error (fail-closed): {ex.Message}";
186+
var failReason = "Gateway error (fail-closed).";
187187

188188
Metrics?.RecordMcpDecision(false, agentId, toolName, sw.Elapsed.TotalMilliseconds, "error");
189189

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/McpMessageHandler.cs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,11 +133,13 @@ public void RegisterResource(string uriPattern, Dictionary<string, object> resou
133133
}
134134
catch (UnauthorizedAccessException ex)
135135
{
136-
return JsonRpcError(id, -32003, ex.Message);
136+
Logger?.LogWarning(ex, "MCP message denied by governance");
137+
return JsonRpcError(id, -32003, "Access denied by governance policy.");
137138
}
138139
catch (Exception ex)
139140
{
140-
return JsonRpcError(id, -32603, $"Internal error: {ex.Message}");
141+
Logger?.LogError(ex, "MCP message handling failed");
142+
return JsonRpcError(id, -32603, "Internal error.");
141143
}
142144
}
143145

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/McpMessageSigner.cs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,14 +69,14 @@ public sealed class McpMessageSigner : IDisposable
6969
/// <summary>
7070
/// Initializes a new message signer with the given shared secret (HMAC-SHA256).
7171
/// </summary>
72-
/// <param name="signingKey">Shared secret key (minimum 16 bytes, 32 recommended).</param>
72+
/// <param name="signingKey">Shared secret key (minimum 32 bytes).</param>
7373
/// <param name="nonceStore">The nonce store used for replay protection.</param>
7474
/// <param name="timeProvider">The clock used for timestamps and replay-window checks.</param>
7575
public McpMessageSigner(byte[] signingKey, IMcpNonceStore? nonceStore = null, TimeProvider? timeProvider = null)
7676
{
7777
ArgumentNullException.ThrowIfNull(signingKey);
78-
if (signingKey.Length < 16)
79-
throw new ArgumentException("Signing key must be at least 16 bytes.", nameof(signingKey));
78+
if (signingKey.Length < 32)
79+
throw new ArgumentException("Signing key must be at least 32 bytes.", nameof(signingKey));
8080
_signingKey = signingKey;
8181
_nonceStore = nonceStore ?? new InMemoryMcpNonceStore();
8282
_timeProvider = timeProvider ?? TimeProvider.System;
@@ -226,8 +226,8 @@ public McpVerificationResult VerifyMessage(McpSignedEnvelope envelope)
226226
}
227227
catch (Exception ex)
228228
{
229-
// Fail-closed
230-
return McpVerificationResult.Failed($"Verification error (fail-closed): {ex.Message}");
229+
Logger?.LogError(ex, "MCP message verification failed closed");
230+
return McpVerificationResult.Failed("Verification error (fail-closed).");
231231
}
232232
}
233233

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/McpSessionAuthenticator.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,8 @@ public McpSessionAuthenticator(IMcpSessionStore sessionStore, TimeProvider? time
9595
UserId = userId,
9696
CreatedAt = now,
9797
ExpiresAt = now.Add(SessionTtl),
98-
// Composite key for rate limiting: userId:agentId or just agentId
99-
RateLimitKey = userId is not null ? $"{userId}:{agentId}" : agentId
98+
// Composite key for rate limiting: userId|agentId or just agentId
99+
RateLimitKey = userId is not null ? $"{userId}|{agentId}" : agentId
100100
};
101101

102102
if (!TrySetSession(session))

‎packages/agent-governance-dotnet/src/AgentGovernance/Mcp/McpSlidingRateLimiter.cs‎

Lines changed: 79 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,10 @@ public sealed class McpSlidingRateLimiter
2121
{
2222
private readonly IMcpRateLimitStore _rateLimitStore;
2323
private readonly ConcurrentDictionary<string, object> _bucketLocks = new(StringComparer.OrdinalIgnoreCase);
24+
private readonly ConcurrentDictionary<string, DateTimeOffset> _lockLastTouched = new(StringComparer.OrdinalIgnoreCase);
2425
private readonly ConcurrentDictionary<string, byte> _trackedAgents = new(StringComparer.OrdinalIgnoreCase);
2526
private readonly TimeProvider _timeProvider;
27+
private DateTimeOffset _lastLockSweep;
2628

2729
/// <summary>
2830
/// Initializes a new limiter with in-memory persistence and the system clock.
@@ -41,6 +43,7 @@ public McpSlidingRateLimiter(IMcpRateLimitStore rateLimitStore, TimeProvider? ti
4143
{
4244
_rateLimitStore = rateLimitStore ?? throw new ArgumentNullException(nameof(rateLimitStore));
4345
_timeProvider = timeProvider ?? TimeProvider.System;
46+
_lastLockSweep = _timeProvider.GetUtcNow();
4447
}
4548

4649
/// <summary>
@@ -54,6 +57,18 @@ public McpSlidingRateLimiter(IMcpRateLimitStore rateLimitStore, TimeProvider? ti
5457
/// </summary>
5558
public TimeSpan WindowSize { get; init; } = TimeSpan.FromMinutes(5);
5659

60+
/// <summary>
61+
/// Maximum idle time before an unused per-agent lock entry is evicted.
62+
/// Defaults to 15 minutes.
63+
/// </summary>
64+
public TimeSpan LockEntryTtl { get; init; } = TimeSpan.FromMinutes(15);
65+
66+
/// <summary>
67+
/// Minimum time between background sweeps that evict stale per-agent lock entries.
68+
/// Defaults to 5 minutes.
69+
/// </summary>
70+
public TimeSpan LockSweepInterval { get; init; } = TimeSpan.FromMinutes(5);
71+
5772
/// <summary>
5873
/// Optional logger for recording rate limit events.
5974
/// When <c>null</c>, no logging occurs — the limiter operates silently.
@@ -72,10 +87,9 @@ public bool TryAcquire(string agentId)
7287
{
7388
ArgumentException.ThrowIfNullOrWhiteSpace(agentId);
7489

75-
var bucketLock = _bucketLocks.GetOrAdd(agentId, _ => new object());
76-
_trackedAgents[agentId] = 0;
77-
7890
var now = _timeProvider.GetUtcNow();
91+
var bucketLock = GetBucketLock(agentId, now);
92+
_trackedAgents[agentId] = 0;
7993
var cutoff = now - WindowSize;
8094

8195
lock (bucketLock)
@@ -91,6 +105,7 @@ public bool TryAcquire(string agentId)
91105

92106
timestamps.Add(now);
93107
SaveBucket(agentId, timestamps);
108+
MaybeSweepInactiveLocks(now);
94109
return true;
95110
}
96111
}
@@ -105,17 +120,22 @@ public int GetRemainingBudget(string agentId)
105120
{
106121
ArgumentException.ThrowIfNullOrWhiteSpace(agentId);
107122

108-
var bucketLock = _bucketLocks.GetOrAdd(agentId, _ => new object());
123+
var now = _timeProvider.GetUtcNow();
124+
var bucketLock = GetBucketLock(agentId, now);
109125
lock (bucketLock)
110126
{
111127
var timestamps = GetBucketTimestamps(agentId);
112128
if (timestamps.Count == 0)
113129
{
130+
EvictLockIfInactive(agentId, timestamps.Count);
131+
MaybeSweepInactiveLocks(now);
114132
return MaxCallsPerWindow;
115133
}
116134

117-
PruneExpired(timestamps, _timeProvider.GetUtcNow() - WindowSize);
135+
PruneExpired(timestamps, now - WindowSize);
118136
SaveBucket(agentId, timestamps);
137+
EvictLockIfInactive(agentId, timestamps.Count);
138+
MaybeSweepInactiveLocks(now);
119139
return Math.Max(0, MaxCallsPerWindow - timestamps.Count);
120140
}
121141
}
@@ -130,17 +150,22 @@ public int GetCallCount(string agentId)
130150
{
131151
ArgumentException.ThrowIfNullOrWhiteSpace(agentId);
132152

133-
var bucketLock = _bucketLocks.GetOrAdd(agentId, _ => new object());
153+
var now = _timeProvider.GetUtcNow();
154+
var bucketLock = GetBucketLock(agentId, now);
134155
lock (bucketLock)
135156
{
136157
var timestamps = GetBucketTimestamps(agentId);
137158
if (timestamps.Count == 0)
138159
{
160+
EvictLockIfInactive(agentId, timestamps.Count);
161+
MaybeSweepInactiveLocks(now);
139162
return 0;
140163
}
141164

142-
PruneExpired(timestamps, _timeProvider.GetUtcNow() - WindowSize);
165+
PruneExpired(timestamps, now - WindowSize);
143166
SaveBucket(agentId, timestamps);
167+
EvictLockIfInactive(agentId, timestamps.Count);
168+
MaybeSweepInactiveLocks(now);
144169
return timestamps.Count;
145170
}
146171
}
@@ -154,11 +179,14 @@ public void Reset(string agentId)
154179
{
155180
ArgumentException.ThrowIfNullOrWhiteSpace(agentId);
156181

157-
var bucketLock = _bucketLocks.GetOrAdd(agentId, _ => new object());
182+
var now = _timeProvider.GetUtcNow();
183+
var bucketLock = GetBucketLock(agentId, now);
158184
lock (bucketLock)
159185
{
160186
SaveBucket(agentId, []);
161187
_trackedAgents.TryRemove(agentId, out _);
188+
EvictLockIfInactive(agentId, 0);
189+
MaybeSweepInactiveLocks(now);
162190
}
163191
}
164192

@@ -181,12 +209,13 @@ public void ResetAll()
181209
/// <returns>The total number of expired entries removed across all agents.</returns>
182210
public int CleanupExpired()
183211
{
184-
var cutoff = _timeProvider.GetUtcNow() - WindowSize;
212+
var now = _timeProvider.GetUtcNow();
213+
var cutoff = now - WindowSize;
185214
int totalRemoved = 0;
186215

187216
foreach (var agentId in _trackedAgents.Keys.ToArray())
188217
{
189-
var bucketLock = _bucketLocks.GetOrAdd(agentId, _ => new object());
218+
var bucketLock = GetBucketLock(agentId, now);
190219
lock (bucketLock)
191220
{
192221
var timestamps = GetBucketTimestamps(agentId);
@@ -198,13 +227,53 @@ public int CleanupExpired()
198227
if (timestamps.Count == 0)
199228
{
200229
_trackedAgents.TryRemove(agentId, out _);
230+
EvictLockIfInactive(agentId, timestamps.Count);
201231
}
202232
}
203233
}
204234

235+
MaybeSweepInactiveLocks(now);
205236
return totalRemoved;
206237
}
207238

239+
private object GetBucketLock(string agentId, DateTimeOffset now)
240+
{
241+
_lockLastTouched[agentId] = now;
242+
return _bucketLocks.GetOrAdd(agentId, _ => new object());
243+
}
244+
245+
private void EvictLockIfInactive(string agentId, int timestampCount)
246+
{
247+
if (timestampCount > 0 || _trackedAgents.ContainsKey(agentId))
248+
{
249+
return;
250+
}
251+
252+
_bucketLocks.TryRemove(agentId, out _);
253+
_lockLastTouched.TryRemove(agentId, out _);
254+
}
255+
256+
private void MaybeSweepInactiveLocks(DateTimeOffset now)
257+
{
258+
if (now - _lastLockSweep < LockSweepInterval)
259+
{
260+
return;
261+
}
262+
263+
_lastLockSweep = now;
264+
var cutoff = now - LockEntryTtl;
265+
foreach (var (agentId, lastTouched) in _lockLastTouched.ToArray())
266+
{
267+
if (lastTouched > cutoff || _trackedAgents.ContainsKey(agentId))
268+
{
269+
continue;
270+
}
271+
272+
_bucketLocks.TryRemove(agentId, out _);
273+
_lockLastTouched.TryRemove(agentId, out _);
274+
}
275+
}
276+
208277
private List<DateTimeOffset> GetBucketTimestamps(string agentId)
209278
{
210279
return _rateLimitStore.GetBucketAsync(agentId).GetAwaiter().GetResult()?.Timestamps.ToList()

0 commit comments

Comments
 (0)