use an env var for configurable proxy header (#8863)

14 files changed, 79 insertions, 34 deletions

diff --git a/src/state/messages/convo/agent.ts b/src/state/messages/convo/agent.ts
index 2ad4c592e..168002b1f 100644
--- a/src/state/messages/convo/agent.ts
+++ b/src/state/messages/convo/agent.ts
@@ -10,6 +10,7 @@ import EventEmitter from 'eventemitter3'
 import {nanoid} from 'nanoid/non-secure'
 
 import {networkRetry} from '#/lib/async/retry'
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {isNetworkError} from '#/lib/strings/errors'
 import {Logger} from '#/logger'
 import {isNative} from '#/platform/detection'
@@ -33,7 +34,6 @@ import {
 } from '#/state/messages/convo/types'
 import {type MessagesEventBus} from '#/state/messages/events/agent'
 import {type MessagesEventBusError} from '#/state/messages/events/types'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 
 const logger = Logger.create(Logger.Context.ConversationAgent)
 
diff --git a/src/state/messages/events/agent.ts b/src/state/messages/events/agent.ts
index fb3047bf6..e54ea1c77 100644
--- a/src/state/messages/events/agent.ts
+++ b/src/state/messages/events/agent.ts
@@ -3,6 +3,7 @@ import EventEmitter from 'eventemitter3'
 import {nanoid} from 'nanoid/non-secure'
 
 import {networkRetry} from '#/lib/async/retry'
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {isNetworkError} from '#/lib/strings/errors'
 import {Logger} from '#/logger'
 import {
@@ -17,7 +18,6 @@ import {
   type MessagesEventBusParams,
   MessagesEventBusStatus,
 } from '#/state/messages/events/types'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 
 const logger = Logger.create(Logger.Context.DMsAgent)
 
diff --git a/src/state/queries/handle-availability.ts b/src/state/queries/handle-availability.ts
index 9391f5d09..06fc6eebb 100644
--- a/src/state/queries/handle-availability.ts
+++ b/src/state/queries/handle-availability.ts
@@ -1,4 +1,4 @@
-import {Agent, ComAtprotoTempCheckHandleAvailability} from '@atproto/api'
+import {ComAtprotoTempCheckHandleAvailability} from '@atproto/api'
 import {useQuery} from '@tanstack/react-query'
 
 import {
@@ -10,6 +10,7 @@ import {createFullHandle} from '#/lib/strings/handles'
 import {logger} from '#/logger'
 import {useDebouncedValue} from '#/components/live/utils'
 import * as bsky from '#/types/bsky'
+import {Agent} from '../session/agent'
 
 export const RQKEY_handleAvailability = (
   handle: string,
@@ -74,7 +75,7 @@ export async function checkHandleAvailability(
   },
 ) {
   if (serviceDid === BSKY_SERVICE_DID) {
-    const agent = new Agent({service: BSKY_SERVICE})
+    const agent = new Agent(null, {service: BSKY_SERVICE})
     // entryway has a special API for handle availability
     const {data} = await agent.com.atproto.temp.checkHandleAvailability({
       handle,
@@ -109,7 +110,7 @@ export async function checkHandleAvailability(
     }
   } else {
     // 3rd party PDSes won't have this API so just try and resolve the handle
-    const agent = new Agent({service: PUBLIC_BSKY_SERVICE})
+    const agent = new Agent(null, {service: PUBLIC_BSKY_SERVICE})
     try {
       const res = await agent.resolveHandle({
         handle,
diff --git a/src/state/queries/messages/accept-conversation.ts b/src/state/queries/messages/accept-conversation.ts
index 82acb33c8..0c06055b5 100644
--- a/src/state/queries/messages/accept-conversation.ts
+++ b/src/state/queries/messages/accept-conversation.ts
@@ -1,9 +1,12 @@
-import {ChatBskyConvoAcceptConvo, ChatBskyConvoListConvos} from '@atproto/api'
+import {
+  type ChatBskyConvoAcceptConvo,
+  type ChatBskyConvoListConvos,
+} from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {logger} from '#/logger'
 import {useAgent} from '#/state/session'
-import {DM_SERVICE_HEADERS} from './const'
 import {
   RQKEY as CONVO_LIST_KEY,
   RQKEY_ROOT as CONVO_LIST_ROOT_KEY,
diff --git a/src/state/queries/messages/const.ts b/src/state/queries/messages/const.ts
deleted file mode 100644
index 1c5519a63..000000000
--- a/src/state/queries/messages/const.ts
+++ /dev/null
@@ -1,5 +0,0 @@
-import {CHAT_PROXY_DID} from '#/env'
-
-export const DM_SERVICE_HEADERS = {
-  'atproto-proxy': `${CHAT_PROXY_DID}#bsky_chat`,
-}
diff --git a/src/state/queries/messages/conversation.ts b/src/state/queries/messages/conversation.ts
index de5a90571..393bf9e52 100644
--- a/src/state/queries/messages/conversation.ts
+++ b/src/state/queries/messages/conversation.ts
@@ -1,17 +1,17 @@
-import {ChatBskyConvoDefs} from '@atproto/api'
+import {type ChatBskyConvoDefs} from '@atproto/api'
 import {
-  QueryClient,
+  type QueryClient,
   useMutation,
   useQuery,
   useQueryClient,
 } from '@tanstack/react-query'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {STALE} from '#/state/queries'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 import {useOnMarkAsRead} from '#/state/queries/messages/list-conversations'
 import {useAgent} from '#/state/session'
 import {
-  ConvoListQueryData,
+  type ConvoListQueryData,
   getConvoFromQueryData,
   RQKEY_ROOT as LIST_CONVOS_KEY,
 } from './list-conversations'
diff --git a/src/state/queries/messages/get-convo-availability.ts b/src/state/queries/messages/get-convo-availability.ts
index f545c3bba..2392edb09 100644
--- a/src/state/queries/messages/get-convo-availability.ts
+++ b/src/state/queries/messages/get-convo-availability.ts
@@ -1,6 +1,6 @@
 import {useQuery} from '@tanstack/react-query'
 
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {useAgent} from '#/state/session'
 import {STALE} from '..'
 
diff --git a/src/state/queries/messages/get-convo-for-members.ts b/src/state/queries/messages/get-convo-for-members.ts
index 3f45c2328..58c1ab524 100644
--- a/src/state/queries/messages/get-convo-for-members.ts
+++ b/src/state/queries/messages/get-convo-for-members.ts
@@ -1,8 +1,8 @@
-import {ChatBskyConvoGetConvoForMembers} from '@atproto/api'
+import {type ChatBskyConvoGetConvoForMembers} from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {logger} from '#/logger'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 import {useAgent} from '#/state/session'
 import {precacheConvoQuery} from './conversation'
 
diff --git a/src/state/queries/messages/leave-conversation.ts b/src/state/queries/messages/leave-conversation.ts
index b17e515be..986351a07 100644
--- a/src/state/queries/messages/leave-conversation.ts
+++ b/src/state/queries/messages/leave-conversation.ts
@@ -1,13 +1,16 @@
 import {useMemo} from 'react'
-import {ChatBskyConvoLeaveConvo, ChatBskyConvoListConvos} from '@atproto/api'
+import {
+  type ChatBskyConvoLeaveConvo,
+  type ChatBskyConvoListConvos,
+} from '@atproto/api'
 import {
   useMutation,
   useMutationState,
   useQueryClient,
 } from '@tanstack/react-query'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {logger} from '#/logger'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 import {useAgent} from '#/state/session'
 import {RQKEY_ROOT as CONVO_LIST_KEY} from './list-conversations'
 
diff --git a/src/state/queries/messages/list-conversations.tsx b/src/state/queries/messages/list-conversations.tsx
index 3f8252519..c5457d1cb 100644
--- a/src/state/queries/messages/list-conversations.tsx
+++ b/src/state/queries/messages/list-conversations.tsx
@@ -13,10 +13,10 @@ import {
 } from '@tanstack/react-query'
 import throttle from 'lodash.throttle'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {useCurrentConvoId} from '#/state/messages/current-convo-id'
 import {useMessagesEventBus} from '#/state/messages/events'
 import {useModerationOpts} from '#/state/preferences/moderation-opts'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 import {useAgent, useSession} from '#/state/session'
 import {useLeftConvos} from './leave-conversation'
 
diff --git a/src/state/queries/messages/mute-conversation.ts b/src/state/queries/messages/mute-conversation.ts
index da9644145..d668e36cb 100644
--- a/src/state/queries/messages/mute-conversation.ts
+++ b/src/state/queries/messages/mute-conversation.ts
@@ -1,11 +1,15 @@
 import {
-  ChatBskyConvoDefs,
-  ChatBskyConvoListConvos,
-  ChatBskyConvoMuteConvo,
+  type ChatBskyConvoDefs,
+  type ChatBskyConvoListConvos,
+  type ChatBskyConvoMuteConvo,
 } from '@atproto/api'
-import {InfiniteData, useMutation, useQueryClient} from '@tanstack/react-query'
+import {
+  type InfiniteData,
+  useMutation,
+  useQueryClient,
+} from '@tanstack/react-query'
 
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {useAgent} from '#/state/session'
 import {RQKEY as CONVO_KEY} from './conversation'
 import {RQKEY_ROOT as CONVO_LIST_KEY} from './list-conversations'
diff --git a/src/state/queries/messages/update-all-read.ts b/src/state/queries/messages/update-all-read.ts
index 72fa65ee6..3d0fd3a45 100644
--- a/src/state/queries/messages/update-all-read.ts
+++ b/src/state/queries/messages/update-all-read.ts
@@ -1,8 +1,8 @@
-import {ChatBskyConvoListConvos} from '@atproto/api'
+import {type ChatBskyConvoListConvos} from '@atproto/api'
 import {useMutation, useQueryClient} from '@tanstack/react-query'
 
+import {DM_SERVICE_HEADERS} from '#/lib/constants'
 import {logger} from '#/logger'
-import {DM_SERVICE_HEADERS} from '#/state/queries/messages/const'
 import {useAgent} from '#/state/session'
 import {RQKEY as CONVO_LIST_KEY} from './list-conversations'
 
diff --git a/src/state/queries/service.ts b/src/state/queries/service.ts
index 6bfd0b011..e9661db9e 100644
--- a/src/state/queries/service.ts
+++ b/src/state/queries/service.ts
@@ -1,6 +1,7 @@
-import {BskyAgent} from '@atproto/api'
 import {useQuery} from '@tanstack/react-query'
 
+import {Agent} from '../session/agent'
+
 const RQKEY_ROOT = 'service'
 export const RQKEY = (serviceUrl: string) => [RQKEY_ROOT, serviceUrl]
 
@@ -8,7 +9,7 @@ export function useServiceQuery(serviceUrl: string) {
   return useQuery({
     queryKey: RQKEY(serviceUrl),
     queryFn: async () => {
-      const agent = new BskyAgent({service: serviceUrl})
+      const agent = new Agent(null, {service: serviceUrl})
       const res = await agent.com.atproto.server.describeServer()
       return res.data
     },
diff --git a/src/state/session/agent.ts b/src/state/session/agent.ts
index 531e285ab..d063a09a2 100644
--- a/src/state/session/agent.ts
+++ b/src/state/session/agent.ts
@@ -1,8 +1,19 @@
-import {AtpSessionData, AtpSessionEvent, BskyAgent} from '@atproto/api'
+import {
+  Agent as BaseAgent,
+  type AtprotoServiceType,
+  type AtpSessionData,
+  type AtpSessionEvent,
+  BskyAgent,
+  type Did,
+} from '@atproto/api'
+import {type FetchHandler} from '@atproto/api/dist/agent'
+import {type SessionManager} from '@atproto/api/dist/session-manager'
 import {TID} from '@atproto/common-web'
+import {type FetchHandlerOptions} from '@atproto/xrpc'
 
 import {networkRetry} from '#/lib/async/retry'
 import {
+  BLUESKY_PROXY_HEADER,
   BSKY_SERVICE,
   DISCOVER_SAVED_FEED,
   IS_PROD_SERVICE,
@@ -19,12 +30,17 @@ import {
   configureModerationForAccount,
   configureModerationForGuest,
 } from './moderation'
-import {SessionAccount} from './types'
+import {type SessionAccount} from './types'
 import {isSessionExpired, isSignupQueued} from './util'
 
+export type ProxyHeaderValue = `${Did}#${AtprotoServiceType}`
+
 export function createPublicAgent() {
   configureModerationForGuest() // Side effect but only relevant for tests
-  return new BskyAppAgent({service: PUBLIC_BSKY_SERVICE})
+
+  const agent = new BskyAppAgent({service: PUBLIC_BSKY_SERVICE})
+  agent.configureProxy(BLUESKY_PROXY_HEADER)
+  return agent
 }
 
 export async function createAgentAndResume(
@@ -61,6 +77,8 @@ export async function createAgentAndResume(
     }
   }
 
+  agent.configureProxy(BLUESKY_PROXY_HEADER)
+
   return agent.prepare(gates, moderation, onSessionChange)
 }
 
@@ -93,6 +111,9 @@ export async function createAgentAndLogin(
   const account = agentToSessionAccountOrThrow(agent)
   const gates = tryFetchGates(account.did, 'prefer-fresh-gates')
   const moderation = configureModerationForAccount(agent, account)
+
+  agent.configureProxy(BLUESKY_PROXY_HEADER)
+
   return agent.prepare(gates, moderation, onSessionChange)
 }
 
@@ -180,6 +201,8 @@ export async function createAgentAndCreateAccount(
     logger.error(e, {message: `session: failed snoozeEmailConfirmationPrompt`})
   }
 
+  agent.configureProxy(BLUESKY_PROXY_HEADER)
+
   return agent.prepare(gates, moderation, onSessionChange)
 }
 
@@ -234,7 +257,22 @@ export function sessionAccountToSession(
   }
 }
 
+export class Agent extends BaseAgent {
+  constructor(
+    proxyHeader: ProxyHeaderValue | null,
+    options: SessionManager | FetchHandler | FetchHandlerOptions,
+  ) {
+    super(options)
+    if (proxyHeader) {
+      this.configureProxy(proxyHeader)
+    }
+  }
+}
+
 // Not exported. Use factories above to create it.
+// WARN: In the factories above, we _manually set a proxy header_ for the agent after we do whatever it is we are supposed to do.
+// Ideally, we wouldn't be doing this. However, since there is so much logic that requires making calls to the PDS right now, it
+// feels safer to just let those run as-is and set the header afterward.
 let realFetch = globalThis.fetch
 class BskyAppAgent extends BskyAgent {
   persistSessionHandler: ((event: AtpSessionEvent) => void) | undefined =