chore: ai, public course and other ui improvements

apps/api/src/services/agent/agent-tool-schemas.ts

@@ -0,0 +1,239 @@
+import { z } from 'zod';
+import { CoursePlanSchema } from '@cio/ai-assistant';
+import { ZExerciseSectionAfterBehavior } from '@cio/utils/validation/exercise';
+import { ZCourseLandingPageUpdate, ZCourseLandingPageMetadataUpdate } from '@cio/utils/validation/course';
+
+// courseId is NOT a parameter — it's injected from the authenticated request context.
+// This prevents prompt injection from tricking the LLM into targeting another course.
+
+export const emptyParam = z.object({});
+export const lessonReadParam = z.object({ lessonId: z.string(), locale: z.string().default('en') });
+export const exerciseReadParam = z.object({ exerciseId: z.string() });
+export const createSectionParam = z.object({ title: z.string().min(1), order: z.number().int().min(0) });
+export const updateSectionParam = z
+  .object({
+    sectionId: z.string(),
+    title: z.string().min(1).optional(),
+    order: z.number().int().min(0).optional()
+  })
+  .refine((data) => data.title !== undefined || data.order !== undefined, {
+    message: 'Provide at least one field to update'
+  });
+export const createLessonParam = z.object({
+  sectionId: z.string(),
+  title: z.string().min(1),
+  order: z.number().int().min(0)
+});
+export const updateLessonParam = z
+  .object({
+    lessonId: z.string(),
+    title: z.string().min(1).optional(),
+    sectionId: z.string().optional(),
+    order: z.number().int().min(0).optional(),
+    lessonAt: z.string().optional(),
+    callUrl: z.string().optional(),
+    isUnlocked: z.boolean().optional(),
+    public: z.boolean().optional()
+  })
+  .refine(
+    (data) =>
+      data.title !== undefined ||
+      data.sectionId !== undefined ||
+      data.order !== undefined ||
+      data.lessonAt !== undefined ||
+      data.callUrl !== undefined ||
+      data.isUnlocked !== undefined ||
+      data.public !== undefined,
+    {
+      message: 'Provide at least one field to update'
+    }
+  );
+export const updateContentParam = z.object({
+  lessonId: z.string(),
+  locale: z.string().default('en'),
+  content: z.string().min(1)
+});
+
+export const questionSchema = z.object({
+  question: z.string().min(1),
+  questionTypeId: z.number().int().min(1).max(14).default(1),
+  points: z.number().min(0).default(1),
+  order: z.number().int().min(0),
+  options: z.array(z.object({ label: z.string().min(1), isCorrect: z.boolean() }))
+});
+
+export const createExerciseParam = z.object({
+  lessonId: z.string().optional(),
+  sectionId: z.string().optional(),
+  title: z.string().min(1),
+  description: z.string().optional().describe('Optional short description shown to students above the questions.'),
+  order: z
+    .number()
+    .int()
+    .min(0)
+    .optional()
+    .describe(
+      'Display order within the section (0-based). Required when sectionId is provided. Lessons and exercises share the same order space within a section.'
+    ),
+  questions: z.array(questionSchema)
+});
+
+export const updateExerciseParam = z
+  .object({
+    exerciseId: z.string().describe('The exercise to update. Must be a real UUID returned by a prior tool call.'),
+    title: z.string().min(1).optional().describe('New exercise title. Omit to keep unchanged.'),
+    description: z
+      .string()
+      .optional()
+      .describe('New short description shown to students. Pass an empty string to clear it.'),
+    lessonId: z.string().optional().describe('Link the exercise to a lesson in this course.'),
+    sectionId: z.string().optional().describe('Move the exercise to a section in this course.'),
+    order: z.number().int().min(0).optional().describe('New order within the section (0-based).'),
+    dueBy: z
+      .string()
+      .optional()
+      .describe('Due date in ISO 8601 format (e.g. 2026-05-01T23:59:00Z). Omit to keep unchanged.'),
+    isUnlocked: z.boolean().optional().describe('Whether the exercise is unlocked for students.'),
+    allowMultipleAttempts: z.boolean().optional().describe('Whether students can re-take the exercise.')
+  })
+  .refine(
+    (data) =>
+      data.title !== undefined ||
+      data.description !== undefined ||
+      data.lessonId !== undefined ||
+      data.sectionId !== undefined ||
+      data.order !== undefined ||
+      data.dueBy !== undefined ||
+      data.isUnlocked !== undefined ||
+      data.allowMultipleAttempts !== undefined,
+    {
+      message: 'Provide at least one field to update'
+    }
+  );
+
+export const updateExerciseSectionParam = z
+  .object({
+    exerciseId: z
+      .string()
+      .describe('The exercise that contains this section. Must match the exercise you read with get_exercise_details.'),
+    exerciseSectionId: z
+      .string()
+      .describe('The exercise section id from get_exercise_details `sections[].id` — not a course section id.'),
+    title: z.string().min(1).optional().describe('New section heading shown above questions in this block.'),
+    description: z.string().optional().describe('Optional intro HTML for this block. Pass an empty string to clear it.')
+  })
+  .refine((data) => data.title !== undefined || data.description !== undefined, {
+    message: 'Provide at least title or description to update'
+  });
+
+export const createExerciseSectionParam = z.object({
+  exerciseId: z.string().describe('The exercise that will contain this section.'),
+  title: z.string().min(1).describe('Section heading shown above the questions in this block.'),
+  description: z
+    .string()
+    .optional()
+    .describe('Optional intro HTML for this block. Pass an empty string to leave it blank.'),
+  order: z.number().int().min(0).describe('Display order within the exercise (0-based).'),
+  colorTheme: z.enum(['blue', 'green', 'amber', 'rose', 'violet', 'slate']).optional(),
+  afterBehavior: ZExerciseSectionAfterBehavior.optional()
+});
+
+export const addQuestionsParam = z.object({
+  exerciseId: z.string(),
+  exerciseSectionId: z
+    .string()
+    .uuid()
+    .optional()
+    .describe(
+      'In-exercise section id from get_exercise_details `sections[].id` where the new questions should appear (not a course outline section id).'
+    ),
+  questions: z.array(questionSchema)
+});
+
+export const updateOptionSchema = z.object({
+  id: z.number().int().optional(),
+  label: z.string().min(1).optional(),
+  isCorrect: z.boolean().optional(),
+  settings: z.record(z.string(), z.unknown()).optional()
+});
+
+export const updateQuestionPatchSchema = z.object({
+  id: z.number().int(),
+  question: z.string().min(1).optional(),
+  questionTypeId: z.number().int().min(1).max(14).optional(),
+  points: z.number().min(0).optional(),
+  order: z.number().int().min(0).optional(),
+  exerciseSectionId: z
+    .string()
+    .uuid()
+    .nullable()
+    .optional()
+    .describe(
+      'Move the question to this in-exercise block (get_exercise_details sections[].id). Use null to clear section assignment.'
+    ),
+  settings: z.record(z.string(), z.unknown()).optional(),
+  options: z.array(updateOptionSchema).optional()
+});
+
+export const updateQuestionsParam = z.object({
+  exerciseId: z.string(),
+  questions: z.array(updateQuestionPatchSchema).min(1)
+});
+export const coursePlanParam = z.object({ plan: CoursePlanSchema });
+
+// Gemini's tool-schema validator only accepts string enums, so numeric/boolean
+// `z.literal` values must be relaxed to plain types in the schemas the model sees.
+export const agentLessonTabsOrder = z.array(
+  z.object({
+    id: z.number().int().min(1).max(4),
+    name: z.string()
+  })
+);
+export const agentLandingPageMetadataUpdate = ZCourseLandingPageMetadataUpdate.extend({
+  lessonTabsOrder: agentLessonTabsOrder.optional()
+});
+export const updateCourseLandingPageParam = ZCourseLandingPageUpdate.extend({
+  metadata: agentLandingPageMetadataUpdate.optional()
+}).refine(
+  (data) =>
+    data.title !== undefined ||
+    data.description !== undefined ||
+    data.overview !== undefined ||
+    data.cost !== undefined ||
+    data.currency !== undefined ||
+    data.imageUrl !== undefined ||
+    data.generateImage !== undefined ||
+    data.imageQuery !== undefined ||
+    data.metadata !== undefined,
+  {
+    message: 'Provide at least one landing-page field to update'
+  }
+);
+export const goLiveParam = z.object({
+  confirmPublish: z
+    .boolean()
+    .describe('Must be true. Only use this after the teacher explicitly asks to publish or go live.')
+});
+
+export const reorderContentParam = z.object({
+  sections: z
+    .array(
+      z.object({
+        id: z.string().min(1),
+        order: z.number().int().min(0)
+      })
+    )
+    .optional()
+    .describe('Reorder sections by setting new order values'),
+  items: z
+    .array(
+      z.object({
+        id: z.string().min(1),
+        type: z.enum(['LESSON', 'EXERCISE']),
+        order: z.number().int().min(0).optional().describe('New order within the section'),
+        sectionId: z.string().nullable().optional().describe('Move item to a different section')
+      })
+    )
+    .optional()
+    .describe('Reorder or move lessons/exercises')
+});

