Skip to content

Feat/ai chat file mentions clean #468

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Bhagy-Yelleti wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Feat/ai chat file mentions clean #468

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9cbcd75
feat: add file mentions support to AI chat
Bhagy-Yelleti Jun 4, 2026
24b36bb
feat: add file mentions support in AI chat
Bhagy-Yelleti Jun 4, 2026
ec33e65
fix: address file mention review feedback
Bhagy-Yelleti Jun 4, 2026
71ca052
fix: address remaining review feedback
Bhagy-Yelleti Jun 4, 2026
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
368 changes: 220 additions & 148 deletions app/api/chat/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,170 +24,242 @@ WORKFLOW for every request that involves code:

If the user asks you to create a new file, call the edit tool with the full content immediately. Do NOT tell the user what code to write - write it yourself using the tool.`;


const MAX_MENTIONED_FILES_TOTAL_SIZE = 100_000;

const RequestBodySchema = z.object({
messages: z.array(z.any()).max(100),
provider: z.enum(["gemini", "groq", "mistral"]).optional().default("gemini"),
fileTree: z.string().max(50_000).optional(),
userApiKey: z.string().max(256).optional(),
messages: z.array(z.any()).max(100),
provider: z.enum(["gemini", "groq", "mistral"]).optional().default("gemini"),
fileTree: z.string().max(50_000).optional(),
userApiKey: z.string().max(256).optional(),
mentionedFiles: z
.array(
z.object({
path: z.string().max(500),
content: z.string().max(20_000),
}),
)
.max(10)
.optional(),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
});

/**
* HTTP POST handler for the AI chat endpoint. Validates request body,
* enforces rate limits, selects model provider, and streams model output.
*/
export async function POST(request: NextRequest) {
try {

// Rate limiting: 20 requests per minute per IP
const ip = getClientIp(request);
const { allowed, remaining } = await rateLimit(ip, 20, 60_000);

if (!allowed) {
return NextResponse.json(
{ success: false, error: "Rate limit exceeded. Please wait before sending more messages." },
{
status: 429,
headers: {
"Retry-After": "60",
"X-RateLimit-Remaining": String(remaining),
},
}
);
}
try {
// Rate limiting: 20 requests per minute per IP
const ip = getClientIp(request);
const { allowed, remaining } = await rateLimit(ip, 20, 60_000);

const session = await auth();
const isAuthenticated = !!session?.user;

const body = await request.json();
const result = RequestBodySchema.safeParse(body);

if (!result.success) {
return NextResponse.json(
{ success: false, error: "Invalid request", details: result.error.issues },
{ status: 400 }
);
}
if (!allowed) {
return NextResponse.json(
{
success: false,
error:
"Rate limit exceeded. Please wait before sending more messages.",
},
{
status: 429,
headers: {
"Retry-After": "60",
"X-RateLimit-Remaining": String(remaining),
},
},
);
}

const { messages, provider, fileTree, userApiKey } = result.data;
const session = await auth();
const isAuthenticated = !!session?.user;

if (!session?.user?.id && (!userApiKey || userApiKey.trim() === "")) {
return NextResponse.json(
{ success: false, error: "Unauthorized: Please log in or provide your own API key in settings." },
{ status: 401 }
);
}
const body = await request.json();
const result = RequestBodySchema.safeParse(body);

const systemInstruction = fileTree
? `${SYSTEM_PROMPT}\n\nProject file tree:\n${fileTree}`
: SYSTEM_PROMPT;

let model;

if (provider === "gemini") {
const apiKey = userApiKey || (isAuthenticated ? process.env.GEMINI_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Gemini API key not configured. Add your key in AI settings."
: "Unauthorized",
},
{ status: isAuthenticated ? 400 : 401 }
);
}
const google = createGoogleGenerativeAI({ apiKey });
model = google("gemini-2.0-flash");
} else if (provider === "groq") {
const apiKey = userApiKey || (isAuthenticated ? process.env.GROQ_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Groq API key not configured. Add your key in AI settings."
: "Unauthorized"
},
{ status: isAuthenticated ? 400 : 401 }
);
}
const groq = createGroq({ apiKey });
model = groq("llama-3.1-70b-versatile");
} else if (provider === "mistral") {
const apiKey = userApiKey || (isAuthenticated ? process.env.MISTRAL_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Mistral API key not configured. Add your key in AI settings."
: "Unauthorized"
},
{ status: isAuthenticated ? 400 : 401 }
);
}
const mistral = createMistral({ apiKey });
model = mistral("mistral-small-latest");
} else {
return NextResponse.json(
{ success: false, error: "Invalid provider" },
{ status: 400 }
);
}
if (!result.success) {
return NextResponse.json(
{
success: false,
error: "Invalid request",
details: result.error.issues,
},
{ status: 400 },
);
}

type MessagePart = { type: string; text: string };

const validRoles = ["system", "user", "assistant", "data", "tool"];
const sanitizedMessages: Omit<UIMessage, 'id'>[] = [];
for (const raw of messages) {
if (!raw || typeof raw !== "object") {
return NextResponse.json(
{ success: false, error: "Invalid request: each message must be an object" },
{ status: 400 }
);
}

const role = (raw as Record<string, unknown>).role;
if (typeof role !== "string" || !validRoles.includes(role)) {
return NextResponse.json(
{ success: false, error: "Invalid request: each message must have a valid role" },
{ status: 400 }
);
}

const m = raw as { role: "system" | "user" | "assistant" | "data" | "tool"; content?: string; parts?: MessagePart[] };
if (Array.isArray(m.parts)) {
// Ensure each part has at least a type property
if (!m.parts.every(p => p && typeof p === "object" && "type" in p)) {
return NextResponse.json(
{ success: false, error: "Invalid request: malformed parts" },
{ status: 400 }
);
}
sanitizedMessages.push({ role: m.role as UIMessage['role'], parts: m.parts as UIMessage['parts'] });
continue;
}
sanitizedMessages.push({
role: m.role as UIMessage['role'],
parts: typeof m.content === "string" && m.content.trim()
? [{ type: "text" as const, text: m.content }]
: [],
});
}
const { messages, provider, fileTree, userApiKey, mentionedFiles } =
result.data;

if (!session?.user?.id && (!userApiKey || userApiKey.trim() === "")) {
return NextResponse.json(
{
success: false,
error:
"Unauthorized: Please log in or provide your own API key in settings.",
},
{ status: 401 },
);
}

// validate total size of mentioned files
if (mentionedFiles && mentionedFiles.length > 0) {
const totalSize = mentionedFiles.reduce(
(sum, file) => sum + file.content.length,
0,
);
if (totalSize > MAX_MENTIONED_FILES_TOTAL_SIZE) {
return NextResponse.json(
{
success: false,
error:
"Referenced files exceed the maximum total size. Reduce file sizes or number of mentioned files.",
},
{ status: 400 },
);
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

let systemInstruction = SYSTEM_PROMPT;

const resultStream = streamText({
model,
messages: await convertToModelMessages(sanitizedMessages, {
ignoreIncompleteToolCalls: true
}),
system: systemInstruction,
tools,
if (fileTree) {
systemInstruction += `\n\nProject file tree:\n${fileTree}`;
}

let model;

if (provider === "gemini") {
const apiKey =
userApiKey ||
(isAuthenticated ? process.env.GEMINI_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Gemini API key not configured. Add your key in AI settings."
: "Unauthorized",
},
{ status: isAuthenticated ? 400 : 401 },
);
}
const google = createGoogleGenerativeAI({ apiKey });
model = google("gemini-2.0-flash");
} else if (provider === "groq") {
const apiKey =
userApiKey || (isAuthenticated ? process.env.GROQ_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Groq API key not configured. Add your key in AI settings."
: "Unauthorized",
},
{ status: isAuthenticated ? 400 : 401 },
);
}
const groq = createGroq({ apiKey });
model = groq("llama-3.3-70b-versatile");
} else if (provider === "mistral") {
const apiKey =
userApiKey ||
(isAuthenticated ? process.env.MISTRAL_API_KEY : undefined);
if (!apiKey) {
return NextResponse.json(
{
success: false,
error: isAuthenticated
? "Mistral API key not configured. Add your key in AI settings."
: "Unauthorized",
},
{ status: isAuthenticated ? 400 : 401 },
);
}
const mistral = createMistral({ apiKey });
model = mistral("mistral-small-latest");
} else {
return NextResponse.json(
{ success: false, error: "Invalid provider" },
{ status: 400 },
);
}

type MessagePart = { type: string; text: string };

const validRoles = ["system", "user", "assistant", "data", "tool"];
const sanitizedMessages: Omit<UIMessage, "id">[] = [];
for (const raw of messages) {
if (!raw || typeof raw !== "object") {
return NextResponse.json(
{
success: false,
error: "Invalid request: each message must be an object",
},
{ status: 400 },
);
}

const role = (raw as Record<string, unknown>).role;
if (typeof role !== "string" || !validRoles.includes(role)) {
return NextResponse.json(
{
success: false,
error: "Invalid request: each message must have a valid role",
},
{ status: 400 },
);
}

const m = raw as {
role: "system" | "user" | "assistant" | "data" | "tool";
content?: string;
parts?: MessagePart[];
};
if (Array.isArray(m.parts)) {
// Ensure each part has at least a type property
if (!m.parts.every((p) => p && typeof p === "object" && "type" in p)) {
return NextResponse.json(
{ success: false, error: "Invalid request: malformed parts" },
{ status: 400 },
);
}
sanitizedMessages.push({
role: m.role as UIMessage["role"],
parts: m.parts as UIMessage["parts"],
});
continue;
}
sanitizedMessages.push({
role: m.role as UIMessage["role"],
parts:
typeof m.content === "string" && m.content.trim()
? [{ type: "text" as const, text: m.content }]
: [],
});
}

return resultStream.toUIMessageStreamResponse();
} catch (error: unknown) {
return handleApiError(error, "POST /api/chat");
// add referenced files as a user message if they exist
if (mentionedFiles && mentionedFiles.length > 0) {
let referencedFilesContent = "Referenced files:\n\n";
for (const file of mentionedFiles) {
referencedFilesContent += `File: ${file.path}\n\n${file.content}\n\n----------------------------------------\n\n`;
}
sanitizedMessages.push({
role: "user",
parts: [{ type: "text", text: referencedFilesContent }],
});
}

const resultStream = streamText({
model,
messages: await convertToModelMessages(sanitizedMessages, {
ignoreIncompleteToolCalls: true,
}),
system: systemInstruction,
tools,
});

return resultStream.toUIMessageStreamResponse();
} catch (error: unknown) {
return handleApiError(error, "POST /api/chat");
}
}
Loading
Loading