fix(ts): close the model invoke span when the consumer breaks the stream (#3898)

strands-ts/src/agent/__tests__/agent.tracer.test.node.ts

@@ -409,6 +409,20 @@ describe('Agent tracer integration', () => {
   })
 
   describe('model invoke span lifecycle', () => {
+    it('ends the model invoke span once when the consumer breaks mid-model-response', async () => {
+      const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' })
+      const agent = new Agent({ model })
+      const tracer = getLatestTracer()
+
+      for await (const event of agent.stream('Hi')) {
+        if (event.type === 'modelStreamUpdateEvent') break
+      }
+
+      expect(tracer.startModelInvokeSpan).toHaveBeenCalledTimes(1)
+      expect(tracer.endModelInvokeSpan).toHaveBeenCalledTimes(1)
+      expect(tracer.endModelInvokeSpan).toHaveBeenCalledWith({ mock: 'modelSpan' })
+    })
+
     it('starts and ends model span on successful model call', async () => {
       const model = new MockMessageModel().addTurn({ type: 'textBlock', text: 'Hello' })
       const agent = new Agent({ model })
@@ -435,6 +449,7 @@ describe('Agent tracer integration', () => {
 
       await expect(agent.invoke('Hi')).rejects.toThrow()
 
+      expect(tracer.endModelInvokeSpan).toHaveBeenCalledTimes(1)
       expect(tracer.endModelInvokeSpan).toHaveBeenCalledWith(
         { mock: 'modelSpan' },
         expect.objectContaining({ error: expect.any(Error) })

strands-ts/src/agent/agent.ts

@@ -2321,6 +2321,7 @@ export class Agent implements LocalAgent, InvokableAgent {
           ...(ctx.systemPrompt !== undefined && { systemPrompt: ctx.systemPrompt }),
         })
 
+        let modelSpanEnded = false
         try {
           // Wrap the snapshot into a StateStore for the model provider, which expects
           // get/set methods.
@@ -2349,11 +2350,19 @@ export class Agent implements LocalAgent, InvokableAgent {
             ...(usage && { usage }),
             ...(metrics && { metrics }),
           })
+          modelSpanEnded = true
 
           return { result: iterResult.value }
         } catch (error) {
           self._tracer.endModelInvokeSpan(modelSpan, { error: normalizeError(error) })
+          modelSpanEnded = true
           throw error
+        } finally {
+          // A consumer break closes this generator via .return(): finally runs but catch does not,
+          // so neither end call above fires and the span would stay open.
+          if (!modelSpanEnded) {
+            self._tracer.endModelInvokeSpan(modelSpan)
+          }
         }
       }
     )