Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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.',
};
});

Expand Down Expand Up @@ -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`,
}),
Expand All @@ -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.'),
});
Expand Down Expand Up @@ -167,8 +188,8 @@ export function PolicyContentManager({
<div className="space-y-4">
<Card>
<CardContent className="p-4">
<div className="flex gap-4">
<div className="flex-1 min-w-0">
<div className="flex gap-4 h-[60vh]">
<div className="flex-1 min-w-0 h-full overflow-y-auto">
<Tabs
defaultValue={displayFormat}
onValueChange={(format) =>
Expand Down Expand Up @@ -219,22 +240,24 @@ export function PolicyContentManager({
</div>

{showAiAssistant && isAiPolicyAssistantEnabled && (
<div className="w-80 shrink-0 min-h-[400px] self-stretch flex flex-col">
<div className="w-80 shrink-0 self-stretch flex flex-col overflow-hidden">
<PolicyAiAssistant
messages={messages}
status={status}
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"
Expand All @@ -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}`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand All @@ -32,6 +43,8 @@ export function PolicyAiAssistant({
errorMessage,
sendMessage,
close,
onScrollToDiff,
hasActiveProposal,
}: PolicyAiAssistantProps) {
const [input, setInput] = useState('');

Expand All @@ -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'),
),
);

Expand All @@ -52,7 +69,7 @@ export function PolicyAiAssistant({
};

return (
<div className="flex h-full flex-col border-l bg-background">
<div className="flex h-full min-h-0 flex-col border-l bg-background">
<div className="flex items-center justify-between border-b px-3 py-2">
<span className="text-sm font-medium">AI Assistant</span>
{close && (
Expand All @@ -62,7 +79,7 @@ export function PolicyAiAssistant({
)}
</div>

<Conversation>
<Conversation className="min-h-0" aria-label="Policy AI assistant conversation">
<ConversationContent className="gap-3 p-3">
{messages.length === 0 ? (
<div className="text-sm text-muted-foreground py-4">
Expand Down Expand Up @@ -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>
);
Expand Down
Original file line number Diff line number Diff line change
@@ -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>>;
Original file line number Diff line number Diff line change
@@ -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(),
Expand Down
Loading