Vault: auto-link 2FA codes to their logins

apps/desktop/src/components/vault/totp-autolink-panel.tsx

@@ -0,0 +1,166 @@
+import { useMemo, useState } from "react";
+import { motion } from "motion/react";
+import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
+import { toast } from "sonner";
+import { ArrowRight, Link2, X } from "lucide-react";
+import type { Credential, TotpEntry } from "@godmode/shared";
+import { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Spinner } from "@/components/ui/spinner";
+import { api } from "@/lib/api";
+import { qk } from "@/lib/queryKeys";
+import { cn } from "@/lib/utils";
+import { Favicon } from "./favicon";
+import { issuerDomain } from "./use-totp-codes";
+import { domainFromUrl, toastApiError } from "./vault-utils";
+
+const DISMISS_KEY = "godmode-totp-autolink-dismissed";
+
+const pairKey = (p: { credentialId: string; totpId: string }) => `${p.credentialId}:${p.totpId}`;
+
+/** "Stripe:dan@x.de" without an issuer reads as service "Stripe", account "dan@x.de". */
+function entryLabel(t: TotpEntry): { service: string; account: string } {
+  if (t.issuer) return { service: t.issuer, account: t.accountName };
+  const i = t.accountName.indexOf(":");
+  return i > 0 ? { service: t.accountName.slice(0, i), account: t.accountName.slice(i + 1) } : { service: t.accountName, account: "" };
+}
+
+function loadDismissed(): Set<string> {
+  try {
+    return new Set(JSON.parse(localStorage.getItem(DISMISS_KEY) ?? "[]") as string[]);
+  } catch {
+    return new Set();
+  }
+}
+
+/** Logins and 2FA codes that clearly belong together, linked in one click. Agents also link them on their own once a guessed code works. */
+export function TotpAutoLinkPanel({ entries, credById }: { entries: TotpEntry[]; credById: Map<string, Credential> }) {
+  const qc = useQueryClient();
+  const [dismissed, setDismissed] = useState(loadDismissed);
+  const [unchecked, setUnchecked] = useState<Set<string>>(new Set());
+  const [checkedLikely, setCheckedLikely] = useState<Set<string>>(new Set());
+  const suggestions = useQuery({ queryKey: qk.totpAutolink, queryFn: api.totp.autolinkSuggestions, enabled: entries.length > 0 });
+
+  const totpById = useMemo(() => new Map(entries.map((t) => [t.id, t])), [entries]);
+  const rows = useMemo(
+    () =>
+      (suggestions.data ?? [])
+        .filter((s) => !dismissed.has(pairKey(s)))
+        .map((s) => ({ ...s, key: pairKey(s), totp: totpById.get(s.totpId), cred: credById.get(s.credentialId) }))
+        .filter((r): r is typeof r & { totp: TotpEntry; cred: Credential } => !!r.totp && !!r.cred),
+    [suggestions.data, dismissed, totpById, credById],
+  );
+  const isChecked = (r: { key: string; sure: boolean }) => (unchecked.has(r.key) ? false : r.sure || checkedLikely.has(r.key));
+  const selected = rows.filter(isChecked);
+
+  const toggle = (r: { key: string; sure: boolean }, on: boolean) => {
+    const set = (s: Set<string>, add: boolean) => {
+      const next = new Set(s);
+      if (add) next.add(r.key);
+      else next.delete(r.key);
+      return next;
+    };
+    if (r.sure) setUnchecked((s) => set(s, !on));
+    else setCheckedLikely((s) => set(s, on));
+  };
+
+  const dismiss = (keys: string[]) => {
+    setDismissed((d) => {
+      const next = new Set([...d, ...keys]);
+      try {
+        localStorage.setItem(DISMISS_KEY, JSON.stringify([...next].slice(-500)));
+      } catch {
+        /* ignore */
+      }
+      return next;
+    });
+  };
+
+  const link = useMutation({
+    mutationFn: () => api.totp.autolink(selected.map(({ credentialId, totpId }) => ({ credentialId, totpId }))),
+    onSuccess: (res) => {
+      void qc.invalidateQueries({ queryKey: qk.totp });
+      void qc.invalidateQueries({ queryKey: qk.credentials });
+      void qc.invalidateQueries({ queryKey: qk.bootstrap });
+      const n = res.linked.length;
+      if (n) toast.success(`${n} 2FA ${n === 1 ? "code" : "codes"} linked`, { description: "Agents fill them in by themselves on the next sign-in." });
+      else toast.info("Nothing to link anymore", { description: "These logins got a 2FA code in the meantime." });
+    },
+    onError: (e) => toastApiError(e, "Could not link the 2FA codes", qc),
+  });
+
+  if (!rows.length) return null;
+  const sureCount = rows.filter((r) => r.sure).length;
+
+  return (
+    <motion.section
+        initial={{ opacity: 0, y: -6 }}
+        animate={{ opacity: 1, y: 0 }}
+        exit={{ opacity: 0, y: -6 }}
+        aria-label="2FA codes that match a login"
+        className="overflow-hidden rounded-xl border border-brand/25 bg-card shadow-card"
+      >
+        <div className="flex items-start gap-3 px-4 pt-4 pb-3">
+          <span className="grid size-9 shrink-0 place-items-center rounded-lg bg-brand/10 text-brand-strong">
+            <Link2 className="size-4" />
+          </span>
+          <div className="min-w-0 flex-1">
+            <p className="text-sm font-medium tracking-[-0.01em]">
+              {rows.length === 1 ? "1 code matches a login" : `${rows.length} codes match a login`}
+            </p>
+            <p className="mt-0.5 text-xs text-muted-foreground">
+              Matched by service and account. Link them and agents finish two-factor on their own.
+              {sureCount < rows.length ? " Likely matches stay unchecked; agents also try them and link the ones that work." : ""}
+            </p>
+          </div>
+          <Button variant="ghost" size="icon-sm" className="-mt-1 -mr-1 text-muted-foreground" aria-label="Hide these suggestions" onClick={() => dismiss(rows.map((r) => r.key))}>
+            <X />
+          </Button>
+        </div>
+
+        <ul className="max-h-72 divide-y overflow-y-auto border-y">
+          {rows.map((r) => {
+            const checked = isChecked(r);
+            const credDomain = r.cred.domains[0] ?? domainFromUrl(r.cred.url);
+            const { service, account } = entryLabel(r.totp);
+            return (
+              <li key={r.key}>
+                <label className={cn("flex cursor-pointer items-center gap-3 px-4 py-2.5 transition-colors hover:bg-accent/50", !checked && "opacity-75")}>
+                  <Checkbox checked={checked} onCheckedChange={(v) => toggle(r, v === true)} aria-label={`Link ${service} to ${r.cred.name}`} />
+                  <span className="flex min-w-0 flex-1 items-center gap-2.5">
+                    <Favicon domain={credDomain || issuerDomain(r.totp.issuer)} name={service} size="sm" />
+                    <span className="min-w-0 flex-1">
+                      <span className="block truncate text-sm font-medium">{service}</span>
+                      {account && <span className="block truncate text-xs text-muted-foreground">{account}</span>}
+                    </span>
+                  </span>
+                  <ArrowRight className="size-3.5 shrink-0 text-muted-foreground" />
+                  <span className="min-w-0 flex-1">
+                    <span className="block truncate text-sm">{r.cred.name}</span>
+                    <span className="block truncate text-xs text-muted-foreground">{r.cred.username || credDomain}</span>
+                  </span>
+                  <span
+                    className={cn(
+                      "hidden shrink-0 rounded-[5px] px-1.5 py-0.5 text-[10px] font-medium tabular-nums sm:inline",
+                      r.sure ? "bg-success/10 text-success" : "bg-muted text-muted-foreground",
+                    )}
+                  >
+                    {r.sure ? "Match" : "Likely"}
+                  </span>
+                </label>
+              </li>
+            );
+          })}
+        </ul>
+
+        <div className="flex items-center justify-end gap-2 bg-paper-2 px-4 py-3">
+          <Button variant="ghost" size="sm" onClick={() => dismiss(rows.map((r) => r.key))} disabled={link.isPending}>
+            Not now
+          </Button>
+          <Button size="sm" onClick={() => link.mutate()} disabled={!selected.length || link.isPending}>
+            {link.isPending ? <Spinner /> : <Link2 />} Link {selected.length || ""}
+          </Button>
+        </div>
+    </motion.section>
+  );
+}

