From 9de3a32ceb3d7fce4126a8a13aabe421bda858c9 Mon Sep 17 00:00:00 2001
From: Mythie <me@lucasjamessmith.me>
Date: Mon, 20 Jan 2025 09:30:36 +1100
Subject: [PATCH] fix: pass team id to vanilla trpc client

---
 .../documents/data-table-action-button.tsx    | 13 ++++++++---
 packages/trpc/client/index.ts                 | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/apps/web/src/app/(dashboard)/documents/data-table-action-button.tsx b/apps/web/src/app/(dashboard)/documents/data-table-action-button.tsx
index 7d8f4b6a5..1194dfd01 100644
--- a/apps/web/src/app/(dashboard)/documents/data-table-action-button.tsx
+++ b/apps/web/src/app/(dashboard)/documents/data-table-action-button.tsx
@@ -50,9 +50,16 @@ export const DataTableActionButton = ({ row, team }: DataTableActionButtonProps)
   const onDownloadClick = async () => {
     try {
       const document = !recipient
-        ? await trpcClient.document.getDocumentById.query({
-            documentId: row.id,
-          })
+        ? await trpcClient.document.getDocumentById.query(
+            {
+              documentId: row.id,
+            },
+            {
+              context: {
+                teamId: team?.id?.toString(),
+              },
+            },
+          )
         : await trpcClient.document.getDocumentByToken.query({
             token: recipient.token,
           });
diff --git a/packages/trpc/client/index.ts b/packages/trpc/client/index.ts
index 51cda4406..5f0e37ab6 100644
--- a/packages/trpc/client/index.ts
+++ b/packages/trpc/client/index.ts
@@ -12,10 +12,32 @@ export const trpc = createTRPCClient<AppRouter>({
       true: httpLink({
         url: `${getBaseUrl()}/api/trpc`,
         transformer: SuperJSON,
+        headers: (opts) => {
+          if (typeof opts.op.context.teamId === 'string') {
+            return {
+              'x-team-id': opts.op.context.teamId,
+            };
+          }
+
+          return {};
+        },
       }),
       false: httpBatchLink({
         url: `${getBaseUrl()}/api/trpc`,
         transformer: SuperJSON,
+        headers: (opts) => {
+          const operationWithTeamId = opts.opList.find(
+            (op) => op.context.teamId && typeof op.context.teamId === 'string',
+          );
+
+          if (operationWithTeamId && typeof operationWithTeamId.context.teamId === 'string') {
+            return {
+              'x-team-id': operationWithTeamId.context.teamId,
+            };
+          }
+
+          return {};
+        },
       }),
     }),
   ],