99import com .sun .jna .Pointer ;
1010
1111import java .nio .file .Path ;
12+ import java .util .Map ;
13+ import java .util .concurrent .ConcurrentHashMap ;
1214import java .util .concurrent .atomic .AtomicInteger ;
1315import java .util .logging .Logger ;
1416
@@ -58,18 +60,33 @@ interface CopilotRuntimeLibrary extends Library {
5860 /** Corresponds to {@code copilot_runtime_host_start}. */
5961 int copilot_runtime_host_start (byte [] argvJson , int argvJsonLen , byte [] envJson , int envJsonLen );
6062
61- /** Corresponds to {@code copilot_runtime_host_shutdown}. */
62- boolean copilot_runtime_host_shutdown (int serverId );
63+ /**
64+ * Corresponds to {@code copilot_runtime_host_shutdown}.
65+ *
66+ * <p>
67+ * Returns {@code byte} (not Java {@code boolean}) because the Rust ABI exports
68+ * a one-byte {@code bool}. JNA maps Java {@code boolean} as a 32-bit C
69+ * {@code int}, which would read three extra bytes.
70+ */
71+ byte copilot_runtime_host_shutdown (int serverId );
6372
6473 /** Corresponds to {@code copilot_runtime_connection_open}. */
6574 int copilot_runtime_connection_open (int serverId , OutboundCallback callback , Pointer userData , byte [] extSource ,
6675 int extSourceLen , byte [] extName , int extNameLen , byte [] connToken , int connTokenLen );
6776
68- /** Corresponds to {@code copilot_runtime_connection_write}. */
69- boolean copilot_runtime_connection_write (int connectionId , byte [] data , int dataLen );
70-
71- /** Corresponds to {@code copilot_runtime_connection_close}. */
72- boolean copilot_runtime_connection_close (int connectionId );
77+ /**
78+ * Corresponds to {@code copilot_runtime_connection_write}.
79+ *
80+ * @see #copilot_runtime_host_shutdown for why this returns {@code byte}
81+ */
82+ byte copilot_runtime_connection_write (int connectionId , byte [] data , int dataLen );
83+
84+ /**
85+ * Corresponds to {@code copilot_runtime_connection_close}.
86+ *
87+ * @see #copilot_runtime_host_shutdown for why this returns {@code byte}
88+ */
89+ byte copilot_runtime_connection_close (int connectionId );
7390 }
7491
7592 // -------------------------------------------------------------------------
@@ -105,6 +122,12 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi
105122 */
106123 final AtomicInteger activeCallbacks = new AtomicInteger (0 );
107124
125+ /**
126+ * Tracked callback wrappers keyed by connection handle. Prevents GC of the JNA
127+ * callback function pointer while native code still holds it.
128+ */
129+ private final Map <Integer , OutboundCallback > trackedCallbacks = new ConcurrentHashMap <>();
130+
108131 // -------------------------------------------------------------------------
109132 // Constructors
110133 // -------------------------------------------------------------------------
@@ -123,7 +146,11 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi
123146 synchronized (LOAD_LOCK ) {
124147 if (loadedLib == null ) {
125148 LOG .fine (() -> "Loading native library from: " + absPath );
126- loadedLib = Native .load (absPath .toString (), CopilotRuntimeLibrary .class );
149+ try {
150+ loadedLib = Native .load (absPath .toString (), CopilotRuntimeLibrary .class );
151+ } catch (UnsatisfiedLinkError e ) {
152+ throw new IllegalStateException ("Failed to load native library from '" + absPath + "'" , e );
153+ }
127154 loadedPath = absPath ;
128155 LOG .fine (() -> "Native library loaded: " + absPath );
129156 } else if (!absPath .equals (loadedPath )) {
@@ -161,7 +188,7 @@ public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJs
161188
162189 @ Override
163190 public boolean hostShutdown (int serverId ) {
164- return lib .copilot_runtime_host_shutdown (serverId );
191+ return lib .copilot_runtime_host_shutdown (serverId ) != 0 ;
165192 }
166193
167194 @ Override
@@ -176,18 +203,27 @@ public int connectionOpen(int serverId, OutboundCallback callback, Pointer userD
176203 activeCallbacks .decrementAndGet ();
177204 }
178205 };
179- return lib .copilot_runtime_connection_open (serverId , tracked , userData , extSource , extSourceLen , extName ,
180- extNameLen , connToken , connTokenLen );
206+ int connectionId = lib .copilot_runtime_connection_open (serverId , tracked , userData , extSource , extSourceLen ,
207+ extName , extNameLen , connToken , connTokenLen );
208+ if (connectionId != 0 ) {
209+ // Hold a strong reference to prevent GC of the JNA function pointer.
210+ trackedCallbacks .put (connectionId , tracked );
211+ }
212+ return connectionId ;
181213 }
182214
183215 @ Override
184216 public boolean connectionWrite (int connectionId , byte [] data , int dataLen ) {
185- return lib .copilot_runtime_connection_write (connectionId , data , dataLen );
217+ return lib .copilot_runtime_connection_write (connectionId , data , dataLen ) != 0 ;
186218 }
187219
188220 @ Override
189221 public boolean connectionClose (int connectionId ) {
190- return lib .copilot_runtime_connection_close (connectionId );
222+ try {
223+ return lib .copilot_runtime_connection_close (connectionId ) != 0 ;
224+ } finally {
225+ trackedCallbacks .remove (connectionId );
226+ }
191227 }
192228
193229 // -------------------------------------------------------------------------
0 commit comments