@payloadcms/plugin-mcp 4.0.0-internal.a1317b2 → 4.0.0-internal.a2d053c

package/src/endpoint/oauthLoginToken.ts

@@ -0,0 +1,91 @@
+import type { PayloadRequest } from 'payload'
+
+import { APIError, createAuthToken, executeAuthStrategies, redeemRefreshToken } from 'payload'
+
+import type { MCPOAuthTokens } from '../types.js'
+
+import { seal, unseal } from '../utils/seal.js'
+
+/**
+ * The default `issueToken`: a Payload login token, sealed so it only works for MCP, where
+ * `authenticateLoginToken` unseals it. Otherwise an MCP client could use it as the user's admin
+ * panel login. When the collection uses sessions, the client also gets a refresh token.
+ */
+export const issueLoginToken = async ({
+  req,
+}: {
+  req: PayloadRequest
+}): Promise<MCPOAuthTokens> => {
+  const collectionConfig = req.payload.collections[req.user!.collection]!.config
+
+  if (collectionConfig.auth.disableLocalStrategy) {
+    throw new APIError(
+      `MCP clients can't get Payload login tokens for the "${collectionConfig.slug}" collection because its local strategy is disabled. Pass oauth.issueToken to the MCP plugin.`,
+    )
+  }
+
+  // The full document, like the one a password login starts its session on
+  const user = (await req.payload.db.findOne<NonNullable<PayloadRequest['user']>>({
+    collection: collectionConfig.slug,
+    req,
+    where: { id: { equals: req.user!.id } },
+  }))!
+
+  const { refreshToken, token } = await createAuthToken({
+    collectionConfig,
+    email: user.email!,
+    req,
+    shouldCreateRefreshToken: true,
+    user,
+  })
+
+  return {
+    accessToken: sealLoginToken({ req, token }),
+    expiresIn: collectionConfig.auth.tokenExpiration,
+    refreshToken,
+  }
+}
+
+/** The default `refreshToken`. Payload logs out the session when a refresh token is reused. */
+export const refreshLoginToken = async ({
+  refreshToken,
+  req,
+}: {
+  refreshToken: string
+  req: PayloadRequest
+}): Promise<MCPOAuthTokens | null> => {
+  const refreshed = await redeemRefreshToken({ refreshToken, req })
+
+  return (
+    refreshed && {
+      accessToken: sealLoginToken({ req, token: refreshed.token }),
+      expiresIn: refreshed.collectionConfig.auth.tokenExpiration,
+      refreshToken: refreshed.refreshToken,
+    }
+  )
+}
+
+/** Finds the user of an MCP request's access token from `issueLoginToken`. */
+export const authenticateLoginToken = async ({
+  req,
+  token,
+}: {
+  req: PayloadRequest
+  token: string
+}): Promise<PayloadRequest['user']> => {
+  const loginToken = unseal<{ token: string }>({ type: 'access', req, value: token })
+
+  if (!loginToken) {
+    return null
+  }
+
+  const headers = new Headers({
+    Authorization: `JWT ${loginToken.token}`,
+    DisableAutologin: 'true',
+  })
+
+  return (await executeAuthStrategies({ headers, payload: req.payload, req })).user
+}
+
+const sealLoginToken = ({ req, token }: { req: PayloadRequest; token: string }): string =>
+  seal({ type: 'access', data: { token }, req })

package/src/endpoint/oauthServer.ts

