How is it proven that Key Exchange implies OWFs?
Your question suggests the following tentative reduction to obtain a OWF from a secure KE: Given an input, interpret it as the private random coins of two simulated parties to the KE protocol. Based on their private randomness, there will be some sequence of public messages between them, and at some point they will stop and agree on a private shared key (with some high probability). Then the public transcript of their communications is the output of the OWF.

You are right to be worried about this approach in the case when Alice and Bob may not agree with probability 1. In fact, this construction is NOT necessarily a OWF in that case. Rather, it is a weaker primitive known as a "distributionally one-way function," which is, informally, a function for which it is hard to generate a uniformly random preimage of $f(x)$, when $x$ is selected at random. It is known (but non-obvious) that one can construct a true one-way function using a distributionally one-way function. See e.g. Exercise 17 on page 96 of Goldreich "Foundations of Cryptography: Basic Tools," which is available for limited viewing on Google preview.

To directly see why the above reduction does not suffice, we can use the common "make a stupid modification" technique. Namely, suppose there exists a secure KE protocol where Alice and Bob agree with probability $1 - \epsilon$. Consider a stupidly-modified KE protocol that works as follows: Alice and Bob first each flip 100 coins. For each of Alice, Bob individually: If all 100 coins are zero, their remaining randomness is interpreted as encoding a sequence of messages which they send to the other party, ignoring whatever the other one says; then they just output a random "shared key." (Otherwise they follow the original protocol.) Given any public transcript, it is possible (with negligible probability) that it was generated by an Alice and Bob who BOTH flipped 100 zeroes and happened to have their remaining randomness specify that exact sequence of messages.

This is still a secure KE protocol, where Alice and Bob will agree with probability at least $1 - \epsilon - 2^{-100}$. On the other hand, it utterly fails to be a OWF in the hoped-for manner, since an attacker who sees some public transcript can always trivially invert by choosing 100-zeroes randomness for each of Alice and Bob.