apps/api/src/services/agent/chat-context.ts

@@ -0,0 +1,160 @@
+import { AppError } from '@api/utils/errors';
+import { getDocumentText } from '@api/services/agent/document';
+import { redis } from '@api/utils/redis/redis';
+import { getCourseSectionBinding, getExerciseCourseBinding, getLessonCourseBinding } from '@cio/db/queries/agent';
+import { z } from 'zod';
+import { CoursePlanSchema } from '@cio/ai-assistant';
+
+type ResourceOwnershipRow = {
+  courseId: string | null;
+  title?: string | null;
+};
+
+function buildResourceOwnershipError(params: {
+  resourceType: 'Lesson' | 'Exercise' | 'Section';
+  resourceId: string;
+  courseId: string;
+  resource?: ResourceOwnershipRow;
+}) {
+  const { resourceType, resourceId, resource } = params;
+
+  if (!resource) {
+    return new AppError(
+      `${resourceType} ${resourceId} does not exist in this course. The ID may have been hallucinated — call get_course_structure to fetch real IDs and retry.`,
+      'RESOURCE_NOT_IN_COURSE',
+      403
+    );
+  }
+
+  const titleSuffix = resource.title ? ` (${resource.title})` : '';
+
+  return new AppError(
+    `${resourceType} ${resourceId}${titleSuffix} belongs to a different course. Call get_course_structure to fetch IDs for the current course and retry.`,
+    'RESOURCE_NOT_IN_COURSE',
+    403
+  );
+}
+
+const UUID_REGEX = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
+
+export function assertValidUuid(
+  resourceType: 'Lesson' | 'Exercise' | 'Section' | 'ExerciseSection',
+  value: string
+): void {
+  if (!UUID_REGEX.test(value)) {
+    throw new AppError(
+      `${resourceType} id "${value}" is not a valid UUID. Call get_course_structure to fetch real IDs and try again — never invent or guess UUIDs.`,
+      'INVALID_RESOURCE_ID',
+      400
+    );
+  }
+}
+
+export async function verifyLessonBelongsToCourse(lessonId: string, courseId: string): Promise<void> {
+  assertValidUuid('Lesson', lessonId);
+
+  const lesson = await getLessonCourseBinding(lessonId);
+
+  if (!lesson || lesson.courseId !== courseId) {
+    throw buildResourceOwnershipError({
+      resourceType: 'Lesson',
+      resourceId: lessonId,
+      courseId,
+      resource: lesson
+    });
+  }
+}
+
+export async function verifyExerciseBelongsToCourse(exerciseId: string, courseId: string): Promise<void> {
+  assertValidUuid('Exercise', exerciseId);
+
+  const exercise = await getExerciseCourseBinding(exerciseId);
+
+  if (!exercise || exercise.courseId !== courseId) {
+    throw buildResourceOwnershipError({
+      resourceType: 'Exercise',
+      resourceId: exerciseId,
+      courseId,
+      resource: exercise
+    });
+  }
+}
+
+export async function verifySectionBelongsToCourse(sectionId: string, courseId: string): Promise<void> {
+  assertValidUuid('Section', sectionId);
+
+  const section = await getCourseSectionBinding(sectionId);
+
+  if (!section || section.courseId !== courseId) {
+    throw buildResourceOwnershipError({
+      resourceType: 'Section',
+      resourceId: sectionId,
+      courseId,
+      resource: section
+    });
+  }
+}
+
+type AttachedMessage = { metadata?: { attachment?: { documentId?: string } } };
+
+export function collectDocumentIds(messages: unknown[], currentDocumentId?: string): string[] {
+  const ids = new Set<string>();
+
+  for (const msg of messages as AttachedMessage[]) {
+    const id = msg?.metadata?.attachment?.documentId;
+
+    if (id) ids.add(id);
+  }
+
+  if (currentDocumentId) ids.add(currentDocumentId);
+
+  return Array.from(ids);
+}
+
+export async function loadDocumentsText(documentIds: string[], userId: string): Promise<string | undefined> {
+  const loaded = await Promise.all(
+    documentIds.map(async (id) => {
+      const text = await getDocumentText(id, userId, redis);
+
+      return text ? { id, text } : null;
+    })
+  );
+
+  const sections = loaded
+    .filter((d): d is { id: string; text: string } => d !== null)
+    .map((d, i) => `--- Document ${i + 1} (id: ${d.id}) ---\n${d.text}`);
+
+  return sections.length > 0 ? sections.join('\n\n') : undefined;
+}
+
+type PlanMetadataMessage = {
+  role?: string;
+  metadata?: {
+    plan?: {
+      action?: string;
+      payload?: unknown;
+    };
+  };
+};
+
+export function getLatestImplementationPlan(messages: unknown[]): z.infer<typeof CoursePlanSchema> | undefined {
+  for (let index = messages.length - 1; index >= 0; index--) {
+    const message = messages[index] as PlanMetadataMessage;
+
+    if (message?.role !== 'user') {
+      continue;
+    }
+
+    if (message?.metadata?.plan?.action !== 'implement_course_plan') {
+      continue;
+    }
+
+    const parsedPlan = CoursePlanSchema.safeParse(message.metadata.plan.payload);
+
+    if (parsedPlan.success) {
+      return parsedPlan.data;
+    }
+  }
+
+  return undefined;
+}

apps/api/src/services/agent/chat-tools.ts

