fix

Author: icecrasher321

apps/sim/app/api/mcp/tools/execute/route.test.ts

@@ -151,6 +151,34 @@ describe('MCP tool execution private secret provenance', () => {
     expect(mockExecuteTool.mock.calls[0]?.[5]).toBeUndefined()
   })
 
+  it('preserves MCP error status and message when attaching private provenance', async () => {
+    mockExecuteTool.mockResolvedValueOnce({
+      isError: true,
+      content: [{ type: 'text', text: 'Provider rejected the request' }],
+    })
+    const request = createRequest({
+      'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
+    })
+
+    const response = await POST(request, {})
+    const body = (await response.json()) as Record<string, unknown>
+
+    expect(response.status).toBe(400)
+    expect(response.headers.get('x-sim-private-tool-metadata')).toBe(
+      'resolved-secret-provenance-v1'
+    )
+    expect(body).toMatchObject({
+      success: false,
+      error: 'Provider rejected the request',
+      __resolvedSecretTraceProvenance: {
+        version: 1,
+        complete: true,
+        entries: [],
+        scope: { userId: 'user-1', workspaceId: 'workspace-1' },
+      },
+    })
+  })
+
   it('attaches private provenance without imposing a second functional response limit', async () => {
     const largeText = 'x'.repeat(10 * 1024 * 1024 + 1)
     mockExecuteTool.mockResolvedValueOnce({
@@ -160,7 +188,16 @@ describe('MCP tool execution private secret provenance', () => {
       'x-sim-request-private-tool-metadata': 'resolved-secret-provenance-v1',
     })
 
-    const response = await POST(request, {})
+    const response = await (async () => {
+      const responseJsonSpy = vi.spyOn(Response.prototype, 'json')
+      try {
+        const result = await POST(request, {})
+        expect(responseJsonSpy).not.toHaveBeenCalled()
+        return result
+      } finally {
+        responseJsonSpy.mockRestore()
+      }
+    })()
     const body = (await response.json()) as Record<string, unknown>
 
     expect(response.status).toBe(200)

apps/sim/app/api/mcp/tools/execute/route.ts

@@ -15,10 +15,9 @@ import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { SIM_VIA_HEADER } from '@/lib/execution/call-chain'
 import { parseRemainingExecutionDeadlineMs } from '@/lib/execution/execution-deadline-header'
 import {
-  PRIVATE_TOOL_METADATA_RESPONSE_HEADER,
-  RESOLVED_SECRET_PROVENANCE_FIELD,
   RESOLVED_SECRET_PROVENANCE_METADATA_V1,
   requestsPrivateToolMetadata,
+  serializePrivateToolMetadataResponseEnvelope,
 } from '@/lib/execution/private-tool-metadata'
 import {
   mcpBodyReadErrorResponse,
@@ -33,7 +32,7 @@ import {
   type McpToolCall,
   type McpToolResult,
 } from '@/lib/mcp/types'
-import { categorizeError, createMcpErrorResponse, createMcpSuccessResponse } from '@/lib/mcp/utils'
+import { categorizeError } from '@/lib/mcp/utils'
 import {
   assertPermissionsAllowed,
   McpToolsNotAllowedError,
@@ -66,23 +65,21 @@ function hasType(prop: unknown): prop is SchemaProperty {
   return typeof prop === 'object' && prop !== null && 'type' in prop
 }
 
-async function attachPrivateProvenance(
-  response: NextResponse,
-  provenance: ResolvedSecretTraceProvenanceAccumulator
-): Promise<NextResponse> {
-  const parsed: unknown = await response.json()
-  if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
-    throw new Error('MCP response is not a JSON object')
+function createToolExecutionResponse(
+  body: Record<string, unknown>,
+  status: number,
+  provenance: ResolvedSecretTraceProvenanceAccumulator | undefined
+): NextResponse {
+  if (!provenance) {
+    return NextResponse.json(body, { status })
   }
-  const payload = parsed as Record<string, unknown>
-
-  const headers = new Headers(response.headers)
-  headers.delete('content-length')
-  headers.set(PRIVATE_TOOL_METADATA_RESPONSE_HEADER, RESOLVED_SECRET_PROVENANCE_METADATA_V1)
-  return NextResponse.json(
-    { ...payload, [RESOLVED_SECRET_PROVENANCE_FIELD]: provenance.exportProvenance() },
-    { status: response.status, headers }
+
+  const envelope = serializePrivateToolMetadataResponseEnvelope(
+    body,
+    RESOLVED_SECRET_PROVENANCE_METADATA_V1,
+    provenance.exportProvenance()
   )
+  return NextResponse.json(envelope.body, { status, headers: envelope.headers })
 }
 
 /**
@@ -103,13 +100,22 @@ export const POST = withRouteHandler(
             resolvedSecretTraceProvenance.record(provenance)
           }
         : undefined
-      const response = await (async (): Promise<NextResponse> => {
+      const errorResponse = (message: string, status: number): NextResponse =>
+        createToolExecutionResponse(
+          { success: false, error: message },
+          status,
+          resolvedSecretTraceProvenance
+        )
+      const successResponse = <T>(data: T, status = 200): NextResponse =>
+        createToolExecutionResponse({ success: true, data }, status, resolvedSecretTraceProvenance)
+
+      return (async (): Promise<NextResponse> => {
         try {
           const rawBody = await readMcpJsonBodyWithLimit(request)
           const parsedBody = mcpToolExecutionBodySchema.safeParse(rawBody)
 
           if (!parsedBody.success) {
-            return createMcpErrorResponse(parsedBody.error, 'Invalid request format', 400)
+            return errorResponse('Invalid request format', 400)
           }
 
           const body = parsedBody.data
@@ -136,7 +142,7 @@ export const POST = withRouteHandler(
             })
           } catch (err) {
             if (err instanceof McpToolsNotAllowedError) {
-              return createMcpErrorResponse(err, err.message, 403)
+              return errorResponse(err.message, 403)
             }
             throw err
           }
@@ -160,11 +166,7 @@ export const POST = withRouteHandler(
               logger.warn(`[${requestId}] Tool ${toolName} not found on server ${serverId}`, {
                 availableTools: tools.map((t) => t.name),
               })
-              return createMcpErrorResponse(
-                new Error('Tool not found'),
-                'Tool not found on the specified server',
-                404
-              )
+              return errorResponse('Tool not found on the specified server', 404)
             }
 
             if (tool.inputSchema?.properties) {
@@ -229,11 +231,7 @@ export const POST = withRouteHandler(
             const validationError = validateToolArguments(tool, args)
             if (validationError) {
               logger.warn(`[${requestId}] Tool validation failed: ${validationError}`)
-              return createMcpErrorResponse(
-                new Error(`Invalid arguments for tool ${toolName}: ${validationError}`),
-                'Invalid tool arguments',
-                400
-              )
+              return errorResponse('Invalid tool arguments', 400)
             }
           }
 
@@ -305,11 +303,7 @@ export const POST = withRouteHandler(
             logger.warn(
               `[${requestId}] Tool execution returned error for ${toolName} on ${serverId}`
             )
-            return createMcpErrorResponse(
-              transformedResult,
-              transformedResult.error || 'Tool execution failed',
-              400
-            )
+            return errorResponse(transformedResult.error || 'Tool execution failed', 400)
           }
           logger.info(`[${requestId}] Successfully executed tool ${toolName} on server ${serverId}`)
 
@@ -330,7 +324,7 @@ export const POST = withRouteHandler(
             })
           }
 
-          return createMcpSuccessResponse(transformedResult)
+          return successResponse(transformedResult)
         } catch (error) {
           if (getErrorMessage(error) === 'Tool execution timeout') {
             resolvedSecretTraceProvenance?.markIncomplete()
@@ -347,27 +341,24 @@ export const POST = withRouteHandler(
             logger.warn(`[${requestId}] OAuth re-authorization required for MCP tool execution`, {
               serverId: errorServerId,
             })
-            return NextResponse.json(
+            return createToolExecutionResponse(
               {
                 success: false,
                 error: 'OAuth re-authorization required',
                 code: 'reauth_required',
                 serverId: errorServerId,
               },
-              { status: 401 }
+              401,
+              resolvedSecretTraceProvenance
             )
           }
 
           logger.error(`[${requestId}] Error executing MCP tool:`, error)
 
           const { message, status } = categorizeError(error)
-          return createMcpErrorResponse(new Error(message), message, status)
+          return errorResponse(message, status)
         }
       })()
-
-      return resolvedSecretTraceProvenance
-        ? attachPrivateProvenance(response, resolvedSecretTraceProvenance)
-        : response
     }
   )
 )

apps/sim/ee/workspace-forking/lib/copy/copy-resources.test.ts

@@ -281,6 +281,14 @@ describe('copyForkResourceContent', () => {
       workspaceId: 'child-ws',
       originalName: 'report.pdf',
     })
+    expect(mockRecordKnowledgeBaseFileOwnership).toHaveBeenNthCalledWith(1, {
+      key: uploadArg.customKey,
+      userId: 'user-1',
+      workspaceId: 'child-ws',
+      originalName: 'report.pdf',
+      contentType: 'application/pdf',
+      size: 321,
+    })
     expect(mockRecordKnowledgeBaseFileOwnership).toHaveBeenCalledWith(
       {
         key: uploadArg.customKey,
@@ -292,6 +300,9 @@ describe('copyForkResourceContent', () => {
       },
       expect.anything()
     )
+    expect(mockRecordKnowledgeBaseFileOwnership.mock.invocationCallOrder[0]).toBeLessThan(
+      storageServiceMockFns.mockUploadFile.mock.invocationCallOrder[0]
+    )
     expect(mockRecordKnowledgeBaseFileOwnership.mock.invocationCallOrder[0]).toBeLessThan(
       mockIncrementStorageUsageInTx.mock.invocationCallOrder[0]
     )
@@ -494,6 +505,42 @@ describe('copyForkResourceContent', () => {
     )
   })
 
+  it('leaves a discoverable ownership reservation when a copied KB upload fails', async () => {
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([sourceDoc])
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([sourceDoc])
+    storageServiceMockFns.mockUploadFile.mockRejectedValueOnce(new Error('upload failed'))
+
+    const result = await copyForkResourceContent({
+      contentPlan: basePlan({
+        knowledgeBases: [
+          {
+            sourceId: 'src-kb',
+            childId: 'child-kb',
+            documentIdMap: { 'doc-1': 'child-doc-1' },
+          },
+        ],
+      }),
+      requestId: 'test',
+    })
+
+    const targetKey = `kb/fork-child-doc-1-${sha256Hex(Buffer.from('blob-bytes'))}`
+    expect(result.failed).toBe(1)
+    expect(mockRecordKnowledgeBaseFileOwnership).toHaveBeenCalledWith({
+      key: targetKey,
+      userId: 'user-1',
+      workspaceId: 'child-ws',
+      originalName: 'report.pdf',
+      contentType: 'application/pdf',
+      size: 321,
+    })
+    expect(mockRecordKnowledgeBaseFileOwnership.mock.invocationCallOrder[0]).toBeLessThan(
+      storageServiceMockFns.mockUploadFile.mock.invocationCallOrder[0]
+    )
+    expect(storageServiceMockFns.mockDeleteFile).not.toHaveBeenCalled()
+  })
+
   it('does not resolve KB billing context for an empty document page', async () => {
     const result = await copyForkResourceContent({
       contentPlan: basePlan({

apps/sim/ee/workspace-forking/lib/copy/copy-resources.ts

@@ -1796,10 +1796,12 @@ async function copyDocumentEmbeddings(
  * copy fails. A stored source blob is required to copy successfully; callers keep the target
  * placeholder archived and report the existing resource failure. The content digest in the key
  * makes reuse safe for identical retries and prevents a later source snapshot from adopting or
- * overwriting bytes left by an earlier failed attempt.
+ * overwriting bytes left by an earlier failed attempt. Ownership is recorded before storage I/O,
+ * matching the presigned-upload lifecycle: successful finalization reuses the immutable binding,
+ * while the existing orphan-binding sweep eventually reclaims an abandoned object or reservation.
  */
 async function copyKbDocumentBlob(
-  doc: { storageKey: string | null; filename: string; mimeType: string },
+  doc: { storageKey: string | null; filename: string; mimeType: string; fileSize: number },
   childWorkspaceId: string,
   userId: string,
   childDocumentId: string
@@ -1811,6 +1813,14 @@ async function copyKbDocumentBlob(
     maxBytes: MAX_FILE_SIZE,
   })
   const targetKey = deriveKbDocumentStorageKey(childDocumentId, sha256Hex(buffer))
+  await recordKnowledgeBaseFileOwnership({
+    key: targetKey,
+    userId,
+    workspaceId: childWorkspaceId,
+    originalName: doc.filename,
+    contentType: doc.mimeType,
+    size: doc.fileSize,
+  })
   const existing = await headObject(targetKey, 'knowledge-base')
   if (!existing) {
     await uploadFile({