@tuturuuu/ai 0.0.10

package/src/chat/google/route.ts

@@ -0,0 +1,632 @@
+import { google } from '@ai-sdk/google';
+import { createAdminClient } from '@tuturuuu/supabase/next/server';
+import {
+  normalizeWorkspaceId,
+  verifyWorkspaceMembershipType,
+} from '@tuturuuu/utils/workspace-helper';
+import {
+  consumeStream,
+  type ModelMessage,
+  smoothStream,
+  stepCountIs,
+  streamText,
+} from 'ai';
+import { type NextRequest, NextResponse } from 'next/server';
+import { normalizeStableModelId } from '../../credits/model-mapping';
+import {
+  PlanModelResolutionError,
+  resolvePlanModel,
+} from '../../credits/resolve-plan-model';
+import { withAiMemory } from '../../memory';
+import type { CreditSource as SharedCreditSource } from '../credit-source';
+import {
+  shouldForceGoogleSearchForLatestUserMessage,
+  shouldForceRenderUiForLatestUserMessage,
+  shouldForceWorkspaceMembersForLatestUserMessage,
+  shouldPreferMarkdownTablesForLatestUserMessage,
+  shouldResolveWorkspaceContextForLatestUserMessage,
+  shouldUseParallelChecksForLatestUserMessage,
+} from '../mira-render-ui-policy';
+import { ChatRequestBodySchema, mapToUIMessages } from './chat-request-schema';
+import { systemInstruction } from './default-system-instruction';
+import { prepareMiraToolStep } from './mira-step-preparation';
+import {
+  type AiRouteAuthResult,
+  isInternalTuturuuuAiUser,
+  resolveAiRouteAuth,
+} from './route-auth';
+import {
+  moveTempFilesToThread,
+  resolveChatIdForUser,
+} from './route-chat-resolution';
+import { performCreditPreflight } from './route-credits';
+import {
+  persistLatestUserMessage,
+  prepareProcessedMessages,
+} from './route-message-preparation';
+import { prepareMiraRuntime } from './route-mira-runtime';
+import {
+  buildAbortedStreamFinishResponse,
+  persistAssistantResponse,
+} from './stream-finish-persistence';
+
+type ThinkingMode = 'fast' | 'thinking';
+
+function splitSystemMessages(messages: ModelMessage[]) {
+  const systemMessages: string[] = [];
+  const nonSystemMessages: ModelMessage[] = [];
+
+  for (const message of messages) {
+    if (message.role === 'system') {
+      const systemMessage = getTextContent(message.content);
+      if (systemMessage) systemMessages.push(systemMessage);
+      continue;
+    }
+
+    nonSystemMessages.push(message);
+  }
+
+  return {
+    messages: nonSystemMessages,
+    system: systemMessages.join('\n\n').trim(),
+  };
+}
+
+function getTextContent(content: ModelMessage['content']) {
+  if (typeof content === 'string') return content.trim();
+  if (!Array.isArray(content)) return '';
+
+  return content
+    .map((part) => (part.type === 'text' ? part.text.trim() : ''))
+    .filter(Boolean)
+    .join('\n')
+    .trim();
+}
+
+function mergeSystemInstructions(...instructions: Array<string | null>) {
+  return instructions
+    .map((instruction) => instruction?.trim())
+    .filter(Boolean)
+    .join('\n\n');
+}
+
+export function createPOST(
+  _options: {
+    serverAPIKeyFallback?: boolean;
+    /** Gateway provider prefix for bare model names (e.g., 'openai', 'anthropic', 'vertex'). Defaults to 'google'. */
+    defaultProvider?: string;
+    resolveAuth?: (request: NextRequest) => Promise<AiRouteAuthResult | null>;
+  } = {}
+) {
+  const defaultProvider = _options.defaultProvider ?? 'google';
+
+  // Higher-order function that returns the actual request handler
+  return async function handler(req: NextRequest): Promise<Response> {
+    try {
+      const sbAdmin = await createAdminClient();
+      let requestBody: unknown;
+      try {
+        requestBody = await req.json();
+      } catch (error) {
+        return NextResponse.json(
+          {
+            error: 'Invalid JSON payload',
+            message:
+              error instanceof Error ? error.message : 'Malformed JSON body',
+          },
+          { status: 400 }
+        );
+      }
+
+      const parsedBody = ChatRequestBodySchema.safeParse(requestBody);
+      if (!parsedBody.success) {
+        return NextResponse.json(
+          {
+            error: 'Invalid request body',
+            issues: parsedBody.error.issues,
+          },
+          { status: 400 }
+        );
+      }
+
+      const {
+        id,
+        model,
+        messages,
+        wsId,
+        workspaceContextId,
+        isMiraMode,
+        timezone,
+        thinkingMode: rawThinkingMode,
+        creditSource: requestedCreditSourceRaw,
+        creditWsId: rawCreditWsId,
+        observabilityContext,
+        taskBoardContext,
+      } = parsedBody.data;
+      const thinkingMode: ThinkingMode =
+        rawThinkingMode === 'thinking' ? 'thinking' : 'fast';
+
+      if (!messages) {
+        console.error('Missing messages');
+        return new Response('Missing messages', { status: 400 });
+      }
+
+      const auth =
+        (await _options.resolveAuth?.(req)) ?? (await resolveAiRouteAuth(req));
+      if (!auth.ok) return auth.response;
+      const { supabase, user } = auth;
+
+      if (isMiraMode && !(await isInternalTuturuuuAiUser(auth))) {
+        return NextResponse.json(
+          { error: 'Mira mode is limited to @tuturuuu.com accounts' },
+          { status: 403 }
+        );
+      }
+
+      // Normalize both workspace identifiers so slugs like 'personal' resolve to UUIDs.
+      let normalizedWsId: string | null = null;
+      let requestedCreditWsId: string | undefined;
+      try {
+        normalizedWsId = wsId
+          ? await normalizeWorkspaceId(wsId, supabase, req)
+          : null;
+        requestedCreditWsId = rawCreditWsId
+          ? await normalizeWorkspaceId(rawCreditWsId, supabase, req)
+          : undefined;
+      } catch (normError) {
+        console.error(
+          'Workspace ID normalization failed:',
+          normError instanceof Error ? normError.message : normError
+        );
+        return NextResponse.json(
+          { error: 'Invalid workspace identifier' },
+          { status: 422 }
+        );
+      }
+
+      if (normalizedWsId) {
+        const contextMembership = await verifyWorkspaceMembershipType({
+          wsId: normalizedWsId,
+          userId: user.id,
+          supabase: sbAdmin,
+          requiredType: 'MEMBER',
+        });
+
+        if (contextMembership.error === 'membership_lookup_failed') {
+          console.error('DB error checking workspace membership');
+          return NextResponse.json(
+            { error: 'Internal error verifying workspace access' },
+            { status: 500 }
+          );
+        }
+
+        if (!contextMembership.ok) {
+          return NextResponse.json(
+            { error: 'Workspace access denied' },
+            { status: 403 }
+          );
+        }
+      }
+
+      const requestedCreditSource: SharedCreditSource =
+        requestedCreditSourceRaw ?? 'workspace';
+      let billingWsId: string | null = normalizedWsId ?? null;
+
+      if (requestedCreditSource === 'personal') {
+        const { data: personalWorkspace, error: personalWorkspaceError } =
+          await sbAdmin
+            .from('workspaces')
+            .select('id, workspace_members!inner(user_id)')
+            .eq('personal', true)
+            .eq('workspace_members.user_id', user.id)
+            .maybeSingle();
+
+        if (personalWorkspaceError) {
+          console.error(
+            'DB error looking up personal workspace:',
+            personalWorkspaceError.message
+          );
+          return NextResponse.json(
+            { error: 'Internal error resolving personal workspace' },
+            { status: 500 }
+          );
+        }
+
+        if (!personalWorkspace?.id) {
+          return NextResponse.json(
+            {
+              error:
+                'Personal workspace not found. Please ensure your account has a personal workspace.',
+              code: 'PERSONAL_WORKSPACE_NOT_FOUND',
+            },
+            { status: 403 }
+          );
+        }
+
+        if (
+          requestedCreditWsId &&
+          requestedCreditWsId !== personalWorkspace.id
+        ) {
+          return NextResponse.json(
+            {
+              error:
+                'Invalid credit workspace for personal credit source selection.',
+              code: 'INVALID_CREDIT_SOURCE',
+            },
+            { status: 403 }
+          );
+        }
+
+        billingWsId = personalWorkspace.id;
+      } else if (requestedCreditWsId) {
+        if (normalizedWsId && requestedCreditWsId !== normalizedWsId) {
+          return NextResponse.json(
+            {
+              error: 'Invalid credit workspace for workspace source selection.',
+              code: 'INVALID_CREDIT_SOURCE',
+            },
+            { status: 403 }
+          );
+        }
+
+        if (!normalizedWsId) {
+          const billingMembership = await verifyWorkspaceMembershipType({
+            wsId: requestedCreditWsId,
+            userId: user.id,
+            supabase: sbAdmin,
+            requiredType: 'MEMBER',
+          });
+
+          if (billingMembership.error === 'membership_lookup_failed') {
+            console.error('Failed to check billing workspace membership');
+            return NextResponse.json(
+              { error: 'Internal server error' },
+              { status: 500 }
+            );
+          }
+
+          if (!billingMembership.ok) {
+            return NextResponse.json(
+              {
+                error: 'Workspace access denied for selected credit workspace.',
+                code: 'INVALID_CREDIT_SOURCE',
+              },
+              { status: 403 }
+            );
+          }
+
+          billingWsId = requestedCreditWsId;
+        }
+      }
+
+      let resolvedModelId: string;
+      try {
+        if (billingWsId) {
+          const resolvedPlanModel = await resolvePlanModel({
+            capability: 'language',
+            requestedModel: model,
+            wsId: billingWsId,
+          });
+          resolvedModelId = resolvedPlanModel.modelId;
+        } else if (model) {
+          resolvedModelId = normalizeStableModelId(
+            model.includes('/') ? model : `${defaultProvider}/${model}`
+          );
+        } else {
+          return NextResponse.json(
+            {
+              error: 'Model is required when billing workspace is unavailable.',
+            },
+            { status: 400 }
+          );
+        }
+      } catch (error) {
+        if (error instanceof PlanModelResolutionError) {
+          const status =
+            error.code === 'WORKSPACE_ID_REQUIRED'
+              ? 400
+              : error.code === 'NO_ALLOCATION'
+                ? 503
+                : 500;
+          return NextResponse.json(
+            { error: error.message, code: error.code },
+            { status }
+          );
+        }
+
+        console.error('Failed to resolve chat model:', error);
+        return NextResponse.json(
+          { error: 'Failed to resolve chat model' },
+          { status: 500 }
+        );
+      }
+
+      const resolvedChatId = await resolveChatIdForUser(id, () =>
+        sbAdmin
+          .from('ai_chats')
+          .select('id')
+          .eq('creator_id', user.id)
+          .order('created_at', { ascending: false })
+          .limit(1)
+          .single()
+      );
+      if ('error' in resolvedChatId) {
+        return resolvedChatId.error;
+      }
+      const chatId = resolvedChatId.chatId;
+
+      const moveFilesError = await moveTempFilesToThread({
+        loadThread: () =>
+          sbAdmin
+            .from('ai_chat_messages')
+            .select('role, ai_chats!chat_id!inner(creator_id)')
+            .eq('chat_id', chatId)
+            .eq('ai_chats.creator_id', user.id),
+        listFiles: (tempStoragePath) =>
+          sbAdmin.storage.from('workspaces').list(tempStoragePath),
+        moveFile: (fromPath, toPath) =>
+          sbAdmin.storage.from('workspaces').move(fromPath, toPath),
+        wsId: normalizedWsId ?? undefined,
+        chatId,
+        userId: user.id,
+      });
+      if (moveFilesError) {
+        return moveFilesError;
+      }
+
+      const normalizedMessages = mapToUIMessages(messages);
+      const preparedMessages = await prepareProcessedMessages(
+        normalizedMessages,
+        normalizedWsId ?? undefined,
+        chatId,
+        req,
+        {
+          attachYoutubeVideoInput:
+            resolvedModelId.toLowerCase().startsWith('google/') ||
+            resolvedModelId.toLowerCase().startsWith('google-vertex/'),
+        }
+      );
+      if ('error' in preparedMessages) {
+        return preparedMessages.error;
+      }
+      const { processedMessages } = preparedMessages;
+      const promptMessages = splitSystemMessages(processedMessages);
+
+      const persistUserMessageError = await persistLatestUserMessage({
+        processedMessages,
+        chatId,
+        insertChatMessage: async (args) => {
+          const { error } = await supabase.rpc(
+            'insert_ai_chat_message' as never,
+            args as never
+          );
+          return { error };
+        },
+        source: isMiraMode ? 'Mira' : 'Rewise',
+      });
+      if (persistUserMessageError) {
+        return persistUserMessageError;
+      }
+
+      const creditPreflight = await performCreditPreflight({
+        wsId: billingWsId ?? normalizedWsId ?? undefined,
+        model: resolvedModelId,
+        userId: user.id,
+        sbAdmin,
+      });
+      if ('error' in creditPreflight) {
+        return creditPreflight.error;
+      }
+      const { cappedMaxOutput } = creditPreflight;
+
+      // Mutable ref so the render_ui preprocessor can read current steps
+      // at Zod-validation time (before the execute handler runs).
+      const stepsRef: { current: unknown[] } = { current: [] };
+
+      const { miraSystemPrompt, miraTools } = await prepareMiraRuntime({
+        isMiraMode,
+        wsId: normalizedWsId ?? undefined,
+        workspaceContextId,
+        creditWsId: billingWsId ?? normalizedWsId ?? undefined,
+        request: req,
+        user,
+        userId: user.id,
+        chatId,
+        supabase,
+        toolSupabase: sbAdmin as typeof supabase,
+        timezone,
+        taskBoardContext,
+        getSteps: () => stepsRef.current,
+      });
+
+      const effectiveSource = isMiraMode ? 'Mira' : 'Rewise';
+
+      const resolvedGatewayModel = await withAiMemory({
+        customId: chatId,
+        model: google(resolvedModelId.split('/').slice(-1)[0]!),
+        product: isMiraMode ? 'mira' : 'rewise',
+        source: effectiveSource,
+        surface: 'shared_chat',
+        userId: user.id,
+        wsId: billingWsId ?? normalizedWsId,
+      });
+
+      // Reasoning mode: default to fast unless the client explicitly requests thinking.
+      const modelLower = resolvedModelId.toLowerCase();
+      const supportsThinking =
+        modelLower.includes('gemini-2.5') || modelLower.includes('gemini-3');
+      const thinkingConfig = supportsThinking
+        ? thinkingMode === 'thinking'
+          ? { thinkingConfig: { includeThoughts: true } }
+          : {
+              thinkingConfig: {
+                thinkingBudget: 0,
+                includeThoughts: false,
+              },
+            }
+        : {};
+      const forceRenderUi = shouldForceRenderUiForLatestUserMessage(
+        promptMessages.messages
+      );
+      const forceGoogleSearch = shouldForceGoogleSearchForLatestUserMessage(
+        promptMessages.messages
+      );
+      const preferMarkdownTables =
+        shouldPreferMarkdownTablesForLatestUserMessage(promptMessages.messages);
+      const needsWorkspaceContextResolution =
+        shouldResolveWorkspaceContextForLatestUserMessage(
+          promptMessages.messages
+        );
+      const needsWorkspaceMembersTool =
+        shouldForceWorkspaceMembersForLatestUserMessage(
+          promptMessages.messages
+        );
+      const needsParallelChecks = shouldUseParallelChecksForLatestUserMessage(
+        promptMessages.messages
+      );
+
+      // Provider-native Google Search tool for non-Mira mode.
+      const googleSearchTool = {
+        google_search: google.tools.googleSearch({}),
+      };
+
+      type PrepareStep = NonNullable<
+        NonNullable<Parameters<typeof streamText>[0]>['prepareStep']
+      >;
+      const prepareStep: PrepareStep = ({ steps }) => {
+        // Keep the mutable ref in sync so the render_ui preprocessor can
+        // read current steps during Zod validation.
+        stepsRef.current = steps;
+        return prepareMiraToolStep({
+          steps,
+          forceGoogleSearch,
+          forceRenderUi,
+          needsParallelChecks,
+          needsWorkspaceContextResolution,
+          needsWorkspaceMembersTool,
+          preferMarkdownTables,
+        });
+      };
+
+      let assistantResponsePersisted = false;
+      const persistChatAssistantResponse = async (
+        response: Parameters<typeof persistAssistantResponse>[0]['response']
+      ) => {
+        if (assistantResponsePersisted) return;
+        assistantResponsePersisted = true;
+        await persistAssistantResponse({
+          response,
+          sbAdmin,
+          chatId,
+          userId: user.id,
+          model: resolvedModelId,
+          effectiveSource,
+          wsId: billingWsId ?? normalizedWsId ?? undefined,
+          observabilityContext,
+        });
+      };
+
+      const result = streamText({
+        abortSignal: req.signal,
+        experimental_transform: smoothStream({ delayInMs: null }),
+        maxRetries: 0,
+        model: resolvedGatewayModel,
+        messages: promptMessages.messages,
+        system: mergeSystemInstructions(
+          isMiraMode && miraSystemPrompt ? miraSystemPrompt : systemInstruction,
+          promptMessages.system
+        ),
+        ...(cappedMaxOutput ? { maxOutputTokens: cappedMaxOutput } : {}),
+        ...(miraTools
+          ? {
+              tools: { ...miraTools, ...googleSearchTool } as NonNullable<
+                Parameters<typeof streamText>[0]
+              >['tools'],
+              stopWhen: stepCountIs(25),
+              toolChoice: 'auto' as const,
+              prepareStep: prepareStep as NonNullable<
+                NonNullable<Parameters<typeof streamText>[0]>['prepareStep']
+              >,
+            }
+          : {
+              tools: googleSearchTool,
+            }),
+        providerOptions: {
+          google: {
+            ...thinkingConfig,
+            safetySettings: [
+              {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                threshold: 'BLOCK_NONE',
+              },
+            ],
+          },
+          vertex: {
+            ...thinkingConfig,
+            safetySettings: [
+              {
+                category: 'HARM_CATEGORY_HARASSMENT',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_HATE_SPEECH',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                threshold: 'BLOCK_NONE',
+              },
+              {
+                category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                threshold: 'BLOCK_NONE',
+              },
+            ],
+          },
+        },
+        onAbort: async ({ steps }) =>
+          persistChatAssistantResponse(buildAbortedStreamFinishResponse(steps)),
+        onFinish: persistChatAssistantResponse,
+        timeout: {
+          chunkMs: 20_000,
+          stepMs: 45_000,
+          totalMs: 120_000,
+        },
+      });
+
+      // Per https://ai-sdk.dev/docs/advanced/stopping-streams: consumeSseStream ensures
+      // the stream is consumed on abort so cleanup can run; use onFinish in toUIMessageStreamResponse
+      // to handle isAborted when needed.
+      return result.toUIMessageStreamResponse({
+        consumeSseStream: consumeStream,
+        sendReasoning: true,
+        sendSources: true,
+      });
+    } catch (error) {
+      if (error instanceof Error) {
+        console.log(error.message);
+        return NextResponse.json(
+          {
+            message: `## Edge API Failure\nCould not complete the request. Please view the **Stack trace** below.\n\`\`\`bash\n${error instanceof Error ? error.stack : 'Unknown error'}\n\`\`\``,
+          },
+          {
+            status: 500,
+          }
+        );
+      }
+      console.log(error);
+      return new Response('Internal Server Error', { status: 500 });
+    }
+  };
+}