@@ -0,0 +1,655 @@
+import { tool } from 'ai';
+import { CoursePlanSchema } from '@cio/ai-assistant';
+import { AppError } from '@api/utils/errors';
+import { trackAgentEvent, AgentEvent } from '@api/utils/tinybird';
+import { getCourseContentItems } from '@cio/db/queries/course/content';
+import { getExerciseSectionsByExerciseId } from '@cio/db/queries/exercise';
+import type { TCourseLandingPageUpdate } from '@cio/utils/validation/course';
+import { listCourseSections, createCourseSection, updateCourseSectionService } from '@api/services/course/section';
+import { createLesson, getLesson, updateLessonService } from '@api/services/lesson/lesson';
+import { upsertLessonLanguageService } from '@api/services/lesson-language';
+import {
+  createExercise,
+  getExercise,
+  createExerciseSectionService,
+  updateExerciseService,
+  updateExerciseSectionMetadataService
+} from '@api/services/exercise/exercise';
+import { reorderCourseContent } from '@api/services/course/content';
+import { normalizeAgentLessonContent } from '@api/services/agent/lesson-content';
+import { buildUpdatedQuestions } from '@api/services/agent/question-update';
+import { updateCourseLandingPageService } from '@api/services/course/landing-page';
+import { getCourseGoLiveReadiness, publishCourseWhenReady } from '@api/services/course/go-live-readiness';
+import {
+  assertValidUuid,
+  verifyExerciseBelongsToCourse,
+  verifyLessonBelongsToCourse,
+  verifySectionBelongsToCourse
+} from '@api/services/agent/chat-context';
+import {
+  addQuestionsParam,
+  coursePlanParam,
+  createExerciseParam,
+  createExerciseSectionParam,
+  createLessonParam,
+  createSectionParam,
+  emptyParam,
+  exerciseReadParam,
+  goLiveParam,
+  lessonReadParam,
+  reorderContentParam,
+  updateContentParam,
+  updateCourseLandingPageParam,
+  updateExerciseParam,
+  updateExerciseSectionParam,
+  updateLessonParam,
+  updateQuestionsParam,
+  updateSectionParam
+} from '@api/services/agent/agent-tool-schemas';
+
+function summarizeAgentDebugValue(value: unknown, depth = 0): unknown {
+  if (value == null) return value;
+
+  if (typeof value === 'string') {
+    return value.length > 200 ? `${value.slice(0, 200)}… (${value.length} chars)` : value;
+  }
+
+  if (typeof value === 'number' || typeof value === 'boolean') {
+    return value;
+  }
+
+  if (Array.isArray(value)) {
+    if (depth >= 2) {
+      return `[Array(${value.length})]`;
+    }
+
+    return {
+      count: value.length,
+      items: value.slice(0, 3).map((item) => summarizeAgentDebugValue(item, depth + 1))
+    };
+  }
+
+  if (typeof value === 'object') {
+    if (depth >= 2) {
+      return '[Object]';
+    }
+
+    return Object.fromEntries(
+      Object.entries(value as Record<string, unknown>).map(([key, nestedValue]) => [
+        key,
+        summarizeAgentDebugValue(nestedValue, depth + 1)
+      ])
+    );
+  }
+
+  return String(value);
+}
+
+function logAgentToolDebug(
+  phase: 'start' | 'success' | 'error',
+  toolName: string,
+  details: { courseId: string; userId: string; args?: unknown; result?: unknown; error?: unknown }
+) {
+  const base = `[agent-tool:${phase}] ${toolName}`;
+
+  if (phase === 'error') {
+    console.error(base, {
+      courseId: details.courseId,
+      userId: details.userId,
+      args: summarizeAgentDebugValue(details.args),
+      error:
+        details.error instanceof Error
+          ? {
+              message: details.error.message,
+              stack: details.error.stack
+            }
+          : summarizeAgentDebugValue(details.error)
+    });
+    return;
+  }
+
+  console.info(base, {
+    courseId: details.courseId,
+    userId: details.userId,
+    args: summarizeAgentDebugValue(details.args),
+    result: summarizeAgentDebugValue(details.result)
+  });
+}
+
+function sanitizeToolError(toolName: string, error: unknown): Error {
+  if (error instanceof AppError) {
+    return error;
+  }
+
+  const raw = error instanceof Error ? error.message : String(error);
+  const looksLikeSqlLeak =
+    /^Failed query:/i.test(raw) ||
+    /\bparams:\s/i.test(raw) ||
+    /\binvalid input syntax for type uuid\b/i.test(raw) ||
+    /\bselect\s+.+\bfrom\s+/i.test(raw);
+
+  if (looksLikeSqlLeak) {
+    return new Error(
+      `${toolName} failed because of an invalid or unknown ID. Call get_course_structure to fetch the current IDs and retry — never invent UUIDs.`
+    );
+  }
+
+  return error instanceof Error ? error : new Error(raw);
+}
+
+async function executeAgentTool<TArgs, TResult>(
+  toolName: string,
+  params: { orgId: string; userId: string; courseId: string; args?: TArgs },
+  execute: () => Promise<TResult>
+): Promise<TResult> {
+  trackAgentEvent(AgentEvent.TOOL_CALLED, {
+    orgId: params.orgId,
+    userId: params.userId,
+    courseId: params.courseId,
+    toolName
+  });
+  logAgentToolDebug('start', toolName, {
+    courseId: params.courseId,
+    userId: params.userId,
+    args: params.args
+  });
+
+  try {
+    const result = await execute();
+
+    trackAgentEvent(AgentEvent.TOOL_COMPLETED, {
+      orgId: params.orgId,
+      userId: params.userId,
+      courseId: params.courseId,
+      toolName,
+      success: true
+    });
+    logAgentToolDebug('success', toolName, {
+      courseId: params.courseId,
+      userId: params.userId,
+      args: params.args,
+      result
+    });
+
+    return result;
+  } catch (error) {
+    logAgentToolDebug('error', toolName, {
+      courseId: params.courseId,
+      userId: params.userId,
+      args: params.args,
+      error
+    });
+
+    throw sanitizeToolError(toolName, error);
+  }
+}
+
+export function buildAgentTools(orgId: string, userId: string, courseId: string) {
+  return {
+    get_course_structure: tool({
+      description:
+        'Get the full course structure including sections, lessons, and exercises as a tree. The courseId is automatically set — do not pass it.',
+      inputSchema: emptyParam,
+      execute: async () => {
+        return executeAgentTool('get_course_structure', { orgId, userId, courseId }, async () => {
+          const [items, sections] = await Promise.all([getCourseContentItems(courseId), listCourseSections(courseId)]);
+          return { sections, items };
+        });
+      }
+    }),
+
+    get_lesson_content: tool({
+      description: 'Get the HTML content of a specific lesson in this course.',
+      inputSchema: lessonReadParam,
+      execute: async (args) => {
+        return executeAgentTool('get_lesson_content', { orgId, userId, courseId, args }, async () => {
+          await verifyLessonBelongsToCourse(args.lessonId, courseId);
+          const lesson = await getLesson(args.lessonId);
+          const lessonWithLangs = lesson as {
+            id: string;
+            title: string;
+            lessonLanguages?: Array<{ locale: string; content: string | null }>;
+          };
+          const langContent = lessonWithLangs.lessonLanguages?.find((ll) => ll.locale === args.locale);
+          return { id: lesson.id, title: lesson.title, content: langContent?.content || null, locale: args.locale };
+        });
+      }
+    }),
+
+    get_exercise_details: tool({
+      description:
+        'Get an exercise with sections, questions, and answer options from this course. Each sections entry has an id — use it with update_exercise_section to change that question block title or description.',
+      inputSchema: exerciseReadParam,
+      execute: async (args) => {
+        return executeAgentTool('get_exercise_details', { orgId, userId, courseId, args }, async () => {
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+          return getExercise(args.exerciseId);
+        });
+      }
+    }),
+
+    create_section: tool({
+      description: 'Create a new section in this course.',
+      inputSchema: createSectionParam,
+      execute: async (args) => {
+        return executeAgentTool('create_section', { orgId, userId, courseId, args }, async () => {
+          const section = await createCourseSection(courseId, { title: args.title, courseId, order: args.order });
+          return { id: section.id, title: section.title, order: section.order };
+        });
+      }
+    }),
+
+    update_section: tool({
+      description:
+        'Update metadata for an existing section in this course. Use this to rename a section or change its order instead of creating a new section.',
+      inputSchema: updateSectionParam,
+      execute: async (args) => {
+        return executeAgentTool('update_section', { orgId, userId, courseId, args }, async () => {
+          await verifySectionBelongsToCourse(args.sectionId, courseId);
+          const section = await updateCourseSectionService(args.sectionId, {
+            ...(args.title !== undefined ? { title: args.title } : {}),
+            ...(args.order !== undefined ? { order: args.order } : {})
+          });
+
+          return { id: section.id, title: section.title, updated: true };
+        });
+      }
+    }),
+
+    create_lesson: tool({
+      description:
+        'Create a new lesson within a section of this course. Use update_lesson_content after to add content.',
+      inputSchema: createLessonParam,
+      execute: async (args) => {
+        return executeAgentTool('create_lesson', { orgId, userId, courseId, args }, async () => {
+          await verifySectionBelongsToCourse(args.sectionId, courseId);
+          const lesson = await createLesson(courseId, {
+            title: args.title,
+            courseId,
+            sectionId: args.sectionId,
+            order: args.order
+          });
+          return { id: lesson.id, title: lesson.title, order: lesson.order };
+        });
+      }
+    }),
+
+    update_lesson: tool({
+      description:
+        'Update metadata for an existing lesson in this course. Use this to rename, move, reorder, schedule, or change visibility/unlock settings instead of creating a new lesson.',
+      inputSchema: updateLessonParam,
+      execute: async (args) => {
+        return executeAgentTool('update_lesson', { orgId, userId, courseId, args }, async () => {
+          if (args.sectionId) {
+            await verifySectionBelongsToCourse(args.sectionId, courseId);
+          }
+
+          await verifyLessonBelongsToCourse(args.lessonId, courseId);
+          const lesson = await updateLessonService(args.lessonId, {
+            ...(args.title !== undefined ? { title: args.title } : {}),
+            ...(args.sectionId !== undefined ? { sectionId: args.sectionId } : {}),
+            ...(args.order !== undefined ? { order: args.order } : {}),
+            ...(args.lessonAt !== undefined ? { lessonAt: args.lessonAt } : {}),
+            ...(args.callUrl !== undefined ? { callUrl: args.callUrl } : {}),
+            ...(args.isUnlocked !== undefined ? { isUnlocked: args.isUnlocked } : {}),
+            ...(args.public !== undefined ? { public: args.public } : {})
+          });
+
+          return { id: lesson.id, title: lesson.title, updated: true };
+        });
+      }
+    }),
+
+    update_lesson_content: tool({
+      description:
+        'Update the text content of a lesson in this course. Replaces full lesson HTML for the given locale. For lesson HTML, put only the lesson body in the content. Do not include the lesson title. Do not use h1 or h2 anywhere in lesson HTML. Start headings at h3 because that is the highest heading level allowed in lesson content.',
+      inputSchema: updateContentParam,
+      execute: async (args) => {
+        return executeAgentTool('update_lesson_content', { orgId, userId, courseId, args }, async () => {
+          await verifyLessonBelongsToCourse(args.lessonId, courseId);
+          const lesson = await getLesson(args.lessonId);
+          const normalizedContent = normalizeAgentLessonContent(args.content, lesson.title);
+
+          await upsertLessonLanguageService(args.lessonId, {
+            locale: args.locale as 'en',
+            content: normalizedContent
+          });
+          return {
+            lessonId: args.lessonId,
+            locale: args.locale,
+            contentLength: normalizedContent.length,
+            updated: true
+          };
+        });
+      }
+    }),
+
+    create_exercise: tool({
+      description: 'Create a new exercise with questions and answer options in this course.',
+      inputSchema: createExerciseParam,
+      execute: async (args) => {
+        return executeAgentTool('create_exercise', { orgId, userId, courseId, args }, async () => {
+          if (args.lessonId) {
+            await verifyLessonBelongsToCourse(args.lessonId, courseId);
+          }
+
+          if (args.sectionId) {
+            await verifySectionBelongsToCourse(args.sectionId, courseId);
+          }
+
+          const exercise = await createExercise({
+            title: args.title,
+            description: args.description,
+            courseId,
+            lessonId: args.lessonId,
+            sectionId: args.sectionId,
+            order: args.order,
+            questions: args.questions.map((q, i) => ({
+              question: q.question,
+              questionTypeId: q.questionTypeId,
+              points: q.points,
+              order: q.order ?? i,
+              options: q.options.map((o) => ({ label: o.label, isCorrect: o.isCorrect }))
+            }))
+          });
+          return { id: exercise.id, title: exercise.title, questionCount: args.questions.length };
+        });
+      }
+    }),
+
+    update_exercise: tool({
+      description:
+        "Update an existing exercise's metadata (title, description, linked lesson, section, order, due date, lock state, allow-multiple-attempts). Use this for editing the exercise itself — not for changing its questions. To edit questions, use update_questions; to add questions, use add_questions.",
+      inputSchema: updateExerciseParam,
+      execute: async (args) => {
+        return executeAgentTool('update_exercise', { orgId, userId, courseId, args }, async () => {
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+
+          if (args.lessonId) {
+            await verifyLessonBelongsToCourse(args.lessonId, courseId);
+          }
+
+          if (args.sectionId) {
+            await verifySectionBelongsToCourse(args.sectionId, courseId);
+          }
+
+          const updated = await updateExerciseService(args.exerciseId, {
+            title: args.title,
+            description: args.description,
+            lessonId: args.lessonId,
+            sectionId: args.sectionId,
+            order: args.order,
+            dueBy: args.dueBy,
+            isUnlocked: args.isUnlocked,
+            allowMultipleAttempts: args.allowMultipleAttempts
+          });
+
+          return {
+            id: updated.id,
+            title: updated.title,
+            description: updated.description ?? null,
+            dueBy: updated.dueBy ?? null,
+            updated: true
+          };
+        });
+      }
+    }),
+
+    update_exercise_section: tool({
+      description:
+        'Update the title and/or description of one section inside an exercise (question groups). This is not the same as update_section, which edits course outline sections. Always use ids from get_exercise_details `sections` for this tool.',
+      inputSchema: updateExerciseSectionParam,
+      execute: async (args) => {
+        return executeAgentTool('update_exercise_section', { orgId, userId, courseId, args }, async () => {
+          assertValidUuid('Exercise', args.exerciseId);
+          assertValidUuid('ExerciseSection', args.exerciseSectionId);
+
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+
+          const updated = await updateExerciseSectionMetadataService(args.exerciseId, args.exerciseSectionId, {
+            ...(args.title !== undefined ? { title: args.title } : {}),
+            ...(args.description !== undefined ? { description: args.description } : {})
+          });
+
+          return { ...updated, updated: true };
+        });
+      }
+    }),
+
+    create_exercise_section: tool({
+      description:
+        'Create a new section inside an existing exercise. Use this to add a question block before adding or moving questions into it.',
+      inputSchema: createExerciseSectionParam,
+      execute: async (args) => {
+        return executeAgentTool('create_exercise_section', { orgId, userId, courseId, args }, async () => {
+          assertValidUuid('Exercise', args.exerciseId);
+
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+
+          const created = await createExerciseSectionService(args.exerciseId, {
+            title: args.title,
+            description: args.description,
+            order: args.order,
+            colorTheme: args.colorTheme,
+            afterBehavior: args.afterBehavior
+          });
+
+          return { ...created, created: true };
+        });
+      }
+    }),
+
+    add_questions: tool({
+      description:
+        'Add questions to an existing exercise in this course. When get_exercise_details lists in-exercise sections, pass exerciseSectionId so new questions are added to the correct block.',
+      inputSchema: addQuestionsParam,
+      execute: async (args) => {
+        return executeAgentTool('add_questions', { orgId, userId, courseId, args }, async () => {
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+
+          if (args.exerciseSectionId !== undefined) {
+            assertValidUuid('ExerciseSection', args.exerciseSectionId);
+
+            const sections = await getExerciseSectionsByExerciseId(args.exerciseId);
+            const belongs = sections.some((section) => section.id === args.exerciseSectionId);
+
+            if (!belongs) {
+              throw new AppError(
+                'That exercise section was not found on this exercise. Call get_exercise_details and use a section id from the sections array.',
+                'VALIDATION_ERROR',
+                404
+              );
+            }
+          }
+
+          const existingExercise = await getExercise(args.exerciseId);
+          const existingQuestions = existingExercise.questions || [];
+          const nextOrder = existingQuestions.length;
+
+          const newQuestions = args.questions.map((q, i) => ({
+            question: q.question,
+            questionTypeId: q.questionTypeId,
+            points: q.points,
+            order: q.order ?? nextOrder + i,
+            ...(args.exerciseSectionId !== undefined ? { exerciseSectionId: args.exerciseSectionId } : {}),
+            options: q.options.map((o) => ({ label: o.label, isCorrect: o.isCorrect }))
+          }));
+
+          const allQuestions = [
+            ...existingQuestions.map((eq) => ({
+              id: Number(eq.id),
+              question: eq.title,
+              questionTypeId: eq.questionTypeId,
+              points: eq.points,
+              order: eq.order,
+              options: eq.options.map((o) => ({ id: Number(o.id), label: o.label || '', isCorrect: o.isCorrect }))
+            })),
+            ...newQuestions
+          ];
+
+          await updateExerciseService(args.exerciseId, { questions: allQuestions });
+          return { exerciseId: args.exerciseId, addedCount: args.questions.length, totalCount: allQuestions.length };
+        });
+      }
+    }),
+
+    update_questions: tool({
+      description:
+        'Update existing questions in an exercise. Pass only fields you want to change; `id` identifies the question. Optional `exerciseSectionId` moves the question to another in-exercise block (use get_exercise_details section ids), or null to unassign. For NUMERIC, the correct answer is `settings.correctValue` (number) — do NOT add options to NUMERIC questions. For STAR use `settings.correctValue`. For WORD_BANK use `settings.correctAnswers` and `settings.template`. RADIO/CHECKBOX/TRUE_FALSE use `options[].isCorrect` (include option `id` to edit, omit `id` to add). `settings` is shallow-merged with existing settings. Omit `options` entirely to leave existing options untouched.',
+      inputSchema: updateQuestionsParam,
+      execute: async (args) => {
+        return executeAgentTool('update_questions', { orgId, userId, courseId, args }, async () => {
+          await verifyExerciseBelongsToCourse(args.exerciseId, courseId);
+
+          const sectionIdsInPatches = new Set<string>();
+          for (const patch of args.questions) {
+            if (patch.exerciseSectionId === undefined || patch.exerciseSectionId === null) {
+              continue;
+            }
+
+            assertValidUuid('ExerciseSection', patch.exerciseSectionId);
+            sectionIdsInPatches.add(patch.exerciseSectionId);
+          }
+
+          if (sectionIdsInPatches.size > 0) {
+            const sections = await getExerciseSectionsByExerciseId(args.exerciseId);
+            const validIds = new Set(sections.map((section) => section.id));
+
+            for (const sectionId of sectionIdsInPatches) {
+              if (!validIds.has(sectionId)) {
+                throw new AppError(
+                  'That exercise section was not found on this exercise. Call get_exercise_details and use a section id from the sections array.',
+                  'VALIDATION_ERROR',
+                  404
+                );
+              }
+            }
+          }
+
+          const existing = await getExercise(args.exerciseId);
+
+          const merged = buildUpdatedQuestions(
+            (existing.questions || []).map((q) => ({
+              id: q.id,
+              title: q.title,
+              questionTypeId: q.questionTypeId,
+              points: q.points,
+              order: q.order,
+              exerciseSectionId: q.exerciseSectionId,
+              settings: q.settings,
+              options: q.options
+            })),
+            args.questions,
+            args.exerciseId
+          );
+
+          await updateExerciseService(args.exerciseId, { questions: merged });
+
+          return { exerciseId: args.exerciseId, updatedCount: merged.length };
+        });
+      }
+    }),
+
+    reorder_content: tool({
+      description:
+        'Reorder sections, lessons, or exercises in this course. Can change the order of sections, change the order of lessons/exercises within a section, or move items between sections. Use get_course_structure first to see current order, and copy item IDs exactly from that response. Never rewrite or guess UUIDs.',
+      inputSchema: reorderContentParam,
+      execute: async (args) => {
+        return executeAgentTool('reorder_content', { orgId, userId, courseId, args }, async () => {
+          if (args.sections) {
+            for (const section of args.sections) {
+              await verifySectionBelongsToCourse(section.id, courseId);
+            }
+          }
+
+          if (args.items) {
+            for (const item of args.items) {
+              if (item.type === 'LESSON') {
+                await verifyLessonBelongsToCourse(item.id, courseId);
+              } else {
+                await verifyExerciseBelongsToCourse(item.id, courseId);
+              }
+
+              if (item.sectionId) {
+                await verifySectionBelongsToCourse(item.sectionId, courseId);
+              }
+            }
+          }
+
+          return reorderCourseContent(courseId, {
+            sections: args.sections,
+            items: args.items
+          });
+        });
+      }
+    }),
+
+    update_course_landing_page: tool({
+      description:
+        'Update course-level landing page fields for this course, including public title, description, overview, goals, requirements, instructor metadata, pricing, and banner image. The courseId is automatically set — do not pass it.',
+      inputSchema: updateCourseLandingPageParam,
+      execute: async (args) => {
+        return executeAgentTool('update_course_landing_page', { orgId, userId, courseId, args }, async () => {
+          const result = await updateCourseLandingPageService(courseId, args as TCourseLandingPageUpdate);
+
+          return {
+            courseId: result.course.id,
+            title: result.course.title,
+            description: result.course.description,
+            courseUrl: result.courseUrl,
+            bannerImageUrl: result.bannerImageUrl,
+            updated: true
+          };
+        });
+      }
+    }),
+
+    check_course_go_live_readiness: tool({
+      description:
+        'Check whether this course is ready to go live. Returns blockers, warnings, suggested fixes, and the public course URL if available. This tool does not publish the course.',
+      inputSchema: emptyParam,
+      execute: async () => {
+        return executeAgentTool('check_course_go_live_readiness', { orgId, userId, courseId }, async () => {
+          return getCourseGoLiveReadiness(courseId);
+        });
+      }
+    }),
+
+    go_live_course: tool({
+      description:
+        'Publish this course only after the teacher explicitly asks to go live. Runs the readiness checklist first, generates a slug when needed, and fails with blockers if the course is not ready.',
+      inputSchema: goLiveParam,
+      execute: async (args) => {
+        return executeAgentTool('go_live_course', { orgId, userId, courseId, args }, async () => {
+          const result = await publishCourseWhenReady(courseId);
+
+          return {
+            courseId: result.course.id,
+            title: result.course.title,
+            slug: result.course.slug,
+            isPublished: result.course.isPublished,
+            readiness: result.readiness
+          };
+        });
+      }
+    }),
+
+    generate_course_plan: tool({
+      description:
+        'Generate a structured course plan with sections and lessons. Always use this when asked to design or plan a course.',
+      inputSchema: coursePlanParam,
+      execute: async (args) => {
+        return executeAgentTool('generate_course_plan', { orgId, userId, courseId, args }, async () => {
+          const envelope = coursePlanParam.parse(args);
+          const plan = CoursePlanSchema.parse(envelope.plan);
+          const sectionCount = plan.sections.length;
+          const itemCount = plan.sections.reduce((sum, section) => sum + section.items.length, 0);
+
+          trackAgentEvent(AgentEvent.PLAN_GENERATED, { orgId, userId, courseId, sectionCount, itemCount });
+
+          return plan;
+        });
+      }
+    })
+  };
+}

