Improve abort tests across all SDKs; add Go unsubscribe tests (#48)

* Fix session event handler unsubscription and add tests

The unsubscribe function was failing due to invalid function pointer comparisons.
Refactored handler registration to use unique IDs for reliable cleanup.

Tests verify:
- Multiple handlers can be registered and all receive events
- Unsubscribing one handler doesn't affect others
- Calling unsubscribe multiple times is safe
- Handlers are called in registration order
- Concurrent subscribe/unsubscribe is safe

Co-authored-by: nathfavour <116535483+nathfavour@users.noreply.github.com>

* Fix abort test to use non-blocking send()

* Clean up "should abort a session"

* Add equivalent abort test improvements to Go, Python, and .NET

* Formatting

* More lint/format

* Update test_session.py

* Fix risk of flakiness

---------

Co-authored-by: nathfavour <116535483+nathfavour@users.noreply.github.com>

Author: SteveSandersonMS

dotnet/test/Harness/TestHelper.cs

@@ -73,4 +73,29 @@ async void CheckExistingMessages()
 
         return null;
     }
+
+    public static async Task<T> GetNextEventOfTypeAsync<T>(
+        CopilotSession session,
+        TimeSpan? timeout = null) where T : SessionEvent
+    {
+        var tcs = new TaskCompletionSource<T>();
+        using var cts = new CancellationTokenSource(timeout ?? TimeSpan.FromSeconds(60));
+
+        using var subscription = session.On(evt =>
+        {
+            if (evt is T matched)
+            {
+                tcs.TrySetResult(matched);
+            }
+            else if (evt is SessionErrorEvent error)
+            {
+                tcs.TrySetException(new Exception(error.Data.Message ?? "session error"));
+            }
+        });
+
+        cts.Token.Register(() => tcs.TrySetException(
+            new TimeoutException($"Timeout waiting for event of type '{typeof(T).Name}'")));
+
+        return await tcs.Task;
+    }
 }

dotnet/test/SessionTests.cs

@@ -201,23 +201,32 @@ public async Task Should_Abort_A_Session()
     {
         var session = await Client.CreateSessionAsync();
 
+        // Set up wait for tool execution to start BEFORE sending
+        var toolStartTask = TestHelper.GetNextEventOfTypeAsync<ToolExecutionStartEvent>(session);
+        var sessionIdleTask = TestHelper.GetNextEventOfTypeAsync<SessionIdleEvent>(session);
+
         // Send a message that will take some time to process
-        await session.SendAsync(new MessageOptions { Prompt = "What is 1+1?" });
+        await session.SendAsync(new MessageOptions
+        {
+            Prompt = "run the shell command 'sleep 100' (note this works on both bash and PowerShell)"
+        });
+
+        // Wait for tool execution to start
+        await toolStartTask;
 
-        // Abort the session immediately
+        // Abort the session
         await session.AbortAsync();
+        await sessionIdleTask;
 
         // The session should still be alive and usable after abort
         var messages = await session.GetMessagesAsync();
         Assert.NotEmpty(messages);
 
-        // TODO: We should do something to verify it really did abort (e.g., is there an abort event we can see,
-        // or can we check that the session became idle without receiving an assistant message?). Right now
-        // I'm not seeing any evidence that it actually does abort.
+        // Verify an abort event exists in messages
+        Assert.Contains(messages, m => m is AbortEvent);
 
         // We should be able to send another message
-        await session.SendAsync(new MessageOptions { Prompt = "What is 2+2?" });
-        var answer = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var answer = await session.SendAndWaitAsync(new MessageOptions { Prompt = "What is 2+2?" });
         Assert.NotNull(answer);
         Assert.Contains("4", answer!.Data.Content ?? string.Empty);
     }

go/client.go

@@ -39,8 +39,6 @@ import (
 	"strings"
 	"sync"
 	"time"
-
-	"github.com/github/copilot-sdk/go/generated"
 )
 
 // Client manages the connection to the Copilot CLI server and provides session management.
@@ -923,7 +921,7 @@ func (c *Client) setupNotificationHandler() {
 				return
 			}
 
-			event, err := generated.UnmarshalSessionEvent(eventJSON)
+			event, err := UnmarshalSessionEvent(eventJSON)
 			if err != nil {
 				return
 			}

go/e2e/session_test.go

@@ -472,18 +472,57 @@ func TestSession(t *testing.T) {
 			t.Fatalf("Failed to create session: %v", err)
 		}
 
