evOo-local/packages/adapters/src/coach-tools.ts

import type { ConnectorTool } from "@rakazo/adapter-kit";
import {
  CHECK_IN_NOTE_MAX_LENGTH,
  COACH_TITLE_MAX_LENGTH,
  DAILY_BRIEF_MAX_LENGTH,
  GOAL_DETAIL_MAX_LENGTH,
  GOAL_STATUSES,
  GoalStatusSchema,
  HABIT_CADENCES,
  HabitCadenceSchema,
  LocalDateSchema,
} from "@rakazo/contracts";
import type { PrismaClient } from "@rakazo/db";
import {
  CoachNotFoundError,
  createGoal,
  createHabit,
  InvalidCoachInputError,
  isPersonalMainBot,
  logCheckIn,
  updateGoal,
  updateHabit,
  upsertDailyBrief,
} from "@rakazo/db";
import { z } from "zod";

/** The main bot's coach tools, in the order they appear in its tool list. */
export const COACH_TOOL_NAMES = [
  "update_goal",
  "update_habit",
  "log_check_in",
  "write_daily_brief",
] as const;
export type CoachToolName = (typeof COACH_TOOL_NAMES)[number];

export function isCoachToolName(name: string): name is CoachToolName {
  return (COACH_TOOL_NAMES as readonly string[]).includes(name);
}

const DATE_HINT = "The user's local date, YYYY-MM-DD.";

export const coachTools: ConnectorTool[] = [
  {
    name: "update_goal",
    description:
      "Track a goal the user committed to. Without id it creates one; with id it changes it, and status done or dropped closes it. Track only what the user stated or agreed to.",
    inputSchema: {
      type: "object",
      properties: {
        id: { type: "string", description: "Id or exact title of the goal to change." },
        title: { type: "string", description: 'A short outcome, e.g. "Run a 10k in November".' },
        detail: { type: "string", description: "Optional context, such as why it matters." },
        status: { type: "string", enum: [...GOAL_STATUSES] },
        target_date: { type: "string", description: `${DATE_HINT} An empty string clears it.` },
      },
    },
  },
  {
    name: "update_habit",
    description:
      "Track a habit the user wants to keep. Without id it creates one; with id it changes it, and active false stops tracking it.",
    inputSchema: {
      type: "object",
      properties: {
        id: { type: "string", description: "Id or exact title of the habit to change." },
        title: { type: "string", description: 'A short practice, e.g. "Walk 20 minutes".' },
        cadence: { type: "string", enum: [...HABIT_CADENCES] },
        goal: {
          type: "string",
          description: "Id or exact title of the goal it serves. An empty string unlinks it.",
        },
        active: { type: "boolean" },
      },
    },
  },
  {
    name: "log_check_in",
    description:
      "Record what the user reported for a day: a habit kept or missed, progress on a goal, or, with neither, how the day went. Logging the same day again updates it.",
    inputSchema: {
      type: "object",
      properties: {
        habit: { type: "string", description: "Id or exact title of the habit." },
        goal: { type: "string", description: "Id or exact title of the goal." },
        date: { type: "string", description: `${DATE_HINT} Defaults to today.` },
        done: { type: "boolean", description: "Whether the habit was kept." },
        mood: { type: "integer", minimum: 1, maximum: 5 },
        energy: { type: "integer", minimum: 1, maximum: 5 },
        note: { type: "string", description: "The user's words, briefly." },
      },
    },
  },
  {
    name: "write_daily_brief",
    description:
      "Save the user's brief for the day, shown at the top of their Today page. Plain text, two to five short lines. Writing again replaces it.",
    inputSchema: {
      type: "object",
      properties: {
        text: { type: "string" },
        date: { type: "string", description: `${DATE_HINT} Defaults to today.` },
      },
      required: ["text"],
    },
  },
];

const Ref = z.string().trim().min(1).max(COACH_TITLE_MAX_LENGTH);
const ClearableDate = z.union([LocalDateSchema, z.literal("")]);

const UpdateGoalArgs = z.object({
  id: Ref.optional(),
  title: z.string().trim().min(1).max(COACH_TITLE_MAX_LENGTH).optional(),
  detail: z.string().trim().max(GOAL_DETAIL_MAX_LENGTH).optional(),
  status: GoalStatusSchema.optional(),
  target_date: ClearableDate.optional(),
});

