The proof you give is fine as a proof for modules, but it is not a valid category theoretical argument, because you are ignoring the role of the codomain in the definition of an arrow.
Recall that a category $\mathcal{C}$ is a collection $\mathbf{Ob}(\mathcal{C})$ of “objects”, a collection $\mathbf{Arr}(\mathcal{C})$ of “arrows”, a function $\mathrm{domain}\colon\mathbf{Arr}(\mathcal{C})\to\mathbf{Ob}(\mathcal{C})$, the “domain of the arrow”, a function $\mathrm{codomain}\colon\mathbf{Arr}(\mathcal{C})\to\mathbf{Ob}(\mathcal{C})$, the “codomain of the arrow”, an associative partial operation $\circ$ on the arrows, where $g\circ f$ is defined if $\mathrm{codomain}(f)=\mathrm{domain}(g)$, with $\mathrm{domain}(g\circ f) = \mathrm{domain}(f)$ and $\mathrm{codomain}(g\circ f) = \mathrm{codomain}(g)$, with “identity functions” associated to each object which act like one-sided identities to $\circ$. Given objects $A$ and $B$, we let $\mathcal{C}(A,B)$ denote the collection of all arrows $r$ with the property that $\mathrm{domain}(r)=A$ and $\mathrm{codomain}(r)=B$. (So $\circ$, “composition”, restricts to a union of functions of the form $\mathcal{C}(B,C)\times\mathcal{C}(A,B)\to\mathcal{C}(A,C)$ )
In a category, for two arrows to be equal they must have the same domain and the same codomain. In concrete categories (categories where the objects are sets and the arrows are functions between sets) you have to keep track of both the domain or the codomain. You are in the situation of Bourbaki, where a function is not merely a set of ordered pairs, but rather a function is really a triple: $(A,B,f)$, where $f$ is the set of pairs, $A$ is the domain, and $B$ is the specified codomain. (See this answer, for instance, for a discussion of the two different ways to define a function).
The problem in your argument, purely from the Category Theory point of view, is that you using $g$ to denote two different arrows: one is an arrow with underlying function $g$, domain $X$ and codomain $K$; and the other is an arrow with underlying function $g$, domain $X$ and codomain $A$. If $A\neq K$, then these are not the same arrow.
Instead, you would define a new arrow $\mathfrak{g}\colon X\to K$, defined by letting $\mathfrak{g}(x)=g(x)$ for each $x\in X$; this arrow has the property that $i\mathfrak{g}=g$.
The uniqueness will now follow by the following (categorical) argument: if $h\colon X\to K$ with the property that $ih = g$, then $i\mathfrak{g} = g =ih$. Now, since $i$ is a monomorphism, it can be cancelled on the left, and thus from $i\mathfrak{g}=ih$ you conclude that $\mathfrak{g}=h$, giving uniqueness.
These issues are mainly technical, and not conceptual. But if you are trying to figure out exactly what assumptions you need in your category to have the argument go through (as you seem to be doing), then you need to be careful with the technical issues.