apps/desktop/src/lib/api.ts

@@ -203,6 +203,8 @@ import type { DeviceGroup, DeviceGroupPatch, DeviceInvite, HandoffInput, Handoff
   TotpImportInput,
   TotpImportResult,
   TotpInput,
+  TotpAutoLinkResult,
+  TotpAutoLinkSuggestion,
   UpdateReport,
   UsageSummary,
   VaultStatus,
@@ -498,6 +500,8 @@ export const api = {
     update: (id: string, input: Partial<TotpInput>) => patch<TotpEntry>(`/api/totp/${id}`, input),
     delete: (id: string) => del<{ ok: true }>(`/api/totp/${id}`),
     import: (input: TotpImportInput) => post<TotpImportResult>("/api/totp/import", input),
+    autolinkSuggestions: () => get<TotpAutoLinkSuggestion[]>("/api/totp/autolink"),
+    autolink: (pairs?: { credentialId: string; totpId: string }[]) => post<TotpAutoLinkResult>("/api/totp/autolink", { pairs }),
   },
 
   workspaces: {

apps/desktop/src/lib/queryKeys.ts

@@ -29,6 +29,7 @@ export const qk = {
   totp: ["totp"] as unknown[],
   totpList: (workspaceId: string) => ["totp", "list", workspaceId],
   totpCodes: ["totp", "codes"] as unknown[],
+  totpAutolink: ["totp", "autolink"] as unknown[],
   cards: ["cards"] as unknown[],
   cardList: (workspaceId: string) => ["cards", "list", workspaceId],
   cardPurchases: ["cards", "purchases"] as unknown[],

apps/desktop/src/lib/realtime.ts

@@ -49,7 +49,7 @@ const ENTITY_KEYS: Record<EntityName, readonly unknown[][]> = {
   workspaces: [qk.workspaces],
   agents: [qk.agents, qk.spend, qk.bootstrap],
   routines: [qk.routines, qk.bootstrap],
-  credentials: [qk.credentials],
+  credentials: [qk.credentials, qk.totpAutolink],
   totp: [qk.totp],
   passkeys: [qk.passkeys],
   "payment-cards": [qk.cards],

apps/desktop/src/pages/vault/totp-page.tsx

@@ -11,6 +11,7 @@ import { EmptyState, PageBody, PageHeader } from "@/components/common";
 import { TotpCard, TotpCardSkeleton } from "@/components/vault/totp-card";
 import { TotpEditDialog } from "@/components/vault/totp-edit-dialog";
 import { TotpImportDialog } from "@/components/vault/totp-import-dialog";
+import { TotpAutoLinkPanel } from "@/components/vault/totp-autolink-panel";
 import { ConfirmDeleteDialog } from "@/components/vault/confirm-dialog";
 import { ScopeFilterSelect } from "@/components/vault/workspace-select";
 import { useTotpCodes } from "@/components/vault/use-totp-codes";
@@ -148,6 +149,8 @@ export default function TotpPage() {
           </div>
         )}
 
+        {hasEntries && creds.data && <TotpAutoLinkPanel entries={list.data ?? []} credById={credById} />}
+
         {list.isLoading ? (
           <div className="@container">
             <div className="grid gap-3 @xl:grid-cols-2 @4xl:grid-cols-3">

packages/core/src/browser/fill.ts

@@ -97,6 +97,24 @@ interface LocateResult {
   focusInFrame?: boolean;
 }
 
+/** Is a visible 2FA code input (or a row of one-digit boxes) on this document, and does it say a code was wrong? */
+const OTP_PROBE = String.raw`() => {
+  const OTP_RE = /one[-_ ]?time|otp|totp|2fa|mfa|two[-_ ]?factor|verification|verify|security[-_ ]?code|auth(entication)?[-_ ]?code|\bcode\b|token|\bpin\b/;
+  const visible = (el) => {
+    const st = getComputedStyle(el);
+    if (st.display === "none" || st.visibility === "hidden") return false;
+    const r = el.getBoundingClientRect();
+    return r.width >= 2 && r.height >= 2;
+  };
+  const attrText = (el) => ["name", "id", "autocomplete", "placeholder", "aria-label", "data-testid"].map((a) => el.getAttribute(a) || "").join(" ").toLowerCase();
+  const inputs = [...document.querySelectorAll("input")].filter((el) => !el.disabled && ["", "text", "tel", "number", "password"].includes((el.getAttribute("type") || "").toLowerCase()) && visible(el));
+  const field = inputs.some((el) => el.getAttribute("autocomplete") === "one-time-code" || (OTP_RE.test(attrText(el)) && !/user|e-?mail|search/.test(attrText(el))));
+  const boxes = inputs.filter((el) => el.maxLength === 1).length >= 4;
+  const text = (document.body && document.body.innerText || "").slice(0, 20000);
+  const error = /incorrect|invalid|wrong code|didn.t work|not valid|expired|too many|try again|falsch|ung(\u00fc|ue)ltig|abgelaufen|zu viele|erneut versuchen|nicht korrekt/i.test(text);
+  return { otp: field || boxes, error };
+}`;
+
 /* ------------------------------------------------------------------ */
 /* In-page scripts (run in an isolated world; state kept on its global) */
 /* ------------------------------------------------------------------ */
@@ -688,3 +706,44 @@ export async function fillIntoActivePage(
     return { ok: false, url: "", detail: scrub(`Could not fill the field: ${err instanceof Error ? err.message : String(err)}`) };
   }
 }
+
+export interface OtpProbe {
+  url: string;
+  otpField: boolean;
+  /** The page says a code was wrong, expired or tried too often. */
+  error: boolean;
+}
+
+/** Whether a 2FA code field is still showing on the page or one of its frames — gone means the site took the code. */
+export async function otpFieldOnPage(page: PageSession): Promise<OtpProbe> {
+  await page.waitForReady(3000).catch(() => {});
+  const url = await page.evaluate<string>("location.href", { timeoutMs: 3000 }).catch(() => "");
+  const contexts: FrameContext[] = [{ session: page, contextId: await page.isolatedWorld(), isChildFrame: false }, ...(await frameContexts(page).catch(() => []))];
+  const out: OtpProbe = { url, otpField: false, error: false };
+  for (const ctx of contexts) {
+    const res = await call<{ otp: boolean; error: boolean }>(ctx, OTP_PROBE, {}).catch(() => null);
+    out.otpField ||= !!res?.otp;
+    out.error ||= !!res?.error;
+  }
+  return out;
+}
+
+/** otpFieldOnPage for the page fillIntoActivePage would pick; null when there is no page to look at. */
+export async function probeOtpInActivePage(browser: {
+  client: CdpClient;
+  port?: number;
+  chatPage?: (urlContains?: string) => PageTarget | null;
+}): Promise<OtpProbe | null> {
+  try {
+    const target = browser.chatPage ? browser.chatPage() : await pickActivePage(browser.client, { port: browser.port });
+    if (!target) return null;
+    const page = await attachToPage(browser.client, target.targetId);
+    try {
+      return await otpFieldOnPage(page);
+    } finally {
+      await page.detach();
+    }
+  } catch {
+    return null;
+  }
+}

packages/core/src/browser/manager.ts

@@ -19,7 +19,7 @@ import { resolveUvx, toolPath } from "../services/doctor";
 import { hasBrowserSubscribers, hasBrowserWatchers } from "../server/ws";
 import { CdpClient, attachToPage, getCookies, pickActivePage, probeCdp, isUserPage, type CdpCookie, type CdpCookieParam, type PageSession, type PageTarget } from "./cdp";
 import { clearLaunchMarker, findChrome, isProcessAlive, launchChrome, readLaunchMarker, writeLaunchMarker, type ChromeProcess } from "./chrome";
-import { fillIntoActivePage, fillPrecheck, type FillKind } from "./fill";
+import { fillIntoActivePage, fillPrecheck, probeOtpInActivePage, type OtpProbe, type FillKind } from "./fill";
 import { solveInBrowser, type SolveOptions, type SolveResult } from "./captcha";
 import { browserUseCommand, browserUseEnv, writeBrowserUseConfig } from "./browserUse";
 import { stealthArgs, stopProbes, windowedUserAgent } from "./stealth";
@@ -822,6 +822,13 @@ export async function fillIntoPage(
   }
 }
 
+/** Whether the chat's page still shows a 2FA code field (see probeOtpInActivePage); null when the browser or tab is gone. */
+export async function probeOtpPage(profileId: string, conversationId?: string): Promise<OtpProbe | null> {
+  const rb = getRunning(profileId);
+  if (!rb) return null;
+  return probeOtpInActivePage({ client: rb.client, port: rb.port, ...(conversationId ? { chatPage: () => chatPage(rb, conversationId) } : {}) });
+}
+
 /** Solve the CAPTCHA on the chat's page (without a chat: the active page) through the selected solving service. */
 export async function solveCaptchaInPage(profileId: string, opts: SolveOptions & { conversationId?: string }): Promise<SolveResult> {
   requireRow(profileId);