Skip to content

Commit c57e7a3

Browse files
Copilotbrunoborges
andcommitted
Add closed-session guard implementation
Co-authored-by: brunoborges <129743+brunoborges@users.noreply.github.com>
1 parent e265cf4 commit c57e7a3

2 files changed

Lines changed: 381 additions & 1 deletion

File tree

src/main/java/com/github/copilot/sdk/CopilotSession.java

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,9 @@ public final class CopilotSession implements AutoCloseable {
9797
private final AtomicReference<SessionHooks> hooksHandler = new AtomicReference<>();
9898
private volatile EventErrorHandler eventErrorHandler;
9999
private volatile EventErrorPolicy eventErrorPolicy = EventErrorPolicy.PROPAGATE_AND_LOG_ERRORS;
100+
101+
/** Tracks whether this session instance has been terminated via close(). */
102+
private volatile boolean isTerminated = false;
100103

101104
/**
102105
* Creates a new session with the given ID and RPC client.
@@ -186,11 +189,13 @@ public String getWorkspacePath() {
186189
* @param handler
187190
* the error handler, or {@code null} to use only the default logging
188191
* behavior
192+
* @throws IllegalStateException if this session has been terminated
189193
* @see EventErrorHandler
190194
* @see #setEventErrorPolicy(EventErrorPolicy)
191195
* @since 1.0.8
192196
*/
193197
public void setEventErrorHandler(EventErrorHandler handler) {
198+
ensureNotTerminated();
194199
this.eventErrorHandler = handler;
195200
}
196201

@@ -224,11 +229,13 @@ public void setEventErrorHandler(EventErrorHandler handler) {
224229
* @param policy
225230
* the error policy (default is
226231
* {@link EventErrorPolicy#PROPAGATE_AND_LOG_ERRORS})
232+
* @throws IllegalStateException if this session has been terminated
227233
* @see EventErrorPolicy
228234
* @see #setEventErrorHandler(EventErrorHandler)
229235
* @since 1.0.8
230236
*/
231237
public void setEventErrorPolicy(EventErrorPolicy policy) {
238+
ensureNotTerminated();
232239
if (policy == null) {
233240
throw new NullPointerException("policy must not be null");
234241
}
@@ -244,9 +251,11 @@ public void setEventErrorPolicy(EventErrorPolicy policy) {
244251
* @param prompt
245252
* the message text to send
246253
* @return a future that resolves with the message ID assigned by the server
254+
* @throws IllegalStateException if this session has been terminated
247255
* @see #send(MessageOptions)
248256
*/
249257
public CompletableFuture<String> send(String prompt) {
258+
ensureNotTerminated();
250259
return send(new MessageOptions().setPrompt(prompt));
251260
}
252261

@@ -260,9 +269,11 @@ public CompletableFuture<String> send(String prompt) {
260269
* the message text to send
261270
* @return a future that resolves with the final assistant message event, or
262271
* {@code null} if no assistant message was received
272+
* @throws IllegalStateException if this session has been terminated
263273
* @see #sendAndWait(MessageOptions)
264274
*/
265275
public CompletableFuture<AssistantMessageEvent> sendAndWait(String prompt) {
276+
ensureNotTerminated();
266277
return sendAndWait(new MessageOptions().setPrompt(prompt));
267278
}
268279

@@ -275,10 +286,12 @@ public CompletableFuture<AssistantMessageEvent> sendAndWait(String prompt) {
275286
* @param options
276287
* the message options containing the prompt and attachments
277288
* @return a future that resolves with the message ID assigned by the server
289+
* @throws IllegalStateException if this session has been terminated
278290
* @see #sendAndWait(MessageOptions)
279291
* @see #send(String)
280292
*/
281293
public CompletableFuture<String> send(MessageOptions options) {
294+
ensureNotTerminated();
282295
var request = new SendMessageRequest();
283296
request.setSessionId(sessionId);
284297
request.setPrompt(options.getPrompt());
@@ -304,10 +317,12 @@ public CompletableFuture<String> send(MessageOptions options) {
304317
* {@code null} if no assistant message was received. The future
305318
* completes exceptionally with a TimeoutException if the timeout
306319
* expires.
320+
* @throws IllegalStateException if this session has been terminated
307321
* @see #sendAndWait(MessageOptions)
308322
* @see #send(MessageOptions)
309323
*/
310324
public CompletableFuture<AssistantMessageEvent> sendAndWait(MessageOptions options, long timeoutMs) {
325+
ensureNotTerminated();
311326
var future = new CompletableFuture<AssistantMessageEvent>();
312327
var lastAssistantMessage = new AtomicReference<AssistantMessageEvent>();
313328

@@ -365,9 +380,11 @@ public CompletableFuture<AssistantMessageEvent> sendAndWait(MessageOptions optio
365380
* the message options containing the prompt and attachments
366381
* @return a future that resolves with the final assistant message event, or
367382
* {@code null} if no assistant message was received
383+
* @throws IllegalStateException if this session has been terminated
368384
* @see #sendAndWait(MessageOptions, long)
369385
*/
370386
public CompletableFuture<AssistantMessageEvent> sendAndWait(MessageOptions options) {
387+
ensureNotTerminated();
371388
return sendAndWait(options, 60000);
372389
}
373390

@@ -397,11 +414,13 @@ public CompletableFuture<AssistantMessageEvent> sendAndWait(MessageOptions optio
397414
* @param handler
398415
* a callback to be invoked when a session event occurs
399416
* @return a Closeable that, when closed, unsubscribes the handler
417+
* @throws IllegalStateException if this session has been terminated
400418
* @see #on(Class, Consumer)
401419
* @see AbstractSessionEvent
402420
* @see #setEventErrorPolicy(EventErrorPolicy)
403421
*/
404422
public Closeable on(Consumer<AbstractSessionEvent> handler) {
423+
ensureNotTerminated();
405424
eventHandlers.add(handler);
406425
return () -> eventHandlers.remove(handler);
407426
}
@@ -447,10 +466,12 @@ public Closeable on(Consumer<AbstractSessionEvent> handler) {
447466
* @param handler
448467
* a callback invoked when events of this type occur
449468
* @return a Closeable that unsubscribes the handler when closed
469+
* @throws IllegalStateException if this session has been terminated
450470
* @see #on(Consumer)
451471
* @see AbstractSessionEvent
452472
*/
453473
public <T extends AbstractSessionEvent> Closeable on(Class<T> eventType, Consumer<T> handler) {
474+
ensureNotTerminated();
454475
Consumer<AbstractSessionEvent> wrapper = event -> {
455476
if (eventType.isInstance(event)) {
456477
handler.accept(eventType.cast(event));
@@ -708,9 +729,11 @@ CompletableFuture<Object> handleHooksInvoke(String hookType, JsonNode input) {
708729
* assistant responses, tool invocations, and other session events.
709730
*
710731
* @return a future that resolves with a list of all session events
732+
* @throws IllegalStateException if this session has been terminated
711733
* @see AbstractSessionEvent
712734
*/
713735
public CompletableFuture<List<AbstractSessionEvent>> getMessages() {
736+
ensureNotTerminated();
714737
return rpc.invoke("session.getMessages", Map.of("sessionId", sessionId), GetMessagesResponse.class)
715738
.thenApply(response -> {
716739
var events = new ArrayList<AbstractSessionEvent>();
@@ -737,20 +760,40 @@ public CompletableFuture<List<AbstractSessionEvent>> getMessages() {
737760
* continuing to generate a response.
738761
*
739762
* @return a future that completes when the abort is acknowledged
763+
* @throws IllegalStateException if this session has been terminated
740764
*/
741765
public CompletableFuture<Void> abort() {
766+
ensureNotTerminated();
742767
return rpc.invoke("session.abort", Map.of("sessionId", sessionId), Void.class);
743768
}
769+
770+
/**
771+
* Verifies that this session has not yet been terminated.
772+
*
773+
* @throws IllegalStateException if close() has already been invoked
774+
*/
775+
private void ensureNotTerminated() {
776+
if (isTerminated) {
777+
throw new IllegalStateException("Session is closed");
778+
}
779+
}
744780

745781
/**
746782
* Disposes the session and releases all associated resources.
747783
* <p>
748784
* This destroys the session on the server, clears all event handlers, and
749785
* releases tool and permission handlers. After calling this method, the session
750-
* cannot be used again.
786+
* cannot be used again. Subsequent calls to this method have no effect.
751787
*/
752788
@Override
753789
public void close() {
790+
synchronized (this) {
791+
if (isTerminated) {
792+
return; // Already terminated - no-op
793+
}
794+
isTerminated = true;
795+
}
796+
754797
try {
755798
rpc.invoke("session.destroy", Map.of("sessionId", sessionId), Void.class).get(5, TimeUnit.SECONDS);
756799
} catch (Exception e) {

0 commit comments

Comments
 (0)