TUS: Fix stem uploads by adding a new publishStems to SDK

packages/common/src/api/tan-query/upload/usePublishStems.ts

@@ -1,11 +1,19 @@
 import { HashId, Id, type UploadResponse } from '@audius/sdk'
-import { mutationOptions, useMutation } from '@tanstack/react-query'
+import {
+  mutationOptions,
+  useMutation,
+  useQueryClient
+} from '@tanstack/react-query'
 
-import { trackMetadataForUploadToSdk } from '~/adapters'
-import { StemCategory, Name } from '~/models'
+import {
+  StemCategory,
+  Name,
+  type StemUpload,
+  type TrackMetadata
+} from '~/models'
 import { ProgressStatus, uploadActions } from '~/store'
-import type { TrackMetadataForUpload } from '~/store'
 
+import { getStemsQueryKey } from '../tracks/useStems'
 import { useCurrentUserId } from '../users/account/useCurrentUserId'
 import { useQueryContext, type QueryContextType } from '../utils'
 
@@ -21,9 +29,11 @@ type PublishStemsContext = Pick<
 type PublishStemsParams = {
   clientId: string
   parentTrackId: number
-  metadata: TrackMetadataForUpload
-  imageUploadResponse: UploadResponse
-  stemsUploadResponses: UploadResponse[]
+  parentMetadata: Omit<TrackMetadata, 'artwork' | 'track_id'>
+  stems: {
+    metadata: StemUpload
+    audioUploadResponse: UploadResponse
+  }[]
 }
 
 export const publishStems = async (
@@ -43,26 +53,16 @@ export const publishStems = async (
 
   const sdk = await audiusSdk()
   return await Promise.all(
-    (params.metadata.stems ?? []).map(async (stem, index) => {
+    (params.stems ?? []).map(async (stem, index) => {
       try {
-        const stemUploadResponse = params.stemsUploadResponses?.[index]
-        if (!stemUploadResponse) {
-          throw new Error(`No upload response found for stem ${index}`)
-        }
         const metadata = {
-          ...stem.metadata,
-          genre: params.metadata.genre,
-          is_downloadable: true,
-          stem_of: {
-            category: stem.category ?? StemCategory.OTHER,
-            parent_track_id: params.parentTrackId
-          }
+          category: stem.metadata.category ?? StemCategory.OTHER,
+          parentTrackId: Id.parse(params.parentTrackId)
         }
-        const stemRes = await sdk.tracks.publishTrack({
+        const stemRes = await sdk.tracks.publishStem({
           userId: Id.parse(userId),
-          metadata: trackMetadataForUploadToSdk(metadata),
-          audioUploadResponse: stemUploadResponse,
-          imageUploadResponse: params.imageUploadResponse
+          metadata,
+          audioUploadResponse: stem.audioUploadResponse
         })
         dispatch(
           updateProgress({
@@ -77,7 +77,7 @@ export const publishStems = async (
             eventName: Name.STEM_COMPLETE_UPLOAD,
             id: HashId.parse(stemRes.trackId),
             parent_track_id: params.parentTrackId,
-            category: stem.category ?? StemCategory.OTHER
+            category: stem.metadata.category ?? StemCategory.OTHER
           })
         )
         return { trackId: stemRes.trackId, error: null }
@@ -109,10 +109,16 @@ export const usePublishStems = (
   }
 ) => {
   const context = useQueryContext()
+  const queryClient = useQueryClient()
   const { data: userId } = useCurrentUserId()
 
   return useMutation({
     ...options,
-    ...getPublishStemsOptions({ ...context, userId: userId! })
+    ...getPublishStemsOptions({ ...context, userId: userId! }),
+    onSuccess: (_, params) => {
+      queryClient.invalidateQueries({
+        queryKey: getStemsQueryKey(params.parentTrackId)
+      })
+    }
   })
 }

packages/common/src/api/tan-query/upload/usePublishTracks.ts

@@ -136,9 +136,12 @@ export const publishTracks = async (
         publishParentTrack(),
         publishStems(context, {
           clientId: param.clientId,
-          metadata: param.metadata,
-          imageUploadResponse: param.imageUploadResponse,
-          stemsUploadResponses: param.stemsUploadResponses ?? [],
+          parentMetadata: param.metadata,
+          stems:
+            param.metadata.stems?.map((stem, index) => ({
+              metadata: stem,
+              audioUploadResponse: param.stemsUploadResponses?.[index]!
+            })) ?? [],
           parentTrackId: trackId
         })
       ])

packages/discovery-provider/src/tasks/entity_manager/entities/track.py

@@ -287,7 +287,7 @@ def populate_track_record_metadata(track_record: Track, track_metadata, handle,
                 track_record.remix_of = track_metadata["remix_of"]
 
         elif key == "route_id":
-            if "title" in track_metadata:
+            if "title" in track_metadata and track_metadata["title"] is not None:
                 track_record.route_id = helpers.create_track_route_id(
                     track_metadata["title"], handle
                 )

packages/sdk/src/sdk/api/tracks/TracksApi.ts

@@ -58,7 +58,9 @@ import {
   ShareTrackSchema,
   ShareTrackRequest,
   type PublishTrackRequest,
-  PublishTrackSchema
+  PublishTrackSchema,
+  type PublishStemRequest,
+  PublishStemSchema
 } from './types'
 
 // Extend that new class
@@ -231,6 +233,53 @@ export class TracksApi extends GeneratedTracksApi {
     )
   }
 
+  /** @hidden
+   * Publishes a stem that was uploaded using storage node uploadFileV2 uploads.
+   */
+  async publishStem(
+    params: PublishStemRequest,
+    advancedOptions?: AdvancedOptions
+  ) {
+    const {
+      userId,
+      metadata: parsedMetadata,
+      audioUploadResponse
+    } = await parseParams('publishStem', PublishStemSchema)(params)
+
+    const trackMetadata = {
+      title: audioUploadResponse.orig_filename || 'Untitled Stem',
+      isStreamGated: false,
+      streamConditions: undefined,
+      isUnlisted: false,
+      fieldVisibility: {
+        genre: false,
+        mood: false,
+        tags: false,
+        share: false,
+        playCount: false
+      },
+      isDownloadable: true,
+      stemOf: parsedMetadata
+    }
+
+    const metadata = this.trackUploadHelper.transformTrackUploadMetadata(
+      trackMetadata,
+      userId
+    )
+
+    const populatedMetadata =
+      this.trackUploadHelper.populateTrackMetadataWithUploadResponse(
+        metadata,
+        audioUploadResponse
+      )
+
+    return this.writeTrackToChain(
+      params.userId,
+      populatedMetadata,
+      advancedOptions
+    )
+  }
+
   /** @hidden
    * Write track upload to chain
    */

packages/sdk/src/sdk/api/tracks/types.ts

@@ -86,6 +86,13 @@ export const USDCPurchaseConditions = z
   })
   .strict()
 
+export const UploadStemMetadataSchema = z.object({
+  category: z
+    .enum(Object.values(StemCategory) as [StemCategory, ...StemCategory[]])
+    .default(StemCategory.OTHER),
+  parentTrackId: HashId
+})
+
 export const UploadTrackMetadataSchema = z.object({
   trackId: z.optional(HashId),
   aiAttributionUserId: z.optional(HashId),
@@ -152,14 +159,7 @@ export const UploadTrackMetadataSchema = z.object({
       })
       .strict()
   ),
-  stemOf: z.optional(
-    z.object({
-      category: z
-        .enum(Object.values(StemCategory) as [StemCategory, ...StemCategory[]])
-        .default(StemCategory.OTHER),
-      parentTrackId: HashId
-    })
-  ),
+  stemOf: z.optional(UploadStemMetadataSchema.strict()),
   tags: z.optional(z.string()),
   title: z.string({
     required_error: messages.titleRequiredError
@@ -425,9 +425,18 @@ export const PublishTrackSchema = z
     userId: HashId,
     metadata: UploadTrackMetadataSchema.strict(),
     audioUploadResponse: UploadResponseSchema,
-    imageUploadResponse: UploadResponseSchema,
-    stemsUploadResponses: z.array(UploadResponseSchema).optional()
+    imageUploadResponse: UploadResponseSchema
   })
   .strict()
 
 export type PublishTrackRequest = z.input<typeof PublishTrackSchema>
+
+export const PublishStemSchema = z
+  .object({
+    userId: HashId,
+    metadata: UploadStemMetadataSchema.strict(),
+    audioUploadResponse: UploadResponseSchema
+  })
+  .strict()
+
+export type PublishStemRequest = z.input<typeof PublishStemSchema>

packages/web/src/store/application/ui/stemsUpload/sagas.ts

@@ -1,11 +1,14 @@
-import { getStemsQueryKey } from '@audius/common/api'
+import {
+  getStemsQueryKey,
+  queryCurrentUserId,
+  queryTrack
+} from '@audius/common/api'
 import { Name, StemCategory } from '@audius/common/models'
+import { publishStems } from '@audius/common/src/api/tan-query/upload/usePublishStems'
 import { getContext, stemsUploadActions } from '@audius/common/store'
 import { takeEvery, put, call } from 'typed-redux-saga'
 
 import { make } from 'common/store/analytics/actions'
-import { handleUploads } from 'common/store/upload/sagas'
-import { createStemMetadata } from 'pages/upload-page/store/utils/stems'
 
 const { startStemUploads, stemUploadsSucceeded } = stemsUploadActions
 
@@ -14,28 +17,62 @@ function* watchUploadStems() {
     startStemUploads.type,
     function* (action: ReturnType<typeof startStemUploads>) {
       const { uploads, parentId, batchUID } = action.payload
-      const stemTracks = uploads.map((u) => {
-        const metadata = createStemMetadata({
-          parentTrackId: parentId,
-          track: u.metadata,
-          stemCategory: u.category ?? StemCategory.OTHER
+
+      const parentTrack = yield* call(queryTrack, parentId)
+      if (!parentTrack) {
+        throw new Error(`Parent track with ID ${parentId} not found`)
+      }
+
+      const audiusSdk = yield* getContext('audiusSdk')
+      const dispatch = yield* getContext('dispatch')
+      const reportToSentry = yield* getContext('reportToSentry')
+      const analytics = yield* getContext('analytics')
+      const userId = yield* call(queryCurrentUserId)
+      if (!userId) {
+        throw new Error('No user ID found for stem upload')
+      }
+
+      const results = yield* call(async () => {
+        const sdk = await audiusSdk()
+        const uploadHandles = uploads.map((stem, index) => {
+          return sdk.tracks.uploadTrackFiles({
+            audioFile: stem.file
+          })
         })
-        return {
-          ...u,
-          metadata
-        }
-      })
-      const trackIds = yield* call(handleUploads, {
-        tracks: stemTracks,
-        kind: 'stems'
+        const uploadResponses = await Promise.all(
+          uploadHandles.map((handle) => handle.start())
+        )
+        const res = await publishStems(
+          {
+            audiusSdk,
+            dispatch,
+            userId,
+            reportToSentry,
+            analytics
+          },
+          {
+            clientId: batchUID,
+            parentTrackId: parentId,
+            parentMetadata: parentTrack,
+            stems: uploadResponses.map((s, i) => ({
+              metadata: uploads[i],
+              audioUploadResponse: s.audioUploadResponse!
+            }))
+          }
+        )
+        return res
       })
 
       yield* put(stemUploadsSucceeded({ parentId, batchUID }))
 
-      if (trackIds) {
-        for (let i = 0; i < trackIds.length; i += 1) {
-          const trackId = trackIds[i]
-          const category = stemTracks[i].metadata.stem_of?.category
+      if (results) {
+        for (let i = 0; i < results.length; i += 1) {
+          const { trackId, error } = results[i]
+          if (error) {
+            console.error(`Error uploading stem ${i}:`, error)
+            continue
+          }
+          const category = uploads[i].category ?? StemCategory.OTHER
           const recordEvent = make(Name.STEM_COMPLETE_UPLOAD, {
             id: trackId,
             parent_track_id: parentId,