Merge pull request #341 from rumpl/stream-start-stop

Add session and agent info in the stream start/stop events

Author: rumpl

pkg/runtime/client.go

@@ -361,9 +361,9 @@ func (c *Client) runAgentWithAgentName(ctx context.Context, sessionID, agent, ag
 				case "error":
 					eventChan <- Error(event["error"].(string))
 				case "stream_started":
-					eventChan <- StreamStarted()
+					eventChan <- StreamStarted(sessionID, event["agent_name"].(string))
 				case "stream_stopped":
-					eventChan <- StreamStopped()
+					eventChan <- StreamStopped(sessionID, event["agent_name"].(string))
 				case "authorization_required":
 					eventChan <- AuthorizationRequired(event["server_url"].(string), event["server_type"].(string), event["confirmation"].(string))
 				case "session_compaction":

pkg/runtime/event.go

@@ -103,17 +103,21 @@ func ToolCallResponse(toolCall tools.ToolCall, response, agentName string) Event
 func (e *ToolCallResponseEvent) isEvent() {}
 
 type StreamStartedEvent struct {
-	Type string `json:"type"`
+	Type      string `json:"type"`
+	SessionID string `json:"session_id,omitempty"`
+	AgentContext
 }
 
-func StreamStarted() Event {
+func StreamStarted(sessionID, agentName string) Event {
 	return &StreamStartedEvent{
-		Type: "stream_started",
+		Type:         "stream_started",
+		SessionID:    sessionID,
+		AgentContext: AgentContext{AgentName: agentName},
 	}
 }
 
 func (e *StreamStartedEvent) GetAgentName() string {
-	return ""
+	return e.AgentName
 }
 
 func (e *StreamStartedEvent) isEvent() {}
@@ -253,17 +257,21 @@ func SessionCompaction(sessionID, status string) Event {
 func (e *SessionCompactionEvent) isEvent() {}
 
 type StreamStoppedEvent struct {
-	Type string `json:"type"`
+	Type      string `json:"type"`
+	SessionID string `json:"session_id,omitempty"`
+	AgentContext
 }
 
-func StreamStopped() Event {
+func StreamStopped(sessionID, agentName string) Event {
 	return &StreamStoppedEvent{
-		Type: "stream_stopped",
+		Type:         "stream_stopped",
+		SessionID:    sessionID,
+		AgentContext: AgentContext{AgentName: agentName},
 	}
 }
 
 func (e *StreamStoppedEvent) GetAgentName() string {
-	return ""
+	return e.AgentName
 }
 
 func (e *StreamStoppedEvent) isEvent() {}

pkg/runtime/runtime.go

@@ -148,7 +148,7 @@ func (r *runtime) handleOAuthAuthorizationFlow(ctx context.Context, sess *sessio
 func (r *runtime) finalizeEventChannel(ctx context.Context, sess *session.Session, events chan Event) {
 	defer close(events)
 
-	events <- StreamStopped()
+	events <- StreamStopped(sess.ID, r.currentAgent)
 
 	telemetry.RecordSessionEnd(ctx)
 
@@ -172,7 +172,7 @@ func (r *runtime) RunStream(ctx context.Context, sess *session.Session) <-chan E
 			events <- UserMessage(messages[len(messages)-1].Content)
 		}
 
-		events <- StreamStarted()
+		events <- StreamStarted(sess.ID, a.Name())
 
 		defer r.finalizeEventChannel(ctx, sess, events)
 

pkg/runtime/runtime_test.go

@@ -181,10 +181,10 @@ func TestSimple(t *testing.T) {
 
 	expectedEvents := []Event{
 		UserMessage("Hi"),
-		StreamStarted(),
+		StreamStarted(sess.ID, "root"),
 		AgentChoice("root", "Hello"),
 		TokenUsage(3, 2, 5, 0, 0),
-		StreamStopped(),
+		StreamStopped(sess.ID, "root"),
 	}
 
 	require.Equal(t, expectedEvents, events)
@@ -206,14 +206,14 @@ func TestMultipleContentChunks(t *testing.T) {
 
 	expectedEvents := []Event{
 		UserMessage("Please greet me"),
-		StreamStarted(),
+		StreamStarted(sess.ID, "root"),
 		AgentChoice("root", "Hello "),
 		AgentChoice("root", "there, "),
 		AgentChoice("root", "how "),
 		AgentChoice("root", "are "),
 		AgentChoice("root", "you?"),
 		TokenUsage(8, 12, 20, 0, 0),
-		StreamStopped(),
+		StreamStopped(sess.ID, "root"),
 	}
 
 	require.Equal(t, expectedEvents, events)
@@ -233,12 +233,12 @@ func TestWithReasoning(t *testing.T) {
 
 	expectedEvents := []Event{
 		UserMessage("Hi"),
-		StreamStarted(),
+		StreamStarted(sess.ID, "root"),
 		AgentChoiceReasoning("root", "Let me think about this..."),
 		AgentChoiceReasoning("root", " I should respond politely."),
 		AgentChoice("root", "Hello, how can I help you?"),
 		TokenUsage(10, 15, 25, 0, 0),
-		StreamStopped(),
+		StreamStopped(sess.ID, "root"),
 	}
 
 	require.Equal(t, expectedEvents, events)
@@ -259,13 +259,13 @@ func TestMixedContentAndReasoning(t *testing.T) {
 
 	expectedEvents := []Event{
 		UserMessage("Hi there"),
-		StreamStarted(),
+		StreamStarted(sess.ID, "root"),
 		AgentChoiceReasoning("root", "The user wants a greeting"),
 		AgentChoice("root", "Hello!"),
 		AgentChoiceReasoning("root", " I should be friendly"),
 		AgentChoice("root", " How can I help you today?"),
 		TokenUsage(15, 20, 35, 0, 0),
-		StreamStopped(),
+		StreamStopped(sess.ID, "root"),
 	}
 
 	require.Equal(t, expectedEvents, events)
@@ -318,66 +318,6 @@ func TestErrorEvent(t *testing.T) {
 	require.Contains(t, errorEvent.Error, "simulated error")
 }
 
-func TestRuntimeRunStream_TableDriven(t *testing.T) {
-	tests := []struct {
-		name           string
-		streamBuilder  func() *streamBuilder
-		userMessage    string
-		expectedEvents []Event
-	}{
-		{
-			name: "single_word_response",
-			streamBuilder: func() *streamBuilder {
-				return newStreamBuilder().AddContent("Yes").AddStopWithUsage(2, 1)
-			},
-			userMessage: "Confirm",
-			expectedEvents: []Event{
-				UserMessage("Confirm"),
-				StreamStarted(),
-				AgentChoice("root", "Yes"),
-				TokenUsage(2, 1, 3, 0, 0),
-				StreamStopped(),
-			},
-		},
-		{
-			name: "reasoning_only_response",
-			streamBuilder: func() *streamBuilder {
-				return newStreamBuilder().AddReasoning("Thinking...").AddStopWithUsage(5, 3)
-			},
-			userMessage: "Think about this",
-			expectedEvents: []Event{
-				UserMessage("Think about this"),
-				StreamStarted(),
-				AgentChoiceReasoning("root", "Thinking..."),
-				TokenUsage(5, 3, 8, 0, 0),
-				StreamStopped(),
-			},
-		},
-		{
-			name: "zero_token_response",
-			streamBuilder: func() *streamBuilder {
-				return newStreamBuilder().AddStopWithUsage(0, 0)
-			},
-			userMessage: "Empty",
-			expectedEvents: []Event{
-				UserMessage("Empty"),
-				StreamStarted(),
-				TokenUsage(0, 0, 0, 0, 0),
-				StreamStopped(),
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			stream := tt.streamBuilder().Build()
-			sess := session.New(session.WithUserMessage("", tt.userMessage))
-			events := runSession(t, sess, stream)
-			require.Equal(t, tt.expectedEvents, events)
-		})
-	}
-}
-
 func TestContextCancellation(t *testing.T) {
 	stream := newStreamBuilder().
 		AddContent("This should not complete").