-
Notifications
You must be signed in to change notification settings - Fork 7
Expand file tree
/
Copy pathJsonRpcClient.java.html
More file actions
362 lines (321 loc) · 22.7 KB
/
Copy pathJsonRpcClient.java.html
File metadata and controls
362 lines (321 loc) · 22.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
<?xml version="1.0" encoding="UTF-8"?><!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-strict.dtd"><html xmlns="http://www.w3.org/1999/xhtml" lang=""><head><meta http-equiv="Content-Type" content="text/html;charset=UTF-8"/><link rel="stylesheet" href="../jacoco-resources/report.css" type="text/css"/><link rel="shortcut icon" href="../jacoco-resources/report.gif" type="image/gif"/><title>JsonRpcClient.java</title><link rel="stylesheet" href="../jacoco-resources/prettify.css" type="text/css"/><script type="text/javascript" src="../jacoco-resources/prettify.js"></script></head><body onload="window['PR_TAB_WIDTH']=4;prettyPrint()"><div class="breadcrumb" id="breadcrumb"><span class="info"><a href="../jacoco-sessions.html" class="el_session">Sessions</a></span><a href="../index.html" class="el_report">GitHub Copilot SDK :: Java</a> > <a href="index.source.html" class="el_package">com.github.copilot</a> > <span class="el_source">JsonRpcClient.java</span></div><h1>JsonRpcClient.java</h1><pre class="source lang-java linenums">/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
*--------------------------------------------------------------------------------------------*/
package com.github.copilot;
import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiConsumer;
import java.util.logging.Level;
import java.util.logging.Logger;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.github.copilot.rpc.JsonRpcError;
import com.github.copilot.rpc.JsonRpcRequest;
import com.github.copilot.rpc.JsonRpcResponse;
/**
* JSON-RPC 2.0 client implementation for communicating with the Copilot CLI.
*
* @since 1.0.0
*/
class JsonRpcClient implements AutoCloseable {
<span class="nc" id="L42"> private static final Logger LOG = Logger.getLogger(JsonRpcClient.class.getName());</span>
<span class="nc" id="L43"> private static final ObjectMapper MAPPER = createObjectMapper();</span>
private final InputStream inputStream;
private final OutputStream outputStream;
private final Socket socket;
private final Process process;
<span class="nc" id="L49"> private final AtomicLong requestIdCounter = new AtomicLong(0);</span>
<span class="nc" id="L50"> private final Map<Long, CompletableFuture<JsonNode>> pendingRequests = new ConcurrentHashMap<>();</span>
<span class="nc" id="L51"> private final Map<String, BiConsumer<String, JsonNode>> notificationHandlers = new ConcurrentHashMap<>();</span>
private final ExecutorService readerExecutor;
<span class="nc" id="L53"> private volatile boolean running = true;</span>
<span class="nc" id="L55"> private JsonRpcClient(InputStream inputStream, OutputStream outputStream, Socket socket, Process process) {</span>
<span class="nc" id="L56"> this.inputStream = inputStream;</span>
<span class="nc" id="L57"> this.outputStream = outputStream;</span>
<span class="nc" id="L58"> this.socket = socket;</span>
<span class="nc" id="L59"> this.process = process;</span>
<span class="nc" id="L60"> this.readerExecutor = Executors.newSingleThreadExecutor(r -> {</span>
<span class="nc" id="L61"> Thread t = new Thread(r, "jsonrpc-reader");</span>
<span class="nc" id="L62"> t.setDaemon(true);</span>
<span class="nc" id="L63"> return t;</span>
});
<span class="nc" id="L65"> startReader();</span>
<span class="nc" id="L66"> }</span>
static ObjectMapper createObjectMapper() {
<span class="nc" id="L69"> var mapper = new ObjectMapper();</span>
<span class="nc" id="L70"> mapper.registerModule(new JavaTimeModule());</span>
<span class="nc" id="L71"> mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);</span>
<span class="nc" id="L72"> mapper.configure(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, false);</span>
<span class="nc" id="L73"> mapper.setDefaultPropertyInclusion(JsonInclude.Include.NON_NULL);</span>
<span class="nc" id="L74"> return mapper;</span>
}
public static ObjectMapper getObjectMapper() {
<span class="nc" id="L78"> return MAPPER;</span>
}
/**
* Creates a JSON-RPC client using stdio with a process.
*/
public static JsonRpcClient fromProcess(Process process) {
<span class="nc" id="L85"> return new JsonRpcClient(process.getInputStream(), process.getOutputStream(), null, process);</span>
}
/**
* Creates a JSON-RPC client using TCP socket.
*/
public static JsonRpcClient fromSocket(Socket socket) throws IOException {
<span class="nc" id="L92"> return new JsonRpcClient(socket.getInputStream(), socket.getOutputStream(), socket, null);</span>
}
/**
* Registers a handler for JSON-RPC method calls (requests/notifications from
* server).
*/
public void registerMethodHandler(String method, BiConsumer<String, JsonNode> handler) {
<span class="nc" id="L100"> notificationHandlers.put(method, handler);</span>
<span class="nc" id="L101"> }</span>
/**
* Sends a JSON-RPC request and waits for the response.
*/
public <T> CompletableFuture<T> invoke(String method, Object params, Class<T> responseType) {
<span class="nc" id="L107"> long timingNanos = System.nanoTime();</span>
<span class="nc" id="L108"> long id = requestIdCounter.incrementAndGet();</span>
<span class="nc" id="L109"> var future = new CompletableFuture<JsonNode>();</span>
<span class="nc" id="L110"> pendingRequests.put(id, future);</span>
<span class="nc" id="L112"> var request = new JsonRpcRequest();</span>
<span class="nc" id="L113"> request.setJsonrpc("2.0");</span>
<span class="nc" id="L114"> request.setId(id);</span>
<span class="nc" id="L115"> request.setMethod(method);</span>
<span class="nc" id="L116"> request.setParams(params);</span>
try {
<span class="nc" id="L119"> sendMessage(request);</span>
<span class="nc" id="L120"> } catch (IOException e) {</span>
<span class="nc" id="L121"> pendingRequests.remove(id);</span>
<span class="nc" id="L122"> future.completeExceptionally(e);</span>
<span class="nc" id="L123"> }</span>
<span class="nc" id="L125"> return future.thenApply(result -> {</span>
try {
<span class="nc" id="L127"> T value = null;</span>
<span class="nc bnc" id="L128" title="All 4 branches missed."> if (responseType != Void.class && responseType != void.class) {</span>
<span class="nc" id="L129"> value = MAPPER.treeToValue(result, responseType);</span>
}
<span class="nc" id="L131"> LoggingHelpers.logTiming(LOG, Level.FINE,</span>
"JsonRpc.invoke JSON-RPC request finished. Elapsed={Elapsed}, Method=" + method + ", RequestId="
+ id + ", Status=Succeeded",
timingNanos);
<span class="nc" id="L135"> return value;</span>
<span class="nc" id="L136"> } catch (JsonProcessingException e) {</span>
<span class="nc" id="L137"> throw new CompletionException(e);</span>
}
<span class="nc" id="L139"> }).exceptionally(ex -> {</span>
<span class="nc" id="L140"> LoggingHelpers.logTiming(LOG, Level.WARNING, ex,</span>
"JsonRpc.invoke JSON-RPC request finished. Elapsed={Elapsed}, Method=" + method + ", RequestId="
+ id + ", Status=Failed",
timingNanos);
<span class="nc bnc" id="L144" title="All 2 branches missed."> throw ex instanceof RuntimeException re ? re : new RuntimeException(ex);</span>
});
}
/**
* Sends a JSON-RPC notification (no response expected).
*/
public void notify(String method, Object params) throws IOException {
<span class="nc" id="L152"> var notification = new JsonRpcRequest();</span>
<span class="nc" id="L153"> notification.setJsonrpc("2.0");</span>
<span class="nc" id="L154"> notification.setMethod(method);</span>
<span class="nc" id="L155"> notification.setParams(params);</span>
<span class="nc" id="L156"> sendMessage(notification);</span>
<span class="nc" id="L157"> }</span>
/**
* Sends a JSON-RPC response to a server request.
*/
public void sendResponse(Object id, Object result) throws IOException {
<span class="nc" id="L163"> var response = new JsonRpcResponse();</span>
<span class="nc" id="L164"> response.setJsonrpc("2.0");</span>
<span class="nc" id="L165"> response.setId(id);</span>
<span class="nc" id="L166"> response.setResult(result);</span>
<span class="nc" id="L167"> sendMessage(response);</span>
<span class="nc" id="L168"> }</span>
/**
* Sends a JSON-RPC error response to a server request.
*/
public void sendErrorResponse(Object id, int code, String message) throws IOException {
<span class="nc" id="L174"> var response = new JsonRpcResponse();</span>
<span class="nc" id="L175"> response.setJsonrpc("2.0");</span>
<span class="nc" id="L176"> response.setId(id);</span>
<span class="nc" id="L177"> var error = new JsonRpcError();</span>
<span class="nc" id="L178"> error.setCode(code);</span>
<span class="nc" id="L179"> error.setMessage(message);</span>
<span class="nc" id="L180"> response.setError(error);</span>
<span class="nc" id="L181"> sendMessage(response);</span>
<span class="nc" id="L182"> }</span>
private synchronized void sendMessage(Object message) throws IOException {
<span class="nc" id="L185"> String json = MAPPER.writeValueAsString(message);</span>
<span class="nc" id="L186"> byte[] content = json.getBytes(StandardCharsets.UTF_8);</span>
<span class="nc" id="L187"> String header = "Content-Length: " + content.length + "\r\n\r\n";</span>
<span class="nc" id="L189"> outputStream.write(header.getBytes(StandardCharsets.UTF_8));</span>
<span class="nc" id="L190"> outputStream.write(content);</span>
<span class="nc" id="L191"> outputStream.flush();</span>
<span class="nc" id="L193"> LOG.fine("Sent: " + json);</span>
<span class="nc" id="L194"> }</span>
private void startReader() {
<span class="nc" id="L197"> readerExecutor.submit(() -> {</span>
try {
// We need to read bytes because Content-Length specifies bytes, not characters.
// Using BufferedReader would cause issues with multi-byte UTF-8 characters.
<span class="nc" id="L201"> var bis = new BufferedInputStream(inputStream);</span>
<span class="nc bnc" id="L203" title="All 2 branches missed."> while (running) {</span>
// Read headers line by line
<span class="nc" id="L205"> int contentLength = -1;</span>
<span class="nc" id="L206"> var headerLine = new StringBuilder();</span>
<span class="nc" id="L207"> boolean lastWasCR = false;</span>
<span class="nc" id="L208"> boolean inHeaders = true;</span>
<span class="nc bnc" id="L210" title="All 2 branches missed."> while (inHeaders) {</span>
<span class="nc" id="L211"> int b = bis.read();</span>
<span class="nc bnc" id="L212" title="All 2 branches missed."> if (b == -1) {</span>
<span class="nc" id="L213"> return;</span>
}
<span class="nc bnc" id="L216" title="All 2 branches missed."> if (b == '\r') {</span>
<span class="nc" id="L217"> lastWasCR = true;</span>
<span class="nc bnc" id="L218" title="All 2 branches missed."> } else if (b == '\n') {</span>
<span class="nc" id="L219"> String line = headerLine.toString();</span>
<span class="nc" id="L220"> headerLine.setLength(0);</span>
<span class="nc" id="L221"> lastWasCR = false;</span>
<span class="nc bnc" id="L223" title="All 2 branches missed."> if (line.isEmpty()) {</span>
// End of headers (blank line)
<span class="nc" id="L225"> inHeaders = false;</span>
<span class="nc bnc" id="L226" title="All 2 branches missed."> } else if (line.toLowerCase().startsWith("content-length:")) {</span>
<span class="nc" id="L227"> contentLength = Integer.parseInt(line.substring(15).trim());</span>
}
<span class="nc" id="L229"> } else {</span>
<span class="nc bnc" id="L230" title="All 2 branches missed."> if (lastWasCR) {</span>
<span class="nc" id="L231"> headerLine.append('\r');</span>
<span class="nc" id="L232"> lastWasCR = false;</span>
}
<span class="nc" id="L234"> headerLine.append((char) b);</span>
}
<span class="nc" id="L236"> }</span>
<span class="nc bnc" id="L238" title="All 2 branches missed."> if (contentLength <= 0) {</span>
<span class="nc" id="L239"> continue;</span>
}
// Read content as bytes (Content-Length specifies bytes, not characters)
<span class="nc" id="L243"> byte[] buffer = new byte[contentLength];</span>
<span class="nc" id="L244"> int read = 0;</span>
<span class="nc bnc" id="L245" title="All 2 branches missed."> while (read < contentLength) {</span>
<span class="nc" id="L246"> int result = bis.read(buffer, read, contentLength - read);</span>
<span class="nc bnc" id="L247" title="All 2 branches missed."> if (result == -1) {</span>
<span class="nc" id="L248"> return;</span>
}
<span class="nc" id="L250"> read += result;</span>
<span class="nc" id="L251"> }</span>
<span class="nc" id="L253"> String content = new String(buffer, StandardCharsets.UTF_8);</span>
<span class="nc" id="L254"> LOG.fine("Received: " + content);</span>
<span class="nc" id="L256"> handleMessage(content);</span>
<span class="nc" id="L257"> }</span>
<span class="nc" id="L258"> } catch (Exception e) {</span>
<span class="nc bnc" id="L259" title="All 2 branches missed."> if (running) {</span>
<span class="nc" id="L260"> LOG.log(Level.SEVERE, "Error in JSON-RPC reader", e);</span>
}
<span class="nc" id="L262"> }</span>
<span class="nc" id="L263"> });</span>
<span class="nc" id="L264"> }</span>
private void handleMessage(String content) {
try {
<span class="nc" id="L268"> JsonNode node = MAPPER.readTree(content);</span>
// Check if this is a response to our request
<span class="nc bnc" id="L271" title="All 8 branches missed."> if (node.has("id") && !node.get("id").isNull() && (node.has("result") || node.has("error"))) {</span>
<span class="nc" id="L272"> long id = node.get("id").asLong();</span>
<span class="nc" id="L273"> CompletableFuture<JsonNode> future = pendingRequests.remove(id);</span>
<span class="nc bnc" id="L274" title="All 2 branches missed."> if (future != null) {</span>
<span class="nc bnc" id="L275" title="All 2 branches missed."> if (node.has("error")) {</span>
<span class="nc" id="L276"> JsonNode errorNode = node.get("error");</span>
<span class="nc bnc" id="L277" title="All 2 branches missed."> String errorMessage = errorNode.has("message")</span>
<span class="nc" id="L278"> ? errorNode.get("message").asText()</span>
<span class="nc" id="L279"> : "Unknown error";</span>
<span class="nc bnc" id="L280" title="All 2 branches missed."> int errorCode = errorNode.has("code") ? errorNode.get("code").asInt() : -1;</span>
<span class="nc" id="L281"> future.completeExceptionally(new JsonRpcException(errorCode, errorMessage));</span>
<span class="nc" id="L282"> } else {</span>
<span class="nc" id="L283"> future.complete(node.get("result"));</span>
}
}
<span class="nc" id="L286"> }</span>
// Check if this is a request from server (has method and id)
<span class="nc bnc" id="L288" title="All 2 branches missed."> else if (node.has("method")) {</span>
<span class="nc" id="L289"> String method = node.get("method").asText();</span>
<span class="nc" id="L290"> JsonNode params = node.get("params");</span>
<span class="nc bnc" id="L291" title="All 4 branches missed."> Object id = node.has("id") && !node.get("id").isNull() ? node.get("id") : null;</span>
<span class="nc" id="L293"> LOG.fine("Received method: " + method);</span>
<span class="nc" id="L295"> BiConsumer<String, JsonNode> handler = notificationHandlers.get(method);</span>
<span class="nc bnc" id="L296" title="All 2 branches missed."> if (handler != null) {</span>
try {
// Create a context that includes the request ID for responses
<span class="nc bnc" id="L299" title="All 2 branches missed."> handler.accept(id != null ? id.toString() : null, params);</span>
<span class="nc" id="L300"> } catch (Exception e) {</span>
<span class="nc" id="L301"> LOG.log(Level.SEVERE, "Error handling method " + method, e);</span>
<span class="nc bnc" id="L302" title="All 2 branches missed."> if (id != null) {</span>
try {
<span class="nc" id="L304"> sendErrorResponse(id, -32603, e.getMessage());</span>
<span class="nc" id="L305"> } catch (IOException ioe) {</span>
<span class="nc" id="L306"> LOG.log(Level.SEVERE, "Failed to send error response", ioe);</span>
<span class="nc" id="L307"> }</span>
}
<span class="nc" id="L309"> }</span>
} else {
<span class="nc" id="L311"> LOG.fine("No handler for method: " + method);</span>
<span class="nc bnc" id="L312" title="All 2 branches missed."> if (id != null) {</span>
try {
<span class="nc" id="L314"> sendErrorResponse(id, -32601, "Method not found: " + method);</span>
<span class="nc" id="L315"> } catch (IOException ioe) {</span>
<span class="nc" id="L316"> LOG.log(Level.SEVERE, "Failed to send error response", ioe);</span>
<span class="nc" id="L317"> }</span>
}
}
}
<span class="nc" id="L321"> } catch (JsonProcessingException e) {</span>
<span class="nc" id="L322"> LOG.log(Level.SEVERE, "Error parsing JSON-RPC message", e);</span>
<span class="nc" id="L323"> }</span>
<span class="nc" id="L324"> }</span>
@Override
public void close() {
<span class="nc" id="L328"> running = false;</span>
<span class="nc" id="L329"> readerExecutor.shutdownNow();</span>
// Cancel all pending requests
<span class="nc" id="L332"> pendingRequests.forEach((id, future) -> future.completeExceptionally(new IOException("Client closed")));</span>
<span class="nc" id="L333"> pendingRequests.clear();</span>
try {
<span class="nc bnc" id="L336" title="All 2 branches missed."> if (socket != null) {</span>
<span class="nc" id="L337"> socket.close();</span>
}
<span class="nc" id="L339"> } catch (IOException e) {</span>
<span class="nc" id="L340"> LOG.log(Level.FINE, "Error closing socket", e);</span>
<span class="nc" id="L341"> }</span>
<span class="nc bnc" id="L343" title="All 2 branches missed."> if (process != null) {</span>
<span class="nc" id="L344"> process.destroy();</span>
}
<span class="nc" id="L346"> }</span>
public boolean isConnected() {
<span class="nc bnc" id="L349" title="All 2 branches missed."> if (socket != null) {</span>
<span class="nc bnc" id="L350" title="All 4 branches missed."> return socket.isConnected() && !socket.isClosed();</span>
}
<span class="nc bnc" id="L352" title="All 2 branches missed."> if (process != null) {</span>
<span class="nc" id="L353"> return process.isAlive();</span>
}
<span class="nc" id="L355"> return false;</span>
}
public Process getProcess() {
<span class="nc" id="L359"> return process;</span>
}
}
</pre><div class="footer"><span class="right">Created with <a href="http://www.jacoco.org/jacoco">JaCoCo</a> 0.8.14.202510111229</span></div></body></html>