@@ -0,0 +1,396 @@
+import type {
+  OAuthClientInformationFull,
+  OAuthTokens,
+  OpenIdProviderDiscoveryMetadata,
+} from '@modelcontextprotocol/server'
+import type { Endpoint, PayloadHandler, PayloadRequest } from 'payload'
+
+import { OAuthErrorCode } from '@modelcontextprotocol/server'
+import { createHash } from 'crypto'
+import { formatAdminURL } from 'payload/shared'
+
+import type { MCPBuiltinOAuthConfig, MCPOAuthTokens } from '../types.js'
+
+import { getPluginConfig } from '../utils/getPluginConfig.js'
+import { seal, unseal } from '../utils/seal.js'
+import { getMCPURL } from './oauth.js'
+import { issueLoginToken, refreshLoginToken } from './oauthLoginToken.js'
+
+type Client = { name: string; redirectURIs: string[] }
+
+type Consent = {
+  clientID: string
+  codeChallenge: string
+  collection: string
+  redirectURI: string
+  state: null | string
+  userID: number | string
+}
+
+type Code = {
+  clientID: string
+  codeChallenge: string
+  redirectURI: string
+  tokens: MCPOAuthTokens
+}
+
+type RefreshToken = {
+  clientID: string
+  refreshToken: string
+}
+
+/**
+ * Turns Payload into the OAuth authorization server for MCP clients when `oauth.type` is
+ * `'builtin'`. Clients register, the user allows them after logging in to the admin panel, and
+ * the client receives the tokens `issueToken` creates.
+ *
+ * Clients, consent forms, authorization codes, and refresh tokens are encrypted and handed to the
+ * client instead of being stored. Only the default tokens write to the database, to the user's
+ * login sessions.
+ */
+export const oauthServerEndpoints: Endpoint[] = [
+  {
+    disableAuth: true,
+    handler: (req) => withIssuer(req, (issuer) => metadataResponse(issuer)),
+    method: 'get',
+    // OpenID discovery is the only metadata location MCP clients check below the issuer's path,
+    // which lets this work under /api without a route at the site root
+    path: '/mcp/oauth/.well-known/openid-configuration',
+  },
+  {
+    disableAuth: true,
+    handler: (req) => withIssuer(req, () => Response.json({ keys: [] })),
+    method: 'get',
+    path: '/mcp/oauth/jwks',
+  },
+  {
+    disableAuth: true,
+    handler: (req) => withIssuer(req, () => register(req)),
+    method: 'post',
+    path: '/mcp/oauth/register',
+  },
+  {
+    handler: (req) => withIssuer(req, () => authorize(req)),
+    method: 'get',
+    path: '/mcp/oauth/authorize',
+  },
+  {
+    handler: (req) => withIssuer(req, () => approve(req)),
+    method: 'post',
+    path: '/mcp/oauth/authorize',
+  },
+  {
+    disableAuth: true,
+    handler: (req) => withIssuer(req, () => token(req)),
+    method: 'post',
+    path: '/mcp/oauth/token',
+  },
+]
+
+/** The authorization server URL, or `undefined` when Payload isn't the authorization server. */
+const getIssuer = (req: PayloadRequest): string | undefined =>
+  getPluginConfig({ config: req.payload.config }).oauth?.type === 'builtin'
+    ? `${getMCPURL({ config: req.payload.config, url: req.url! })}/oauth`
+    : undefined
+
+const withIssuer = (
+  req: PayloadRequest,
+  handler: (issuer: string) => Promise<Response> | Response,
+): Promise<Response> | Response => {
+  const issuer = getIssuer(req)
+  return issuer ? handler(issuer) : new Response(null, { status: 404 })
+}
+
+/**
+ * The user logged in to this browser, who can allow MCP clients. Tokens and API keys in the
+ * `Authorization` header can't, or an MCP client could use its access token to approve itself a
+ * new, separate login.
+ */
+const getBrowserUser = (req: PayloadRequest): PayloadRequest['user'] =>
+  req.headers.has('Authorization') ? null : req.user
+
+const metadataResponse = (issuer: string): Response =>
+  Response.json(
+    {
+      authorization_endpoint: `${issuer}/authorize`,
+      code_challenge_methods_supported: ['S256'],
+      grant_types_supported: ['authorization_code', 'refresh_token'],
+      // Required by OpenID discovery. No ID tokens are issued.
+      id_token_signing_alg_values_supported: ['RS256'],
+      issuer,
+      jwks_uri: `${issuer}/jwks`,
+      registration_endpoint: `${issuer}/register`,
+      response_types_supported: ['code'],
+      subject_types_supported: ['public'],
+      token_endpoint: `${issuer}/token`,
+      token_endpoint_auth_methods_supported: ['none'],
+    } satisfies OpenIdProviderDiscoveryMetadata,
+    // Not cached, since without `serverURL` its URLs come from the request
+    { headers: { 'Access-Control-Allow-Origin': '*', 'Cache-Control': 'no-store' } },
+  )
+
+const register: PayloadHandler = async (req) => {
+  const { client_name, redirect_uris } = await req.json!().catch(() => ({}))
+
+  if (
+    !Array.isArray(redirect_uris) ||
+    !redirect_uris.length ||
+    !redirect_uris.every(
+      (uri) => typeof uri === 'string' && URL.canParse(uri) && !uri.includes('#'),
+    )
+  ) {
+    return Response.json({ error: OAuthErrorCode.InvalidRedirectUri }, { status: 400 })
+  }
+
+  const client: Client = {
+    name: typeof client_name === 'string' ? client_name : 'An MCP client',
+    redirectURIs: redirect_uris,
+  }
+
+  return Response.json(
+    {
+      client_id: seal({ type: 'client', data: client, req }),
+      client_name: client.name,
+      grant_types: ['authorization_code', 'refresh_token'],
+      redirect_uris: client.redirectURIs,
+      response_types: ['code'],
+      token_endpoint_auth_method: 'none',
+    } satisfies OAuthClientInformationFull,
+    { status: 201 },
+  )
+}
+
+const authorize: PayloadHandler = (req) => {
+  const url = new URL(req.url!)
+  const params = url.searchParams
+  const clientID = params.get('client_id')
+  const client = unseal<Client>({ type: 'client', req, value: clientID })
+  const redirectURI = params.get('redirect_uri')
+
+  // Never redirect to an address the client didn't register
+  if (!client || !redirectURI || !client.redirectURIs.includes(redirectURI)) {
+    return page({
+      content: '<p>Start the sign-in again from your MCP client.</p>',
+      status: 400,
+      title: 'This sign-in link is invalid',
+    })
+  }
+
+  const state = params.get('state')
+  const codeChallenge = params.get('code_challenge')
+
+  if (
+    params.get('response_type') !== 'code' ||
+    params.get('code_challenge_method') !== 'S256' ||
+    !codeChallenge
+  ) {
+    return redirect({ params: { error: OAuthErrorCode.InvalidRequest, state }, redirectURI })
+  }
+
+  const user = getBrowserUser(req)
+
+  if (!user) {
+    const loginURL = formatAdminURL({ adminRoute: req.payload.config.routes.admin, path: '/login' })
+
+    return new Response(null, {
+      headers: {
+        Location: `${loginURL}?redirect=${encodeURIComponent(url.pathname + url.search)}`,
+      },
+      status: 302,
+    })
+  }
+
+  const consent: Consent = {
+    clientID: clientID!,
+    codeChallenge,
+    collection: user.collection,
+    redirectURI,
+    state,
+    userID: user.id,
+  }
+
+  return page({
+    content: `<p>It will be able to use this app's MCP tools as ${escapeHTML(String(user.email ?? user.id))}, and you will return to ${escapeHTML(new URL(redirectURI).host || redirectURI)}.</p>
+<form method="post">
+<input name="consent" type="hidden" value="${seal({ type: 'consent', data: consent, expiresIn: 20 * 60 * 60, req })}">
+<button name="decision" type="submit" value="allow">Allow</button>
+<button name="decision" type="submit" value="deny">Deny</button>
+</form>`,
+    status: 200,
+    title: `Allow ${client.name} to use your account?`,
+  })
+}
+
+const approve: PayloadHandler = async (req) => {
+  const form = new URLSearchParams(await req.text!())
+  const consent = unseal<Consent>({ type: 'consent', req, value: form.get('consent') })
+  const user = getBrowserUser(req)
+
+  // The form is only valid for the user it was shown to
+  if (!consent || user?.id !== consent.userID || user.collection !== consent.collection) {
+    return page({
+      content: '<p>Start the sign-in again from your MCP client.</p>',
+      status: 400,
+      title: 'This sign-in request expired',
+    })
+  }
+
+  const { clientID, codeChallenge, redirectURI, state } = consent
+
+  if (form.get('decision') !== 'allow') {
+    return redirect({ params: { error: OAuthErrorCode.AccessDenied, state }, redirectURI })
+  }
+
+  const code: Code = {
+    clientID,
+    codeChallenge,
+    redirectURI,
+    tokens: await getTokenHooks(req).issueToken({ req }),
+  }
+
+  return redirect({
+    params: { code: seal({ type: 'code', data: code, expiresIn: 60, req }), state },
+    redirectURI,
+  })
+}
+
+const token: PayloadHandler = async (req) => {
+  const form = new URLSearchParams(await req.text!())
+  const clientID = form.get('client_id')
+  const tokens =
+    clientID &&
+    (form.get('grant_type') === 'refresh_token'
+      ? await refresh({ clientID, form, req })
+      : redeemCode({ clientID, form, req }))
+
+  if (!clientID || !tokens) {
+    return Response.json({ error: OAuthErrorCode.InvalidGrant }, { status: 400 })
+  }
+
+  return Response.json(
+    {
+      access_token: tokens.accessToken,
+      expires_in: tokens.expiresIn,
+      refresh_token:
+        tokens.refreshToken &&
+        seal({
+          type: 'refresh',
+          data: { clientID, refreshToken: tokens.refreshToken } satisfies RefreshToken,
+          req,
+        }),
+      token_type: 'Bearer',
+    } satisfies OAuthTokens,
+    { headers: { 'Cache-Control': 'no-store' } },
+  )
+}
+
+const redeemCode = ({
+  clientID,
+  form,
+  req,
+}: {
+  clientID: string
+  form: URLSearchParams
+  req: PayloadRequest
+}): MCPOAuthTokens | undefined => {
+  const code = unseal<Code>({ type: 'code', req, value: form.get('code') })
+
+  if (
+    form.get('grant_type') === 'authorization_code' &&
+    code?.clientID === clientID &&
+    code.redirectURI === form.get('redirect_uri') &&
+    // PKCE: only the client that started the sign-in knows the verifier
+    createHash('sha256')
+      .update(form.get('code_verifier') ?? '')
+      .digest('base64url') === code.codeChallenge
+  ) {
+    return code.tokens
+  }
+}
+
+const refresh = async ({
+  clientID,
+  form,
+  req,
+}: {
+  clientID: string
+  form: URLSearchParams
+  req: PayloadRequest
+}): Promise<MCPOAuthTokens | null | undefined> => {
+  const refreshToken = unseal<RefreshToken>({
+    type: 'refresh',
+    req,
+    value: form.get('refresh_token'),
+  })
+
+  // Only the client a refresh token was issued to can use it
+  if (refreshToken?.clientID === clientID) {
+    return await getTokenHooks(req).refreshToken?.({
+      refreshToken: refreshToken.refreshToken,
+      req,
+    })
+  }
+}
+
+/** `issueToken` and `refreshToken`, or the defaults for Payload login tokens. */
+const getTokenHooks = (req: PayloadRequest) => {
+  const oauth = getPluginConfig({ config: req.payload.config }).oauth as MCPBuiltinOAuthConfig
+
+  return oauth.issueToken
+    ? { issueToken: oauth.issueToken, refreshToken: oauth.refreshToken }
+    : { issueToken: issueLoginToken, refreshToken: refreshLoginToken }
+}
+
+const redirect = ({
+  params,
+  redirectURI,
+}: {
+  params: Record<string, null | string>
+  redirectURI: string
+}): Response => {
+  const location = new URL(redirectURI)
+
+  for (const [key, value] of Object.entries(params)) {
+    if (value) {
+      location.searchParams.set(key, value)
+    }
+  }
+
+  return new Response(null, { headers: { Location: location.href }, status: 302 })
+}
+
+const page = ({
+  content,
+  status,
+  title,
+}: {
+  content: string
+  status: number
+  title: string
+}): Response =>
+  new Response(
+    `<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>${escapeHTML(title)}</title>
+<style>
+body { font-family: system-ui, sans-serif; margin: 4rem auto; max-width: 32rem; overflow-wrap: anywhere; padding: 0 1rem }
+button { font: inherit; margin-inline-end: 0.5rem; padding: 0.5rem 1rem }
+</style>
+</head>
+<body><h1>${escapeHTML(title)}</h1>${content}</body>
+</html>`,
+    {
+      headers: {
+        // Stop other sites from framing the page to trick users into clicking Allow
+        'Content-Security-Policy': "frame-ancestors 'none'",
+        'Content-Type': 'text/html; charset=utf-8',
+      },
+      status,
+    },
+  )
+
+const escapeHTML = (value: string): string =>
+  value.replace(/[&<>"']/g, (character) => `&#${character.charCodeAt(0)};`)

package/src/index.ts

@@ -1,8 +1,18 @@
 import { definePlugin } from 'payload'
 
-import type { AuthorizedMCP, MCPPluginConfig, SanitizedMCPPluginConfig } from './types.js'
+import type {
+  AuthorizedMCP,
+  MCPBuiltinOAuthConfig,
+  MCPExternalOAuthConfig,
+  MCPOAuthConfig,
+  MCPOAuthTokens,
+  MCPPluginConfig,
+  SanitizedMCPPluginConfig,
+} from './types.js'
 
 import { mcpEndpoint } from './endpoint/index.js'
+import { oauthProtectedResourceEndpoint } from './endpoint/oauth.js'
+import { oauthServerEndpoints } from './endpoint/oauthServer.js'
 import { sanitizeMCPConfig } from './mcp/sanitizeMCPConfig.js'
 
 declare module 'payload' {
@@ -13,11 +23,27 @@ declare module 'payload' {
     /** After the plugin's `plugin` callback runs, `options` holds the sanitized config. */
     '@payloadcms/plugin-mcp': MCPPluginConfig
   }
+  export interface RequestContext {
+    /** Set on MCP requests before they're authenticated. */
+    mcp?: {
+      /** The MCP endpoint's address, which OAuth access tokens for MCP are issued for. */
+      resource: string
+    }
+  }
 }
 
-export type { AuthorizedMCP, MCPPluginConfig, SanitizedMCPPluginConfig }
+export type {
+  AuthorizedMCP,
+  MCPBuiltinOAuthConfig,
+  MCPExternalOAuthConfig,
+  MCPOAuthConfig,
+  MCPOAuthTokens,
+  MCPPluginConfig,
+  SanitizedMCPPluginConfig,
+}
 
 export { defineCollectionTool, defineGlobalTool, definePrompt, defineTool } from './defineTool.js'
+export { oauthProtectedResourceHandler } from './endpoint/oauth.js'
 
 export const mcpPlugin = definePlugin<MCPPluginConfig>({
   slug: '@payloadcms/plugin-mcp',
@@ -29,7 +55,7 @@ export const mcpPlugin = definePlugin<MCPPluginConfig>({
     const registered = plugins['@payloadcms/plugin-mcp']
     if (registered) {
       // @ts-expect-error
-      registered.sanitizedOptions = pluginConfig as unknown as typeof registered.options
+      registered.sanitizedOptions = pluginConfig
     }
 
     if (pluginConfig.disabled) {
@@ -40,15 +66,24 @@ export const mcpPlugin = definePlugin<MCPPluginConfig>({
       ...config,
       endpoints: [
         ...(config.endpoints ?? []),
-        // Payload prefixes /api, so the full path is /api/mcp.
-        { handler: mcpEndpoint, method: 'post', path: '/mcp' },
+        // Payload prefixes /api, so the full path is /api/mcp. getAuthorizedMCP authenticates the
+        // request with dev auto-login turned off, so Payload doesn't authenticate it first.
+        { disableAuth: true, handler: mcpEndpoint, method: 'post', path: '/mcp' },
         // Streamable HTTP's optional server=>client GET stream. We don't offer
         // one, so answer 405 per spec-— clients then skip it cleanly
         {
+          disableAuth: true,
           handler: () => new Response(null, { headers: { Allow: 'POST' }, status: 405 }),
           method: 'get',
           path: '/mcp',
         },
+        {
+          disableAuth: true,
+          handler: oauthProtectedResourceEndpoint,
+          method: 'get',
+          path: '/mcp/oauth-protected-resource',
+        },
+        ...oauthServerEndpoints,
       ],
     }
   },

package/src/mcp/buildMcpServer.ts

@@ -17,8 +17,6 @@ import type {
 import { getLogger } from '../utils/getLogger.js'
 import { toStandardSchema } from '../utils/toStandardSchema.js'
 
-z.config(z.locales.en())
-
 /**
  * Serving-entry-agnostic core: registers every authorized MCP item onto a fresh
  * `McpServer` and returns it. The HTTP and stdio entry point callers provide fresh
@@ -27,7 +25,7 @@ z.config(z.locales.en())
  *
  * `req` is the request context handlers see. For HTTP it's the live
  * `PayloadRequest` derived from the incoming HTTP request; for stdio it's a
- * synthesized one built via `createLocalReq`.
+ * synthesized one built via `createPayloadRequest`.
  */
 export const buildMcpServer = ({
   authorizedMCP,
@@ -38,6 +36,8 @@ export const buildMcpServer = ({
   pluginConfig: SanitizedMCPPluginConfig
   req: PayloadRequest
 }): McpServer => {
+  z.config(z.locales.en())
+
   const serverOptions = pluginConfig.mcp?.serverOptions || {}
   const server = new McpServer(
     { name: 'Payload MCP Server', version: '1.0.0', ...serverOptions.serverInfo },
@@ -250,7 +250,7 @@ const withSlugInput = ({ input }: { input?: ToolInputSchema }): ToolInputSchema
   const slugSchema = z.string().check(z.describe(description))
 
   if (!input) {
-    return strictObject({ slug: slugSchema }) as unknown as ToolInputSchema
+    return strictObject({ slug: slugSchema })
   }
 
   if (typeof input === 'object' && input !== null && '~standard' in input) {
@@ -261,7 +261,7 @@ const withSlugInput = ({ input }: { input?: ToolInputSchema }): ToolInputSchema
       return input
     }
 
-    return addSlugToJSONSchema({ description, schema: schema as JsonSchemaType })
+    return addSlugToJSONSchema({ description, schema })
   }
 
   return addSlugToJSONSchema({ description, schema: input })

package/src/mcp/builtin/collections/authTools.ts

@@ -1,4 +1,4 @@
-import { createLocalReq, strictObject, z } from 'payload'
+import { createPayloadRequest, strictObject, z } from 'payload'
 
 import type { MCPToolResponse } from '../../../types.js'
 
@@ -53,7 +53,10 @@ export const authCollectionTool = defineCollectionTool({
     if (input.headers) {
       authHeaders = new Headers(input.headers)
     }
-    const authReq = await createLocalReq({ req: { headers: authHeaders } }, req.payload)
+    const authReq = await createPayloadRequest({
+      payload: req.payload,
+      req: { headers: authHeaders },
+    })
     const result = await req.payload.auth({ headers: authHeaders, req: authReq })
 
     if (result.user) {
@@ -232,7 +235,7 @@ export const unlockCollectionTool = defineCollectionTool({
           text: `# Account Unlocked\n\n**Collection:** ${slug}\n**Email:** ${input.email}\n**Result:** ${result ? 'Success' : 'Failed'}`,
         },
       ],
-      doc: { result } as Record<string, unknown>,
+      doc: { result },
     }
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : 'Unknown error'
@@ -267,7 +270,7 @@ export const verifyCollectionTool = defineCollectionTool({
           text: `# Email Verification Successful\n\n**Collection:** ${slug}\n**Result:** ${result ? 'Success' : 'Failed'}`,
         },
       ],
-      doc: { result } as Record<string, unknown>,
+      doc: { result },
     }
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : 'Unknown error'

package/src/mcp/builtin/collections/createTool.ts

@@ -32,8 +32,17 @@ export const createDocumentsTool = defineCollectionTool({
   const payload = req.payload
   const collectionConfig = payload.collections[slug]?.config
   const logger = getLogger({ payload })
-  const { depth, documents, draft, fallbackLocale, locale, populate, publishAllLocales, select } =
-    input
+  const {
+    depth,
+    documents,
+    draft,
+    fallbackLocale,
+    locale,
+    populate,
+    publishAllLocales,
+    returning,
+    select,
+  } = input
   const shouldUsePartialSchema =
     draft === true && collectionConfig !== undefined && !hasDraftValidationEnabled(collectionConfig)
 
@@ -41,7 +50,9 @@ export const createDocumentsTool = defineCollectionTool({
 
   try {
     const virtualFieldNames = getCollectionVirtualFieldNames(payload.config, slug)
-    const docs: Array<{ doc: Record<string, unknown>; index: number }> = []
+    const docs: Array<
+      { doc: Record<string, unknown>; index: number } | { id: number | string; index: number }
+    > = []
     const errors: Array<{ index: number; message: string }> = []
     let validationSchema: Record<string, unknown> | undefined
 
@@ -69,10 +80,10 @@ export const createDocumentsTool = defineCollectionTool({
           ...(file ? { file } : {}),
           ...(locale ? { locale } : {}),
           ...(fallbackLocale !== undefined ? { fallbackLocale } : {}),
-          ...(select ? { select } : {}),
+          select: returning ? select : { id: true },
         })
 
-        docs.push({ doc: result as Record<string, unknown>, index })
+        docs.push(returning ? { doc: result, index } : { id: result.id, index })
       } catch (error) {
         const message = error instanceof Error ? error.message : 'Unknown error'
         const formattedError = formatEntityError({