refactor(policy): update policy details and AI assistant components for improved functionality

Author: Itsnotaka

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

@@ -10,22 +10,16 @@ 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 { DefaultChatTransport } from 'ai';
 import { structuredPatch } from 'diff';
 import { CheckCircle, Loader2, Sparkles, X } from 'lucide-react';
 import { useAction } from 'next-safe-action/hooks';
-import { useFeatureFlagEnabled } from 'posthog-js/react';
-import { useMemo, useState } from 'react';
+import { useCallback, useMemo, useRef, useState } from 'react';
 import { toast } from 'sonner';
 import { switchPolicyDisplayFormatAction } from '../../actions/switch-policy-display-format';
 import { PdfViewer } from '../../components/PdfViewer';
 import { updatePolicy } from '../actions/update-policy';
+import type { PolicyChatUIMessage } from '../types';
 import { markdownToTipTapJSON } from './ai/markdown-utils';
 import { PolicyAiAssistant } from './ai/policy-ai-assistant';
 
@@ -49,24 +43,32 @@ interface LatestProposal {
  key: string;
  content: string;
  summary: string;
+ title: string;
+ detail: string;
+ reviewHint: string;
 }
 
-function getLatestProposedPolicy(messages: UIMessage[]): LatestProposal | null {
+function getLatestProposedPolicy(messages: PolicyChatUIMessage[]): 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 (part.type !== 'tool-proposePolicy') return;
+ if (part.state === 'input-streaming' || part.state === 'output-error') return;
+ const input = part.input;
  if (!input?.content) return;
 
  latest = {
  key: `${lastAssistantMessage.id}:${index}`,
  content: input.content,
  summary: input.summary ?? 'Proposing policy changes',
+ title: input.title ?? input.summary ?? 'Policy updates ready for your review',
+ detail:
+ input.detail ??
+ 'I have prepared an updated version of this policy based on your instructions.',
+ reviewHint: input.reviewHint ?? 'Review the proposed changes below before applying them.',
  };
  });
 
