fix: improve authentication SSE

src/controllers/auth/anilist/getWatchingTitles.ts

@@ -45,6 +45,7 @@ const GetWatchingTitlesQuery = graphql(`
         currentPage
         hasNextPage
         perPage
+        total
       }
     }
   }
@@ -53,21 +54,29 @@ const GetWatchingTitlesQuery = graphql(`
 export function getWatchingTitles(
   username: string,
   page: number,
+  aniListToken: string,
   executionCtx: ExecutionContext,
 ) {
   const client = new GraphQLClient("https://graphql.anilist.co/");
 
   return client
-    .request(GetWatchingTitlesQuery, { userName: username, page })
+    .request(
+      GetWatchingTitlesQuery,
+      { userName: username, page },
+      { Authorization: `Bearer ${aniListToken}` },
+    )
     .then((data) => data?.Page!)
     .catch((err) => {
+      console.error("Failed to get watching titles");
+      console.error(err);
+
       const response = err.response;
       if (response.status === 429) {
         console.log("429, retrying in", response.headers.get("Retry-After"));
         executionCtx.waitUntil(
           sleep(Number(response.headers.get("Retry-After")!) * 1000),
         );
-        return getWatchingTitles(username, page, executionCtx);
+        return getWatchingTitles(username, page, aniListToken, executionCtx);
       }
 
       throw err;

src/controllers/auth/anilist/index.ts

@@ -41,6 +41,21 @@ const route = createRoute({
       "x-anilist-token": z.string(),
       "x-aniplay-device-id": z.string(),
     }),
+    // Uncomment when testing locally
+    // headers: z.object({
+    //   "x-anilist-token":
+    //     process.env.NODE_ENV === "production"
+    //       ? z.string()
+    //       : z.string().optional(),
+    //   "x-aniplay-device-id":
+    //     process.env.NODE_ENV === "production"
+    //       ? z.string()
+    //       : z.string().optional(),
+    // }),
+    // query: z.object({
+    //   aniListToken: z.string().optional(),
+    //   deviceId: z.string().optional(),
+    // }),
   },
   responses: {
     200: {
@@ -76,24 +91,38 @@ const route = createRoute({
 const app = new OpenAPIHono<Env>();
 
 app.openapi(route, async (c) => {
-  const deviceId = await c.req.header("X-Aniplay-Device-Id");
+  const deviceId =
-  const aniListToken = await c.req.header("X-AniList-Token");
+    c.req.header("X-Aniplay-Device-Id") ?? c.req.query("deviceId");
+  const aniListToken =
+    c.req.header("X-AniList-Token") ?? c.req.query("aniListToken");
 
   if (!aniListToken) {
     return c.json(ErrorResponse, { status: 401 });
   }
 
+  let user: Awaited<ReturnType<typeof getUser>>;
   try {
-    const user = await getUser(aniListToken);
+    user = await getUser(aniListToken);
     if (!user) {
       return c.json(ErrorResponse, { status: 401 });
     }
+  } catch (error) {
+    console.error(new Error("Failed to authenticate with AniList"));
+    console.error(error);
+    return c.json(ErrorResponse, { status: 500 });
+  }
 
+  try {
     await associateDeviceIdWithUsername(
       env(c, "workerd"),
       deviceId!,
       user.name!,
     );
+  } catch (error) {
+    console.error(new Error("Failed to associate device"));
+    console.error(error);
+    return c.json(ErrorResponse, { status: 500 });
+  }
 
   return streamSSE(
     c,
@@ -105,14 +134,22 @@ app.openapi(route, async (c) => {
 
       do {
         const { mediaList, pageInfo } = await getWatchingTitles(
-            user.name,
+          user.name!,
           currentPage++,
+          aniListToken,
           c.executionCtx,
         );
         if (!mediaList) {
           break;
         }
 
+        if (!(pageInfo?.hasNextPage ?? false) && (pageInfo?.total ?? 0) > 0) {
+          stream.writeSSE({
+            event: "count",
+            data: pageInfo!.total.toString(),
+          });
+        }
+
         for (const mediaObj of mediaList) {
           const media = mediaObj?.media!;
           if (!media) {
@@ -153,7 +190,7 @@ app.openapi(route, async (c) => {
           await fetchEpisodes(
             media.id,
             readEnvVariable<boolean>(c.env, "ENABLE_ANIFY"),
-            ).then(({ result: episodes }) => {
+          ).then(({ result: { episodes } }) => {
             stream.writeSSE({
               event: "title",
               data: JSON.stringify({ title: media, episodes }),
@@ -162,23 +199,22 @@ app.openapi(route, async (c) => {
           });
         }
 
+        hasNextPage = pageInfo?.hasNextPage ?? false;
+        hasNextPage = pageInfo?.hasNextPage ?? false;
+        console.log(hasNextPage);
         hasNextPage = pageInfo?.hasNextPage ?? false;
         console.log(hasNextPage);
       } while (hasNextPage);
 
-        await stream.close();
+      // send end event instead of closing the connection to let the client know that the stream didn't end abruptly
+      await stream.writeSSE({ event: "end", data: "end" });
     },
     async (err, stream) => {
-        stream.writeln("An error occurred!");
+      console.error("Error occurred in SSE");
       console.error(err);
+      stream.writeln("An error occurred");
     },
   );
-  } catch (error) {
-    console.error(
-      new Error("Failed to authenticate with AniList", { cause: error }),
-    );
-    return c.json(ErrorResponse, { status: 500 });
-  }
 });
 
 export default app;