Refactor to be purely event-driven in order to support direct Rpc use

Author: stephentoub

dotnet/src/CopilotTelemetry.cs

@@ -54,7 +54,7 @@ public partial class CopilotSession : IAsyncDisposable
     private event SessionEventHandler? _eventHandlers;
     private readonly Dictionary<string, AIFunction> _toolHandlers = new();
     private readonly JsonRpc _rpc;
-    private readonly CopilotTelemetry.AgentTurnTracker? _turnTracker;
+    private readonly CopilotTelemetry.AgentTurnTracker? _telemetryTracker;
     private volatile PermissionRequestHandler? _permissionHandler;
     private volatile UserInputHandler? _userInputHandler;
     private SessionHooks? _hooks;
@@ -82,12 +82,12 @@ public partial class CopilotSession : IAsyncDisposable
     /// </value>
     public string? WorkspacePath { get; internal set; }
 
-    internal string TelemetryProviderName => _turnTracker?.ProviderName ?? OpenTelemetryConsts.DefaultProviderName;
-    internal string? TelemetryServerAddress => _turnTracker?.ServerAddress;
-    internal int? TelemetryServerPort => _turnTracker?.ServerPort;
-    internal ActivityContext TelemetryActivityContext => _turnTracker?.GetActivityContext() ?? default;
+    internal string TelemetryProviderName => _telemetryTracker?.ProviderName ?? OpenTelemetryConsts.DefaultProviderName;
+    internal string? TelemetryServerAddress => _telemetryTracker?.ServerAddress;
+    internal int? TelemetryServerPort => _telemetryTracker?.ServerPort;
+    internal ActivityContext TelemetryActivityContext => _telemetryTracker?.GetActivityContext() ?? default;
     internal ActivityContext GetTelemetryToolCallParentContext(string toolCallId) =>
-        _turnTracker?.GetToolCallParentContext(toolCallId) ?? default;
+        _telemetryTracker?.GetToolCallParentContext(toolCallId) ?? default;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="CopilotSession"/> class.
@@ -119,7 +119,7 @@ internal CopilotSession(
     {
         SessionId = sessionId;
         _rpc = rpc;
-        _turnTracker = telemetry is not null ? new CopilotTelemetry.AgentTurnTracker(telemetry, sessionId, model, provider, systemMessage, tools, streaming, agentName, agentDescription) : null;
+        _telemetryTracker = telemetry is not null ? new CopilotTelemetry.AgentTurnTracker(telemetry, sessionId, model, provider, systemMessage, tools, streaming, agentName, agentDescription) : null;
         WorkspacePath = workspacePath;
     }
 
@@ -156,27 +156,18 @@ private Task<T> InvokeRpcAsync<T>(string method, object?[]? args, CancellationTo
     /// </example>
     public async Task<string> SendAsync(MessageOptions options, CancellationToken cancellationToken = default)
     {
-        _turnTracker?.BeginSend(options.Prompt);
-        try
+        var request = new SendMessageRequest
         {
-            var request = new SendMessageRequest
-            {
-                SessionId = SessionId,
-                Prompt = options.Prompt,
-                Attachments = options.Attachments,
-                Mode = options.Mode
-            };
+            SessionId = SessionId,
+            Prompt = options.Prompt,
+            Attachments = options.Attachments,
+            Mode = options.Mode
+        };
 
-            var response = await InvokeRpcAsync<SendMessageResponse>(
-                "session.send", [request], cancellationToken);
+        var response = await InvokeRpcAsync<SendMessageResponse>(
+            "session.send", [request], cancellationToken);
 
-            return response.MessageId;
-        }
-        catch (Exception ex) when (_turnTracker is { } tracker)
-        {
-            tracker.CompleteTurnWithError(ex);
-            throw;
-        }
+        return response.MessageId;
     }
 
     /// <summary>
@@ -238,28 +229,18 @@ void Handler(SessionEvent evt)
 
         await SendAsync(options, cancellationToken);
 
-        try
-        {
-            using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-            cts.CancelAfter(effectiveTimeout);
-
-            using var registration = cts.Token.Register(() =>
-            {
-                if (cancellationToken.IsCancellationRequested)
-                    tcs.TrySetCanceled(cancellationToken);
-                else
-                    tcs.TrySetException(new TimeoutException($"SendAndWaitAsync timed out after {effectiveTimeout}"));
-            });
+        using var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        cts.CancelAfter(effectiveTimeout);
 
-            return await tcs.Task;
-        }
-        catch (Exception ex) when (_turnTracker is { } tracker)
+        using var registration = cts.Token.Register(() =>
         {
-            // If timeout/cancellation occurs before DispatchEvent handles the turn-ending event,
-            // complete the telemetry span with the error (idempotent if already completed).
-            tracker.CompleteTurnWithError(ex);
-            throw;
-        }
+            if (cancellationToken.IsCancellationRequested)
+                tcs.TrySetCanceled(cancellationToken);
+            else
+                tcs.TrySetException(new TimeoutException($"SendAndWaitAsync timed out after {effectiveTimeout}"));
+        });
+
+        return await tcs.Task;
     }
 
     /// <summary>
