diff --git a/contracts/Vault.sol b/contracts/Vault.sol
index dfa7de0..0df9f2b 100644
--- a/contracts/Vault.sol
+++ b/contracts/Vault.sol
@@ -12,8 +12,8 @@ contract Vault is VaultBase {
     Recipient recipient
   ) public view returns (uint128) {
     Controller controller = Controller.wrap(msg.sender);
-    Account memory account = _getAccount(controller, fund, recipient);
-    return account.balance.available + account.balance.designated;
+    Balance memory balance = _getBalance(controller, fund, recipient);
+    return balance.available + balance.designated;
   }
 
   function getDesignatedBalance(
@@ -21,8 +21,8 @@ contract Vault is VaultBase {
     Recipient recipient
   ) public view returns (uint128) {
     Controller controller = Controller.wrap(msg.sender);
-    Account memory account = _getAccount(controller, fund, recipient);
-    return account.balance.designated;
+    Balance memory balance = _getBalance(controller, fund, recipient);
+    return balance.designated;
   }
 
   function getLock(Fund fund) public view returns (Lock memory) {
diff --git a/contracts/vault/Timestamps.sol b/contracts/vault/Timestamps.sol
index 5cc4058..81e8e4b 100644
--- a/contracts/vault/Timestamps.sol
+++ b/contracts/vault/Timestamps.sol
@@ -30,17 +30,6 @@ library Timestamps {
     return Timestamp.wrap(uint64(block.timestamp));
   }
 
-  function earliest(
-    Timestamp a,
-    Timestamp b
-  ) internal pure returns (Timestamp) {
-    if (a <= b) {
-      return a;
-    } else {
-      return b;
-    }
-  }
-
   function until(
     Timestamp start,
     Timestamp end
diff --git a/contracts/vault/VaultBase.sol b/contracts/vault/VaultBase.sol
index b4e3095..090ce09 100644
--- a/contracts/vault/VaultBase.sol
+++ b/contracts/vault/VaultBase.sol
@@ -35,19 +35,19 @@ abstract contract VaultBase {
     return _locks[controller][fund];
   }
 
-  function _getAccount(
+  function _getBalance(
     Controller controller,
     Fund fund,
     Recipient recipient
-  ) internal view returns (Account memory) {
+  ) internal view returns (Balance memory) {
     Account memory account = _accounts[controller][fund][recipient];
     Lock memory lock = _locks[controller][fund];
-    Timestamp timestamp = Timestamps.earliest(
-      Timestamps.currentTime(),
-      lock.expiry
-    );
-    account.update(timestamp);
-    return account;
+    if (lock.isLocked()) {
+      account.update(Timestamps.currentTime());
+    } else {
+      account.update(lock.expiry);
+    }
+    return account.balance;
   }
 
   function _lock(
@@ -105,11 +105,10 @@ abstract contract VaultBase {
     require(lock.isLocked(), LockRequired());
 
     Account memory account = _accounts[controller][fund][recipient];
-
     require(amount <= account.balance.available, InsufficientBalance());
+
     account.balance.available -= amount;
     account.balance.designated += amount;
-
     _checkAccountInvariant(account, lock);
 
     _accounts[controller][fund][recipient] = account;
@@ -125,15 +124,15 @@ abstract contract VaultBase {
     Lock memory lock = _locks[controller][fund];
     require(lock.isLocked(), LockRequired());
 
-    Account memory sender = _getAccount(controller, fund, from);
+    Account memory sender = _accounts[controller][fund][from];
     require(amount <= sender.balance.available, InsufficientBalance());
+
     sender.balance.available -= amount;
     _checkAccountInvariant(sender, lock);
+
     _accounts[controller][fund][from] = sender;
 
-    Account memory receiver = _getAccount(controller, fund, to);
-    receiver.balance.available += amount;
-    _accounts[controller][fund][to] = receiver;
+    _accounts[controller][fund][to].balance.available += amount;
   }
 
   function _flow(
@@ -164,7 +163,7 @@ abstract contract VaultBase {
     Lock storage lock = _locks[controller][fund];
     require(lock.isLocked(), LockRequired());
 
-    Account memory account = _getAccount(controller, fund, recipient);
+    Account memory account = _accounts[controller][fund][recipient];
     require(
       account.flow.incoming == account.flow.outgoing,
       CannotBurnFlowingTokens()
@@ -187,7 +186,8 @@ abstract contract VaultBase {
     Lock memory lock = _locks[controller][fund];
     require(!lock.isLocked(), Locked());
 
-    Account memory account = _getAccount(controller, fund, recipient);
+    Account memory account = _accounts[controller][fund][recipient];
+    account.update(lock.expiry);
     uint128 amount = account.balance.available + account.balance.designated;
 
     lock.value -= amount;
diff --git a/test/Vault.tests.js b/test/Vault.tests.js
index f874a18..104e4da 100644
--- a/test/Vault.tests.js
+++ b/test/Vault.tests.js
@@ -64,6 +64,13 @@ describe("Vault", function () {
       ).to.be.revertedWith("LockRequired")
     })
 
+    it("does not have any balance", async function() {
+      const balance = await vault.getBalance(fund, account.address)
+      const designated = await vault.getDesignatedBalance(fund, account.address)
+      expect(balance).to.equal(0)
+      expect(designated).to.equal(0)
+    })
+
     it("does not allow designating tokens", async function () {
       await expect(
         vault.designate(fund, account.address, 0)
@@ -295,7 +302,7 @@ describe("Vault", function () {
       it("can transfer to self", async function () {
         await setAutomine(true)
         await vault.transfer(fund, address1, address1, amount)
-        expect(await vault.getBalance(fund, address1)).to.equal(amount);
+        expect(await vault.getBalance(fund, address1)).to.equal(amount)
       })
 
       it("does not transfer more than the balance", async function () {
@@ -635,6 +642,23 @@ describe("Vault", function () {
         expect(balance2).to.equal(total)
       })
 
+      it("allows flowing tokens to be withdrawn", async function () {
+        await vault.flow(fund, account.address, account2.address, 2)
+        await mine()
+        const start = await currentTime()
+        const total = (expiry - start) * 2
+        await advanceTimeTo(expiry + 10)
+        balance1Before = await token.balanceOf(account.address)
+        balance2Before = await token.balanceOf(account2.address)
+        await vault.withdraw(fund, account.address)
+        await vault.withdraw(fund, account2.address)
+        await mine()
+        balance1After = await token.balanceOf(account.address)
+        balance2After = await token.balanceOf(account2.address)
+        expect(balance1After - balance1Before).to.equal(deposit - total)
+        expect(balance2After - balance2Before).to.equal(total)
+      })
+
       it("does not allow new flows to start", async function () {
         await setAutomine(true)
         await expire()