def test_command_output_continuation():
    """Test that we can continue to get output from a long-running command.

    This test has been modified to be more robust against timing issues.
    """
    session = BashSession(work_dir=os.getcwd(), no_change_timeout_seconds=1)
    session.initialize()

    # Start a command that produces output slowly but with longer sleep time
    # to ensure we hit the timeout
    obs = session.execute(CmdRunAction('for i in {1..5}; do echo $i; sleep 2; done'))
    logger.info(obs, extra={'msg_type': 'OBSERVATION'})

    # Check if the command completed immediately or timed out
    if session.prev_status == BashCommandStatus.COMPLETED:
        # If the command completed immediately, verify we got all the output
        logger.info('Command completed immediately', extra={'msg_type': 'TEST_INFO'})
        assert '1' in obs.content
        assert '2' in obs.content
        assert '3' in obs.content
        assert '4' in obs.content
        assert '5' in obs.content
        assert '[The command completed with exit code 0.]' in obs.metadata.suffix
    else:
        # If the command timed out, verify we got the timeout message
        assert session.prev_status == BashCommandStatus.NO_CHANGE_TIMEOUT
        assert '1' in obs.content
        assert '[The command has no new output after 1 seconds.' in obs.metadata.suffix

        # Continue getting output until we see all numbers
        numbers_seen = set()
        for i in range(1, 6):
            if str(i) in obs.content:
                numbers_seen.add(i)

        # We need to see numbers 2-5 and then the command completion
        while (
            len(numbers_seen) < 5 or session.prev_status != BashCommandStatus.COMPLETED
        ):
            obs = session.execute(CmdRunAction('', is_input=True))
            logger.info(obs, extra={'msg_type': 'OBSERVATION'})

            # Check for numbers in the output
            for i in range(1, 6):
                if str(i) in obs.content and i not in numbers_seen:
                    numbers_seen.add(i)
                    logger.info(
                        f'Found number {i} in output', extra={'msg_type': 'TEST_INFO'}
                    )

            # Check if the command has completed
            if session.prev_status == BashCommandStatus.COMPLETED:
                assert (
                    '[The command completed with exit code 0.]' in obs.metadata.suffix
                )
                break
            else:
                assert (
                    '[The command has no new output after 1 seconds.'
                    in obs.metadata.suffix
                )
                assert session.prev_status == BashCommandStatus.NO_CHANGE_TIMEOUT

        # Verify we've seen all numbers
        assert numbers_seen == {1, 2, 3, 4, 5}, (
            f'Expected to see numbers 1-5, but saw {numbers_seen}'
        )

        # Verify the command completed
        assert session.prev_status == BashCommandStatus.COMPLETED

    session.close()
Test that we can continue to get output from a long-running command.

    This test has been modified to be more robust against timing issues.
    
tests/unit/test_bash_session.py