From eb9a77ea644739beefd4bfbe8acfe036c74f7e9a Mon Sep 17 00:00:00 2001
From: Mike Cao <mike@mikecao.com>
Date: Mon, 20 Jul 2026 17:15:52 -0700
Subject: [PATCH] Return 404 when session is not found and add test coverage
 for GET session route

---
 .../sessions/[sessionId]/route.test.ts        | 53 +++++++++++++++++++
 .../[websiteId]/sessions/[sessionId]/route.ts | 53 +++++++++----------
 2 files changed, 77 insertions(+), 29 deletions(-)
 create mode 100644 src/app/api/websites/[websiteId]/sessions/[sessionId]/route.test.ts

diff --git a/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.test.ts b/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.test.ts
new file mode 100644
index 000000000..7a50b1770
--- /dev/null
+++ b/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.test.ts
@@ -0,0 +1,53 @@
+import { beforeEach, expect, test, vi } from 'vitest';
+import { parseRequest } from '@/lib/request';
+import { canViewWebsiteSection } from '@/permissions';
+import { getLinkedDistinctIds, getLinkedSessionIds, getWebsiteSession } from '@/queries/sql';
+import { GET } from './route';
+
+vi.mock('@/lib/request', () => ({
+  parseRequest: vi.fn(),
+}));
+
+vi.mock('@/permissions', () => ({
+  canViewWebsiteSection: vi.fn(),
+}));
+
+vi.mock('@/queries/sql', () => ({
+  getLinkedDistinctIds: vi.fn(),
+  getLinkedSessionIds: vi.fn(),
+  getWebsiteSession: vi.fn(),
+}));
+
+const parseRequestMock = vi.mocked(parseRequest);
+const canViewWebsiteSectionMock = vi.mocked(canViewWebsiteSection);
+const getLinkedDistinctIdsMock = vi.mocked(getLinkedDistinctIds);
+const getLinkedSessionIdsMock = vi.mocked(getLinkedSessionIds);
+const getWebsiteSessionMock = vi.mocked(getWebsiteSession);
+
+beforeEach(() => {
+  parseRequestMock.mockReset();
+  canViewWebsiteSectionMock.mockReset();
+  getLinkedDistinctIdsMock.mockReset();
+  getLinkedSessionIdsMock.mockReset();
+  getWebsiteSessionMock.mockReset();
+});
+
+test('GET returns not found when the session does not exist', async () => {
+  parseRequestMock.mockResolvedValue({ auth: {}, error: undefined });
+  canViewWebsiteSectionMock.mockResolvedValue(true);
+  getWebsiteSessionMock.mockResolvedValue(undefined);
+
+  const response = await GET(
+    new Request('http://localhost/api/websites/website-1/sessions/missing-session'),
+    {
+      params: Promise.resolve({ websiteId: 'website-1', sessionId: 'missing-session' }),
+    },
+  );
+
+  expect(response.status).toBe(404);
+  await expect(response.json()).resolves.toMatchObject({
+    error: { code: 'not-found', status: 404 },
+  });
+  expect(getLinkedDistinctIdsMock).not.toHaveBeenCalled();
+  expect(getLinkedSessionIdsMock).not.toHaveBeenCalled();
+});
diff --git a/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.ts b/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.ts
index 6f1a08d40..2c32d52cc 100644
--- a/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.ts
+++ b/src/app/api/websites/[websiteId]/sessions/[sessionId]/route.ts
@@ -1,5 +1,5 @@
 import { parseRequest } from '@/lib/request';
-import { json, unauthorized } from '@/lib/response';
+import { json, notFound, unauthorized } from '@/lib/response';
 import { canViewWebsiteSection } from '@/permissions';
 import { getLinkedDistinctIds, getLinkedSessionIds, getWebsiteSession } from '@/queries/sql';
 
@@ -16,42 +16,37 @@ export async function GET(
   const { websiteId, sessionId } = await params;
 
   if (
-    !(await canViewWebsiteSection(auth, websiteId, [
-      'sessions',
-      'events',
-      'realtime',
-      'revenue',
-    ]))
+    !(await canViewWebsiteSection(auth, websiteId, ['sessions', 'events', 'realtime', 'revenue']))
   ) {
     return unauthorized();
   }
 
   const data = await getWebsiteSession(websiteId, sessionId);
 
-  let stitchedSessionCount = 0;
-
-  if (data) {
-    let sessionIds = [sessionId];
-    const distinctIds = data.distinctId
-      ? [data.distinctId]
-      : await getLinkedDistinctIds(websiteId, sessionId);
-
-    if (!data.distinctId && distinctIds.length === 1) {
-      data.distinctId = distinctIds[0];
-    }
-
-    if (distinctIds.length) {
-      const links = await Promise.all(
-        distinctIds.map(distinctId => getLinkedSessionIds(websiteId, distinctId)),
-      );
-      const linkedIds = links.flatMap(group => group.map(link => link.sessionId));
-
-      sessionIds = Array.from(new Set([sessionId, ...linkedIds]));
-    }
-
-    stitchedSessionCount = sessionIds.length;
+  if (!data) {
+    return notFound();
   }
 
+  let sessionIds = [sessionId];
+  const distinctIds = data.distinctId
+    ? [data.distinctId]
+    : await getLinkedDistinctIds(websiteId, sessionId);
+
+  if (!data.distinctId && distinctIds.length === 1) {
+    data.distinctId = distinctIds[0];
+  }
+
+  if (distinctIds.length) {
+    const links = await Promise.all(
+      distinctIds.map(distinctId => getLinkedSessionIds(websiteId, distinctId)),
+    );
+    const linkedIds = links.flatMap(group => group.map(link => link.sessionId));
+
+    sessionIds = Array.from(new Set([sessionId, ...linkedIds]));
+  }
+
+  const stitchedSessionCount = sessionIds.length;
+
   return json({
     ...data,
     stitchedSessionCount,