From 10a6a5fd248ed5be52ab6e925f52d299af3f4bf9 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Thu, 15 Aug 2024 01:19:34 +0100
Subject: [PATCH] Add logoutCurrentAccount()

---
 src/state/session/__tests__/session-test.ts | 97 +++++++++++++++++++++
 src/state/session/index.tsx                 | 26 +++++-
 src/state/session/reducer.ts                | 19 ++++
 src/state/session/types.ts                  |  3 +
 4 files changed, 144 insertions(+), 1 deletion(-)

diff --git a/src/state/session/__tests__/session-test.ts b/src/state/session/__tests__/session-test.ts
index 6f6a0ad15a..3e22c262cb 100644
--- a/src/state/session/__tests__/session-test.ts
+++ b/src/state/session/__tests__/session-test.ts
@@ -674,6 +674,103 @@ describe('session', () => {
     expect(state.currentAgentState.did).toBe(undefined)
   })
 
+  it('can log out of the current account', () => {
+    let state = getInitialState([])
+
+    const agent1 = new BskyAgent({service: 'https://alice.com'})
+    agent1.sessionManager.session = {
+      active: true,
+      did: 'alice-did',
+      handle: 'alice.test',
+      accessJwt: 'alice-access-jwt-1',
+      refreshJwt: 'alice-refresh-jwt-1',
+    }
+    state = run(state, [
+      {
+        type: 'switched-to-account',
+        newAgent: agent1,
+        newAccount: agentToSessionAccountOrThrow(agent1),
+      },
+    ])
+    expect(state.accounts.length).toBe(1)
+    expect(state.accounts[0].accessJwt).toBe('alice-access-jwt-1')
+    expect(state.accounts[0].refreshJwt).toBe('alice-refresh-jwt-1')
+    expect(state.currentAgentState.did).toBe('alice-did')
+
+    const agent2 = new BskyAgent({service: 'https://bob.com'})
+    agent2.sessionManager.session = {
+      active: true,
+      did: 'bob-did',
+      handle: 'bob.test',
+      accessJwt: 'bob-access-jwt-1',
+      refreshJwt: 'bob-refresh-jwt-1',
+    }
+    state = run(state, [
+      {
+        type: 'switched-to-account',
+        newAgent: agent2,
+        newAccount: agentToSessionAccountOrThrow(agent2),
+      },
+    ])
+    expect(state.accounts.length).toBe(2)
+    expect(state.accounts[0].accessJwt).toBe('bob-access-jwt-1')
+    expect(state.accounts[0].refreshJwt).toBe('bob-refresh-jwt-1')
+    expect(state.currentAgentState.did).toBe('bob-did')
+
+    state = run(state, [
+      {
+        type: 'logged-out-current-account',
+      },
+    ])
+    expect(state.accounts.length).toBe(2)
+    expect(state.accounts[0].accessJwt).toBe(undefined)
+    expect(state.accounts[0].refreshJwt).toBe(undefined)
+    expect(state.accounts[1].accessJwt).toBe('alice-access-jwt-1')
+    expect(state.accounts[1].refreshJwt).toBe('alice-refresh-jwt-1')
+    expect(state.currentAgentState.did).toBe(undefined)
+    expect(printState(state)).toMatchInlineSnapshot(`
+      {
+        "accounts": [
+          {
+            "accessJwt": undefined,
+            "active": true,
+            "did": "bob-did",
+            "email": undefined,
+            "emailAuthFactor": false,
+            "emailConfirmed": false,
+            "handle": "bob.test",
+            "pdsUrl": undefined,
+            "refreshJwt": undefined,
+            "service": "https://bob.com/",
+            "signupQueued": false,
+            "status": undefined,
+          },
+          {
+            "accessJwt": "alice-access-jwt-1",
+            "active": true,
+            "did": "alice-did",
+            "email": undefined,
+            "emailAuthFactor": false,
+            "emailConfirmed": false,
+            "handle": "alice.test",
+            "pdsUrl": undefined,
+            "refreshJwt": "alice-refresh-jwt-1",
+            "service": "https://alice.com/",
+            "signupQueued": false,
+            "status": undefined,
+          },
+        ],
+        "currentAgentState": {
+          "agent": {
+            "service": "https://public.api.bsky.app/",
+          },
+          "did": undefined,
+        },
+        "needsPersist": true,
+      }
+    `)
+  })
+
   it('updates stored account with refreshed tokens', () => {
     let state = getInitialState([])
 
diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 86919ca138..21fe7f75b9 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -35,6 +35,7 @@ const AgentContext = React.createContext<BskyAgent | null>(null)
 const ApiContext = React.createContext<SessionApiContext>({
   createAccount: async () => {},
   login: async () => {},
+  logoutCurrentAccount: async () => {},
   logoutEveryAccount: async () => {},
   resumeSession: async () => {},
   removeAccount: () => {},
@@ -115,6 +116,21 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     [onAgentSessionChange, cancelPendingTask],
   )
 
+  const logoutCurrentAccount = React.useCallback<
+    SessionApiContext['logoutEveryAccount']
+  >(
+    logContext => {
+      addSessionDebugLog({type: 'method:start', method: 'logout'})
+      cancelPendingTask()
+      dispatch({
+        type: 'logged-out-current-account',
+      })
+      logEvent('account:loggedOut', {logContext, scope: 'current'})
+      addSessionDebugLog({type: 'method:end', method: 'logout'})
+    },
+    [cancelPendingTask],
+  )
+
   const logoutEveryAccount = React.useCallback<
     SessionApiContext['logoutEveryAccount']
   >(
@@ -232,11 +248,19 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
     () => ({
       createAccount,
       login,
+      logoutCurrentAccount,
       logoutEveryAccount,
       resumeSession,
       removeAccount,
     }),
-    [createAccount, login, logoutEveryAccount, resumeSession, removeAccount],
+    [
+      createAccount,
+      login,
+      logoutCurrentAccount,
+      logoutEveryAccount,
+      resumeSession,
+      removeAccount,
+    ],
   )
 
   // @ts-ignore
diff --git a/src/state/session/reducer.ts b/src/state/session/reducer.ts
index d21e228db5..22ba47162a 100644
--- a/src/state/session/reducer.ts
+++ b/src/state/session/reducer.ts
@@ -41,6 +41,9 @@ export type Action =
       type: 'removed-account'
       accountDid: string
     }
+  | {
+      type: 'logged-out-current-account'
+    }
   | {
       type: 'logged-out-every-account'
     }
@@ -138,6 +141,22 @@ let reducer = (state: State, action: Action): State => {
         needsPersist: true,
       }
     }
+    case 'logged-out-current-account': {
+      const {currentAgentState} = state
+      return {
+        accounts: state.accounts.map(a =>
+          a.did === currentAgentState.did
+            ? {
+                ...a,
+                refreshJwt: undefined,
+                accessJwt: undefined,
+              }
+            : a,
+        ),
+        currentAgentState: createPublicAgentState(),
+        needsPersist: true,
+      }
+    }
     case 'logged-out-every-account': {
       return {
         accounts: state.accounts.map(a => ({
diff --git a/src/state/session/types.ts b/src/state/session/types.ts
index e2528a7c4b..d32259de9d 100644
--- a/src/state/session/types.ts
+++ b/src/state/session/types.ts
@@ -29,6 +29,9 @@ export type SessionApiContext = {
     },
     logContext: LogEvents['account:loggedIn']['logContext'],
   ) => Promise<void>
+  logoutCurrentAccount: (
+    logContext: LogEvents['account:loggedOut']['logContext'],
+  ) => void
   logoutEveryAccount: (
     logContext: LogEvents['account:loggedOut']['logContext'],
   ) => void