-		// Send a message that will take some time to process
-		_, err = session.Send(copilot.MessageOptions{Prompt: "What is 1+1?"})
+		// Set up event listeners BEFORE sending to avoid race conditions
+		toolStartCh := make(chan *copilot.SessionEvent, 1)
+		toolStartErrCh := make(chan error, 1)
+		go func() {
+			evt, err := testharness.GetNextEventOfType(session, copilot.ToolExecutionStart, 60*time.Second)
+			if err != nil {
+				toolStartErrCh <- err
+			} else {
+				toolStartCh <- evt
+			}
+		}()
+
+		sessionIdleCh := make(chan *copilot.SessionEvent, 1)
+		sessionIdleErrCh := make(chan error, 1)
+		go func() {
+			evt, err := testharness.GetNextEventOfType(session, copilot.SessionIdle, 60*time.Second)
+			if err != nil {
+				sessionIdleErrCh <- err
+			} else {
+				sessionIdleCh <- evt
+			}
+		}()
+
+		// Send a message that triggers a long-running shell command
+		_, err = session.Send(copilot.MessageOptions{Prompt: "run the shell command 'sleep 100' (note this works on both bash and PowerShell)"})
 		if err != nil {
 			t.Fatalf("Failed to send message: %v", err)
 		}
 
-		// Abort the session immediately
+		// Wait for tool.execution_start
+		select {
+		case <-toolStartCh:
+			// Tool execution has started
+		case err := <-toolStartErrCh:
+			t.Fatalf("Failed waiting for tool.execution_start: %v", err)
+		}
+
+		// Abort the session
 		err = session.Abort()
 		if err != nil {
 			t.Fatalf("Failed to abort session: %v", err)
 		}
 
