[Session] Logging (#4476)

* Add session logging (console.log)

* Hook it up for real

* Send type separately

3 files changed, 178 insertions, 6 deletions

diff --git a/src/state/session/index.tsx b/src/state/session/index.tsx
index 314945bcf..3aac19025 100644
--- a/src/state/session/index.tsx
+++ b/src/state/session/index.tsx
@@ -19,6 +19,7 @@ import {
 import {getInitialState, reducer} from './reducer'
 
 export {isSignupQueued} from './util'
+import {addSessionDebugLog} from './logging'
 export type {SessionAccount} from '#/state/session/types'
 import {SessionApiContext, SessionStateContext} from '#/state/session/types'
 
@@ -40,9 +41,11 @@ const ApiContext = React.createContext<SessionApiContext>({
 
 export function Provider({children}: React.PropsWithChildren<{}>) {
   const cancelPendingTask = useOneTaskAtATime()
-  const [state, dispatch] = React.useReducer(reducer, null, () =>
-    getInitialState(persisted.get('session').accounts),
-  )
+  const [state, dispatch] = React.useReducer(reducer, null, () => {
+    const initialState = getInitialState(persisted.get('session').accounts)
+    addSessionDebugLog({type: 'reducer:init', state: initialState})
+    return initialState
+  })
 
   const onAgentSessionChange = React.useCallback(
     (agent: BskyAgent, accountDid: string, sessionEvent: AtpSessionEvent) => {
@@ -63,6 +66,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
 
   const createAccount = React.useCallback<SessionApiContext['createAccount']>(
     async params => {
+      addSessionDebugLog({type: 'method:start', method: 'createAccount'})
       const signal = cancelPendingTask()
       track('Try Create Account')
       logEvent('account:create:begin', {})
@@ -81,12 +85,14 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       })
       track('Create Account')
       logEvent('account:create:success', {})
+      addSessionDebugLog({type: 'method:end', method: 'createAccount', account})
     },
     [onAgentSessionChange, cancelPendingTask],
   )
 
   const login = React.useCallback<SessionApiContext['login']>(
     async (params, logContext) => {
+      addSessionDebugLog({type: 'method:start', method: 'login'})
       const signal = cancelPendingTask()
       const {agent, account} = await createAgentAndLogin(
         params,
@@ -103,23 +109,31 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       })
       track('Sign In', {resumedSession: false})
       logEvent('account:loggedIn', {logContext, withPassword: true})
+      addSessionDebugLog({type: 'method:end', method: 'login', account})
     },
     [onAgentSessionChange, cancelPendingTask],
   )
 
   const logout = React.useCallback<SessionApiContext['logout']>(
     logContext => {
+      addSessionDebugLog({type: 'method:start', method: 'logout'})
       cancelPendingTask()
       dispatch({
         type: 'logged-out',
       })
       logEvent('account:loggedOut', {logContext})
+      addSessionDebugLog({type: 'method:end', method: 'logout'})
     },
     [cancelPendingTask],
   )
 
   const resumeSession = React.useCallback<SessionApiContext['resumeSession']>(
     async storedAccount => {
+      addSessionDebugLog({
+        type: 'method:start',
+        method: 'resumeSession',
+        account: storedAccount,
+      })
       const signal = cancelPendingTask()
       const {agent, account} = await createAgentAndResume(
         storedAccount,
@@ -134,17 +148,24 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
         newAgent: agent,
         newAccount: account,
       })
+      addSessionDebugLog({type: 'method:end', method: 'resumeSession', account})
     },
     [onAgentSessionChange, cancelPendingTask],
   )
 
   const removeAccount = React.useCallback<SessionApiContext['removeAccount']>(
     account => {
+      addSessionDebugLog({
+        type: 'method:start',
+        method: 'removeAccount',
+        account,
+      })
       cancelPendingTask()
       dispatch({
         type: 'removed-account',
         accountDid: account.did,
       })
+      addSessionDebugLog({type: 'method:end', method: 'removeAccount', account})
     },
     [cancelPendingTask],
   )
@@ -152,18 +173,21 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
   React.useEffect(() => {
     if (state.needsPersist) {
       state.needsPersist = false
-      persisted.write('session', {
+      const persistedData = {
         accounts: state.accounts,
         currentAccount: state.accounts.find(
           a => a.did === state.currentAgentState.did,
         ),
-      })
+      }
+      addSessionDebugLog({type: 'persisted:broadcast', data: persistedData})
+      persisted.write('session', persistedData)
     }
   }, [state])
 
   React.useEffect(() => {
     return persisted.onUpdate(() => {
       const synced = persisted.get('session')
+      addSessionDebugLog({type: 'persisted:receive', data: synced})
       dispatch({
         type: 'synced-accounts',
         syncedAccounts: synced.accounts,
@@ -177,7 +201,14 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
           resumeSession(syncedAccount)
         } else {
           const agent = state.currentAgentState.agent as BskyAgent
+          const prevSession = agent.session
           agent.session = sessionAccountToSession(syncedAccount)
+          addSessionDebugLog({
+            type: 'agent:patch',
+            agent,
+            prevSession,
+            nextSession: agent.session,
+          })
         }
       }
     })
@@ -215,6 +246,7 @@ export function Provider({children}: React.PropsWithChildren<{}>) {
       // Read the previous value and immediately advance the pointer.
       const prevAgent = currentAgentRef.current
       currentAgentRef.current = agent
+      addSessionDebugLog({type: 'agent:switch', prevAgent, nextAgent: agent})
       // We never reuse agents so let's fully neutralize the previous one.
       // This ensures it won't try to consume any refresh tokens.
       prevAgent.session = undefined
diff --git a/src/state/session/logging.ts b/src/state/session/logging.ts
new file mode 100644
index 000000000..16aa66fe7
--- /dev/null
+++ b/src/state/session/logging.ts
@@ -0,0 +1,137 @@
+import {AtpSessionData} from '@atproto/api'
+import {sha256} from 'js-sha256'
+import {Statsig} from 'statsig-react-native-expo'
+
+import {Schema} from '../persisted'
+import {Action, State} from './reducer'
+import {SessionAccount} from './types'
+
+type Reducer = (state: State, action: Action) => State
+
+type Log =
+  | {
+      type: 'reducer:init'
+      state: State
+    }
+  | {
+      type: 'reducer:call'
+      action: Action
+      prevState: State
+      nextState: State
+    }
+  | {
+      type: 'method:start'
+      method:
+        | 'createAccount'
+        | 'login'
+        | 'logout'
+        | 'resumeSession'
+        | 'removeAccount'
+      account?: SessionAccount
+    }
+  | {
+      type: 'method:end'
+      method:
+        | 'createAccount'
+        | 'login'
+        | 'logout'
+        | 'resumeSession'
+        | 'removeAccount'
+      account?: SessionAccount
+    }
+  | {
+      type: 'persisted:broadcast'
+      data: Schema['session']
+    }
+  | {
+      type: 'persisted:receive'
+      data: Schema['session']
+    }
+  | {
+      type: 'agent:switch'
+      prevAgent: object
+      nextAgent: object
+    }
+  | {
+      type: 'agent:patch'
+      agent: object
+      prevSession: AtpSessionData | undefined
+      nextSession: AtpSessionData
+    }
+
+export function wrapSessionReducerForLogging(reducer: Reducer): Reducer {
+  return function loggingWrapper(prevState: State, action: Action): State {
+    const nextState = reducer(prevState, action)
+    addSessionDebugLog({type: 'reducer:call', prevState, action, nextState})
+    return nextState
+  }
+}
+
+let nextMessageIndex = 0
+const MAX_SLICE_LENGTH = 1000
+
+export function addSessionDebugLog(log: Log) {
+  try {
+    if (!Statsig.initializeCalled() || !Statsig.getStableID()) {
+      // Drop these logs for now.
+      return
+    }
+    if (!Statsig.checkGate('debug_session')) {
+      return
+    }
+    const messageIndex = nextMessageIndex++
+    const {type, ...content} = log
+    let payload = JSON.stringify(content, replacer)
+
+    let nextSliceIndex = 0
+    while (payload.length > 0) {
+      const sliceIndex = nextSliceIndex++
+      const slice = payload.slice(0, MAX_SLICE_LENGTH)
+      payload = payload.slice(MAX_SLICE_LENGTH)
+      Statsig.logEvent('session:debug', null, {
+        realmId,
+        messageIndex: String(messageIndex),
+        messageType: type,
+        sliceIndex: String(sliceIndex),
+        slice,
+      })
+    }
+  } catch (e) {
+    console.error(e)
+  }
+}
+
+let agentIds = new WeakMap<object, string>()
+let realmId = Math.random().toString(36).slice(2)
+let nextAgentId = 1
+
+function getAgentId(agent: object) {
+  let id = agentIds.get(agent)
+  if (id === undefined) {
+    id = realmId + '::' + nextAgentId++
+    agentIds.set(agent, id)
+  }
+  return id
+}
+
+function replacer(key: string, value: unknown) {
+  if (typeof value === 'object' && value != null && 'api' in value) {
+    return getAgentId(value)
+  }
+  if (
+    key === 'service' ||
+    key === 'email' ||
+    key === 'emailConfirmed' ||
+    key === 'emailAuthFactor' ||
+    key === 'pdsUrl'
+  ) {
+    return undefined
+  }
+  if (
+    typeof value === 'string' &&
+    (key === 'refreshJwt' || key === 'accessJwt')
+  ) {
+    return sha256(value)
+  }
+  return value
+}
diff --git a/src/state/session/reducer.ts b/src/state/session/reducer.ts
index 7f3080935..0a537b42c 100644
--- a/src/state/session/reducer.ts
+++ b/src/state/session/reducer.ts
@@ -1,6 +1,7 @@
 import {AtpSessionEvent} from '@atproto/api'
 
 import {createPublicAgent} from './agent'
+import {wrapSessionReducerForLogging} from './logging'
 import {SessionAccount} from './types'
 
 // A hack so that the reducer can't read anything from the agent.
@@ -64,7 +65,7 @@ export function getInitialState(persistedAccounts: SessionAccount[]): State {
   }
 }
 
-export function reducer(state: State, action: Action): State {
+let reducer = (state: State, action: Action): State => {
   switch (action.type) {
     case 'received-agent-event': {
       const {agent, accountDid, refreshedAccount, sessionEvent} = action
@@ -166,3 +167,5 @@ export function reducer(state: State, action: Action): State {
     }
   }
 }
+reducer = wrapSessionReducerForLogging(reducer)
+export {reducer}