diff --git a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java index 5303c4f500..299ac18a2b 100644 --- a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java +++ b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java @@ -46,6 +46,7 @@ class JsonRpcClient implements AutoCloseable { private final OutputStream outputStream; private final Socket socket; private final Process process; + private final boolean ownsStreams; private final AtomicLong requestIdCounter = new AtomicLong(0); private final Map> pendingRequests = new ConcurrentHashMap<>(); private final Map> notificationHandlers = new ConcurrentHashMap<>(); @@ -53,10 +54,16 @@ class JsonRpcClient implements AutoCloseable { private volatile boolean running = true; private JsonRpcClient(InputStream inputStream, OutputStream outputStream, Socket socket, Process process) { + this(inputStream, outputStream, socket, process, false); + } + + private JsonRpcClient(InputStream inputStream, OutputStream outputStream, Socket socket, Process process, + boolean ownsStreams) { this.inputStream = inputStream; this.outputStream = outputStream; this.socket = socket; this.process = process; + this.ownsStreams = ownsStreams; this.readerExecutor = Executors.newSingleThreadExecutor(r -> { Thread t = new Thread(r, "jsonrpc-reader"); t.setDaemon(true); @@ -93,6 +100,15 @@ public static JsonRpcClient fromSocket(Socket socket) throws IOException { return new JsonRpcClient(socket.getInputStream(), socket.getOutputStream(), socket, null); } + /** + * Creates a JSON-RPC client over arbitrary input/output streams. The client + * takes ownership of the streams and closes them when {@link #close()} is + * called. + */ + public static JsonRpcClient fromStreams(InputStream inputStream, OutputStream outputStream) { + return new JsonRpcClient(inputStream, outputStream, null, null, true); + } + /** * Registers a handler for JSON-RPC method calls (requests/notifications from * server). @@ -344,6 +360,19 @@ public void close() { if (process != null) { process.destroy(); } + + if (ownsStreams) { + try { + inputStream.close(); + } catch (IOException e) { + LOG.log(Level.FINE, "Error closing input stream", e); + } + try { + outputStream.close(); + } catch (IOException e) { + LOG.log(Level.FINE, "Error closing output stream", e); + } + } } public boolean isConnected() { diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/FfiOutputStream.java b/java/sdk/src/main/java/com/github/copilot/ffi/FfiOutputStream.java new file mode 100644 index 0000000000..4198f08f02 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/ffi/FfiOutputStream.java @@ -0,0 +1,63 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +import java.io.IOException; +import java.io.OutputStream; +import java.util.Arrays; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; + +final class FfiOutputStream extends OutputStream { + + private final NativeBinding nativeBinding; + private final AtomicInteger connectionId; + private final AtomicBoolean closing; + private final ReentrantLock operationLock; + + FfiOutputStream(NativeBinding nativeBinding, AtomicInteger connectionId, AtomicBoolean closing, + ReentrantLock operationLock) { + this.nativeBinding = Objects.requireNonNull(nativeBinding, "nativeBinding must not be null"); + this.connectionId = Objects.requireNonNull(connectionId, "connectionId must not be null"); + this.closing = Objects.requireNonNull(closing, "closing must not be null"); + this.operationLock = Objects.requireNonNull(operationLock, "operationLock must not be null"); + } + + @Override + public void write(int b) throws IOException { + write(new byte[]{(byte) b}, 0, 1); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + Objects.requireNonNull(b, "buffer must not be null"); + if (off < 0 || len < 0 || off + len > b.length) { + throw new IndexOutOfBoundsException("Invalid off/len for buffer of length " + b.length); + } + if (len == 0) { + return; + } + + operationLock.lock(); + try { + if (closing.get()) { + throw new IOException("The in-process runtime connection is closed."); + } + int id = connectionId.get(); + if (id == 0) { + throw new IOException("The in-process runtime connection is closed."); + } + + byte[] payload = (off == 0 && len == b.length) ? b : Arrays.copyOfRange(b, off, off + len); + if (!nativeBinding.connectionWrite(id, payload, payload.length)) { + throw new IOException("Failed to write a frame to the in-process runtime connection."); + } + } finally { + operationLock.unlock(); + } + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java b/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java new file mode 100644 index 0000000000..6a71c71431 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/ffi/FfiRuntimeHost.java @@ -0,0 +1,349 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; +import java.util.logging.Level; +import java.util.logging.Logger; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.github.copilot.rpc.CopilotClientMode; +import com.github.copilot.rpc.CopilotClientOptions; +import com.sun.jna.Callback; +import com.sun.jna.Native; +import com.sun.jna.Pointer; + +/** + * Manages the in-process FFI runtime lifecycle. + */ +public final class FfiRuntimeHost implements AutoCloseable { + + private static final Logger LOG = Logger.getLogger(FfiRuntimeHost.class.getName()); + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private final NativeBinding nativeBinding; + private final QueueInputStream receiveStream; + private final AtomicBoolean closing = new AtomicBoolean(false); + private final AtomicBoolean disposed = new AtomicBoolean(false); + private final AtomicInteger serverId = new AtomicInteger(0); + private final AtomicInteger connectionId = new AtomicInteger(0); + private final AtomicInteger activeCallbacks = new AtomicInteger(0); + private final Object callbackDrainMonitor = new Object(); + private final ReentrantLock operationLock = new ReentrantLock(); + private final FfiOutputStream sendStream; + private final String libraryPath; + + private volatile OutboundCallback callbackRef; + + /** + * Creates an FFI runtime host using the resolved bundled native library. + * + * @throws IOException + * if the runtime library cannot be resolved + */ + public FfiRuntimeHost() throws IOException { + this(resolveLibraryPath(), null, new QueueInputStream()); + } + + FfiRuntimeHost(NativeBinding nativeBinding, String libraryPath) { + this(nativeBinding, libraryPath, new QueueInputStream()); + } + + FfiRuntimeHost(NativeBinding nativeBinding, String libraryPath, QueueInputStream receiveStream) { + this.nativeBinding = Objects.requireNonNull(nativeBinding, "nativeBinding must not be null"); + this.receiveStream = Objects.requireNonNull(receiveStream, "receiveStream must not be null"); + this.sendStream = new FfiOutputStream(this.nativeBinding, this.connectionId, this.closing, this.operationLock); + this.libraryPath = libraryPath; + Native.setCallbackExceptionHandler((Callback callback, Throwable throwable) -> LOG.log(Level.WARNING, + "Unhandled exception in FFI callback", throwable)); + } + + private FfiRuntimeHost(Path libraryPath, NativeBinding nativeBinding, QueueInputStream receiveStream) { + this(nativeBinding == null ? new JnaNativeBinding(libraryPath) : nativeBinding, libraryPath.toString(), + receiveStream); + } + + private static Path resolveLibraryPath() throws IOException { + return NativeRuntimeLoader.resolve(); + } + + /** + * Starts the in-process runtime and opens a connection. + * + * @param entrypointPath + * runtime entrypoint path passed in {@code argv_json} + * @param options + * client options used to construct {@code argv_json} and + * {@code env_json} + */ + public void start(String entrypointPath, CopilotClientOptions options) { + Objects.requireNonNull(entrypointPath, "entrypointPath must not be null"); + Objects.requireNonNull(options, "options must not be null"); + if (disposed.get()) { + throw new IllegalStateException("FfiRuntimeHost is already closed."); + } + if (serverId.get() != 0 || connectionId.get() != 0) { + throw new IllegalStateException("FfiRuntimeHost has already been started."); + } + + byte[] argvJson = buildArgvJson(entrypointPath, options); + byte[] envJson = buildEnvJson(options); + int hostHandle = runHostStartOnBlockingThread(argvJson, envJson); + if (hostHandle == 0) { + String lib = libraryPath != null ? libraryPath : ""; + throw new IllegalStateException( + "copilot_runtime_host_start failed (library '" + lib + "', entrypoint '" + entrypointPath + "')."); + } + + // Hold operationLock while publishing handles to serialize with close(). + // Recheck disposed in case close() ran while hostStart was blocking. + operationLock.lock(); + try { + if (disposed.get()) { + try { + nativeBinding.hostShutdown(hostHandle); + } catch (Throwable ignored) { + // Best effort + } + throw new IllegalStateException("FfiRuntimeHost was closed during startup."); + } + serverId.set(hostHandle); + + OutboundCallback callback = createOutboundCallback(); + callbackRef = callback; + int connHandle = nativeBinding.connectionOpen(hostHandle, callback, Pointer.NULL, null, 0, null, 0, null, + 0); + if (connHandle == 0) { + try { + nativeBinding.hostShutdown(hostHandle); + } catch (Throwable ignored) { + // Best effort + } + serverId.set(0); + callbackRef = null; + throw new IllegalStateException("copilot_runtime_connection_open failed."); + } + connectionId.set(connHandle); + LOG.fine(() -> "Started FFI runtime host. Library=" + libraryPath + ", serverId=" + hostHandle + + ", connectionId=" + connHandle); + } finally { + operationLock.unlock(); + } + } + + public InputStream getReceiveStream() { + return receiveStream; + } + + public OutputStream getSendStream() { + return sendStream; + } + + @Override + public void close() { + if (!disposed.compareAndSet(false, true)) { + return; + } + + closing.set(true); + + operationLock.lock(); + try { + int connHandle = connectionId.getAndSet(0); + if (connHandle != 0) { + try { + nativeBinding.connectionClose(connHandle); + } catch (Throwable t) { + LOG.log(Level.FINE, "Failed to close FFI connection", t); + } + } + } finally { + operationLock.unlock(); + } + + drainActiveCallbacks(); + + int hostHandle = serverId.getAndSet(0); + if (hostHandle != 0) { + try { + nativeBinding.hostShutdown(hostHandle); + } catch (Throwable t) { + LOG.log(Level.FINE, "Failed to shut down FFI host", t); + } + } + + try { + receiveStream.close(); + } catch (Throwable ignored) { + // never throw from close + } + + callbackRef = null; + } + + private void drainActiveCallbacks() { + while (activeCallbacks.get() > 0) { + synchronized (callbackDrainMonitor) { + if (activeCallbacks.get() == 0) { + return; + } + try { + callbackDrainMonitor.wait(10L); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + } + } + + private OutboundCallback createOutboundCallback() { + return (userData, data, len) -> { + if (closing.get()) { + return; + } + activeCallbacks.incrementAndGet(); + try { + if (closing.get() || data == null || len <= 0) { + return; + } + byte[] bytes = data.getByteArray(0, len); + if (!closing.get()) { + receiveStream.enqueue(bytes); + } + } catch (Throwable t) { + LOG.log(Level.WARNING, "Exception in FFI outbound callback", t); + } finally { + if (activeCallbacks.decrementAndGet() == 0) { + synchronized (callbackDrainMonitor) { + callbackDrainMonitor.notifyAll(); + } + } + } + }; + } + + private int runHostStartOnBlockingThread(byte[] argvJson, byte[] envJson) { + ReaderThreadFactory readerThreadFactory = new ReaderThreadFactory(); + ExecutorService executor = Executors + .newSingleThreadExecutor(runnable -> readerThreadFactory.create(runnable, "copilot-ffi-host-start")); + try { + Future future = executor.submit(() -> nativeBinding.hostStart(argvJson, argvJson.length, envJson, + envJson == null ? 0 : envJson.length)); + return future.get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted while starting in-process runtime host.", e); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException runtimeException) { + throw runtimeException; + } + throw new IllegalStateException("Failed to start in-process runtime host.", cause); + } finally { + executor.shutdownNow(); + try { + executor.awaitTermination(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + + private static byte[] buildArgvJson(String entrypointPath, CopilotClientOptions options) { + List argv = new ArrayList<>(); + if (entrypointPath.toLowerCase().endsWith(".js")) { + argv.add("node"); + } + argv.add(entrypointPath); + argv.add("--embedded-host"); + argv.add("--no-auto-update"); + + String logLevel = options.getLogLevel(); + if (logLevel != null && !logLevel.isBlank()) { + argv.add("--log-level"); + argv.add(logLevel); + } + + String gitHubToken = options.getGitHubToken(); + if (gitHubToken != null && !gitHubToken.isEmpty()) { + argv.add("--auth-token-env"); + argv.add("COPILOT_SDK_AUTH_TOKEN"); + } + + boolean useLoggedInUser = options.getUseLoggedInUser().orElse(gitHubToken == null || gitHubToken.isEmpty()); + if (!useLoggedInUser) { + argv.add("--no-auto-login"); + } + + if (options.getSessionIdleTimeoutSeconds().isPresent() + && options.getSessionIdleTimeoutSeconds().getAsInt() > 0) { + argv.add("--session-idle-timeout"); + argv.add(String.valueOf(options.getSessionIdleTimeoutSeconds().getAsInt())); + } + + if (options.isRemote()) { + argv.add("--remote"); + } + + String[] cliArgs = options.getCliArgs(); + if (cliArgs != null && cliArgs.length > 0) { + for (String arg : cliArgs) { + if (arg != null && !arg.isBlank()) { + argv.add(arg); + } + } + } + + return jsonBytes(argv); + } + + private static byte[] buildEnvJson(CopilotClientOptions options) { + Map env = new LinkedHashMap<>(); + + String token = options.getGitHubToken(); + if (token != null && !token.isEmpty()) { + env.put("COPILOT_SDK_AUTH_TOKEN", token); + } + String copilotHome = options.getCopilotHome(); + if (copilotHome != null && !copilotHome.isEmpty()) { + env.put("COPILOT_HOME", copilotHome); + } + if (options.getMode() == CopilotClientMode.EMPTY) { + env.put("COPILOT_DISABLE_KEYTAR", "1"); + } + + if (env.isEmpty()) { + return null; + } + return jsonBytes(env); + } + + private static byte[] jsonBytes(Object value) { + try { + return MAPPER.writeValueAsString(value).getBytes(StandardCharsets.UTF_8); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to serialize FFI JSON parameter.", e); + } + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/QueueInputStream.java b/java/sdk/src/main/java/com/github/copilot/ffi/QueueInputStream.java new file mode 100644 index 0000000000..977182d5f4 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/ffi/QueueInputStream.java @@ -0,0 +1,119 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +import java.io.IOException; +import java.io.InputStream; +import java.util.Objects; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * {@link InputStream} backed by a {@link BlockingQueue} of byte-array chunks. + * + *

+ * Used by the in-process FFI transport to bridge native callback frames into + * the JSON-RPC reader. + */ +public class QueueInputStream extends InputStream { + + private static final byte[] EOF_SENTINEL = new byte[0]; + + private final BlockingQueue queue; + private final AtomicBoolean closed = new AtomicBoolean(false); + + private byte[] currentChunk; + private int currentOffset; + private boolean eof; + + /** + * Creates a queue-backed input stream with an unbounded queue. + */ + public QueueInputStream() { + this(new LinkedBlockingQueue<>()); + } + + /** + * Testing constructor that injects a queue implementation. + * + * @param queue + * backing queue + */ + QueueInputStream(BlockingQueue queue) { + this.queue = Objects.requireNonNull(queue, "queue must not be null"); + } + + void enqueue(byte[] bytes) { + if (bytes == null || bytes.length == 0 || closed.get()) { + return; + } + queue.offer(bytes); + } + + @Override + public int read() throws IOException { + byte[] one = new byte[1]; + int read = read(one, 0, 1); + if (read == -1) { + return -1; + } + return one[0] & 0xFF; + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + Objects.requireNonNull(b, "buffer must not be null"); + if (off < 0 || len < 0 || off + len > b.length) { + throw new IndexOutOfBoundsException("Invalid off/len for buffer of length " + b.length); + } + if (len == 0) { + return 0; + } + if (eof) { + return -1; + } + + while (currentChunk == null || currentOffset >= currentChunk.length) { + byte[] next; + try { + next = queue.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while waiting for callback data", e); + } + if (next == EOF_SENTINEL) { + eof = true; + return -1; + } + if (next.length == 0) { + continue; + } + currentChunk = next; + currentOffset = 0; + } + + int available = currentChunk.length - currentOffset; + int toCopy = Math.min(available, len); + System.arraycopy(currentChunk, currentOffset, b, off, toCopy); + currentOffset += toCopy; + return toCopy; + } + + @Override + public int available() { + if (currentChunk == null || currentOffset >= currentChunk.length) { + return 0; + } + return currentChunk.length - currentOffset; + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + queue.offer(EOF_SENTINEL); + } + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/ReaderThreadFactory.java b/java/sdk/src/main/java/com/github/copilot/ffi/ReaderThreadFactory.java new file mode 100644 index 0000000000..b0824fa9a7 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/ffi/ReaderThreadFactory.java @@ -0,0 +1,22 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +/** + * Creates reader threads for FFI queue consumption. + * + *

+ * Baseline (JDK 17) implementation creates a daemon platform thread. The JDK 25 + * multi-release overlay switches this to a virtual thread with the same + * package-private API. + */ +final class ReaderThreadFactory { + + Thread create(Runnable task, String name) { + Thread thread = new Thread(task, name); + thread.setDaemon(true); + return thread; + } +} diff --git a/java/sdk/src/main/java25/com/github/copilot/ffi/ReaderThreadFactory.java b/java/sdk/src/main/java25/com/github/copilot/ffi/ReaderThreadFactory.java new file mode 100644 index 0000000000..a67346b889 --- /dev/null +++ b/java/sdk/src/main/java25/com/github/copilot/ffi/ReaderThreadFactory.java @@ -0,0 +1,15 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +/** + * JDK 25 multi-release variant of {@link ReaderThreadFactory}. + */ +final class ReaderThreadFactory { + + Thread create(Runnable task, String name) { + return Thread.ofVirtual().name(name).unstarted(task); + } +} diff --git a/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java b/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java new file mode 100644 index 0000000000..df57fed42a --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/ffi/FfiRuntimeHostTest.java @@ -0,0 +1,407 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.github.copilot.rpc.CopilotClientMode; +import com.github.copilot.rpc.CopilotClientOptions; +import com.sun.jna.Library; +import com.sun.jna.Memory; +import com.sun.jna.Native; +import com.sun.jna.Pointer; + +class FfiRuntimeHostTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final String TEST_LIB_PATH_PROP = "copilot.test.nativelib.path"; + private static final String SPIKE_LIB_PATH = System.getProperty(TEST_LIB_PATH_PROP, + "../../1917-java-embed-rust-cli-runtime-remove-before-merge" + "/spike-3-4-jna-callback-and-threading" + + "/rust-dll/target/release/libcallback_test.so"); + + interface CallbackTestLib extends Library { + int host_start(); + + byte host_shutdown(int serverHandle); + + int connection_open(int serverHandle, OutboundCallback callback, Pointer userData, int burstCount); + + byte connection_write(int connectionHandle, byte[] data, int len); + + byte connection_close(int connectionHandle); + } + + private static boolean testLibExists() { + return Files.isRegularFile(Path.of(SPIKE_LIB_PATH).toAbsolutePath().normalize()); + } + + private static CallbackTestLib loadTestLib() { + return Native.load(Path.of(SPIKE_LIB_PATH).toAbsolutePath().normalize().toString(), CallbackTestLib.class); + } + + /** + * Integration test that exercises the real JNA callback/lifecycle path against + * a native test library. Skipped in normal CI because the Rust test crate is + * not built as part of the Maven build. To run locally, build the Rust crate + * and set {@code -Dcopilot.test.nativelib.path=}. + */ + @Test + void startWithSpikeLibrarySupportsLifecycleAndDataFlow() throws Exception { + assumeTrue(testLibExists(), "Native test library not found at " + SPIKE_LIB_PATH + + ". Build the Rust test crate or set -D" + TEST_LIB_PATH_PROP + " to run this test."); + CallbackTestLib callbackTestLib = loadTestLib(); + AtomicInteger writes = new AtomicInteger(); + + NativeBinding binding = new NativeBinding() { + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + return callbackTestLib.host_start(); + } + + @Override + public boolean hostShutdown(int serverId) { + return callbackTestLib.host_shutdown(serverId) != 0; + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + return callbackTestLib.connection_open(serverId, callback, userData, 1); + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + writes.incrementAndGet(); + return callbackTestLib.connection_write(connectionId, data, dataLen) != 0; + } + + @Override + public boolean connectionClose(int connectionId) { + return callbackTestLib.connection_close(connectionId) != 0; + } + }; + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "spike-lib"); + host.start("/tmp/runtime.js", new CopilotClientOptions()); + + InputStream in = host.getReceiveStream(); + byte[] buffer = new byte[512]; + int read = in.read(buffer); + assertTrue(read > 0); + String content = new String(buffer, 0, read, StandardCharsets.UTF_8); + assertTrue(content.contains("jsonrpc"), "callback payload should contain JSON-RPC"); + + OutputStream out = host.getSendStream(); + out.write("{\"jsonrpc\":\"2.0\"}".getBytes(StandardCharsets.UTF_8)); + assertEquals(1, writes.get()); + + assertDoesNotThrow(host::close); + } + + @Test + void startBuildsExpectedArgvAndEnvJson() throws Exception { + class RecordingBinding implements NativeBinding { + byte[] argv; + byte[] env; + + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + this.argv = argvJson; + this.env = envJson; + return 11; + } + + @Override + public boolean hostShutdown(int serverId) { + return true; + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + return 21; + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + return true; + } + + @Override + public boolean connectionClose(int connectionId) { + return true; + } + } + + RecordingBinding binding = new RecordingBinding(); + CopilotClientOptions options = new CopilotClientOptions().setLogLevel("debug").setGitHubToken("gh-token") + .setCopilotHome("/tmp/copilot-home").setUseLoggedInUser(false).setSessionIdleTimeoutSeconds(42) + .setRemote(true).setMode(CopilotClientMode.EMPTY).setCliArgs(new String[]{"--extra-flag"}); + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "/tmp/runtime.node"); + host.start("/tmp/entrypoint.js", options); + + List argv = MAPPER.readValue(binding.argv, new TypeReference>() { + }); + assertEquals("node", argv.get(0)); + assertEquals("/tmp/entrypoint.js", argv.get(1)); + assertTrue(argv.contains("--embedded-host")); + assertTrue(argv.contains("--no-auto-update")); + assertTrue(argv.contains("--auth-token-env")); + assertTrue(argv.contains("COPILOT_SDK_AUTH_TOKEN")); + assertTrue(argv.contains("--no-auto-login")); + assertTrue(argv.contains("--session-idle-timeout")); + assertTrue(argv.contains("42")); + assertTrue(argv.contains("--remote")); + assertTrue(argv.contains("--extra-flag")); + + Map env = MAPPER.readValue(binding.env, new TypeReference>() { + }); + assertEquals("gh-token", env.get("COPILOT_SDK_AUTH_TOKEN")); + assertEquals("/tmp/copilot-home", env.get("COPILOT_HOME")); + assertEquals("1", env.get("COPILOT_DISABLE_KEYTAR")); + } + + @Test + void callbackExceptionIsContainedAndDoesNotEscapeAcrossFfiBoundary() { + AtomicBoolean callbackReturned = new AtomicBoolean(false); + NativeBinding binding = new NativeBinding() { + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + return 1; + } + + @Override + public boolean hostShutdown(int serverId) { + return true; + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + Memory mem = new Memory(5); + mem.write(0, "hello".getBytes(StandardCharsets.UTF_8), 0, 5); + callback.invoke(Pointer.NULL, mem, 5); + callbackReturned.set(true); + return 2; + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + return true; + } + + @Override + public boolean connectionClose(int connectionId) { + return true; + } + }; + + QueueInputStream throwingStream = new QueueInputStream() { + @Override + void enqueue(byte[] bytes) { + throw new RuntimeException("boom"); + } + }; + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "test-lib", throwingStream); + assertDoesNotThrow(() -> host.start("/tmp/entrypoint", new CopilotClientOptions())); + assertTrue(callbackReturned.get(), "callback should return normally even when enqueue throws"); + } + + @Test + void closeNeverThrowsEvenWhenNativeCloseFails() { + NativeBinding binding = new NativeBinding() { + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + return 5; + } + + @Override + public boolean hostShutdown(int serverId) { + throw new RuntimeException("shutdown failed"); + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + return 9; + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + return true; + } + + @Override + public boolean connectionClose(int connectionId) { + throw new RuntimeException("close failed"); + } + }; + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "test-lib"); + host.start("/tmp/entrypoint", new CopilotClientOptions()); + assertDoesNotThrow(host::close); + } + + @Test + void writeAndCloseAreSerializedByOperationLock() throws Exception { + CountDownLatch writeStarted = new CountDownLatch(1); + CountDownLatch allowWriteToFinish = new CountDownLatch(1); + AtomicInteger writes = new AtomicInteger(0); + + NativeBinding binding = new NativeBinding() { + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + return 3; + } + + @Override + public boolean hostShutdown(int serverId) { + return true; + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + return 4; + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + writes.incrementAndGet(); + writeStarted.countDown(); + try { + allowWriteToFinish.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return true; + } + + @Override + public boolean connectionClose(int connectionId) { + return true; + } + }; + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "test-lib"); + host.start("/tmp/entrypoint", new CopilotClientOptions()); + + CompletableFuture writer = CompletableFuture.runAsync(() -> { + try { + host.getSendStream().write("ping".getBytes(StandardCharsets.UTF_8)); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + + assertTrue(writeStarted.await(2, TimeUnit.SECONDS)); + CompletableFuture closer = CompletableFuture.runAsync(host::close); + allowWriteToFinish.countDown(); + + writer.get(5, TimeUnit.SECONDS); + closer.get(5, TimeUnit.SECONDS); + assertEquals(1, writes.get()); + assertThrows(IOException.class, () -> host.getSendStream().write("late".getBytes(StandardCharsets.UTF_8))); + } + + @Test + void closeDrainsActiveCallbacksBeforeHostShutdown() throws Exception { + CountDownLatch callbackEntered = new CountDownLatch(1); + CountDownLatch allowCallbackToReturn = new CountDownLatch(1); + AtomicBoolean shutdownObservedAfterCallbackReturn = new AtomicBoolean(false); + AtomicBoolean callbackFinished = new AtomicBoolean(false); + AtomicReference callbackRef = new AtomicReference<>(); + + NativeBinding binding = new NativeBinding() { + @Override + public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJsonLen) { + return 7; + } + + @Override + public boolean hostShutdown(int serverId) { + shutdownObservedAfterCallbackReturn.set(callbackFinished.get()); + return true; + } + + @Override + public int connectionOpen(int serverId, OutboundCallback callback, Pointer userData, byte[] extSource, + int extSourceLen, byte[] extName, int extNameLen, byte[] connToken, int connTokenLen) { + callbackRef.set(callback); + return 8; + } + + @Override + public boolean connectionWrite(int connectionId, byte[] data, int dataLen) { + return true; + } + + @Override + public boolean connectionClose(int connectionId) { + return true; + } + }; + + QueueInputStream blockingStream = new QueueInputStream() { + @Override + void enqueue(byte[] bytes) { + callbackEntered.countDown(); + try { + allowCallbackToReturn.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + callbackFinished.set(true); + super.enqueue(bytes); + } + }; + + FfiRuntimeHost host = new FfiRuntimeHost(binding, "test-lib", blockingStream); + host.start("/tmp/entrypoint", new CopilotClientOptions()); + assertNotNull(callbackRef.get()); + + CompletableFuture callbackFuture = CompletableFuture.runAsync(() -> { + Memory mem = new Memory(1); + mem.setByte(0, (byte) 'x'); + callbackRef.get().invoke(Pointer.NULL, mem, 1); + }); + + assertTrue(callbackEntered.await(2, TimeUnit.SECONDS)); + CompletableFuture closeFuture = CompletableFuture.runAsync(host::close); + Thread.sleep(150); + assertFalse(closeFuture.isDone(), "close should wait for active callback to drain"); + allowCallbackToReturn.countDown(); + callbackFuture.get(5, TimeUnit.SECONDS); + closeFuture.get(5, TimeUnit.SECONDS); + assertTrue(shutdownObservedAfterCallbackReturn.get(), "host_shutdown should run after callback drains"); + } +} diff --git a/java/sdk/src/test/java/com/github/copilot/ffi/QueueInputStreamTest.java b/java/sdk/src/test/java/com/github/copilot/ffi/QueueInputStreamTest.java new file mode 100644 index 0000000000..6fd9616706 --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/ffi/QueueInputStreamTest.java @@ -0,0 +1,81 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot.ffi; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import org.junit.jupiter.api.Test; + +class QueueInputStreamTest { + + @Test + void readReturnsEnqueuedBytesAcrossMultipleChunks() throws Exception { + QueueInputStream stream = new QueueInputStream(); + stream.enqueue("hello ".getBytes(StandardCharsets.UTF_8)); + stream.enqueue("world".getBytes(StandardCharsets.UTF_8)); + + byte[] buffer = new byte[11]; + int first = stream.read(buffer, 0, 6); + int second = stream.read(buffer, 6, 5); + + assertEquals(6, first); + assertEquals(5, second); + assertArrayEquals("hello world".getBytes(StandardCharsets.UTF_8), buffer); + } + + @Test + void readBlocksUntilDataArrives() throws Exception { + QueueInputStream stream = new QueueInputStream(); + + CompletableFuture readFuture = CompletableFuture.supplyAsync(() -> { + try { + return stream.read(); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + + Thread.sleep(100); + stream.enqueue(new byte[]{(byte) 'A'}); + + assertEquals((int) 'A', readFuture.get(2, TimeUnit.SECONDS)); + } + + @Test + void closeSignalsEndOfStream() throws Exception { + QueueInputStream stream = new QueueInputStream(); + stream.enqueue("x".getBytes(StandardCharsets.UTF_8)); + + assertEquals('x', stream.read()); + stream.close(); + assertEquals(-1, stream.read()); + } + + @Test + void closeUnblocksPendingReadWithEof() throws Exception { + QueueInputStream stream = new QueueInputStream(); + + CompletableFuture readFuture = CompletableFuture.supplyAsync(() -> { + try { + return stream.read(); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + + Thread.sleep(100); + stream.close(); + + assertEquals(-1, readFuture.get(2, TimeUnit.SECONDS)); + assertTrue(readFuture.isDone()); + } +}