@@ -100,13 +102,18 @@ export function PolicyContentManager({
  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 diffViewerRef = useRef<HTMLDivElement>(null);
+
+ const isAiPolicyAssistantEnabled = true;
+ const scrollToDiffViewer = useCallback(() => {
+ diffViewerRef.current?.scrollIntoView({ behavior: 'smooth', block: 'start' });
+ }, []);
 
  const {
  messages,
  status,
  sendMessage: baseSendMessage,
- } = useChat({
+ } = useChat<PolicyChatUIMessage>({
  transport: new DefaultChatTransport({
  api: `/api/policies/${policyId}/chat`,
  }),
@@ -128,6 +135,20 @@ export function PolicyContentManager({
 
  const proposedPolicyMarkdown = activeProposal?.content ?? null;
 
+ const hasPendingProposal = useMemo(
+ () =>
+ messages.some(
+ (m) =>
+ m.role === 'assistant' &&
+ m.parts?.some(
+ (part) =>
+ part.type === 'tool-proposePolicy' &&
+ (part.state === 'input-streaming' || part.state === 'input-available'),
+ ),
+ ),
+ [messages],
+ );
+
  const switchFormat = useAction(switchPolicyDisplayFormatAction, {
  onError: () => toast.error('Failed to switch view.'),
  });
@@ -226,15 +247,17 @@ export function PolicyContentManager({
  errorMessage={chatErrorMessage}
  sendMessage={sendMessage}
  close={() => setShowAiAssistant(false)}
+ onScrollToDiff={scrollToDiffViewer}
+ hasActiveProposal={!!activeProposal && !hasPendingProposal}
  />
  </div>
  )}
  </div>
  </CardContent>
  </Card>
 
- {proposedPolicyMarkdown && diffPatch && activeProposal && (
- <div className="space-y-2">
+ {proposedPolicyMarkdown && diffPatch && activeProposal && !hasPendingProposal && (
+ <div ref={diffViewerRef} className="space-y-2">
  <div className="flex items-center justify-end gap-2">
  <Button
  variant="ghost"
@@ -261,7 +284,7 @@ export function PolicyContentManager({
 }
 
 function createGitPatch(fileName: string, oldStr: string, newStr: string): string {
- const patch = structuredPatch(fileName, fileName, oldStr, newStr);
+ const patch = structuredPatch(fileName, fileName, oldStr, newStr, '', '', { context: 1 });
  const lines: string[] = [
  `diff --git a/${fileName} b/${fileName}`,
  `--- a/${fileName}`,

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

@@ -12,18 +12,29 @@ import {
  PromptInputTextarea,
 } from '@comp/ui/ai-elements/prompt-input';
 import { Tool, ToolHeader } from '@comp/ui/ai-elements/tool';
+import { Badge } from '@comp/ui/badge';
 import { Button } from '@comp/ui/button';
 import { cn } from '@comp/ui/cn';
-import { getToolName, isToolUIPart, type ChatStatus, type ToolUIPart, type UIMessage } from 'ai';
-import { X } from 'lucide-react';
+import type { ChatStatus } from 'ai';
+import {
+ ArrowDownIcon,
+ CheckCircleIcon,
+ CircleIcon,
+ ClockIcon,
+ X,
+ XCircleIcon,
+} from 'lucide-react';
 import { useState } from 'react';
+import type { PolicyChatUIMessage } from '../../types';
 
 interface PolicyAiAssistantProps {
- messages: UIMessage[];
+ messages: PolicyChatUIMessage[];
  status: ChatStatus;
  errorMessage?: string | null;
  sendMessage: (payload: { text: string }) => void;
  close?: () => void;
+ onScrollToDiff?: () => void;
+ hasActiveProposal?: boolean;
 }
 
 export function PolicyAiAssistant({
@@ -32,6 +43,8 @@ export function PolicyAiAssistant({
  errorMessage,
  sendMessage,
  close,
+ onScrollToDiff,
+ hasActiveProposal,
 }: PolicyAiAssistantProps) {
  const [input, setInput] = useState('');
 
@@ -41,7 +54,11 @@ export function PolicyAiAssistant({
  (m) =>
  m.role === 'assistant' &&
  m.parts.some(
- (p) => isToolUIPart(p) && (p.state === 'input-streaming' || p.state === 'input-available'),
+ (p) =>
+ p.type === 'tool-proposePolicy' &&
+ (p.state === 'input-streaming' ||
+ p.state === 'output-available' ||
+ p.state === 'output-error'),
  ),
  );
 
@@ -93,18 +110,92 @@ export function PolicyAiAssistant({
  );
  }
 
- if (isToolUIPart(part) && getToolName(part) === 'proposePolicy') {
- const toolPart = part as ToolUIPart;
- const toolInput = toolPart.input as { content?: string; summary?: string };
+ if (part.type === 'tool-proposePolicy') {
+ const toolInput = part.input;
+
+ const isInProgress =
+ part.state === 'input-streaming' || part.state === 'input-available';
+ const isCompleted = part.state === 'output-available';
+
+ const title =
+ (isCompleted
+ ? toolInput?.title || toolInput?.summary
+ : toolInput?.title || 'Configuring policy updates') ||
+ 'Policy updates ready for your review';
+
+ const bodyText = (() => {
+ if (isInProgress) {
+ return (
+ toolInput?.detail ||
+ 'I am preparing an updated version of this policy. Please wait a moment before accepting any changes.'
+ );
+ }
+ if (isCompleted) {
+ return (
+ toolInput?.detail ||
+ 'The updated policy is ready. Review the diff in the editor before applying changes.'
+ );
+ }
+ return (
+ toolInput?.detail ||
+ 'Review the proposed changes in the editor preview below before applying them.'
+ );
+ })();
+
+ const truncatedBodyText =
+ bodyText.length > 180 ? `${bodyText.slice(0, 177)}…` : bodyText;
+
+ type ToolState = typeof part.state;
+ const statusPill = (() => {
+ const labels: Record<ToolState, string> = {
+ 'input-streaming': 'Drafting',
+ 'input-available': 'Running',
+ 'output-available': 'Completed',
+ 'output-error': 'Error',
+ };
+
+ const icons: Record<ToolState, React.ReactNode> = {
+ 'input-streaming': <CircleIcon className="size-3" />,
+ 'input-available': <ClockIcon className="size-3 animate-pulse" />,
+ 'output-available': <CheckCircleIcon className="size-3 text-emerald-600" />,
+ 'output-error': <XCircleIcon className="size-3 text-red-600" />,
+ };
+
+ return (
+ <Badge
+ className="gap-1.5 rounded-full border border-border/60 bg-background/80 px-2 py-0.5 text-[10px] uppercase tracking-[0.16em]"
+ variant="secondary"
+ >
+ {icons[part.state]}
+ {labels[part.state]}
+ </Badge>
+ );
+ })();
+
  return (
  <Tool key={`${message.id}-${index}`} className="mt-2">
  <ToolHeader
- title={toolInput?.summary || 'Proposing policy changes'}
- type={toolPart.type}
- state={toolPart.state}
+ title={title}
+ meta={statusPill}
+ onClick={isCompleted && onScrollToDiff ? onScrollToDiff : undefined}
+ className={
+ isCompleted && onScrollToDiff
+ ? 'cursor-pointer hover:bg-muted/50'
+ : undefined
+ }
  />
- <p className="px-3 pb-2 text-xs text-muted-foreground">
- View the proposed changes in the editor preview
+ <p className="px-3 py-2 text-[10px] text-muted-foreground">
+ {truncatedBodyText}
+ {hasActiveProposal && onScrollToDiff && (
+ <button
+ type="button"
+ onClick={onScrollToDiff}
+ className="flex items-center gap-1.5 text-[11px] text-primary hover:underline"
+ >
+ <ArrowDownIcon className="size-3" />
+ View proposed changes
+ </button>
+ )}
  </p>
  </Tool>
  );

apps/app/src/app/(app)/[orgId]/policies/[policyId]/editor/tools/policy-tools.ts

@@ -0,0 +1,45 @@
+import { type InferUITools, tool } from 'ai';
+import { z } from 'zod';
+
+export function getPolicyTools() {
+ return {
+ proposePolicy: tool({
+ description:
+ 'Propose an updated version of the policy. Use this tool whenever the user asks you to make changes, edits, or improvements to the policy. You must provide the COMPLETE policy content, not just the changes.',
+ inputSchema: z.object({
+ content: z
+ .string()
+ .describe(
+ 'The complete updated policy content in markdown format. Must include the entire policy, not just the changed sections.',
+ ),
+ summary: z
+ .string()
+ .describe('One to two sentences summarizing the changes. No bullet points.'),
+ title: z
+ .string()
+ .describe(
+ 'A short, sentence-case heading (~4–10 words) that clearly states the main change, for use in a small review banner.',
+ ),
+ detail: z
+ .string()
+ .describe(
+ 'One or two plain-text sentences briefly explaining what changed and why, shown in the review banner.',
+ ),
+ reviewHint: z
+ .string()
+ .describe(
+ 'A very short imperative phrase that tells the user to review the updated policy content in the editor below.',
+ ),
+ }),
+ execute: async ({ summary, title, detail, reviewHint }) => ({
+ success: true,
+ summary,
+ title,
+ detail,
+ reviewHint,
+ }),
+ }),
+ };
+}
+
+export type PolicyToolSet = InferUITools<ReturnType<typeof getPolicyTools>>;

apps/app/src/app/(app)/[orgId]/policies/[policyId]/editor/types/index.ts

@@ -1,4 +1,8 @@
+import type { UIMessage } from 'ai';
 import { z } from 'zod';
+import type { PolicyToolSet } from '../tools/policy-tools';
+
+export type PolicyChatUIMessage = UIMessage<never, never, PolicyToolSet>;
 
 export const policyDetailsSchema = z.object({
  id: z.string(),