Fix output handling to match new format.