[dev] [Itsnotaka] daniel/ai-policy-edit (#1838)

* feat(api): add AI chat endpoint for policy editing assistance, initial draft for ai policy edits

* fix: type error

* feat(policy-editor): integrate AI-assisted policy editing with markdown support

* refactor(api): streamline POST function and enhance markdown guidelines

* refactor(policy-editor): improve policy details layout and diff viewer integration

* refactor(policy-editor): simplify policy details component and enhance AI assistant integration

* refactor(policy-editor): remove unused AI assistant logic and simplify component structure

* feat(ui): add new components to package.json for diff viewer and AI elements

* chore: update lockfile

* refactor(tsconfig): reorganize compiler options and update paths

* fix(policies): resolve infinite loop in policy AI assistant

* fix(api): update policy editing assistant instructions and tool usage

---------

Co-authored-by: Daniel Fu <[email protected]>
Co-authored-by: Amp <[email protected]>

Author: 3 people

.husky/commit-msg

@@ -1,4 +1 @@
-#!/usr/bin/env sh
-. "$(dirname "$0")/_/husky.sh"
-
 npx commitlint --edit $1

apps/app/src/app/(app)/[orgId]/policies/[policyId]/editor/components/PolicyDetails.tsx

@@ -1,15 +1,22 @@
 'use client';
 
 import { PolicyEditor } from '@/components/editor/policy-editor';
+import { useChat } from '@ai-sdk/react';
 import { Button } from '@comp/ui/button';
 import { Card, CardContent } from '@comp/ui/card';
-
 import { DiffViewer } from '@comp/ui/diff-viewer';
 import { validateAndFixTipTapContent } from '@comp/ui/editor';
 import '@comp/ui/editor.css';
 import { Tabs, TabsContent, TabsList, TabsTrigger } from '@comp/ui/tabs';
 import type { PolicyDisplayFormat } from '@db';
 import type { JSONContent } from '@tiptap/react';
+import {
+  DefaultChatTransport,
+  getToolName,
+  isToolUIPart,
+  type ToolUIPart,
+  type UIMessage,
+} from 'ai';
 import { structuredPatch } from 'diff';
 import { CheckCircle, Loader2, Sparkles, X } from 'lucide-react';
 import { useAction } from 'next-safe-action/hooks';
@@ -22,6 +29,50 @@ import { updatePolicy } from '../actions/update-policy';
 import { markdownToTipTapJSON } from './ai/markdown-utils';
 import { PolicyAiAssistant } from './ai/policy-ai-assistant';
 
+function mapChatErrorToMessage(error: unknown): string {
+  const e = error as { status?: number };
+  const status = e?.status;
+
+  if (status === 401 || status === 403) {
+    return "You don't have access to this policy's AI assistant.";
+  }
+  if (status === 404) {
+    return 'This policy could not be found. It may have been removed.';
+  }
+  if (status === 429) {
+    return 'Too many requests. Please wait a moment and try again.';
+  }
+  return 'The AI assistant is currently unavailable. Please try again.';
+}
+
+interface LatestProposal {
+  key: string;
+  content: string;
+  summary: string;
+}
+
+function getLatestProposedPolicy(messages: UIMessage[]): LatestProposal | null {
+  const lastAssistantMessage = [...messages].reverse().find((m) => m.role === 'assistant');
+  if (!lastAssistantMessage?.parts) return null;
+
+  let latest: LatestProposal | null = null;
+
+  lastAssistantMessage.parts.forEach((part, index) => {
+    if (!isToolUIPart(part) || getToolName(part) !== 'proposePolicy') return;
+    const toolPart = part as ToolUIPart;
+    const input = toolPart.input as { content?: string; summary?: string } | undefined;
+    if (!input?.content) return;
+
+    latest = {
+      key: `${lastAssistantMessage.id}:${index}`,
+      content: input.content,
+      summary: input.summary ?? 'Proposing policy changes',
+    };
+  });
+
+  return latest;
+}
+
 interface PolicyContentManagerProps {
   policyId: string;
   policyContent: JSONContent | JSONContent[];
@@ -46,10 +97,37 @@ export function PolicyContentManager({
     return formattedContent;
   });
 
-  const [proposedPolicyMarkdown, setProposedPolicyMarkdown] = useState<string | null>(null);
+  const [dismissedProposalKey, setDismissedProposalKey] = useState<string | null>(null);
   const [isApplying, setIsApplying] = useState(false);
+  const [chatErrorMessage, setChatErrorMessage] = useState<string | null>(null);
   const isAiPolicyAssistantEnabled = useFeatureFlagEnabled('is-ai-policy-assistant-enabled');
 
+  const {
+    messages,
+    status,
+    sendMessage: baseSendMessage,
+  } = useChat({
+    transport: new DefaultChatTransport({
+      api: `/api/policies/${policyId}/chat`,
+    }),
+    onError(error) {
+      console.error('Policy AI chat error:', error);
+      setChatErrorMessage(mapChatErrorToMessage(error));
+    },
+  });
+
+  const sendMessage = (payload: { text: string }) => {
+    setChatErrorMessage(null);
+    baseSendMessage(payload);
+  };
+
+  const latestProposal = useMemo(() => getLatestProposedPolicy(messages), [messages]);
+
+  const activeProposal =
+    latestProposal && latestProposal.key !== dismissedProposalKey ? latestProposal : null;
+
+  const proposedPolicyMarkdown = activeProposal?.content ?? null;
+
   const switchFormat = useAction(switchPolicyDisplayFormatAction, {
     onError: () => toast.error('Failed to switch view.'),
   });
@@ -65,15 +143,17 @@ export function PolicyContentManager({
   }, [currentPolicyMarkdown, proposedPolicyMarkdown]);
 
   async function applyProposedChanges() {
-    if (!proposedPolicyMarkdown) return;
+    if (!activeProposal) return;
+
+    const { content, key } = activeProposal;
 
     setIsApplying(true);
     try {
-      const jsonContent = markdownToTipTapJSON(proposedPolicyMarkdown);
+      const jsonContent = markdownToTipTapJSON(content);
       await updatePolicy({ policyId, content: jsonContent });
       setCurrentContent(jsonContent);
       setEditorKey((prev) => prev + 1);
-      setProposedPolicyMarkdown(null);
+      setDismissedProposalKey(key);
       toast.success('Policy updated with AI suggestions');
     } catch (err) {
       console.error('Failed to apply changes:', err);
@@ -141,8 +221,10 @@ export function PolicyContentManager({
             {showAiAssistant && isAiPolicyAssistantEnabled && (
               <div className="w-80 shrink-0 min-h-[400px] self-stretch flex flex-col">
                 <PolicyAiAssistant
-                  policyId={policyId}
-                  onProposedPolicyChange={setProposedPolicyMarkdown}
+                  messages={messages}
+                  status={status}
+                  errorMessage={chatErrorMessage}
+                  sendMessage={sendMessage}
                   close={() => setShowAiAssistant(false)}
                 />
               </div>
@@ -151,10 +233,14 @@ export function PolicyContentManager({
         </CardContent>
       </Card>
 
-      {proposedPolicyMarkdown && diffPatch && (
+      {proposedPolicyMarkdown && diffPatch && activeProposal && (
         <div className="space-y-2">
           <div className="flex items-center justify-end gap-2">
-            <Button variant="ghost" size="sm" onClick={() => setProposedPolicyMarkdown(null)}>
+            <Button
+              variant="ghost"
+              size="sm"
+              onClick={() => setDismissedProposalKey(activeProposal.key)}
+            >
               <X className="h-3 w-3 mr-1" />
               Dismiss
             </Button>

apps/app/src/app/(app)/[orgId]/policies/[policyId]/editor/components/ai/policy-ai-assistant.tsx

@@ -9,58 +9,43 @@ import {
 import { Tool, ToolHeader } from '@comp/ui/ai-elements/tool';
 import { Button } from '@comp/ui/button';
 import { cn } from '@comp/ui/cn';
-import { DefaultChatTransport, getToolName, isToolUIPart } from 'ai';
-import type { ToolUIPart } from 'ai';
-import { useChat } from '@ai-sdk/react';
+import {
+  getToolName,
+  isToolUIPart,
+  type ChatStatus,
+  type ToolUIPart,
+  type UIMessage,
+} from 'ai';
 import { X } from 'lucide-react';
-import { useEffect, useRef, useState } from 'react';
+import { useState } from 'react';
 
 interface PolicyAiAssistantProps {
-  policyId: string;
-  onProposedPolicyChange?: (content: string | null) => void;
+  messages: UIMessage[];
+  status: ChatStatus;
+  errorMessage?: string | null;
+  sendMessage: (payload: { text: string }) => void;
   close?: () => void;
 }
 
 export function PolicyAiAssistant({
-  policyId,
-  onProposedPolicyChange,
+  messages,
+  status,
+  errorMessage,
+  sendMessage,
   close,
 }: PolicyAiAssistantProps) {
   const [input, setInput] = useState('');
-  const lastProcessedToolCallRef = useRef<string | null>(null);
-  
-  const { messages, status, error, sendMessage } = useChat({
-    transport: new DefaultChatTransport({
-      api: `/api/policies/${policyId}/chat`,
-    }),
-  });
 
   const isLoading = status === 'streaming' || status === 'submitted';
 
-  useEffect(() => {
-    const lastAssistantMessage = [...messages].reverse().find((m) => m.role === 'assistant');
-    if (!lastAssistantMessage?.parts) return;
-
-    for (const part of lastAssistantMessage.parts) {
-      if (isToolUIPart(part) && getToolName(part) === 'proposePolicy') {
-        const toolInput = part.input as { content: string; summary: string };
-        
-        if (part.state === 'input-streaming') {
-          onProposedPolicyChange?.(toolInput?.content || '');
-          continue;
-        }
-        
-        if (lastProcessedToolCallRef.current === part.toolCallId) {
-          continue;
-        }
-        
-        if (toolInput?.content) {
-          lastProcessedToolCallRef.current = part.toolCallId;
-          onProposedPolicyChange?.(toolInput.content);
-        }
-      }
-    }
-  }, [messages, onProposedPolicyChange]);
+  const hasActiveTool = messages.some(
+    (m) =>
+      m.role === 'assistant' &&
+      m.parts.some(
+        (p) =>
+          isToolUIPart(p) && (p.state === 'input-streaming' || p.state === 'input-available'),
+      ),
+  );
 
   const handleSubmit = () => {
     if (!input.trim()) return;
@@ -108,10 +93,10 @@ export function PolicyAiAssistant({
                     </div>
                   );
                 }
-                
+
                 if (isToolUIPart(part) && getToolName(part) === 'proposePolicy') {
                   const toolPart = part as ToolUIPart;
-                  const toolInput = part.input as { content: string; summary: string };
+                  const toolInput = toolPart.input as { content?: string; summary?: string };
                   return (
                     <Tool key={`${message.id}-${index}`} className="mt-2">
                       <ToolHeader
@@ -122,20 +107,20 @@ export function PolicyAiAssistant({
                     </Tool>
                   );
                 }
-                
+
                 return null;
               })}
             </div>
           ))
         )}
-        {isLoading && (
+        {isLoading && !hasActiveTool && (
           <div className="text-sm text-muted-foreground">Thinking...</div>
         )}
       </div>
 
-      {error && (
+      {errorMessage && (
         <div className="border-t bg-destructive/10 px-3 py-2">
-          <p className="text-xs text-destructive">{error.message}</p>
+          <p className="text-xs text-destructive">{errorMessage}</p>
         </div>
       )}