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

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,33 @@ 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);
+
         // 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)"
+        });
 
-        // Abort the session immediately
+        // Wait for tool execution to start
+        await toolStartTask;
+
+        // Abort the session
         await session.AbortAsync();
 
+        // Wait for session to become idle after abort
+        await TestHelper.GetNextEventOfTypeAsync<SessionIdleEvent>(session);
+
         // 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,44 @@ 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 wait for tool.execution_start BEFORE sending
+		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
+			}
+		}()
+
+		// 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
+		_, err = testharness.GetNextEventOfType(session, copilot.SessionIdle, 60*time.Second)
+		if err != nil {
+			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 +519,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,40 @@ 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) {
+		if event.Type == eventType {
+			select {
+			case result <- &event:
+			default:
+			}
+		} else if event.Type == 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

@@ -1,11 +1,5 @@
 package copilot
 
-import (
-	"github.com/github/copilot-sdk/go/generated"
-)
-
-type SessionEvent = generated.SessionEvent
-
 // ConnectionState represents the client connection state
 type ConnectionState string
 
@@ -258,13 +252,6 @@ type MessageOptions struct {
 	Mode string
 }
 
-// Attachment represents a file or directory attachment
-type Attachment struct {
-	Type        string `json:"type"` // "file" or "directory"
-	Path        string `json:"path"`
-	DisplayName string `json:"displayName,omitempty"`
-}
-
 // SessionEventHandler is a callback for session events
 type SessionEventHandler func(event SessionEvent)
 

go/types.go

@@ -272,7 +272,7 @@ async function generateGoTypes(schemaPath: string) {
         inputData,
         lang: "go",
         rendererOptions: {
-            package: "generated",
+            package: "copilot",
         },
     });
 
@@ -289,7 +289,7 @@ async function generateGoTypes(schemaPath: string) {
 
 `;
 
-    const outputPath = path.join(__dirname, "../../go/generated/session_events.go");
+    const outputPath = path.join(__dirname, "../../go/generated_session_events.go");
     await fs.mkdir(path.dirname(outputPath), { recursive: true });
     await fs.writeFile(outputPath, banner + generatedCode, "utf-8");
 

nodejs/scripts/generate-session-types.ts

@@ -5,7 +5,7 @@
 from copilot import CopilotClient
 from copilot.types import Tool
 
-from .testharness import E2ETestContext, get_final_assistant_message
+from .testharness import E2ETestContext, get_final_assistant_message, get_next_event_of_type
 
 pytestmark = pytest.mark.asyncio(loop_scope="module")
 
@@ -256,21 +256,39 @@ async def test_should_resume_session_with_custom_provider(self, ctx: E2ETestCont
         assert session2.session_id == session_id
 
     async def test_should_abort_a_session(self, ctx: E2ETestContext):
+        import asyncio
+
         session = await ctx.client.create_session()
 
-        # Send a message that will take some time to process
-        await session.send({"prompt": "What is 1+1?"})
+        # Set up wait for tool.execution_start BEFORE sending
+        wait_for_tool_start = asyncio.create_task(
+            get_next_event_of_type(session, "tool.execution_start", timeout=60.0)
+        )
+
+        # Send a message that will trigger a long-running shell command
+        await session.send(
+            {"prompt": "run the shell command 'sleep 100' (note this works on both bash and PowerShell)"}
+        )
+
+        # Wait for the tool to start executing
+        await wait_for_tool_start
 
-        # Abort the session immediately
+        # Abort the session while the tool is running
         await session.abort()
 
+        # Wait for session to become idle after abort
+        await get_next_event_of_type(session, "session.idle", timeout=30.0)
+
         # The session should still be alive and usable after abort
         messages = await session.get_messages()
         assert len(messages) > 0
 
+        # Verify an abort event exists in messages
+        abort_events = [m for m in messages if m.type.value == "abort"]
+        assert len(abort_events) > 0, "Expected an abort event in messages"
+
         # We should be able to send another message
-        await session.send({"prompt": "What is 2+2?"})
-        answer = await get_final_assistant_message(session)
+        answer = await session.send_and_wait({"prompt": "What is 2+2?"})
         assert "4" in answer.data.content
 
     async def test_should_receive_streaming_delta_events_when_streaming_is_enabled(