Let's start by analyzing the family's problem with the imposition of a tax. Assuming a CRRA utility function: \begin{equation} U = \int_0^\infty e^{-(\rho-n)t} \cdot \left[ \frac{c^{1-\theta}-1}{1-\theta} \right] \operatorname{dt} \hspace{5mm} (1) \end{equation} Agents maximize $(1)$ subject to the assets' restriction: \begin{equation} \dot{a} = w - c + (1-\tau)r \cdot a - n \cdot a \end{equation} The Hamiltonian is given by \begin{equation} \mathcal{H} = e^{-(\rho-n)t} \cdot \left[ \frac{c^{1-\theta}-1}{1-\theta} \right] + \lambda \cdot [ w - c + (1-\tau)r \cdot a - n \cdot a ] \end{equation} The maximum principle conditions are: $$\begin{cases} \frac{\partial \mathcal{H}}{\partial c} = 0 \implies e^{-(\rho-n)t} \cdot c^{-\theta} = \lambda \\ \dot{a} = \frac{\partial \mathcal{H}}{\partial \lambda} \implies \dot{a} = w - c + (1-\tau)r \cdot a - n \cdot a \\ \dot{\lambda} = -\frac{\partial \mathcal{H}}{\partial a} \implies \dot{\lambda} = -\lambda \cdot [(1-\tau)r - n] \\ \lim_{t \to \infty} [\lambda \cdot a] = 0 \end{cases} $$ Is my reasoning correct so far? Or should I have included a transfer from the government to the agents somewhere? Furthermore, whether the government spending is productive or not would also change how we deal with the tax in our AK model, right?
UPDATE: Thinking a little bit more, I've reached the conclusion that the problem should be reestated as:
Maximize $(1)$ subject to assets' restriction $$ \dot{a} = w - c + (1-\tau)r \cdot a - n \cdot a + v $$ where $v = \frac{V}{L} = \frac{\tau \cdot r \cdot a}{L}$ is a per capita lump-sum transfer from the government to the families. Therefore, the Hamiltonian is now given by: \begin{equation} \mathcal{H} = e^{-(\rho-n)t} \cdot \left[ \frac{c^{1-\theta}-1}{1-\theta} \right] + \lambda \cdot [ w - c + (1-\tau)r \cdot a - n \cdot a + v] \end{equation} The maximum principle conditions are: $$\begin{cases} \frac{\partial \mathcal{H}}{\partial c} = 0 \implies e^{-(\rho-n)t} \cdot c^{-\theta} = \lambda \\ \dot{a} = \frac{\partial \mathcal{H}}{\partial \lambda} \implies \dot{a} = w - c + (1-\tau)r \cdot a - n \cdot a + v \\ \dot{\lambda} = -\frac{\partial \mathcal{H}}{\partial a} \implies \dot{\lambda} = -\lambda \cdot [(1-\tau)r - n] \\ \lim_{t \to \infty} [\lambda \cdot a] = 0 \end{cases} $$ Using the equations above, we can find that: $$ \frac{\dot{c}}{c} = \frac{1}{\theta} \cdot [(1-\tau)r - \rho] $$ Considering now a production function in per capita form given by $$ y = Ak, A > 0 $$ we have that the firms profits are (note that I didn't include the tax in $r$ here, per the reasoning that since $r$ enters profits as an expense, in this case it's not taxable, because it's not an asset income as in the representative family's problem) $$ \pi = Ak - \underbrace{w}_{=\frac{\partial y}{\partial L} = 0} - \delta \cdot k - r \cdot k $$ which means the first-order condition of profit maximization will be given by $$ \frac{\partial \pi}{\partial k} = 0 \implies A = f'(k) = \delta - r $$ or $$ r = A - \delta $$ Therefore, $$ \boxed{ \frac{\dot{c}}{c} = \frac{1}{\theta} \cdot [(1-\tau) \cdot (A-\delta) - \rho]} $$ Now, in equilibrium, $a = k$. If we substitute this and $v = \tau \cdot r \cdot k$ in the assets's restriction, we have: \begin{align*} \dot{k} &= -c + (1-\tau) \cdot r \cdot k - n \cdot k + \tau \cdot r \cdot k \\ &= -c - (r+n) \cdot k \end{align*} and, therefore, $$ \boxed{ \dot{k} = -c + (A-\delta-n) \cdot k} $$ Is what I did correct?