Space and time complexity of $L = \{a^nb^{n^2} \mid n≥1\}$
Your confusion stems from a double use of $n$. Let's write $L$ differently:
$$
L = \{ a^i b^{n-i} : n-i = i^2, i \geq 1 \}.
$$
Here $n$ is the input length. It should be clear that the space complexity is $O(n)$.

Your first estimate just translates to a space complexity of $O(i^2) = O(n)$, since $n = i^2 + i$.