Agent refactoring fix