Improve music activity reporting and preserve Comfy host fixes

tests/music-status.test.mjs

@@ -0,0 +1,55 @@
+import test from 'node:test'
+import assert from 'node:assert/strict'
+import { readFileSync } from 'node:fs'
+import ts from 'typescript'
+
+const page = readFileSync(new URL('../pages/music.vue', import.meta.url), 'utf8')
+const source = page.slice(page.indexOf('async function tryRecover()'), page.indexOf('async function attachLiveMusic('))
+const code = ts.transpileModule(source, { compilerOptions: { target: ts.ScriptTarget.ES2022 } }).outputText
+function recovery(fetch, apply, id = { value: 'music-1' }) {
+  return new Function('$fetch', 'applyEvent', 'jobId', 'busy', 'queued', 'settled', `${code}; return tryRecover`)(
+    fetch, apply, id, { value: true }, { value: false }, false
+  )
+}
+
+test('music status recovers without a working event stream, including completion and errors', async () => {
+  for (const status of ['running', 'complete', 'error']) {
+    const snapshot = { status, message: 'YuE Stage A', progress: 20 }
+    const events = []
+    await recovery(async url => {
+      assert.equal(url, '/api/generate/music-1')
+      return snapshot
+    }, event => events.push(event))()
+    assert.deepEqual(events, [snapshot])
+  }
+})
+
+test('late status response cannot overwrite a different music job', async () => {
+  const id = { value: 'music-1' }
+  const events = []
+  await recovery(async () => {
+    id.value = 'music-2'
+    return { status: 'complete' }
+  }, event => events.push(event), id)()
+  assert.deepEqual(events, [])
+})
+
+test('missing live job still falls back to saved audio recovery', async () => {
+  const events = []
+  await recovery(async url => {
+    if (url === '/api/generate/music-1') throw new Error('Job unavailable')
+    assert.equal(url, '/api/generate/recover')
+    return { trackId: 'saved-track' }
+  }, event => events.push(event))()
+  assert.deepEqual(events, [{ trackId: 'saved-track' }])
+})
+
+test('inspecting a queued YuE job uses music controls', () => {
+  const inspect = readFileSync(new URL('../server/utils/queuedInspect.ts', import.meta.url), 'utf8')
+  const fn = inspect.slice(inspect.indexOf('export function inspectStudioJob'), inspect.indexOf('export function inspectLiveMusicJob')).replace('export function', 'function')
+  const js = ts.transpileModule(fn, { compilerOptions: { target: ts.ScriptTarget.ES2022 } }).outputText
+  const result = new Function(`${js}; return inspectStudioJob({kind:'music',payload:{prompt:'rock',musicEngine:'yue'}})`)()
+  assert.equal(result.kind, 'music')
+  assert.equal(result.payload.musicEngine, 'yue')
+  assert.equal(result.payload.tags, 'rock')
+})

tests/test_yue_cancellation.py

@@ -0,0 +1,35 @@
+import ast
+import importlib.util
+import pathlib
+import sys
+import types
+import unittest
+
+spec = importlib.util.spec_from_file_location('patcher', pathlib.Path(__file__).parents[1] / 'scripts/patch_yue_cancellation.py')
+patcher = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(patcher)
+
+class CancellationTests(unittest.TestCase):
+    def test_interrupt_propagates_before_sampling(self):
+        source = 'class BlockTokenRangeProcessor:\n    def __call__(self, input_ids, scores):\n        return scores\n'
+        result = patcher.patch(source)
+        self.assertEqual(result, patcher.patch(result))
+        cancelled = False
+        class Interrupted(Exception):
+            pass
+        def check():
+            if cancelled:
+                raise Interrupted()
+        module = types.ModuleType('comfy.model_management')
+        module.throw_exception_if_processing_interrupted = check
+        sys.modules['comfy.model_management'] = module
+        namespace = {}
+        exec(compile(ast.parse(result), '<patched>', 'exec'), namespace)
+        callback = namespace['BlockTokenRangeProcessor']()
+        self.assertEqual(callback(None, 42), 42)
+        cancelled = True
+        with self.assertRaises(Interrupted):
+            callback(None, 42)
+
+if __name__ == '__main__':
+    unittest.main()