Instead of display: flex on the container, use display: inline-flex.
This switches the container from block-level (which takes the full width of its parent) to inline-level (which takes the width of its content).
This sizing behavior is similar to display: block vs. display: inline-block.
For a related problem and solution see:
- Make flex items take content width, not width of parent container