Forge
search
csharp8f5a77c3
1using System.Collections.Concurrent;
2using System.Text.Json;
3
4namespace AgentClientProtocol;
5
6internal enum JsonRpcErrorCode
7{
8 ParseError = -32700,
9 InvalidRequest = -32600,
10 MethodNotFound = -32601,
11 InvalidParams = -32602,
12 InternalError = -32603
13}
14
15internal sealed class JsonRpcEndpoint(Func<CancellationToken, ValueTask<string?>> readFunc, Func<string, CancellationToken, ValueTask> writeFunc, Func<string, CancellationToken, ValueTask> errorWriteFunc)
16{
17 readonly ConcurrentDictionary<RequestId, TaskCompletionSource<JsonRpcResponse>> pendingRequests = new();
18 readonly ConcurrentDictionary<string, Func<JsonRpcRequest, CancellationToken, ValueTask<JsonRpcResponse>>> requestHandlers = new();
19 readonly ConcurrentDictionary<string, Func<JsonRpcNotification, CancellationToken, ValueTask>> notificationHandlers = new();
20 Func<JsonRpcRequest, CancellationToken, ValueTask<JsonRpcResponse>>? defaultRequestHandler;
21 Func<JsonRpcNotification, CancellationToken, ValueTask>? defaultNotificationHandler;
22 int nextRequestId = 0;
23
24 public void SetRequestHandler(string method, Func<JsonRpcRequest, CancellationToken, ValueTask<JsonRpcResponse>> handler)
25 {
26 requestHandlers.TryAdd(method, handler);
27 }
28
29 public void SetNotificationHandler(string method, Func<JsonRpcNotification, CancellationToken, ValueTask> handler)
30 {
31 notificationHandlers.TryAdd(method, handler);
32 }
33
34 public void SetDefaultRequestHandler(Func<JsonRpcRequest, CancellationToken, ValueTask<JsonRpcResponse>> handler)
35 {
36 defaultRequestHandler = handler;
37 }
38
39 public void SetDefaultNotificationHandler(Func<JsonRpcNotification, CancellationToken, ValueTask> handler)
40 {
41 defaultNotificationHandler = handler;
42 }
43
44 public async Task ReadMessagesAsync(CancellationToken cancellationToken = default)
45 {
46 while (!cancellationToken.IsCancellationRequested)
47 {
48 try
49 {
50 var line = await readFunc(cancellationToken).ConfigureAwait(false);
51 if (string.IsNullOrWhiteSpace(line)) continue;
52
53 var trimmedLineSpan = line.AsSpan().Trim();
54 if (trimmedLineSpan.Length < 2 || trimmedLineSpan[0] != '{' || trimmedLineSpan[^1] != '}') continue; // skip non-json input
55
56 var message = JsonSerializer.Deserialize(line, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()!);
57
58 switch (message)
59 {
60 case JsonRpcRequest request:
61 try
62 {
63 if (requestHandlers.TryGetValue(request.Method, out var requestHandler))
64 {
65 var response = await requestHandler(request, cancellationToken);
66 await writeFunc(JsonSerializer.Serialize(response, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()), cancellationToken);
67 }
68 else if (defaultRequestHandler != null)
69 {
70 var response = await defaultRequestHandler(request, cancellationToken);
71 await writeFunc(JsonSerializer.Serialize(response, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()), cancellationToken);
72 }
73 else
74 {
75 await writeFunc(JsonSerializer.Serialize(new JsonRpcResponse
76 {
77 Id = request.Id,
78 Error = new()
79 {
80 Code = (int)JsonRpcErrorCode.MethodNotFound,
81 Message = $"Method '{request.Method}' is not available",
82 }
83 }, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()), cancellationToken);
84 }
85 }
86 catch (NotImplementedException)
87 {
88 await writeFunc(JsonSerializer.Serialize(new JsonRpcResponse
89 {
90 Id = request.Id,
91 Error = new()
92 {
93 Code = (int)JsonRpcErrorCode.MethodNotFound,
94 Message = $"Method '{request.Method}' is not available",
95 }
96 }, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()), cancellationToken);
97 }
98 catch (AcpException acpException)
99 {
100 await writeFunc(JsonSerializer.Serialize(new JsonRpcResponse
101 {
102 Id = request.Id,
103 Error = new()
104 {
105 Code = acpException.Code,
106 Data = acpException.ErrorData,
107 Message = acpException.Message,
108 }
109 }, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>()), cancellationToken);
110 }
111 break;
112 case JsonRpcResponse response:
113 {
114 if (pendingRequests.TryRemove(response.Id, out var tcs))
115 {
116 tcs.TrySetResult(response);
117 }
118 }
119 break;
120 case JsonRpcNotification notification:
121 if (notificationHandlers.TryGetValue(notification.Method, out var notificationHandler))
122 {
123 await notificationHandler(notification, cancellationToken);
124 }
125 else if (defaultNotificationHandler != null)
126 {
127 await defaultNotificationHandler(notification, cancellationToken);
128 }
129 break;
130 default:
131 throw new AcpException($"Invalid response type: {message?.GetType().Name}", null, (int)JsonRpcErrorCode.InternalError);
132 }
133 }
134 catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
135 {
136 return;
137 }
138 catch (Exception ex)
139 {
140 await errorWriteFunc(ex.ToString(), cancellationToken);
141 }
142 }
143 }
144
145 public async ValueTask SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
146 {
147 if (message is JsonRpcRequest request && !request.Id.IsValid)
148 {
149 request.Id = Interlocked.Increment(ref nextRequestId);
150 }
151
152 var json = JsonSerializer.Serialize(message, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcMessage>());
153 await writeFunc(json, cancellationToken).ConfigureAwait(false);
154 }
155
156 public async ValueTask<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, CancellationToken cancellationToken = default)
157 {
158 if (!request.Id.IsValid) request.Id = Interlocked.Increment(ref nextRequestId);
159
160 var json = JsonSerializer.Serialize(request, AcpJsonSerializerContext.Default.Options.GetTypeInfo<JsonRpcRequest>());
161
162 var tcs = new TaskCompletionSource<JsonRpcResponse>();
163 pendingRequests.TryAdd(request.Id, tcs);
164
165 await writeFunc(json, cancellationToken).ConfigureAwait(false);
166
167 return await tcs.Task;
168 }
169}
View only · write via MCP/CIDE