apps/api/src/services/agent/context-window.ts

@@ -1,27 +1,31 @@
-/**
- * Trims conversation message history to stay within token budget.
- *
- * Strategy:
- * - Always keep the first message (often contains the initial request/document context)
- * - Always keep the last N messages (recent context is most relevant)
- * - Drop middle messages when the array is too long
- *
- * This is a message-count heuristic, not a token counter. Each message averages
- * ~200-500 tokens, so capping at ~40 messages keeps us under ~20K tokens of
- * history, leaving room for the system prompt + document text + LLM response.
- */
+const MAX_MESSAGES = 20;
+const KEEP_RECENT = 16;
+const MAX_TOOL_OUTPUT_CHARS = 2000;
 
-const MAX_MESSAGES = 40;
-const KEEP_RECENT = 30;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function trimToolOutputs(messages: any[]): any[] {
+  return messages.map((message) => {
+    if (!Array.isArray(message.parts)) return message;
+
+    const trimmedParts = message.parts.map((part: any) => {
+      if (part.state !== 'output-available' || part.output == null) return part;
+
+      const serialized = typeof part.output === 'string' ? part.output : JSON.stringify(part.output);
+      if (serialized.length <= MAX_TOOL_OUTPUT_CHARS) return part;
+
+      const truncated = serialized.slice(0, MAX_TOOL_OUTPUT_CHARS) + '…[truncated]';
+      return { ...part, output: truncated };
+    });
+
+    if (trimmedParts === message.parts) return message;
+
+    return { ...message, parts: trimmedParts };
+  });
+}
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export function trimMessageHistory(messages: any[]): any[] {
-  if (messages.length <= MAX_MESSAGES) {
-    return messages;
-  }
+  const countTrimmed = messages.length <= MAX_MESSAGES ? messages : [messages[0], ...messages.slice(-KEEP_RECENT)];
 
-  const firstMessage = messages[0];
-  const recentMessages = messages.slice(-KEEP_RECENT);
-
-  return [firstMessage, ...recentMessages];
+  return trimToolOutputs(countTrimmed);
 }

