Enhance Tests with SafeWriter Closure for Async Writes

- Added calls to CloseSafeWriter in multiple test cases to ensure all asynchronous writes are completed before verifying output.
- Improved test reliability by ensuring that the SafeWriter's buffer is properly flushed, preventing potential race conditions in output verification.

agent/context/jsapi_output_test.go

@@ -771,6 +771,9 @@ func TestJsValueEndBlock(t *testing.T) {
 	}
 	assert.Equal(t, true, result["success"], "EndBlock should work correctly")
 
+	// Close SafeWriter to wait for all async writes to complete
+	cxt.CloseSafeWriter()
+
 	// Verify that block_end event was sent
 	output := mockWriter.buffer.String()
 	assert.Contains(t, output, "block_end", "Output should contain block_end event")
@@ -825,6 +828,9 @@ func TestJsValueSendStream(t *testing.T) {
 	}
 	assert.Equal(t, true, result["success"], "SendStream should work correctly")
 
+	// Close SafeWriter to wait for all async writes to complete
+	cxt.CloseSafeWriter()
+
 	// Verify message_start was sent but NOT message_end
 	output := mockWriter.buffer.String()
 	assert.Contains(t, output, "message_start", "Output should contain message_start event")
@@ -873,6 +879,9 @@ func TestJsValueSendStreamWithBlockID(t *testing.T) {
 	}
 	assert.Equal(t, true, result["success"], "SendStream with blockId should succeed")
 
+	// Close SafeWriter to wait for all async writes to complete
+	cxt.CloseSafeWriter()
+
 	// Verify block_start was also sent
 	output := mockWriter.buffer.String()
 	assert.Contains(t, output, "block_start", "Output should contain block_start event")
@@ -922,6 +931,9 @@ func TestJsValueEnd(t *testing.T) {
 	}
 	assert.Equal(t, true, result["success"], "End should work correctly")
 
+	// Close SafeWriter to wait for all async writes to complete
+	cxt.CloseSafeWriter()
+
 	// Verify message_end was sent
 	output := mockWriter.buffer.String()
 	assert.Contains(t, output, "message_end", "Output should contain message_end event after End()")
@@ -971,6 +983,9 @@ func TestJsValueEndWithFinalContent(t *testing.T) {
 	}
 	assert.Equal(t, true, result["success"], "End with final content should work correctly")
 
+	// Close SafeWriter to wait for all async writes to complete
+	cxt.CloseSafeWriter()
+
 	// Verify message_end was sent
 	output := mockWriter.buffer.String()
 	assert.Contains(t, output, "message_end", "Output should contain message_end event")

agent/context/message_events_test.go

@@ -39,6 +39,10 @@ func TestMessageLifecycleEvents(t *testing.T) {
 	// Flush to ensure all messages are written
 	ctx.Flush()
 
+	// Close SafeWriter to wait for all async writes to complete
+	// SafeWriter uses a channel-based queue, so we must close it before reading buffer
+	ctx.CloseSafeWriter()
+
 	// Parse output to find events
 	output := buf.String()
 	t.Logf("Output:\n%s", output)