const UpdateHabitArgs = z.object({
  id: Ref.optional(),
  title: z.string().trim().min(1).max(COACH_TITLE_MAX_LENGTH).optional(),
  cadence: HabitCadenceSchema.optional(),
  goal: z.union([Ref, z.literal("")]).optional(),
  active: z.boolean().optional(),
});

const Rating = z.number().int().min(1).max(5);

const LogCheckInArgs = z.object({
  habit: Ref.optional(),
  goal: Ref.optional(),
  date: LocalDateSchema.optional(),
  done: z.boolean().optional(),
  mood: Rating.optional(),
  energy: Rating.optional(),
  note: z.string().trim().max(CHECK_IN_NOTE_MAX_LENGTH).optional(),
});

const WriteDailyBriefArgs = z.object({
  text: z.string().trim().min(1).max(DAILY_BRIEF_MAX_LENGTH),
  date: LocalDateSchema.optional(),
});

export interface CoachToolContext {
  userId: string;
  /** The space the run happens in. */
  spaceId: string;
  bot: { id: string; kind: string };
}

type ToolResult = Record<string, unknown>;

function invalid(name: string, error: z.ZodError): ToolResult {
  const issues = error.issues
    .map((issue) => `${issue.path.join(".") || "arguments"}: ${issue.message}`)
    .join("; ");
  return { error: `Invalid ${name} arguments: ${issues}` };
}

/** An id first, then an exact title (any case) among the user's items. */
async function resolveRef(
  find: (where: object) => Promise<Array<{ id: string }>>,
  ref: string,
  what: "goal" | "habit",
): Promise<string | ToolResult> {
  const byId = await find({ id: ref });
  if (byId[0]) return byId[0].id;
  const byTitle = await find({ title: { equals: ref, mode: "insensitive" } });
  if (byTitle.length === 1) return byTitle[0]!.id;
  return byTitle.length > 1
    ? { error: `Several ${what}s are titled "${ref}". Use the id.` }
    : { error: `No ${what} "${ref}".` };
}

function goalRef(prisma: PrismaClient, userId: string, ref: string) {
  return resolveRef(
    (where) =>
      prisma.userGoal.findMany({
        where: { userId, ...where },
        orderBy: [{ status: "asc" }, { createdAt: "desc" }],
        select: { id: true },
        take: 2,
      }),
    ref,
    "goal",
  );
}

function habitRef(prisma: PrismaClient, userId: string, ref: string) {
  return resolveRef(
    (where) =>
      prisma.habit.findMany({
        where: { userId, ...where },
        orderBy: [{ active: "desc" }, { createdAt: "desc" }],
        select: { id: true },
        take: 2,
      }),
    ref,
    "habit",
  );
}

async function updateGoalFromTool(
  prisma: PrismaClient,
  context: CoachToolContext,
  rawArgs: Record<string, unknown>,
): Promise<ToolResult> {
  const parsed = UpdateGoalArgs.safeParse(rawArgs);
  if (!parsed.success) return invalid("update_goal", parsed.error);
  const args = parsed.data;
  const targetDate = args.target_date === undefined ? undefined : args.target_date || null;
  if (!args.id) {
    if (!args.title) return { error: "A new goal needs a title." };
    const goal = await createGoal(prisma, context.userId, {
      title: args.title,
      detail: args.detail,
      targetDate,
      source: "chat",
      createdByBotId: context.bot.id,
    });
    return { ok: true, goal: { id: goal.id, title: goal.title, status: goal.status } };
  }
  const goalId = await goalRef(prisma, context.userId, args.id);
  if (typeof goalId !== "string") return goalId;
  const goal = await updateGoal(prisma, context.userId, goalId, {
    title: args.title,
    detail: args.detail,
    status: args.status,
    targetDate,
  });
  return {
    ok: true,
    goal: { id: goal.id, title: goal.title, status: goal.status, targetDate: goal.targetDate },
  };
}