apps/api/src/services/agent/credit-purchase.ts

@@ -0,0 +1,21 @@
+import { findCreditPurchaseByOrderId, insertCreditPurchaseAndApplyBalance } from '@cio/db/queries/agent';
+import type { CreditPurchaseInput } from '@cio/db/queries/agent';
+import { AppError } from '@api/utils/errors';
+
+/**
+ * Record a Polar token-pack purchase.
+ * Idempotent on `providerOrderId` — Polar may retry the webhook.
+ */
+export async function recordCreditPurchase(input: CreditPurchaseInput) {
+  if (input.tokens <= 0) {
+    throw new AppError('Tokens must be positive', 'INVALID_CREDIT_AMOUNT', 400);
+  }
+
+  const existing = await findCreditPurchaseByOrderId(input.providerOrderId);
+
+  if (existing) {
+    return existing;
+  }
+
+  return insertCreditPurchaseAndApplyBalance(input);
+}

apps/api/src/services/agent/document.ts

@@ -107,6 +107,7 @@ export async function parseAndStoreDocument(
       text: extractedText,
       fileName: file.name,
       mimeType,
+      userId,
       uploadedAt: new Date().toISOString()
     }),
     { EX: DOCUMENT_REDIS_TTL }
@@ -150,16 +151,18 @@ export async function parseAndStoreDocument(
  * Retrieve stored document text. Tries the Redis hot cache first; on miss
  * falls back to Postgres and rehydrates Redis for next time.
  */
-export async function getDocumentText(documentId: string, redis: RedisClient): Promise<string | null> {
+export async function getDocumentText(documentId: string, userId: string, redis: RedisClient): Promise<string | null> {
   const raw = await redis.get(agentDocumentKey(documentId));
 
   if (raw) {
-    const parsed = JSON.parse(raw) as { text: string };
+    const parsed = JSON.parse(raw) as { text: string; userId?: string };
+
+    if (parsed.userId && parsed.userId !== userId) return null;
 
     return parsed.text;
   }
 
-  const record = await getChatDocument(documentId);
+  const record = await getChatDocument(documentId, userId);
 
   if (!record) return null;
 
@@ -169,6 +172,7 @@ export async function getDocumentText(documentId: string, redis: RedisClient): P
       text: record.text,
       fileName: record.fileName,
       mimeType: record.mimeType,
+      userId: record.userId,
       uploadedAt: record.createdAt
     }),
     { EX: DOCUMENT_REDIS_TTL }

apps/api/src/services/agent/question-update.ts

@@ -6,6 +6,7 @@ export type CurrentQuestionForPatch = {
   questionTypeId: number;
   points: number;
   order: number;
+  exerciseSectionId?: string | null;
   settings?: Record<string, unknown>;
   options: Array<{
     id: number | string;
@@ -21,6 +22,7 @@ export type QuestionPatch = {
   questionTypeId?: number;
   points?: number;
   order?: number;
+  exerciseSectionId?: string | null;
   settings?: Record<string, unknown>;
   options?: Array<{
     id?: number;
@@ -43,6 +45,7 @@ export type MergedQuestion = {
   questionTypeId: number;
   points: number;
   order: number;
+  exerciseSectionId?: string | null;
   settings?: Record<string, unknown>;
   options: MergedOption[];
 };
@@ -97,7 +100,7 @@ export function buildUpdatedQuestions(
       });
     }
 
-    return {
+    const merged: MergedQuestion = {
       id: patch.id,
       question: patch.question ?? current.title ?? '',
       questionTypeId: patch.questionTypeId ?? current.questionTypeId,
@@ -106,5 +109,11 @@ export function buildUpdatedQuestions(
       settings: mergedSettings,
       options: mergedOptions
     };
+
+    if (patch.exerciseSectionId !== undefined) {
+      merged.exerciseSectionId = patch.exerciseSectionId;
+    }
+
+    return merged;
   });
 }

apps/api/src/services/agent/title-generation.ts

@@ -1,4 +1,5 @@
-import { generateText } from 'ai';
+import { generateText, generateObject } from 'ai';
+import { z } from 'zod';
 import { AIProvider, type AIProviderConfig } from '@cio/ai-assistant';
 import { createOpenAI } from '@ai-sdk/openai';
 import { createAnthropic } from '@ai-sdk/anthropic';
@@ -9,8 +10,9 @@ import { createGoogleGenerativeAI } from '@ai-sdk/google';
  */
 const TITLE_MODELS: Record<AIProvider, string> = {
   [AIProvider.OPENAI]: 'gpt-4o-mini',
-  [AIProvider.ANTHROPIC]: 'claude-haiku-4-20250414',
-  [AIProvider.GOOGLE]: 'gemini-2.5-flash'
+  [AIProvider.ANTHROPIC]: 'claude-haiku-4-5-20251001',
+  [AIProvider.GOOGLE]: 'gemini-2.5-flash',
+  [AIProvider.MOONSHOT]: 'kimi-k2.6'
 };
 
 function createTitleModel(config: AIProviderConfig) {
@@ -29,6 +31,10 @@ function createTitleModel(config: AIProviderConfig) {
       const google = createGoogleGenerativeAI({ apiKey: config.apiKey });
       return google(modelName);
     }
+    case AIProvider.MOONSHOT: {
+      const moonshot = createOpenAI({ apiKey: config.apiKey, baseURL: 'https://api.moonshot.cn/v1' });
+      return moonshot(modelName);
+    }
     default:
       throw new Error(`Unsupported AI provider: ${config.provider}`);
   }
@@ -44,13 +50,58 @@ export async function generateConversationTitle(
 ): Promise<string> {
   const model = createTitleModel(providerConfig);
 
-  const { text } = await generateText({
-    model,
-    system:
-      "Generate a short title (max 6 words) for a chat conversation based on the user's first message. Return ONLY the title, no quotes, no punctuation at the end. Be concise and descriptive.",
-    prompt: firstMessageText,
-    maxOutputTokens: 30
-  });
+  const systemPrompt =
+    "Generate a short, descriptive title (2–5 words) for a chat conversation based on the user's first message. The title must be at least 2 words — never a single word. Return ONLY the title, no quotes, no punctuation at the end.";
+
+  const truncatedMessage = firstMessageText.slice(0, 1000);
+
+  // Gemini 2.5 counts thinking toward maxOutputTokens; disable thinking and allow headroom for the title text.
+  const { text } = await generateText(
+    providerConfig.provider === AIProvider.GOOGLE
+      ? {
+          model,
+          system: systemPrompt,
+          prompt: truncatedMessage,
+          maxOutputTokens: 256,
+          maxRetries: 0,
+          providerOptions: {
+            google: { thinkingConfig: { thinkingBudget: 0 } }
+          }
+        }
+      : {
+          model,
+          system: systemPrompt,
+          prompt: truncatedMessage,
+          maxOutputTokens: 96,
+          maxRetries: 0
+        }
+  );
 
   return text.trim().slice(0, 80);
 }
+
+const CourseMeta = z.object({
+  title: z.string().describe('3–6 word course title, no quotes, no trailing punctuation'),
+  description: z.string().describe('20–30 word description of what students will learn')
+});
+
+export async function generateCourseMeta(
+  prompt: string,
+  providerConfig: AIProviderConfig
+): Promise<{ title: string; description: string }> {
+  const model = createTitleModel(providerConfig);
+
+  const { object } = await generateObject({
+    model,
+    schema: CourseMeta,
+    system:
+      "You are helping create an online course. Based on the user's course idea, generate a short title (3–6 words, no quotes, no trailing punctuation) and a concise description (20–30 words) explaining what students will learn.",
+    prompt: prompt.slice(0, 1000),
+    maxRetries: 0
+  });
+
+  return {
+    title: object.title.trim().slice(0, 100),
+    description: object.description.trim().slice(0, 300)
+  };
+}

apps/api/src/services/agent/usage.ts

@@ -1,22 +1,51 @@
-import { db } from '@cio/db';
-import * as schema from '@cio/db/schema';
-import { eq, and, gte, sql } from 'drizzle-orm';
-import { getActiveOrganizationPlan } from '@cio/db/queries/organization';
-import { AppError, ErrorCodes } from '@api/utils/errors';
 import type { TokenBalance, TokenUsage } from '@cio/ai-assistant';
+import {
+  aggregateTokenUsageByUser,
+  countRequests,
+  getDailyTokenUsageHistory,
+  getMonthlyTokenUsage,
+  getOrgCreditBalance,
+  insertTokenUsageAndDrainCredits,
+  summarizePurchases,
+  upsertCreditBalance
+} from '@cio/db/queries/agent';
+import type { UsageLeaderboardRow } from '@cio/db/queries/agent';
+import { getActiveOrganizationPlan } from '@cio/db/queries/organization';
+
+import { AppError } from '@api/utils/errors';
+
+function startOfCurrentMonth(): Date {
+  const date = new Date();
+  date.setDate(1);
+  date.setHours(0, 0, 0, 0);
+
+  return date;
+}
 
 // ─── Plan-Based Token Allowances ─────────────────────────────────────────────
 
 const PLAN_TOKEN_ALLOWANCES: Record<string, number> = {
-  BASIC: 50_000,
-  EARLY_ADOPTER: 500_000,
-  ENTERPRISE: 2_000_000
+  BASIC: 500_000,
+  EARLY_ADOPTER: 3_000_000,
+  ENTERPRISE: 15_000_000
 };
 
-/**
- * Get the monthly token allowance for an org based on its active plan.
- * Also checks if the plan payload has a custom `aiTokenAllowance` override.
- */
+// ─── Model Cost Multipliers ───────────────────────────────────────────────────
+
+// Blended multiplier vs Gemini 2.5 Flash baseline ($0.15/1M); ~80/20 input/output mix; unknown → 1×.
+const MODEL_COST_MULTIPLIER: Record<string, number> = {
+  'gemini-2.5-flash': 1,
+  'gpt-5.4-mini': 3,
+  'claude-sonnet-4-6': 11,
+  'claude-haiku-4-5-20251001': 1,
+  'kimi-k2.6': 3
+};
+
+function computeCostUnits(promptTokens: number, completionTokens: number, model: string): number {
+  const multiplier = MODEL_COST_MULTIPLIER[model] ?? 1;
+  return Math.round((promptTokens + completionTokens) * multiplier);
+}
+
 async function getPlanAllowance(orgId: string): Promise<{ planName: string; allowance: number }> {
   const activePlan = await getActiveOrganizationPlan(orgId);
 
@@ -32,42 +61,16 @@ async function getPlanAllowance(orgId: string): Promise<{ planName: string; allo
   return { planName, allowance };
 }
 
-/**
- * Get the total tokens used by an org in the current month.
- */
-async function getMonthlyUsage(orgId: string): Promise<number> {
-  const startOfMonth = new Date();
-  startOfMonth.setDate(1);
-  startOfMonth.setHours(0, 0, 0, 0);
-
-  const [result] = await db
-    .select({
-      totalTokens: sql<number>`COALESCE(SUM(${schema.aiTokenUsage.promptTokens} + ${schema.aiTokenUsage.completionTokens}), 0)`
-    })
-    .from(schema.aiTokenUsage)
-    .where(and(eq(schema.aiTokenUsage.orgId, orgId), gte(schema.aiTokenUsage.createdAt, startOfMonth.toISOString())));
-
-  return Number(result?.totalTokens ?? 0);
+export async function getMonthlyUsage(orgId: string): Promise<number> {
+  return getMonthlyTokenUsage(orgId, startOfCurrentMonth());
 }
 
-/**
- * Get the credit balance for an org.
- */
-async function getCreditBalance(orgId: string): Promise<number> {
-  const [result] = await db
-    .select({ balance: schema.aiCreditBalance.balance })
-    .from(schema.aiCreditBalance)
-    .where(eq(schema.aiCreditBalance.orgId, orgId))
-    .limit(1);
-
-  return result?.balance ?? 0;
+export async function getCreditBalance(orgId: string): Promise<number> {
+  return getOrgCreditBalance(orgId);
 }
 
-/**
- * Get the full token balance for an org: used, allowance, credits, remaining.
- */
 export async function getTokenBalance(orgId: string): Promise<TokenBalance> {
-  const [{ planName, allowance }, monthlyUsage, creditBalance] = await Promise.all([
+  const [{ allowance }, monthlyUsage, creditBalance] = await Promise.all([
     getPlanAllowance(orgId),
     getMonthlyUsage(orgId),
     getCreditBalance(orgId)
@@ -84,9 +87,6 @@ export async function getTokenBalance(orgId: string): Promise<TokenBalance> {
   };
 }
 
-/**
- * Check if the org has sufficient tokens. Throws 402 if exhausted.
- */
 export async function enforceTokenBalance(orgId: string): Promise<TokenBalance> {
   const balance = await getTokenBalance(orgId);
 
@@ -97,9 +97,17 @@ export async function enforceTokenBalance(orgId: string): Promise<TokenBalance>
   return balance;
 }
 
-/**
- * Record token usage after an LLM call.
- */
+export function computePurchasedTokenOverflow(params: {
+  allowance: number;
+  monthlyUsageBefore: number;
+  requestTokens: number;
+}): number {
+  const allowanceRemainingBefore = Math.max(0, params.allowance - params.monthlyUsageBefore);
+
+  return Math.max(0, params.requestTokens - allowanceRemainingBefore);
+}
+
+/** Record token usage after an LLM call; drains purchased credits when usage exceeds plan allowance. */
 export async function recordTokenUsage(
   orgId: string,
   userId: string,
@@ -107,29 +115,28 @@ export async function recordTokenUsage(
   usage: TokenUsage,
   model: string
 ): Promise<void> {
-  await db.insert(schema.aiTokenUsage).values({
+  const costUnits = computeCostUnits(usage.promptTokens, usage.completionTokens, model);
+  const { allowance } = await getPlanAllowance(orgId);
+
+  await insertTokenUsageAndDrainCredits({
     orgId,
     userId,
     courseId,
     promptTokens: usage.promptTokens,
     completionTokens: usage.completionTokens,
-    model
+    costUnits,
+    model,
+    planAllowance: allowance,
+    since: startOfCurrentMonth()
   });
 }
 
-/**
- * Check if an org's plan supports a specific feature (e.g., document upload).
- * Returns the plan name so callers can decide what to gate.
- */
 export async function getOrgPlanName(orgId: string): Promise<string> {
   const { planName } = await getPlanAllowance(orgId);
 
   return planName;
 }
 
-/**
- * Check if an org is on a paid plan (EARLY_ADOPTER or above).
- */
 export async function isOrgOnPaidPlan(orgId: string): Promise<boolean> {
   const planName = await getOrgPlanName(orgId);
 
@@ -143,67 +150,76 @@ export interface DailyUsage {
   tokens: number;
 }
 
-/**
- * Get daily token usage for the current month, grouped by day.
- */
 export async function getUsageHistory(orgId: string): Promise<DailyUsage[]> {
-  const startOfMonth = new Date();
-  startOfMonth.setDate(1);
-  startOfMonth.setHours(0, 0, 0, 0);
-
-  const rows = await db
-    .select({
-      date: sql<string>`DATE(${schema.aiTokenUsage.createdAt})`,
-      tokens: sql<number>`SUM(${schema.aiTokenUsage.promptTokens} + ${schema.aiTokenUsage.completionTokens})`
-    })
-    .from(schema.aiTokenUsage)
-    .where(and(eq(schema.aiTokenUsage.orgId, orgId), gte(schema.aiTokenUsage.createdAt, startOfMonth.toISOString())))
-    .groupBy(sql`DATE(${schema.aiTokenUsage.createdAt})`)
-    .orderBy(sql`DATE(${schema.aiTokenUsage.createdAt})`);
-
-  return rows.map((row) => ({
-    date: String(row.date),
-    tokens: Number(row.tokens)
-  }));
+  return getDailyTokenUsageHistory(orgId, startOfCurrentMonth());
 }
 
-/**
- * Get full detailed usage stats for the usage endpoint.
- */
-export async function getDetailedUsage(orgId: string) {
-  const [balance, history] = await Promise.all([getTokenBalance(orgId), getUsageHistory(orgId)]);
+export interface PurchasedSummary {
+  totalPurchasedTokens: number;
+  totalSpentCents: number;
+  currency: 'USD';
+  currentBalance: number;
+  lastPurchaseAt: string | null;
+}
 
-  return { ...balance, history };
+export async function getPurchasedSummary(orgId: string): Promise<PurchasedSummary> {
+  const [summary, currentBalance] = await Promise.all([summarizePurchases(orgId), getCreditBalance(orgId)]);
+
+  return {
+    totalPurchasedTokens: summary.totalPurchasedTokens,
+    totalSpentCents: summary.totalSpentCents,
+    currency: 'USD',
+    currentBalance,
+    lastPurchaseAt: summary.lastPurchaseAt
+  };
+}
+
+export interface LeaderboardEntry {
+  userId: string;
+  fullname: string | null;
+  email: string | null;
+  avatarUrl: string | null;
+  tokens: number;
+  requests: number;
+  favoriteModel: string | null;
+  percentage: number;
+}
+
+export async function getTeamLeaderboard(orgId: string): Promise<{ entries: LeaderboardEntry[]; totalTokens: number }> {
+  const rows: UsageLeaderboardRow[] = await aggregateTokenUsageByUser(orgId, startOfCurrentMonth());
+  const totalTokens = rows.reduce((sum: number, row: UsageLeaderboardRow) => sum + row.tokens, 0);
+
+  const entries: LeaderboardEntry[] = rows.map((row: UsageLeaderboardRow) => ({
+    userId: row.userId,
+    fullname: row.fullname,
+    email: row.email,
+    avatarUrl: row.avatarUrl,
+    tokens: row.tokens,
+    requests: row.requests,
+    favoriteModel: row.favoriteModel,
+    percentage: totalTokens > 0 ? row.tokens / totalTokens : 0
+  }));
+
+  return { entries, totalTokens };
+}
+
+export async function getDetailedUsage(orgId: string) {
+  const start = startOfCurrentMonth();
+  const [balance, history, requestsThisMonth] = await Promise.all([
+    getTokenBalance(orgId),
+    getDailyTokenUsageHistory(orgId, start),
+    countRequests(orgId, start)
+  ]);
+
+  return { ...balance, history, requestsThisMonth };
 }
 
 // ─── Credit Purchases ────────────────────────────────────────────────────────
 
-/**
- * Add credits to an org's balance.
- * Creates the balance row if it doesn't exist (upsert).
- */
 export async function addCredits(orgId: string, amount: number): Promise<number> {
   if (amount <= 0) {
     throw new AppError('Credit amount must be positive', 'INVALID_CREDIT_AMOUNT', 400);
   }
 
-  const existing = await db
-    .select({ id: schema.aiCreditBalance.id, balance: schema.aiCreditBalance.balance })
-    .from(schema.aiCreditBalance)
-    .where(eq(schema.aiCreditBalance.orgId, orgId))
-    .limit(1);
-
-  if (existing.length > 0) {
-    const newBalance = existing[0].balance + amount;
-    await db
-      .update(schema.aiCreditBalance)
-      .set({ balance: newBalance, updatedAt: new Date().toISOString() })
-      .where(eq(schema.aiCreditBalance.orgId, orgId));
-
-    return newBalance;
-  }
-
-  await db.insert(schema.aiCreditBalance).values({ orgId, balance: amount });
-
-  return amount;
+  return upsertCreditBalance(orgId, amount);
 }