@@ -309,7 +290,7 @@ public IDisposable On(SessionEventHandler handler)
     /// </remarks>
     internal void DispatchEvent(SessionEvent sessionEvent)
     {
-        _turnTracker?.ProcessEvent(sessionEvent);
+        _telemetryTracker?.ProcessEvent(sessionEvent);
 
         // Reading the field once gives us a snapshot; delegates are immutable.
         _eventHandlers?.Invoke(sessionEvent);
@@ -620,7 +601,7 @@ await InvokeRpcAsync<object>(
 
         _eventHandlers = null;
         _toolHandlers.Clear();
-        _turnTracker?.CompleteOnDispose();
+        _telemetryTracker?.CompleteOnDispose();
 
         _permissionHandler = null;
     }

dotnet/src/Session.cs

@@ -1354,7 +1354,7 @@ func (c *Client) executeToolCall(
 
 	if c.telemetry != nil {
 		toolSpanCtx := context.Background()
-		if session.turnTracker != nil {
+		if session.telemetryTracker != nil {
 			if tCtx := session.getToolCallParentContext(toolCallID); tCtx != nil {
 				toolSpanCtx = tCtx
 			}
@@ -1380,10 +1380,10 @@ func (c *Client) executeToolCall(
 			providerName := otelDefaultProviderName
 			var serverAddress string
 			var serverPort int
-			if session.turnTracker != nil {
-				providerName = session.turnTracker.getProviderName()
-				serverAddress = session.turnTracker.getServerAddress()
-				serverPort = session.turnTracker.getServerPort()
+			if session.telemetryTracker != nil {
+				providerName = session.telemetryTracker.getProviderName()
+				serverAddress = session.telemetryTracker.getServerAddress()
+				serverPort = session.telemetryTracker.getServerPort()
 			}
 			ctx := spanCtx
 			if ctx == nil {

go/client.go

@@ -573,8 +573,6 @@ type agentTurnTracker struct {
 	turnID             *string
 	turnInteractionID  *string
 
-	// Stashed user prompt for the first chat turn
-	pendingUserPrompt string
 }
 
 func newAgentTurnTracker(telemetry *copilotTelemetry, sessionID string, model string, provider *ProviderConfig, systemMessage *SystemMessageConfig, tools []Tool, streaming bool, agentName string, agentDescription string) *agentTurnTracker {
@@ -641,45 +639,34 @@ func (t *agentTurnTracker) completeOnDispose() {
 	}
 }
 
-// beginSend is called at the start of Send() to start a span and record the user message.
-func (t *agentTurnTracker) beginSend(ctx context.Context, prompt string) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	if t.agentSpan == nil {
-		spanCtx, span := t.telemetry.startInvokeAgentSpan(
-			ctx,
-			t.sessionID,
-			t.requestModel,
-			t.providerName,
-			t.serverAddress,
-			t.serverPort,
-			t.agentName,
-			t.agentDescription,
-		)
-		t.agentSpan = span
-		t.agentSpanCtx = spanCtx
-		t.agentStartTime = time.Now()
-		t.agentInputMsgs = nil
-		t.agentOutputMsgs = nil
-	}
-
-	// Agent-level input = what the caller sent (all user prompts).
-	if prompt != "" {
-		t.agentInputMsgs = append(t.agentInputMsgs, otelMsg{
-			Role:  "user",
-			Parts: []otelPart{{Type: "text", Content: prompt}},
-		})
-	}
-
-	// Stash user prompt for the first chat turn's input messages.
-	t.pendingUserPrompt = prompt
-}
-
 // processEvent handles telemetry enrichment for dispatched events.
 func (t *agentTurnTracker) processEvent(event SessionEvent) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
+	// A user.message event starts a new invoke_agent span (if not already
+	// active) and records the user prompt.
+	if event.Type == UserMessage {
+		prompt := ""
+		if event.Data.Content != nil {
+			prompt = *event.Data.Content
+		}
+		t.ensureAgentSpan()
+
+		if prompt != "" {
+			t.agentInputMsgs = append(t.agentInputMsgs, otelMsg{
+				Role:  "user",
+				Parts: []otelPart{{Type: "text", Content: prompt}},
+			})
+			t.turnInputMsgs = append(t.turnInputMsgs, otelMsg{
+				Role:  "user",
+				Parts: []otelPart{{Type: "text", Content: prompt}},
+			})
+		}
+
+		return
+	}
+
 	// Route subagent events by parentToolCallId.
 	parentToolCallID := getParentToolCallID(event)
 	if parentToolCallID != "" {
@@ -968,23 +955,43 @@ func (t *agentTurnTracker) processEvent(event SessionEvent) {
 	}
 }
 
-// completeTurnWithError completes the current turn with an error.
-func (t *agentTurnTracker) completeTurnWithError(err error) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	t.completeChatTurnLocked(err)
-	t.completeAgentTurnLocked(err)
-}
-
 // ============================================================================
 // Chat turn lifecycle
 // ============================================================================
 
 // beginChatTurnLocked starts a new chat child span for an LLM turn. Caller must hold mu.
+// ensureAgentSpan ensures the invoke_agent span exists, creating it on demand
+// if needed. Called from both the user.message handler and beginChatTurnLocked
+// so that RPC-initiated turns (no user.message) still get an agent span.
+// Caller must hold mu.
+func (t *agentTurnTracker) ensureAgentSpan() {
+	if t.agentSpan == nil {
+		spanCtx, span := t.telemetry.startInvokeAgentSpan(
+			context.Background(),
+			t.sessionID,
+			t.requestModel,
+			t.providerName,
+			t.serverAddress,
+			t.serverPort,
+			t.agentName,
+			t.agentDescription,
+		)
+		t.agentSpan = span
+		t.agentSpanCtx = spanCtx
+		t.agentStartTime = time.Now()
+		t.agentInputMsgs = nil
+		t.agentOutputMsgs = nil
+	}
+}
+
 func (t *agentTurnTracker) beginChatTurnLocked() {
 	// If there's already an active turn, complete it first.
 	t.completeChatTurnLocked(nil)
 
+	// Ensure the parent agent span exists — covers RPC-initiated turns
+	// where no user.message event preceded the assistant.turn_start.
+	t.ensureAgentSpan()
+
 	t.turnResponseModel = ""
 	t.turnResponseID = ""
 	t.turnInputTokens = 0
@@ -993,7 +1000,6 @@ func (t *agentTurnTracker) beginChatTurnLocked() {
 	t.turnCacheCreationTokens = 0
 	t.firstOutputChunkRecorded = false
 	t.lastOutputChunkTime = time.Time{}
-	t.turnInputMsgs = nil
 	t.turnOutputMsgs = nil
 	t.turnCost = nil
 	t.turnServerDuration = nil
@@ -1002,15 +1008,6 @@ func (t *agentTurnTracker) beginChatTurnLocked() {
 	t.turnID = nil
 	t.turnInteractionID = nil
 
-	// Add stashed user prompt as input message for the first turn.
-	if t.pendingUserPrompt != "" {
-		t.turnInputMsgs = append(t.turnInputMsgs, otelMsg{
-			Role:  "user",
-			Parts: []otelPart{{Type: "text", Content: t.pendingUserPrompt}},
-		})
-		t.pendingUserPrompt = ""
-	}
-
 	parentCtx := t.agentSpanCtx
 	if parentCtx == nil {
 		parentCtx = context.Background()
@@ -1222,7 +1219,6 @@ func (t *agentTurnTracker) completeAgentTurnLocked(err error) {
 	t.agentSpan = nil
 	t.agentSpanCtx = nil
 	t.agentStartTime = time.Time{}
-	t.pendingUserPrompt = ""
 	t.agentInputMsgs = nil
 	t.agentOutputMsgs = nil