async function updateHabitFromTool(
  prisma: PrismaClient,
  context: CoachToolContext,
  rawArgs: Record<string, unknown>,
): Promise<ToolResult> {
  const parsed = UpdateHabitArgs.safeParse(rawArgs);
  if (!parsed.success) return invalid("update_habit", parsed.error);
  const args = parsed.data;
  let goalId: string | null | undefined;
  if (args.goal) {
    const resolved = await goalRef(prisma, context.userId, args.goal);
    if (typeof resolved !== "string") return resolved;
    goalId = resolved;
  } else if (args.goal === "") {
    goalId = null;
  }
  if (!args.id) {
    if (!args.title) return { error: "A new habit needs a title." };
    const habit = await createHabit(prisma, context.userId, {
      title: args.title,
      cadence: args.cadence,
      goalId,
      source: "chat",
      createdByBotId: context.bot.id,
    });
    return { ok: true, habit: { id: habit.id, title: habit.title, cadence: habit.cadence } };
  }
  const habitId = await habitRef(prisma, context.userId, args.id);
  if (typeof habitId !== "string") return habitId;
  const habit = await updateHabit(prisma, context.userId, habitId, {
    title: args.title,
    cadence: args.cadence,
    goalId,
    active: args.active,
  });
  return {
    ok: true,
    habit: { id: habit.id, title: habit.title, cadence: habit.cadence, active: habit.active },
  };
}

async function logCheckInFromTool(
  prisma: PrismaClient,
  context: CoachToolContext,
  rawArgs: Record<string, unknown>,
): Promise<ToolResult> {
  const parsed = LogCheckInArgs.safeParse(rawArgs);
  if (!parsed.success) return invalid("log_check_in", parsed.error);
  const { habit, goal, ...fields } = parsed.data;
  if (habit && goal) return { error: "Log a habit and a goal in separate calls." };
  const habitId = habit ? await habitRef(prisma, context.userId, habit) : undefined;
  if (habitId !== undefined && typeof habitId !== "string") return habitId;
  const goalId = goal ? await goalRef(prisma, context.userId, goal) : undefined;
  if (goalId !== undefined && typeof goalId !== "string") return goalId;
  const checkIn = await logCheckIn(prisma, context.userId, {
    kind: habitId ? "habit" : goalId ? "goal" : "daily",
    habitId,
    goalId,
    ...fields,
    createdByBotId: context.bot.id,
  });
  return {
    ok: true,
    checkIn: {
      kind: checkIn.kind,
      date: checkIn.date,
      done: checkIn.done,
      mood: checkIn.mood,
      energy: checkIn.energy,
    },
  };
}

async function writeDailyBriefFromTool(
  prisma: PrismaClient,
  context: CoachToolContext,
  rawArgs: Record<string, unknown>,
): Promise<ToolResult> {
  const parsed = WriteDailyBriefArgs.safeParse(rawArgs);
  if (!parsed.success) return invalid("write_daily_brief", parsed.error);
  const brief = await upsertDailyBrief(prisma, context.userId, {
    text: parsed.data.text,
    date: parsed.data.date,
    botId: context.bot.id,
  });
  return { ok: true, date: brief.date };
}

const HANDLERS: Record<
  CoachToolName,
  (
    prisma: PrismaClient,
    context: CoachToolContext,
    rawArgs: Record<string, unknown>,
  ) => Promise<ToolResult>
> = {
  update_goal: updateGoalFromTool,
  update_habit: updateHabitFromTool,
  log_check_in: logCheckInFromTool,
  write_daily_brief: writeDailyBriefFromTool,
};

/** Runs one coach tool for the user's personal main agent; every other bot is refused. */
export async function executeCoachTool(
  prisma: PrismaClient,
  context: CoachToolContext,
  name: CoachToolName,
  rawArgs: Record<string, unknown>,
): Promise<ToolResult> {
  if (!(await isPersonalMainBot(prisma, { kind: context.bot.kind, spaceId: context.spaceId }))) {
    return { error: "Only the user's main bot can coach the user." };
  }
  try {
    return await HANDLERS[name](prisma, context, rawArgs);
  } catch (error) {
    if (error instanceof CoachNotFoundError || error instanceof InvalidCoachInputError) {
      return { error: error.message };
    }
    throw error;
  }
}