Feature: Add (opt-in) support for cookies for treaty client

package.json

@@ -59,9 +59,11 @@
         }
     },
     "devDependencies": {
+        "@elysiajs/cookie": "^0.8.0",
         "@elysiajs/cors": "0.7.0",
         "@sinclair/typebox": "^0.31.6",
         "@types/node": "^18.15.5",
+        "@types/tough-cookie": "^4.0.5",
         "bun-types": "^1.0.1",
         "elysia": "0.8.0",
         "esbuild": "^0.19.3",
@@ -77,6 +79,7 @@
         "trailingComma": "none"
     },
     "dependencies": {
-        "superjson": "^2.2.1"
+        "superjson": "^2.2.1",
+        "tough-cookie": "^4.1.3"
     }
 }

src/treaty/cookies.ts

@@ -0,0 +1,16 @@
+import { Cookie, CookieJar } from 'tough-cookie'
+
+export const addCookiesToJar = (
+    jar: CookieJar,
+    setCookieHeaders: string[],
+    requestUrl: string
+) => {
+    for (const setCookieHeader of setCookieHeaders) {
+        const cookie = Cookie.parse(setCookieHeader, { loose: true })
+        if (!cookie) {
+            continue
+        }
+
+        jar.setCookieSync(cookie, requestUrl)
+    }
+}

src/treaty/index.ts

@@ -7,6 +7,9 @@ import type { EdenTreaty } from './types'
 
 export type { EdenTreaty } from './types'
 
+import { CookieJar } from 'tough-cookie'
+import { addCookiesToJar } from './cookies'
+
 // @ts-ignore
 const isServer = typeof FileList === 'undefined'
 
@@ -159,11 +162,17 @@ export class EdenWS<Schema extends InputSchema<any> = InputSchema> {
 const createProxy = (
     domain: string,
     path = '',
-    config: EdenTreaty.Config
+    config: EdenTreaty.Config,
+    cookieJar: CookieJar
 ): Record<string, unknown> =>
     new Proxy(() => {}, {
         get(target, key, value) {
-            return createProxy(domain, `${path}/${key.toString()}`, config)
+            return createProxy(
+                domain,
+                `${path}/${key.toString()}`,
+                config,
+                cookieJar
+            )
         },
         // @ts-ignore
         apply(
@@ -245,7 +254,10 @@ const createProxy = (
                     ...config.$fetch?.headers,
                     ...$fetch?.headers,
                     ...options.headers,
-                    ...$headers
+                    ...$headers,
+                    ...(config.persistCookies
+                        ? { Cookie: cookieJar.getCookieStringSync(url) }
+                        : {})
                 } as Record<string, string>
 
                 if (method !== 'GET' && method !== 'HEAD') {
@@ -318,6 +330,14 @@ const createProxy = (
                     headers
                 })
 
+                if (config.persistCookies) {
+                    addCookiesToJar(
+                        cookieJar,
+                        response.headers.getSetCookie(),
+                        url
+                    )
+                }
+
                 let data
 
                 if (modifiers.getRaw) return response as any
@@ -369,12 +389,15 @@ export const edenTreaty = <App extends Elysia<any, any, any, any, any, any>>(
     config: EdenTreaty.Config = {
         fetcher: fetch
     }
-): EdenTreaty.Create<App> =>
-    new Proxy(
+): EdenTreaty.Create<App> => {
+    const jar = new CookieJar()
+
+    return new Proxy(
         {},
         {
             get(target, key) {
-                return createProxy(domain, key as string, config)
+                return createProxy(domain, key as string, config, jar)
             }
         }
     ) as any
+}

src/treaty/types.ts

@@ -81,6 +81,7 @@ export namespace EdenTreaty {
         $fetch?: RequestInit
         fetcher?: typeof fetch
         transform?: Transform
+        persistCookies?: boolean
     }
 
     export type DetailedResponse = {
@@ -216,4 +217,4 @@ export namespace EdenTreaty {
         K extends keyof WebSocketEventMap,
         Data = unknown
     > = K extends 'message' ? OnMessage<Data> : WebSocketEventMap[K]
-}
+}

test/treaty.test.ts

@@ -1,4 +1,5 @@
 import { Elysia, t } from 'elysia'
+import { cookie } from '@elysiajs/cookie'
 import { edenTreaty } from '../src'
 
 import { beforeEach, describe, expect, it, spyOn, mock } from 'bun:test'
@@ -13,6 +14,7 @@ const prefix =
         app.get(`${prefix}/prefixed`, () => 'hi')
 
 const app = new Elysia()
+    .use(cookie)
     .get('/', () => 'hi')
     .use(prefix('/prefix'))
     .post('/', () => 'hi')
@@ -50,6 +52,17 @@ const app = new Elysia()
     .post('/string', ({ body }) => body, {
         body: t.String()
     })
+    .group('/cookie', (app) =>
+        app
+            .get('/set-cookie', ({ setCookie }) => {
+                setCookie('testCookie', 'hello world')
+                return 'hi'
+            })
+            .get('/get-cookie', ({ cookie: { testCookie } }) => {
+                console.log('here')
+                return testCookie ?? 'no cookie'
+            })
+    )
     .listen(8082)
 
 const client = edenTreaty<typeof app>('http://localhost:8082')
@@ -262,4 +275,16 @@ describe('Eden Rest', () => {
                 .method
         ).toBe('PATCH')
     })
+
+    it('supports cookies with `persistCookies` option', async () => {
+        const client = edenTreaty<typeof app>('http://localhost:8082', {
+            persistCookies: true
+        })
+        const { data } = await client.cookie['get-cookie'].get()
+        expect(data).toBe('no cookie')
+        await client.cookie['set-cookie'].get()
+
+        const { data: afterSetCookie } = await client.cookie['get-cookie'].get()
+        expect(afterSetCookie).toBe('hello world')
+    })
 })