+		// Wait for session.idle after abort
+		select {
+		case <-sessionIdleCh:
+			// Session is idle
+		case err := <-sessionIdleErrCh:
+			t.Fatalf("Failed waiting for session.idle after abort: %v", err)
+		}
+
 		// The session should still be alive and usable after abort
 		messages, err := session.GetMessages()
 		if err != nil {
@@ -493,15 +532,22 @@ func TestSession(t *testing.T) {
 			t.Error("Expected messages to exist after abort")
 		}
 
-		// We should be able to send another message
-		_, err = session.Send(copilot.MessageOptions{Prompt: "What is 2+2?"})
-		if err != nil {
-			t.Fatalf("Failed to send message after abort: %v", err)
+		// Verify messages contain an abort event
+		hasAbortEvent := false
+		for _, msg := range messages {
+			if msg.Type == copilot.Abort {
+				hasAbortEvent = true
+				break
+			}
+		}
+		if !hasAbortEvent {
+			t.Error("Expected messages to contain an 'abort' event")
 		}
 
-		answer, err := testharness.GetFinalAssistantMessage(session, 60*time.Second)
+		// We should be able to send another message
+		answer, err := session.SendAndWait(copilot.MessageOptions{Prompt: "What is 2+2?"}, 60*time.Second)
 		if err != nil {
-			t.Fatalf("Failed to get assistant message after abort: %v", err)
+			t.Fatalf("Failed to send message after abort: %v", err)
 		}
 
 		if answer.Data.Content == nil || !strings.Contains(*answer.Data.Content, "4") {

go/e2e/testharness/helper.go

@@ -54,6 +54,41 @@ func GetFinalAssistantMessage(session *copilot.Session, timeout time.Duration) (
 	}
 }
 
+// GetNextEventOfType waits for and returns the next event of the specified type from a session.
+func GetNextEventOfType(session *copilot.Session, eventType copilot.SessionEventType, timeout time.Duration) (*copilot.SessionEvent, error) {
+	result := make(chan *copilot.SessionEvent, 1)
+	errCh := make(chan error, 1)
+
+	unsubscribe := session.On(func(event copilot.SessionEvent) {
+		switch event.Type {
+		case eventType:
+			select {
+			case result <- &event:
+			default:
+			}
+		case copilot.SessionError:
+			msg := "session error"
+			if event.Data.Message != nil {
+				msg = *event.Data.Message
+			}
+			select {
+			case errCh <- errors.New(msg):
+			default:
+			}
+		}
+	})
+	defer unsubscribe()
+
+	select {
+	case evt := <-result:
+		return evt, nil
+	case err := <-errCh:
+		return nil, err
+	case <-time.After(timeout):
+		return nil, errors.New("timeout waiting for event: " + string(eventType))
+	}
+}
+
 func getExistingFinalResponse(session *copilot.Session) (*copilot.SessionEvent, error) {
 	messages, err := session.GetMessages()
 	if err != nil {

go/generated/session_events.go go/generated_session_events.gogo/generated/session_events.go renamed to go/generated_session_events.go

@@ -6,8 +6,6 @@ import (
 	"fmt"
 	"sync"
 	"time"
-
-	"github.com/github/copilot-sdk/go/generated"
 )
 
 type sessionHandler struct {
@@ -159,17 +157,17 @@ func (s *Session) SendAndWait(options MessageOptions, timeout time.Duration) (*S
 
 	unsubscribe := s.On(func(event SessionEvent) {
 		switch event.Type {
-		case generated.AssistantMessage:
+		case AssistantMessage:
 			mu.Lock()
 			eventCopy := event
 			lastAssistantMessage = &eventCopy
 			mu.Unlock()
-		case generated.SessionIdle:
+		case SessionIdle:
 			select {
 			case idleCh <- struct{}{}:
 			default:
 			}
-		case generated.SessionError:
+		case SessionError:
 			errMsg := "session error"
 			if event.Data.Message != nil {
 				errMsg = *event.Data.Message
@@ -387,7 +385,7 @@ func (s *Session) GetMessages() ([]SessionEvent, error) {
 			continue
 		}
 
-		event, err := generated.UnmarshalSessionEvent(eventJSON)
+		event, err := UnmarshalSessionEvent(eventJSON)
 		if err != nil {
 			continue
 		}

go/session.go

@@ -0,0 +1,121 @@
+package copilot
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestSession_On(t *testing.T) {
+	t.Run("multiple handlers all receive events", func(t *testing.T) {
+		session := &Session{
+			handlers: make([]sessionHandler, 0),
+		}
+
+		var received1, received2, received3 bool
+		session.On(func(event SessionEvent) { received1 = true })
+		session.On(func(event SessionEvent) { received2 = true })
+		session.On(func(event SessionEvent) { received3 = true })
+
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		if !received1 || !received2 || !received3 {
+			t.Errorf("Expected all handlers to receive event, got received1=%v, received2=%v, received3=%v",
+				received1, received2, received3)
+		}
+	})
+
+	t.Run("unsubscribing one handler does not affect others", func(t *testing.T) {
+		session := &Session{
+			handlers: make([]sessionHandler, 0),
+		}
+
+		var count1, count2, count3 int
+		session.On(func(event SessionEvent) { count1++ })
+		unsub2 := session.On(func(event SessionEvent) { count2++ })
+		session.On(func(event SessionEvent) { count3++ })
+
+		// First event - all handlers receive it
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		// Unsubscribe handler 2
+		unsub2()
+
+		// Second event - only handlers 1 and 3 should receive it
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		if count1 != 2 {
+			t.Errorf("Expected handler 1 to receive 2 events, got %d", count1)
+		}
+		if count2 != 1 {
+			t.Errorf("Expected handler 2 to receive 1 event (before unsubscribe), got %d", count2)
+		}
+		if count3 != 2 {
+			t.Errorf("Expected handler 3 to receive 2 events, got %d", count3)
+		}
+	})
+
+	t.Run("calling unsubscribe multiple times is safe", func(t *testing.T) {
+		session := &Session{
+			handlers: make([]sessionHandler, 0),
+		}
+
+		var count int
+		unsub := session.On(func(event SessionEvent) { count++ })
+
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		// Call unsubscribe multiple times - should not panic
+		unsub()
+		unsub()
+		unsub()
+
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		if count != 1 {
+			t.Errorf("Expected handler to receive 1 event, got %d", count)
+		}
+	})
+
+	t.Run("handlers are called in registration order", func(t *testing.T) {
+		session := &Session{
+			handlers: make([]sessionHandler, 0),
+		}
+
+		var order []int
+		session.On(func(event SessionEvent) { order = append(order, 1) })
+		session.On(func(event SessionEvent) { order = append(order, 2) })
+		session.On(func(event SessionEvent) { order = append(order, 3) })
+
+		session.dispatchEvent(SessionEvent{Type: "test"})
+
+		if len(order) != 3 || order[0] != 1 || order[1] != 2 || order[2] != 3 {
+			t.Errorf("Expected handlers to be called in order [1,2,3], got %v", order)
+		}
+	})
+
+	t.Run("concurrent subscribe and unsubscribe is safe", func(t *testing.T) {
+		session := &Session{
+			handlers: make([]sessionHandler, 0),
+		}
+
+		var wg sync.WaitGroup
+		for i := 0; i < 100; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				unsub := session.On(func(event SessionEvent) {})
+				unsub()
+			}()
+		}
+		wg.Wait()
+
+		// Should not panic and handlers should be empty
+		session.handlerMutex.RLock()
+		count := len(session.handlers)
+		session.handlerMutex.RUnlock()
+
+		if count != 0 {
+			t.Errorf("Expected 0 handlers after all unsubscribes, got %d", count)
+		}
+	})
+}