From fd1f7fd6f43cd09bd9d088fe9dcd93f7141bd637 Mon Sep 17 00:00:00 2001
From: Lewis Carhart <lewis@flexops.dev>
Date: Mon, 31 Mar 2025 11:03:20 -0400
Subject: [PATCH] Enhance API responses in organization, policies, and risks
 tools

- Updated `findOrganization`, `getPolicies`, and `getRisks` functions to return structured objects with `message` fields for better error handling and clarity.
- Modified `getPolicyContent` and `getRiskById` to provide consistent response formats, including `null` values and descriptive messages when items are not found.
- Improved `useStreamableText` hook with a debouncing mechanism for better performance and user experience during content updates.
---
 apps/app/src/data/tools/organization.ts    |  9 +++-
 apps/app/src/data/tools/policies.ts        | 22 +++++++++-
 apps/app/src/data/tools/risks.ts           | 22 +++++++++-
 apps/app/src/hooks/use-streamable-text.tsx | 48 ++++++++++++++++------
 4 files changed, 83 insertions(+), 18 deletions(-)

diff --git a/apps/app/src/data/tools/organization.ts b/apps/app/src/data/tools/organization.ts
index c62534ebf1..59ec4ddcce 100644
--- a/apps/app/src/data/tools/organization.ts
+++ b/apps/app/src/data/tools/organization.ts
@@ -28,9 +28,14 @@ export const findOrganization = tool({
 		});
 
 		if (!org) {
-			return { error: "Organization not found" };
+			return {
+				organization: null,
+				message: "Organization not found",
+			};
 		}
 
-		return org;
+		return {
+			organization: org,
+		};
 	},
 });
diff --git a/apps/app/src/data/tools/policies.ts b/apps/app/src/data/tools/policies.ts
index a0fff1ecf1..46ff0a64ac 100644
--- a/apps/app/src/data/tools/policies.ts
+++ b/apps/app/src/data/tools/policies.ts
@@ -37,7 +37,16 @@ export const getPolicies = tool({
 			},
 		});
 
-		return policies;
+		if (policies.length === 0) {
+			return {
+				policies: [],
+				message: "No policies found",
+			};
+		}
+
+		return {
+			policies,
+		};
 	},
 });
 
@@ -61,6 +70,15 @@ export const getPolicyContent = tool({
 			},
 		});
 
-		return policy?.content;
+		if (!policy) {
+			return {
+				content: null,
+				message: "Policy not found",
+			};
+		}
+
+		return {
+			content: policy?.content,
+		};
 	},
 });
diff --git a/apps/app/src/data/tools/risks.ts b/apps/app/src/data/tools/risks.ts
index e9feb7bbd9..57d9b7a8d9 100644
--- a/apps/app/src/data/tools/risks.ts
+++ b/apps/app/src/data/tools/risks.ts
@@ -45,7 +45,16 @@ export const getRisks = tool({
 			},
 		});
 
-		return risks;
+		if (risks.length === 0) {
+			return {
+				risks: [],
+				message: "No risks found",
+			};
+		}
+
+		return {
+			risks,
+		};
 	},
 });
 
@@ -119,6 +128,15 @@ export const getRiskById = tool({
 			},
 		});
 
-		return risk;
+		if (!risk) {
+			return {
+				risk: null,
+				message: "Risk not found",
+			};
+		}
+
+		return {
+			risk,
+		};
 	},
 });
diff --git a/apps/app/src/hooks/use-streamable-text.tsx b/apps/app/src/hooks/use-streamable-text.tsx
index 6836de40b2..0f11a89abc 100644
--- a/apps/app/src/hooks/use-streamable-text.tsx
+++ b/apps/app/src/hooks/use-streamable-text.tsx
@@ -1,12 +1,26 @@
 import { type StreamableValue, readStreamableValue } from "ai/rsc";
-import { useEffect, useState } from "react";
+import { useEffect, useState, useCallback, useTransition } from "react";
 
-export const useStreamableText = (
+function useDebounce<T>(value: T, delay: number): T {
+  const [debouncedValue, setDebouncedValue] = useState<T>(value);
+
+  useEffect(() => {
+    const timer = setTimeout(() => setDebouncedValue(value), delay);
+    return () => clearTimeout(timer);
+  }, [value, delay]);
+
+  return debouncedValue;
+}
+
+export function useStreamableText(
   content: string | StreamableValue<string>,
-) => {
+  debounceMs = 200,
+): string {
   const [rawContent, setRawContent] = useState(
     typeof content === "string" ? content : "",
   );
+  const [isPending, startTransition] = useTransition();
+  const debouncedContent = useDebounce(rawContent, debounceMs);
 
   useEffect(() => {
     if (typeof content === "string") {
@@ -14,23 +28,33 @@ export const useStreamableText = (
       return;
     }
 
-    let isMounted = true;
+    const controller = new AbortController();
+    const { signal } = controller;
 
     (async () => {
       let accumulated = "";
-      for await (const delta of readStreamableValue(content)) {
-        if (!isMounted) break;
-        if (typeof delta === "string") {
-          accumulated += delta;
-          setRawContent(accumulated);
+      try {
+        for await (const delta of readStreamableValue(content)) {
+          if (signal.aborted) break;
+          if (typeof delta === "string") {
+            accumulated += delta;
+            startTransition(() => {
+              setRawContent(accumulated);
+            });
+          }
+        }
+      } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+          return;
         }
+        throw error;
       }
     })();
 
     return () => {
-      isMounted = false;
+      controller.abort();
     };
   }, [content]);
 
-  return rawContent;
-};
+  return debouncedContent;
+}