From 811154d3d7e1823286a4a91285a26de298b238ab Mon Sep 17 00:00:00 2001 From: Aryaveer Date: Fri, 21 Aug 2026 20:21:11 +0530 Subject: [PATCH] feat(music): integrate Flow Music & Google Labs MusicFX generation across CLI, API, MCP, and Chrome extension - Chrome Extension: add host permissions and token capture for flowmusic.app, plus SSE stream listener for clip delivery - Flow Engine (t2m.py): implement T2M generator with primary Flow Music pipeline and automatic Google Labs MusicFX fallback - HTTP API: add POST /v1/audio/generations and /v1/music/generations with Idempotency-Key support and history registry - MCP Server: add generate_flow_music tool with duration, loop, count, and seed parameters - Unified CLI: add 'flow music' command with duration, looping, variation counts, and exact output paths - Media Sniffing: extend signature detection for MP3, WAV, M4A, FLAC, OGG, and AAC formats - Tests & Docs: add 13 unit tests for music generator/API/sniffing and update README documentation --- .gitignore | 6 + README.md | 15 +- flow-agent/flow_engine/__init__.py | 17 +- flow-agent/flow_engine/bridge.py | 69 ++- flow-agent/flow_engine/config.py | 10 + flow-agent/flow_engine/generators/__init__.py | 3 + flow-agent/flow_engine/generators/t2m.py | 416 +++++++++++++++ flow-agent/flow_engine/generators/t2v.py | 1 - flow-agent/flow_engine/generators/v2v.py | 1 - flow-agent/flow_server/batch.py | 1 - flow-agent/flow_server/mcp_server.py | 104 ++++ flow-agent/flow_server/media_history.py | 6 +- flow-agent/flow_server/media_types.py | 21 + flow-agent/flow_server/models.py | 11 + flow-agent/flow_server/routes/generation.py | 137 ++++- flow-agent/main.py | 48 +- flow-agent/tests/test_music_generation.py | 319 +++++++++++ flow-agent/tests/test_recaptcha_self_heal.py | 50 ++ flow-agent/uv.lock | 2 +- flow-extension/background.js | 501 +++++++++++------- flow-extension/manifest.json | 12 +- 21 files changed, 1547 insertions(+), 203 deletions(-) create mode 100644 flow-agent/flow_engine/generators/t2m.py create mode 100644 flow-agent/tests/test_music_generation.py create mode 100644 flow-agent/tests/test_recaptcha_self_heal.py diff --git a/.gitignore b/.gitignore index 3218b61..30384ab 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,12 @@ input/ *.png *.jpg *.jpeg +*.mp3 +*.m4a +*.wav +*.ogg +*.flac +*.aac !flow-extension/icon*.png media-id.js diff --git a/README.md b/README.md index 140f1c7..c4e165e 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # Flow Agent CLI, OpenAI-compatible API, Chrome extension bridge, and MCP server for Google -Flow image and video generation. +Flow image, video, and music generation. [![Release](https://img.shields.io/github/v/release/kodelyx/flow-agent)](https://github.com/kodelyx/flow-agent/releases) [![Build](https://github.com/kodelyx/flow-agent/actions/workflows/build.yml/badge.svg)](https://github.com/kodelyx/flow-agent/actions/workflows/build.yml) @@ -13,9 +13,10 @@ HTTP API, and MCP clients all share one backend and one extension bridge. - Text-to-image and reference-image generation - Text-to-video, image-to-video, first/last-frame, reference-to-video, and video editing +- Text-to-music and soundtrack generation with custom duration, looping, and variations - 4, 6, 8, and 10-second video generation - Reusable generated and uploaded media IDs, including after backend restarts -- Exact `--output` paths with real PNG, JPEG, and WebP conversion +- Exact `--output` paths with real PNG, JPEG, WebP, and audio format conversion - Signature-based MIME and extension detection - Persistent idempotency for safe paid-generation retries - OpenAI-compatible HTTP endpoints and pollable video jobs @@ -133,6 +134,14 @@ flow video "animate this generated image" --start GENERATED_IMAGE_MEDIA_ID flow video "use these existing references" --ref GENERATED_ID UPLOADED_ID ``` +### Music + +```bash +flow music "a cinematic ambient synth soundtrack" --duration 30 +flow music "relaxing lo-fi hip hop beat" --loop --output bgm.mp3 +flow music "epic orchestral intro" --duration 50 --count 2 --seed 42 -o ./soundtrack.mp3 +``` + ### Upload and edit ```bash @@ -283,6 +292,7 @@ JSON-RPC messages are sent to `http://127.0.0.1:8001/messages`. - `get_flow_history` — generated and uploaded media history - `generate_flow_image` — text/reference image generation - `generate_flow_video` — text, start-image, and reference video generation +- `generate_flow_music` — text-to-music and soundtrack generation with loops and duration controls - `upload_flow_media` — upload a local path, URL, or base64 media payload - `download_media_from_url` — download media and optionally upload it to Flow - `edit_flow_video` — edit a video by media ID or local video path @@ -300,6 +310,7 @@ Default base URL: `http://127.0.0.1:8001` | `POST /v1/images/generations` | Generate images | | `POST /v1/videos/generations` | Submit video generation | | `GET /v1/videos/generations/{job_id}` | Poll a video job | +| `POST /v1/audio/generations` | Generate music / audio tracks | | `POST /v1/upload` | Upload an image or video reference | | `GET /download/{filename}` | Download a managed media file | | `GET /sse` | MCP over SSE | diff --git a/flow-agent/flow_engine/__init__.py b/flow-agent/flow_engine/__init__.py index 30efdf1..5ae8b92 100644 --- a/flow-agent/flow_engine/__init__.py +++ b/flow-agent/flow_engine/__init__.py @@ -30,12 +30,23 @@ # ─── Public API ────────────────────────────────────────────── from .bridge import ExtensionBridge -from .config import ASPECTS, DEFAULT_PROJECT, ENDPOINTS, CLIENT_CTX, API_KEY, API_BASE +from .config import ( + ASPECTS, + DEFAULT_PROJECT, + ENDPOINTS, + CLIENT_CTX, + API_KEY, + API_BASE, + DEFAULT_MUSIC_DURATION, + MUSIC_DURATIONS, +) from .generators import ( generate_video, edit_video, upload_image, generate_video_i2v, + generate_music, + download_music, poll_status, download_video, build_client_context, @@ -51,6 +62,8 @@ "edit_video", "upload_image", "generate_video_i2v", + "generate_music", + "download_music", "poll_status", "download_video", "build_client_context", @@ -63,6 +76,8 @@ "CLIENT_CTX", "API_KEY", "API_BASE", + "DEFAULT_MUSIC_DURATION", + "MUSIC_DURATIONS", # Media store "media_store", ] diff --git a/flow-agent/flow_engine/bridge.py b/flow-agent/flow_engine/bridge.py index 58b60e5..6e19e55 100644 --- a/flow-agent/flow_engine/bridge.py +++ b/flow-agent/flow_engine/bridge.py @@ -358,12 +358,13 @@ async def _request_flow_tab_for(self, client_id): async def _force_refresh_client(self, client_id): """Force a client to reload its Flow tab and re-capture a fresh token, bypassing the extension's 50-min freshness cache. Used to self-heal a - stale token (Google invalidates via inactivity before the cache expires). + stale token (Google invalidates via inactivity before the cache expires) + or a stale reCAPTCHA session. """ - if client_id not in self._clients: + if client_id not in self._clients and not self.http_registry.is_connected(client_id) and client_id != "__http__": return try: - log.info("Force-refreshing Flow tab for client %s (stale token self-heal)...", client_id) + log.info("Force-refreshing Flow tab for client %s (stale token/captcha self-heal)...", client_id) self._tokens.pop(client_id, None) await self.send_message_to(client_id, {"method": "force_refresh", "force": True}) await asyncio.sleep(6) @@ -583,6 +584,19 @@ async def api_request(self, url_path, body, captcha_action="VIDEO_GENERATION", m if fallback: log.warning("Client %s still 401 after refresh — failing over to %s", client_id, fallback) result = await self._run_api_request(fallback, url_path, body, captcha_action, method, timeout, meta) + + # Self-heal on reCAPTCHA / UNUSUAL_ACTIVITY: force-reload the Flow tab so that + # the enterprise reCAPTCHA session re-initialises and issues a fresh valid token. + elif self._is_recaptcha_failure(result): + log.warning( + "Client %s returned reCAPTCHA / UNUSUAL_ACTIVITY — " + "force-reloading Flow tab and retrying once", + client_id, + ) + await self._force_refresh_client(client_id) + result = await self._run_api_request( + client_id, url_path, body, captcha_action, method, timeout, meta + ) return result @staticmethod @@ -599,6 +613,30 @@ def _is_unauthenticated(result) -> bool: return True return False + @staticmethod + def _is_recaptcha_failure(result) -> bool: + """True if Google or extension returned a reCAPTCHA / UNUSUAL_ACTIVITY failure.""" + if not isinstance(result, dict): + return False + err_msg = str(result.get("error") or "") + if "CAPTCHA" in err_msg.upper() or "UNUSUAL_ACTIVITY" in err_msg: + return True + status = result.get("status") + if status in (400, 403): + data = result.get("data") + text = "" + if isinstance(data, str): + text = data + elif isinstance(data, dict): + err = data.get("error", {}) + if isinstance(err, dict): + text = f"{err.get('message', '')} {err.get('status', '')} {err.get('details', '')}" + else: + text = str(err) + if "UNUSUAL_ACTIVITY" in text or "recaptcha" in text.lower(): + return True + return False + def _select_client_excluding(self, exclude_id) -> str | None: """Pick another connected client with a token, skipping exclude_id.""" candidates = sorted([c for c in self._clients if c != exclude_id and self._tokens.get(c)]) @@ -649,7 +687,20 @@ async def _do_api_request(self, client_id, url_path, body, captcha_action, metho **(meta or {}), }) - url = f"{API_BASE}{url_path}?key={API_KEY}" + if url_path.startswith("http://") or url_path.startswith("https://"): + url = url_path + is_flowmusic = "flowmusic.app" in url + origin = "https://www.flowmusic.app" if is_flowmusic else CLIENT_CTX["origin"] + referer = "https://www.flowmusic.app/session" if is_flowmusic else CLIENT_CTX["origin"] + "/" + content_type = "application/json" if is_flowmusic else "text/plain;charset=UTF-8" + site = "same-origin" if is_flowmusic else "cross-site" + else: + url = f"{API_BASE}{url_path}?key={API_KEY}" + origin = CLIENT_CTX["origin"] + referer = CLIENT_CTX["origin"] + "/" + content_type = "text/plain;charset=UTF-8" + site = "cross-site" + ua = random.choice(USER_AGENTS) platform = '"macOS"' if "Macintosh" in ua else '"Windows"' @@ -675,18 +726,18 @@ async def _do_api_request(self, client_id, url_path, body, captcha_action, metho "method": method, "headers": { "accept": "*/*", - "content-type": "text/plain;charset=UTF-8", - "origin": CLIENT_CTX["origin"], - "referer": CLIENT_CTX["origin"] + "/", + "content-type": content_type, + "origin": origin, + "referer": referer, "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": platform, "sec-fetch-dest": "empty", "sec-fetch-mode": "cors", - "sec-fetch-site": "cross-site", + "sec-fetch-site": site, "user-agent": ua, }, "body": body, - "captchaAction": captcha_action, + "captchaAction": captcha_action if not ("flowmusic.app" in url) else None, }, } diff --git a/flow-agent/flow_engine/config.py b/flow-agent/flow_engine/config.py index d07997f..b68f658 100644 --- a/flow-agent/flow_engine/config.py +++ b/flow-agent/flow_engine/config.py @@ -116,6 +116,13 @@ def _flow_binary_dir() -> str: "generate_fl": "/v1/video:batchAsyncGenerateVideoStartAndEndImage", "generate_r2v": "/v1/video:batchAsyncGenerateVideoReferenceImages", "generate_edit": "/v1/video:batchAsyncGenerateVideoEditVideo", + "flowmusic_conversation": "https://www.flowmusic.app/__api/conversation", + "flowmusic_stream": "https://www.flowmusic.app/__api/messages/{job_id}/stream?last_id=0", + "flowmusic_clips": "https://www.flowmusic.app/__api/clips", + "generate_music": "/v1:runMusicFx", + "generate_music_demo": "/v1:runSoundDemo", + "generate_music_sound": "/v1/sound:generate", + "generate_music_batch": "/v1/music:batchGenerateMusic", "upload_image": "/v1/flow/uploadImage", "poll_status": "/v1/video:batchCheckAsyncVideoGenerationStatus", "get_media": "/v1/media/{media_id}", @@ -136,6 +143,9 @@ def _flow_binary_dir() -> str: DEFAULT_DURATION = 10 MAX_COUNT = 4 +DEFAULT_MUSIC_DURATION = 30 +MUSIC_DURATIONS = [10, 30, 50, 70] + CREDITS_PER_VIDEO = { 4: 7, 6: 10, diff --git a/flow-agent/flow_engine/generators/__init__.py b/flow-agent/flow_engine/generators/__init__.py index b581362..c879d93 100644 --- a/flow-agent/flow_engine/generators/__init__.py +++ b/flow-agent/flow_engine/generators/__init__.py @@ -4,6 +4,7 @@ from .v2v import edit_video from .i2v import upload_image, generate_video_i2v from .t2i import generate_image, download_image, IMAGE_ASPECTS +from .t2m import generate_music, download_music from .common import poll_status, download_video, build_client_context __all__ = [ @@ -13,6 +14,8 @@ "generate_video_i2v", "generate_image", "download_image", + "generate_music", + "download_music", "IMAGE_ASPECTS", "poll_status", "download_video", diff --git a/flow-agent/flow_engine/generators/t2m.py b/flow-agent/flow_engine/generators/t2m.py new file mode 100644 index 0000000..71047fd --- /dev/null +++ b/flow-agent/flow_engine/generators/t2m.py @@ -0,0 +1,416 @@ +"""Flow Engine — Text to Music (T2M) generator for Google Labs / MusicFX. + +Handles prompt formatting, audio duration, looping, multi-endpoint fallback, +and media downloading/decoding. +""" + +from __future__ import annotations + +import base64 +import logging +import os +import re +import time +import urllib.request +import uuid +from typing import Any + +from ..config import ( + CLIENT_CTX, + DEFAULT_MUSIC_DURATION, + DEFAULT_PROJECT, + ENDPOINTS, +) +from .common import resolve_seed +from flow_server.media_types import ensure_correct_extension, sniff_media_type + +log = logging.getLogger("flow_engine.generators.t2m") + +UUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") + + +def build_music_client_context(project_id: str, tool: str = "MUSIC_FX") -> dict[str, Any]: + """Build client context tailored for MusicFX / Google Labs audio requests.""" + return { + "projectId": project_id or DEFAULT_PROJECT, + "tool": tool, + "userPaygateTier": CLIENT_CTX.get("tier", "PAYGATE_TIER_ONE"), + "sessionId": f";{int(time.time() * 1000)}", + "recaptchaContext": { + "applicationType": CLIENT_CTX.get("recaptcha_app_type", "RECAPTCHA_APPLICATION_TYPE_WEB"), + "token": "", + }, + } + + +def _parse_music_results(data: dict[str, Any], prompt: str) -> list[dict[str, Any]]: + """Parse generated audio tracks from varied Google Labs / MusicFX response formats. + + Returns a list of dicts with keys: + media_id: str + audio_url: str + audio_base64: str + revised_prompt: str + """ + results: list[dict[str, Any]] = [] + + # Format 1: "sounds" array (MusicFX / SoundDemo format) + sounds = data.get("sounds", []) + if isinstance(sounds, list) and sounds: + for item in sounds: + if not isinstance(item, dict): + continue + name = item.get("name", "") + audio_b64 = ( + item.get("audio") + or item.get("rawBytes") + or item.get("audioContent") + or item.get("bytes", "") + ) + url = ( + item.get("downloadUrl") + or item.get("fifeUrl") + or item.get("audioUri") + or item.get("url", "") + ) + revised = item.get("revisedPrompt", prompt) + media_id = name if UUID_RE.match(name) else "" + if not media_id and url: + match = UUID_RE.search(url) + if match: + media_id = match.group() + results.append({ + "media_id": media_id or name, + "audio_url": url, + "audio_base64": audio_b64, + "revised_prompt": revised, + }) + + # Format 2.5: "predictions" array + predictions = data.get("predictions", []) + if isinstance(predictions, list) and predictions: + for item in predictions: + if not isinstance(item, dict): + continue + audio_b64 = ( + item.get("bytesBase64Encoded") + or item.get("audio") + or item.get("rawBytes", "") + ) + url = item.get("downloadUrl") or item.get("url", "") + name = item.get("name", "") + results.append({ + "media_id": name or str(uuid.uuid4()), + "audio_url": url, + "audio_base64": audio_b64, + "revised_prompt": prompt, + }) + + # Format 2: "media" array (Pinhole / Flow style) + media_list = data.get("media", []) + if isinstance(media_list, list) and media_list: + for item in media_list: + if not isinstance(item, dict): + continue + name = item.get("name", "") + audio_obj = item.get("audio", {}) + gen = audio_obj.get("generatedAudio", {}) if isinstance(audio_obj, dict) else {} + url = ( + gen.get("downloadUrl") + or gen.get("fifeUrl") + or gen.get("audioUri") + or item.get("downloadUrl", "") + ) + audio_b64 = ( + gen.get("rawBytes") + or gen.get("audioContent") + or item.get("rawBytes", "") + ) + media_id = name if UUID_RE.match(name) else "" + if not media_id and url: + match = UUID_RE.search(url) + if match: + media_id = match.group() + results.append({ + "media_id": media_id or name, + "audio_url": url, + "audio_base64": audio_b64, + "revised_prompt": prompt, + }) + + # Format 3: Direct single audio response + if not results: + direct_b64 = ( + data.get("audioContent") + or data.get("rawBytes") + or data.get("audio", "") + ) + direct_url = data.get("downloadUrl") or data.get("audioUri", "") + if direct_b64 or direct_url: + results.append({ + "media_id": data.get("name", ""), + "audio_url": direct_url, + "audio_base64": direct_b64 if isinstance(direct_b64, str) else "", + "revised_prompt": prompt, + }) + + return results + + +async def _generate_flowmusic( + bridge, + prompt: str, + model: str | None = None, +) -> list[dict[str, Any]] | None: + """Generate music using Flow Music (flowmusic.app).""" + conv_url = ENDPOINTS.get("flowmusic_conversation", "https://www.flowmusic.app/__api/conversation") + payload = { + "parts": [{"content": prompt, "part_kind": "user-prompt"}], + "client_context": { + "song_queue": [], + "selected_model": None, + "lyrics_id_map": {}, + "ghostwriter_version": "standard", + }, + "model_name": model or "producer:standard", + "mode": "standard", + } + + log.info('Submitting prompt to Flow Music (flowmusic.app): "%s"', prompt[:60]) + resp = await bridge.api_request(conv_url, payload, method="POST") + status = resp.get("status", 0) + data = resp.get("data", {}) + + if status not in (200, 201) or not isinstance(data, dict): + err = data if data else resp.get("error", f"HTTP {status}") + log.error("Flow Music conversation endpoint failed: status=%s data=%s", status, err) + raise ValueError(f"Flow Music conversation creation failed (HTTP {status}): {err}") + + # Check if direct results were returned in data + direct_parsed = _parse_music_results(data, prompt) + if direct_parsed: + return direct_parsed + + job_id = data.get("job_id") + if not job_id: + log.error("Flow Music response contained no job_id: %s", data) + raise ValueError(f"Flow Music response missing job_id: {data}") + + log.info("Flow Music job created: %s, streaming events...", job_id) + stream_url = ENDPOINTS.get( + "flowmusic_stream", + "https://www.flowmusic.app/__api/messages/{job_id}/stream?last_id=0" + ).format(job_id=job_id) + + # Await event stream resolution via extension + stream_resp = await bridge.api_request(stream_url, {}, method="GET", timeout=120) + stream_data = stream_resp.get("data", {}) + clip_ids = [] + if isinstance(stream_data, dict): + clip_ids = stream_data.get("clip_ids", []) + + if not clip_ids: + log.error("Flow Music stream returned no clip_ids: %s", stream_resp) + raise ValueError(f"Flow Music stream completed without clip_ids: {stream_resp}") + + log.info("Flow Music produced %d clip(s): %s. Fetching details...", len(clip_ids), clip_ids) + clips_url = ENDPOINTS.get("flowmusic_clips", "https://www.flowmusic.app/__api/clips") + clips_resp = await bridge.api_request(clips_url, {"clip_ids": clip_ids}, method="POST") + + if clips_resp.get("status") != 200: + err = clips_resp.get("data") or clips_resp.get("error", f"HTTP {clips_resp.get('status')}") + log.error("Flow Music clips endpoint failed (HTTP %s): %s", clips_resp.get("status"), err) + raise ValueError(f"Flow Music clips endpoint failed (HTTP {clips_resp.get('status')}): {err}") + + clips_data = clips_resp.get("data", {}) + clips_dict = clips_data.get("clips", {}) if isinstance(clips_data, dict) else {} + + results = [] + for cid, clip in clips_dict.items(): + if not isinstance(clip, dict): + continue + audio_url = clip.get("audio_url") or clip.get("wav_url", "") + wav_url = clip.get("wav_url") or clip.get("audio_url", "") + title = clip.get("title", "") + duration_obj = clip.get("duration", {}) + duration_sec = 0.0 + if isinstance(duration_obj, dict) and duration_obj.get("value"): + try: + duration_sec = float(duration_obj["value"]) + except (ValueError, TypeError): + duration_sec = 0.0 + + results.append({ + "media_id": cid, + "audio_url": audio_url, + "wav_url": wav_url, + "title": title, + "duration": duration_sec, + "audio_base64": "", + "revised_prompt": prompt, + }) + + if results: + log.info("Flow Music generated %d tracks successfully!", len(results)) + return results if results else None + + +async def generate_music( + bridge, + prompt: str, + project_id: str | None = None, + duration: int = DEFAULT_MUSIC_DURATION, + count: int = 1, + loop: bool = False, + seed: int | None = None, + model: str | None = None, +) -> list[dict[str, Any]] | None: + """Submit a music generation request to Flow Music (flowmusic.app) or Google Labs. + + Args: + bridge: ExtensionBridge instance + prompt: Text prompt describing the music, mood, genre, or instruments + project_id: Flow project ID + duration: Desired duration in seconds (e.g. 10, 30, 50, 70) + count: Number of variations (1-4) + loop: Whether to generate a seamless loop + seed: Optional explicit random seed + model: Optional model override (e.g. producer:standard, MUSIC_FX) + + Returns: + List of dicts with keys (media_id, audio_url, audio_base64, revised_prompt) or None on failure + """ + # Primary engine: Flow Music (flowmusic.app) + if model not in ("MUSIC_FX", "MUSICLM_V2", "DEFAULT"): + try: + flowmusic_results = await _generate_flowmusic(bridge, prompt, model=model) + if flowmusic_results: + return flowmusic_results + except Exception as exc: + log.warning("Flow Music generation failed, falling back to Google Labs MusicFX: %s", exc) + + # 2. Fallback: Google Labs AISandbox endpoints + proj_id = project_id or DEFAULT_PROJECT + count = max(1, min(4, count)) + seed_val = resolve_seed(seed) + + primary_body = { + "clientContext": build_music_client_context(proj_id, tool=model or "MUSICLM_V2"), + "generationCount": count, + "input": { + "textInput": prompt, + }, + "soundLengthSeconds": int(duration), + "loop": bool(loop), + "model": "DEFAULT", + "seed": seed_val, + } + alt_body = { + "clientContext": build_music_client_context(proj_id, tool=model or "MUSIC_FX"), + "generationCount": count, + "inputContext": { + "textInput": prompt, + }, + "soundLengthSeconds": int(duration), + "loop": bool(loop), + "seed": seed_val, + } + + log.info('Generating music via fallback: "%s" (%ds, loop=%s, count=%d)', prompt[:50], duration, loop, count) + + candidate_configs = [ + (ENDPOINTS.get("generate_music", "/v1:runMusicFx"), primary_body), + (ENDPOINTS.get("generate_music", "/v1:runMusicFx"), alt_body), + (ENDPOINTS.get("generate_music_demo", "/v1:runSoundDemo"), primary_body), + (ENDPOINTS.get("generate_music_sound", "/v1/sound:generate"), alt_body), + (ENDPOINTS.get("generate_music_batch", "/v1/music:batchGenerateMusic"), alt_body), + ] + + attempt_errors = [] + for endpoint, req_body in candidate_configs: + try: + result = await bridge.api_request(endpoint, req_body, captcha_action="MUSIC_GENERATION") + status = result.get("status", 0) + data = result.get("data", {}) + err_msg = "" + if isinstance(data, dict): + err_msg = data.get("error", {}).get("message") or data.get("message") or str(data) + elif isinstance(data, str): + err_msg = data + err_msg = err_msg or result.get("error", f"HTTP {status}") + + log.info("Music endpoint %s response: status=%s, error=%s", endpoint, status, err_msg[:200]) + + if status == 200: + parsed = _parse_music_results(data if isinstance(data, dict) else {}, prompt) + if parsed: + log.info("Music generated successfully! (%d tracks)", len(parsed)) + return parsed + else: + log.warning("Music endpoint returned 200 but parse found no audio tracks: %s", str(data)[:200]) + attempt_errors.append(f"{endpoint}: 200 OK but no audio in payload") + else: + attempt_errors.append(f"{endpoint} -> HTTP {status}: {err_msg}") + except Exception as exc: + log.warning("Request to %s encountered exception: %s", endpoint, exc) + attempt_errors.append(f"{endpoint} -> Exception: {exc}") + + raise ValueError(f"Google Labs Music API error: {'; '.join(attempt_errors)}") + + +async def download_music( + bridge, + media_id_or_url: str, + output_path: str, + audio_base64: str | None = None, +) -> str | None: + """Save generated audio to disk from base64 content, signed URL, or media ID. + + Returns the absolute path of the saved audio file, or None on failure. + """ + os.makedirs(os.path.dirname(os.path.abspath(output_path)) or ".", exist_ok=True) + destination = os.path.abspath(output_path) + + # Path 1: Direct base64 data + if audio_base64: + try: + raw_payload = audio_base64 + declared_mime = "audio/mpeg" + if raw_payload.startswith("data:") and "," in raw_payload: + meta, raw_payload = raw_payload.split(",", 1) + declared_mime = meta[5:].split(";", 1)[0] + audio_bytes = base64.b64decode("".join(raw_payload.split()), validate=False) + with open(destination, "wb") as f: + f.write(audio_bytes) + final_path = ensure_correct_extension(destination, declared_mime=declared_mime) + return final_path + except Exception as exc: + log.error("Failed to decode base64 audio: %s", exc) + + # Path 2: Download from URL + target_url = media_id_or_url + if not (target_url.startswith("http://") or target_url.startswith("https://")): + # Path 3: Media ID resolution + signed_url = await bridge.request_media_url(media_id_or_url) + if signed_url: + target_url = signed_url + else: + log.error("Could not obtain download URL for media ID: %s", media_id_or_url) + return None + + try: + req = urllib.request.Request( + target_url, + headers={"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)"}, + ) + with urllib.request.urlopen(req, timeout=90) as response: + audio_bytes = response.read() + + mime = sniff_media_type(audio_bytes, filename=destination) + with open(destination, "wb") as f: + f.write(audio_bytes) + final_path = ensure_correct_extension(destination, declared_mime=mime) + return final_path + except Exception as exc: + log.error("Failed to download audio from %s: %s", target_url, exc) + return None diff --git a/flow-agent/flow_engine/generators/t2v.py b/flow-agent/flow_engine/generators/t2v.py index b09506c..c740cd4 100644 --- a/flow-agent/flow_engine/generators/t2v.py +++ b/flow-agent/flow_engine/generators/t2v.py @@ -1,7 +1,6 @@ """Flow Engine — Text to Video (T2V) generator.""" import logging -import random from ..config import ENDPOINTS from .common import build_client_context, build_generation_context, resolve_seed diff --git a/flow-agent/flow_engine/generators/v2v.py b/flow-agent/flow_engine/generators/v2v.py index 028d6af..80ca9ad 100644 --- a/flow-agent/flow_engine/generators/v2v.py +++ b/flow-agent/flow_engine/generators/v2v.py @@ -1,7 +1,6 @@ """Flow Engine — Video to Video (V2V) editor.""" import logging -import random from ..config import ENDPOINTS from .common import build_client_context, build_generation_context, resolve_seed diff --git a/flow-agent/flow_server/batch.py b/flow-agent/flow_server/batch.py index 372689f..0d20ba2 100644 --- a/flow-agent/flow_server/batch.py +++ b/flow-agent/flow_server/batch.py @@ -21,7 +21,6 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from flow_engine.config import MAX_CONCURRENT_REQUESTS, REQUEST_MIN_INTERVAL def generate_single_item( diff --git a/flow-agent/flow_server/mcp_server.py b/flow-agent/flow_server/mcp_server.py index 2549473..8a858c3 100755 --- a/flow-agent/flow_server/mcp_server.py +++ b/flow-agent/flow_server/mcp_server.py @@ -410,6 +410,47 @@ def handle_tools_list(request_id): }, "required": ["prompt"] } + }, + { + "name": "generate_flow_music", + "description": "Generate music or ambient soundscapes from a text prompt using Google Labs / Flow. Returns local audio path and media ID.", + "inputSchema": { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Text description of the music, genre, mood, instruments, or soundscape" + }, + "duration": { + "type": "integer", + "description": "Duration in seconds (e.g. 10, 30, 50, 70). Default: 30", + "default": 30 + }, + "loop": { + "type": "boolean", + "description": "Whether to generate a seamless audio loop. Default: false", + "default": False + }, + "count": { + "type": "integer", + "description": "Number of variations (1-4). Default: 1", + "minimum": 1, + "maximum": 4, + "default": 1 + }, + "output_dir": { + "type": "string", + "description": "Optional local destination directory; defaults to FLOW_OUTPUT_DIR" + }, + "seed": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295, + "description": "Optional generation seed for reproducible output" + } + }, + "required": ["prompt"] + } } ] @@ -669,6 +710,59 @@ def call_edit_flow_video(prompt, media_id=None, video_path=None, aspect="landsca start_media_id=media_id, ref_media_ids=ref_media_ids, is_video=True, ) +def call_generate_flow_music(prompt, duration=30, loop=False, count=1, output_dir=None, seed=None): + """Generate audio/music tracks and return their details.""" + ensure_backend_running() + payload = { + "prompt": prompt, + "duration": duration or 30, + "loop": bool(loop), + "n": count or 1, + } + if seed is not None: + payload["seed"] = int(seed) + + try: + data = _request_json("/v1/audio/generations", payload, method="POST", timeout=300) + except Exception as exc: + return {"error": f"Audio generation request failed: {exc}"} + + items = data.get("data", []) if isinstance(data, dict) else [] + if not items: + return {"error": f"Audio generation returned no tracks: {data}"} + + target_dir = os.path.abspath(os.path.expanduser(output_dir)) if output_dir else OUTPUT_DIR + os.makedirs(target_dir, exist_ok=True) + + results = [] + for item in items: + url = item.get("url", "") + media_id = item.get("media_id", "") + revised = item.get("revised_prompt", prompt) + + filename = url.split("/media/")[-1] if "/media/" in url else os.path.basename(url) + local_src = os.path.join(OUTPUT_DIR, filename) + + if os.path.exists(local_src) and target_dir != OUTPUT_DIR: + dest = os.path.join(target_dir, filename) + shutil.copy2(local_src, dest) + local_path = dest + else: + local_path = local_src + + results.append({ + "media_id": media_id, + "url": url, + "local_path": local_path, + "revised_prompt": revised, + }) + + return { + "success": True, + "tracks": results, + "primary_track": results[0]["local_path"] if results else None, + } + def call_download_media_from_url(url, output_dir=None, filename=None, upload_to_flow=False, max_size_mb=2048): """Download a remote media asset safely and optionally upload it to Flow.""" @@ -1075,6 +1169,16 @@ def handle_tool_call(request_id, tool_name, arguments): arguments.get("ref_media_ids"), ) content = [{"type": "text", "text": text}] + elif tool_name == "generate_flow_music": + result = call_generate_flow_music( + arguments.get("prompt"), + arguments.get("duration", 30), + arguments.get("loop", False), + arguments.get("count", 1), + arguments.get("output_dir"), + arguments.get("seed"), + ) + content = [{"type": "text", "text": json.dumps(result, indent=2)}] else: return { "jsonrpc": "2.0", diff --git a/flow-agent/flow_server/media_history.py b/flow-agent/flow_server/media_history.py index 58a25da..0c4765b 100644 --- a/flow-agent/flow_server/media_history.py +++ b/flow-agent/flow_server/media_history.py @@ -131,7 +131,11 @@ def record_local_media( actual_mime = sniff_media_type(local_path, declared_mime=mime_type) filename = os.path.basename(local_path) record = { - "type": "video" if actual_mime.startswith("video/") else "image", + "type": ( + "video" + if actual_mime.startswith("video/") + else ("audio" if actual_mime.startswith("audio/") else "image") + ), "prompt": prompt, "timestamp": int(time.time()), "media_id": str(media_id) if media_id else None, diff --git a/flow-agent/flow_server/media_types.py b/flow-agent/flow_server/media_types.py index 21a9fc9..833889a 100644 --- a/flow-agent/flow_server/media_types.py +++ b/flow-agent/flow_server/media_types.py @@ -16,6 +16,15 @@ _MIME_EXTENSIONS = { + "audio/aac": ".aac", + "audio/flac": ".flac", + "audio/mp3": ".mp3", + "audio/mp4": ".m4a", + "audio/mpeg": ".mp3", + "audio/ogg": ".ogg", + "audio/wav": ".wav", + "audio/x-m4a": ".m4a", + "audio/x-wav": ".wav", "image/avif": ".avif", "image/bmp": ".bmp", "image/gif": ".gif", @@ -87,17 +96,29 @@ def _mime_from_signature(header: bytes) -> str | None: return "image/bmp" if header.startswith((b"II*\x00", b"MM\x00*")): return "image/tiff" + if header.startswith(b"ID3") or header.startswith((b"\xff\xfb", b"\xff\xf3", b"\xff\xf2", b"\xff\xe3")): + return "audio/mpeg" + if header.startswith(b"fLaC"): + return "audio/flac" + if header.startswith(b"OggS"): + return "audio/ogg" + if header.startswith((b"\xff\xf1", b"\xff\xf9")): + return "audio/aac" if len(header) >= 12 and header[:4] == b"RIFF": if header[8:12] == b"WEBP": return "image/webp" if header[8:12] == b"AVI ": return "video/x-msvideo" + if header[8:12] == b"WAVE": + return "audio/wav" # ISO base media files share the ftyp container. The major/compatible # brand tells image formats apart from MP4 and QuickTime video. if len(header) >= 12 and header[4:8] == b"ftyp": brands = header[8:64] major_brand = header[8:12] + if any(brand in brands for brand in (b"M4A ", b"M4B ", b"m4a ", b"f4a ")): + return "audio/mp4" if any(brand in brands for brand in (b"avif", b"avis")): return "image/avif" if any(brand in brands for brand in (b"heic", b"heix", b"hevc", b"hevx")): diff --git a/flow-agent/flow_server/models.py b/flow-agent/flow_server/models.py index 5b6a0a5..14c0d1c 100644 --- a/flow-agent/flow_server/models.py +++ b/flow-agent/flow_server/models.py @@ -33,6 +33,17 @@ class VideoGenerationRequest(BaseModel): video_model: Optional[str] = Field(None, description="Optional Flow videoModelKey override") +class MusicGenerationRequest(BaseModel): + prompt: str = Field(..., description="The prompt describing the music, mood, genre, or instruments") + duration: int = Field(30, description="Duration in seconds (e.g. 10, 30, 50, 70)") + loop: bool = Field(False, description="Whether to generate a seamless loop") + n: int = Field(1, ge=1, le=4, description="Number of music tracks to generate (1-4)") + model: Optional[str] = Field(None, description="Optional music model/tool name (e.g. MUSIC_FX)") + seed: Optional[int] = Field(None, ge=0, le=4294967295, description="Optional explicit generation seed") + response_format: str = Field("url", description="Return format: url or b64_json") + user: Optional[str] = None + + class GeneratedMedia(BaseModel): url: Optional[str] = None media_id: Optional[str] = None diff --git a/flow-agent/flow_server/routes/generation.py b/flow-agent/flow_server/routes/generation.py index e7120bc..a56ec01 100644 --- a/flow-agent/flow_server/routes/generation.py +++ b/flow-agent/flow_server/routes/generation.py @@ -13,7 +13,12 @@ from fastapi.responses import JSONResponse from flow_server.config import OUTPUT_DIR, map_size_to_aspect -from flow_server.models import ImageGenerationRequest, VideoGenerationRequest, VideoGenerationResult +from flow_server.models import ( + ImageGenerationRequest, + VideoGenerationRequest, + VideoGenerationResult, + MusicGenerationRequest, +) from flow_server.idempotency import get_idempotency_store from flow_server.jobs import get_job_store from flow_server.history import MediaNotFoundError @@ -21,7 +26,7 @@ from flow_server.media_types import ensure_correct_extension, extension_for_media, sniff_media_type from flow_server.state import verify_api_key, get_active_bridge, publish, append_to_history -from flow_engine import DEFAULT_PROJECT +from flow_engine import DEFAULT_PROJECT, generate_music, download_music from flow_engine.bridge import target_client_id_var from flow_engine.config import CREDITS_PER_VIDEO from flow_engine.generators.t2i import generate_image, download_image @@ -648,3 +653,131 @@ async def poll_and_download(media_id: str, index: int): f"(each {req.duration}s video costs {cost_each} credits)." ) return resp + + +@router.post("/v1/audio/generations", dependencies=[Depends(verify_api_key)]) +@router.post("/v1/music/generations", dependencies=[Depends(verify_api_key)]) +async def openai_generate_audio( + req: MusicGenerationRequest, + x_client_id: Optional[str] = Header(None, alias="X-Client-Id"), + idempotency_key: Optional[str] = Header(None, alias="Idempotency-Key"), +): + """Generate audio/music from a prompt with idempotency and history registry support.""" + x_client_id = _header_string(x_client_id) + key = _header_string(idempotency_key) + if not key: + return await _generate_music(req, x_client_id) + + store = get_idempotency_store(OUTPUT_DIR) + claim = await store.claim(key, _request_payload(req, "audio", x_client_id)) + if claim.action == "conflict": + raise HTTPException( + status_code=409, + detail="Idempotency-Key was already used with a different audio generation request.", + ) + if claim.action == "replay": + return claim.record["response"] + if claim.action == "failed": + error = claim.record.get("error", {}) + raise HTTPException( + status_code=int(error.get("status_code", 500)), + detail=error.get("detail", "The original idempotent audio generation failed."), + ) + if claim.action == "processing": + raise HTTPException( + status_code=409, + detail="This idempotent audio generation is already processing; retry with the same key.", + ) + + try: + response = await _generate_music(req, x_client_id) + except HTTPException as exc: + await store.fail(key, exc.status_code, str(exc.detail)) + raise + except Exception as exc: + await store.fail(key, 500, str(exc)) + raise + await store.succeed(key, response) + return response + + +async def _generate_music(req: MusicGenerationRequest, x_client_id: Optional[str] = None): + """Generate audio from a prompt via Google Labs / MusicFX.""" + target_client_id_var.set(x_client_id) + active_bridge = await get_active_bridge() + project_id = os.environ.get("DEFAULT_PROJECT", DEFAULT_PROJECT) + + try: + results = await generate_music( + active_bridge, + prompt=req.prompt, + project_id=project_id, + duration=req.duration, + count=req.n, + loop=req.loop, + seed=req.seed, + model=req.model, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + except Exception as exc: + log.exception("Audio generation request failed") + raise HTTPException(status_code=500, detail=f"Audio generation failed: {exc}") from exc + + if not results: + raise HTTPException(status_code=500, detail="No audio returned from generator.") + + timestamp = int(time.time()) + data_outputs = [] + + for i, item in enumerate(results): + media_id = item.get("media_id") or uuid.uuid4().hex + audio_url = item.get("audio_url", "") + audio_b64 = item.get("audio_base64", "") + revised = item.get("revised_prompt", req.prompt) + + filename = f"flow_audio_{timestamp}_{uuid.uuid4().hex[:6]}_{i+1}.mp3" + out_path = os.path.join(OUTPUT_DIR, filename) + + saved_path = await download_music( + active_bridge, + media_id_or_url=audio_url or media_id, + output_path=out_path, + audio_base64=audio_b64, + ) + + if not saved_path or not os.path.exists(saved_path): + log.warning("Could not download audio track %d", i + 1) + continue + + saved_filename = os.path.basename(saved_path) + served_url, r2_key = await publish(saved_filename, saved_path) + + entry = { + "url": served_url, + "media_id": media_id, + "revised_prompt": revised, + } + if req.response_format == "b64_json": + with open(saved_path, "rb") as af: + entry["b64_json"] = base64.b64encode(af.read()).decode("utf-8") + + data_outputs.append(entry) + + await append_to_history( + "audio", + served_url, + req.prompt, + media_id, + r2_key, + local_path=saved_path, + project_id=project_id, + ) + + if not data_outputs: + raise HTTPException(status_code=500, detail="Failed to download or save generated audio tracks.") + + return { + "created": timestamp, + "data": data_outputs, + } diff --git a/flow-agent/main.py b/flow-agent/main.py index adf1665..dca2529 100644 --- a/flow-agent/main.py +++ b/flow-agent/main.py @@ -273,14 +273,12 @@ def _download_media(url, output_path, *, preserve_explicit=True): sniff_media_type, ) - # Pass an open stream rather than the path so an unknown signature - # cannot silently fall back to the requested filename extension. with open(temporary_path, "rb") as downloaded_file: actual_mime = sniff_media_type(downloaded_file) - if not actual_mime.startswith(("image/", "video/")): + if not actual_mime.startswith(("image/", "video/", "audio/")): mime_hint = f"; server declared {declared_mime}" if declared_mime else "" raise ValueError( - f"downloaded data from {url} is not recognized image/video media " + f"downloaded data from {url} is not recognized image/video/audio media " f"(signature detected {actual_mime}{mime_hint})" ) @@ -705,11 +703,52 @@ def cmd_batch(argv): batch_main(argv) +def cmd_music(argv): + parser = argparse.ArgumentParser( + prog="flow music", + description="Generate music or audio through the Flow backend.", + ) + parser.add_argument("prompt", help="Text prompt describing music, mood, genre, or instruments") + parser.add_argument("--output", "-o", default=None, help="Exact output file path (e.g. soundtrack.mp3)") + parser.add_argument( + "--duration", + "-d", + type=int, + choices=[10, 30, 50, 70], + default=30, + help="Duration in seconds (default: 30)", + ) + parser.add_argument("--count", "-c", type=int, choices=[1, 2, 3, 4], default=1, help="Number of variations") + parser.add_argument("--loop", "-l", action="store_true", help="Generate a seamless loop") + parser.add_argument("--seed", type=int, help="Generation seed for reproducibility") + parser.add_argument("--idempotency-key", help="Reuse a previous request safely") + args = parser.parse_args(argv) + + _wait_for_generation_ready() + output_path, explicit_output = _requested_output_path(args.output, "soundtrack.mp3") + payload = { + "prompt": args.prompt, + "duration": args.duration, + "n": args.count, + "loop": args.loop, + "response_format": "url", + } + if args.seed is not None: + payload["seed"] = args.seed + + idempotency_key = args.idempotency_key or uuid.uuid4().hex + result = _post_generation( + "/v1/audio/generations", payload, idempotency_key, timeout=300 + ) + _save_outputs(result, output_path, "audio", explicit_output=explicit_output) + + COMMANDS = { "mcp": run_mcp, "image": cmd_image, "batch": cmd_batch, "video": cmd_video, + "music": cmd_music, "edit": cmd_edit, "upload": cmd_upload, "credits": cmd_credits, @@ -727,6 +766,7 @@ def _usage(): print(" image Generate images") print(" batch Generate batch images in parallel") print(" video Generate videos") + print(" music Generate music / audio") print(" edit Edit a video") print(" upload Upload media") print(" credits Show Flow credits") diff --git a/flow-agent/tests/test_music_generation.py b/flow-agent/tests/test_music_generation.py new file mode 100644 index 0000000..0c2a877 --- /dev/null +++ b/flow-agent/tests/test_music_generation.py @@ -0,0 +1,319 @@ +"""Comprehensive test suite for Flow Agent music/audio generation capabilities.""" + +import base64 +import os +import tempfile +import uuid +import pytest +from unittest.mock import AsyncMock, MagicMock + +from flow_server.media_types import ( + extension_for_media, + extension_for_mime, + sniff_media_type, +) +from flow_engine.generators.t2m import ( + _parse_music_results, + build_music_client_context, + download_music, + generate_music, +) +from flow_server.models import MusicGenerationRequest + + +# ─── 1. Media Type & Audio Sniffing Tests ───────────────────────── + +def test_audio_mime_extensions(): + assert extension_for_mime("audio/mpeg") == ".mp3" + assert extension_for_mime("audio/mp3") == ".mp3" + assert extension_for_mime("audio/wav") == ".wav" + assert extension_for_mime("audio/x-wav") == ".wav" + assert extension_for_mime("audio/ogg") == ".ogg" + assert extension_for_mime("audio/flac") == ".flac" + assert extension_for_mime("audio/aac") == ".aac" + assert extension_for_mime("audio/mp4") == ".m4a" + assert extension_for_mime("audio/x-m4a") == ".m4a" + + +def test_audio_signature_sniffing(): + # MP3 with ID3 tag + id3_header = b"ID3\x03\x00\x00\x00\x00\x00\x00" + b"\x00" * 50 + assert sniff_media_type(id3_header) == "audio/mpeg" + + # MP3 sync frame + mp3_frame = b"\xff\xfb\x90\x44" + b"\x00" * 50 + assert sniff_media_type(mp3_frame) == "audio/mpeg" + + # WAV header + wav_header = b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00" + assert sniff_media_type(wav_header) == "audio/wav" + + # OGG header + ogg_header = b"OggS\x00\x02\x00\x00\x00\x00\x00\x00" + assert sniff_media_type(ogg_header) == "audio/ogg" + + # FLAC header + flac_header = b"fLaC\x00\x00\x00\x22" + assert sniff_media_type(flac_header) == "audio/flac" + + # AAC sync word + aac_header = b"\xff\xf1\x50\x80" + assert sniff_media_type(aac_header) == "audio/aac" + + # M4A container + m4a_header = b"\x00\x00\x00\x20ftypM4A \x00\x00\x00\x00M4A mp42isom" + assert sniff_media_type(m4a_header) == "audio/mp4" + + +def test_extension_for_audio_media(): + wav_bytes = b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00" + assert extension_for_media(wav_bytes) == ".wav" + + id3_bytes = b"ID3\x03\x00\x00\x00\x00\x00\x00" + assert extension_for_media(id3_bytes) == ".mp3" + + +# ─── 2. Engine Generator Tests ──────────────────────────────────── + +def test_build_music_client_context(): + ctx = build_music_client_context("test-project-123", tool="MUSIC_FX") + assert ctx["projectId"] == "test-project-123" + assert ctx["tool"] == "MUSIC_FX" + assert "sessionId" in ctx + assert "recaptchaContext" in ctx + + +def test_parse_music_results_sounds_format(): + test_id = str(uuid.uuid4()) + data = { + "sounds": [ + { + "name": test_id, + "audio": base64.b64encode(b"ID3mockaudio").decode("utf-8"), + "fifeUrl": f"https://example.com/audio/{test_id}", + "revisedPrompt": "An epic cinematic ambient soundtrack", + } + ] + } + parsed = _parse_music_results(data, "cinematic music") + assert len(parsed) == 1 + assert parsed[0]["media_id"] == test_id + assert parsed[0]["audio_url"] == f"https://example.com/audio/{test_id}" + assert parsed[0]["revised_prompt"] == "An epic cinematic ambient soundtrack" + + +def test_parse_music_results_media_format(): + test_id = str(uuid.uuid4()) + data = { + "media": [ + { + "name": test_id, + "audio": { + "generatedAudio": { + "downloadUrl": f"https://example.com/download/{test_id}", + "rawBytes": base64.b64encode(b"ID3mockbytes").decode("utf-8"), + } + }, + } + ] + } + parsed = _parse_music_results(data, "ambient loop") + assert len(parsed) == 1 + assert parsed[0]["media_id"] == test_id + assert parsed[0]["audio_url"] == f"https://example.com/download/{test_id}" + + +@pytest.mark.asyncio +async def test_generate_music_success(): + mock_bridge = MagicMock() + test_id = str(uuid.uuid4()) + mock_bridge.api_request = AsyncMock( + return_value={ + "status": 200, + "data": { + "sounds": [ + { + "name": test_id, + "audio": base64.b64encode(b"ID3mockaudio").decode("utf-8"), + "downloadUrl": f"https://example.com/{test_id}.mp3", + } + ] + }, + } + ) + + results = await generate_music( + mock_bridge, + prompt="Cyberpunk beat", + project_id="p-1", + duration=30, + count=1, + loop=True, + ) + + assert results is not None + assert len(results) == 1 + assert results[0]["media_id"] == test_id + mock_bridge.api_request.assert_called_once() + + +@pytest.mark.asyncio +async def test_download_music_from_base64(): + mock_bridge = MagicMock() + b64_audio = base64.b64encode(b"ID3mockaudiobytes" + b"\x00" * 30).decode("utf-8") + + with tempfile.TemporaryDirectory() as tmpdir: + out_file = os.path.join(tmpdir, "output.mp3") + saved_path = await download_music( + mock_bridge, + media_id_or_url="test-id", + output_path=out_file, + audio_base64=b64_audio, + ) + + assert saved_path is not None + assert os.path.exists(saved_path) + with open(saved_path, "rb") as f: + content = f.read() + assert content.startswith(b"ID3") + + +# ─── 3. Pydantic Models & API Route Tests ───────────────────────── + +def test_music_generation_request_validation(): + req = MusicGenerationRequest(prompt="Relaxing piano", duration=30, loop=True, n=2) + assert req.prompt == "Relaxing piano" + assert req.duration == 30 + assert req.loop is True + assert req.n == 2 + assert req.response_format == "url" + + +@pytest.mark.asyncio +async def test_api_audio_generations_endpoint(): + from fastapi.testclient import TestClient + from flow_server.api import app + from flow_server.state import set_bridge + + mock_bridge = MagicMock() + mock_bridge.health_check = AsyncMock(return_value=True) + test_id = str(uuid.uuid4()) + b64_audio = base64.b64encode(b"ID3mockaudiobytes" + b"\x00" * 30).decode("utf-8") + mock_bridge.api_request = AsyncMock( + return_value={ + "status": 200, + "data": { + "sounds": [ + { + "name": test_id, + "audio": b64_audio, + "downloadUrl": f"https://example.com/{test_id}.mp3", + "revisedPrompt": "Relaxing piano melody", + } + ] + }, + } + ) + set_bridge(mock_bridge) + + client = TestClient(app) + response = client.post( + "/v1/audio/generations", + json={"prompt": "Relaxing piano", "duration": 30, "loop": False, "n": 1}, + ) + + assert response.status_code == 200 + data = response.json() + assert "data" in data + assert len(data["data"]) == 1 + assert data["data"][0]["media_id"] == test_id + assert data["data"][0]["revised_prompt"] == "Relaxing piano melody" + + +# ─── 4. MCP Tool Tests ──────────────────────────────────────────── + +def test_mcp_server_music_tool_registered(): + from flow_server.mcp_server import handle_tools_list + + response = handle_tools_list("req-1") + tools = response.get("result", {}).get("tools", []) + tool_names = [t["name"] for t in tools] + assert "generate_flow_music" in tool_names + + music_tool = next(t for t in tools if t["name"] == "generate_flow_music") + assert "prompt" in music_tool["inputSchema"]["required"] + assert "duration" in music_tool["inputSchema"]["properties"] + assert "loop" in music_tool["inputSchema"]["properties"] + + +# ─── 5. Flow Music Pipeline & Fallback Tests ────────────────────── + +@pytest.mark.asyncio +async def test_flowmusic_pipeline_success(): + from flow_engine.generators.t2m import _generate_flowmusic + + clip_id = str(uuid.uuid4()) + mock_bridge = MagicMock() + + async def mock_api_request(url, payload, method="POST", timeout=None, **kwargs): + if "conversation" in url: + return {"status": 200, "data": {"job_id": "job-123"}} + if "stream" in url: + return {"status": 200, "data": {"clip_ids": [clip_id]}} + if "clips" in url: + return { + "status": 200, + "data": { + "clips": { + clip_id: { + "audio_url": f"https://flowmusic.app/audio/{clip_id}.m4a", + "title": "Summer Lo-Fi", + "duration": {"value": 30.0}, + } + } + }, + } + return {"status": 404, "error": "NOT_FOUND"} + + mock_bridge.api_request = AsyncMock(side_effect=mock_api_request) + + results = await _generate_flowmusic(mock_bridge, "relaxing lo-fi beat") + assert results is not None + assert len(results) == 1 + assert results[0]["media_id"] == clip_id + assert results[0]["audio_url"] == f"https://flowmusic.app/audio/{clip_id}.m4a" + + +@pytest.mark.asyncio +async def test_generate_music_fallback_on_flowmusic_error(): + mock_bridge = MagicMock() + test_id = str(uuid.uuid4()) + + # Flow music fails, then Google Labs endpoint succeeds + async def mock_api_request(url, payload, **kwargs): + if "flowmusic.app" in url: + return {"status": 500, "error": "SERVICE_UNAVAILABLE"} + return { + "status": 200, + "data": { + "sounds": [ + { + "name": test_id, + "audio": base64.b64encode(b"ID3mockaudio").decode("utf-8"), + "downloadUrl": f"https://example.com/{test_id}.mp3", + } + ] + }, + } + + mock_bridge.api_request = AsyncMock(side_effect=mock_api_request) + + results = await generate_music( + mock_bridge, + prompt="Calm piano", + project_id="test-proj", + duration=30, + ) + assert results is not None + assert len(results) == 1 + assert results[0]["media_id"] == test_id diff --git a/flow-agent/tests/test_recaptcha_self_heal.py b/flow-agent/tests/test_recaptcha_self_heal.py new file mode 100644 index 0000000..14c3179 --- /dev/null +++ b/flow-agent/tests/test_recaptcha_self_heal.py @@ -0,0 +1,50 @@ +"""Unit tests for bridge error classification and reCAPTCHA self-healing detection.""" + +from flow_engine.bridge import ExtensionBridge + + +def test_is_unauthenticated(): + assert ExtensionBridge._is_unauthenticated({"status": 401}) is True + assert ExtensionBridge._is_unauthenticated({"data": {"error": {"code": 401}}}) is True + assert ExtensionBridge._is_unauthenticated({"data": {"error": {"status": "UNAUTHENTICATED"}}}) is True + assert ExtensionBridge._is_unauthenticated({"status": 200}) is False + assert ExtensionBridge._is_unauthenticated({"status": 400}) is False + + +def test_is_recaptcha_failure(): + # 400 with UNUSUAL_ACTIVITY + err_400 = { + "status": 400, + "data": { + "error": { + "code": 400, + "message": "reCAPTCHA evaluation failed (PUBLIC_ERROR_UNUSUAL_ACTIVITY)", + "status": "INVALID_ARGUMENT", + } + } + } + assert ExtensionBridge._is_recaptcha_failure(err_400) is True + + # 403 / 400 string error + err_str = { + "status": 400, + "data": "Generation failed: reCAPTCHA verification error" + } + assert ExtensionBridge._is_recaptcha_failure(err_str) is True + + # Extension direct captcha failed error + err_ext = { + "status": 403, + "error": "CAPTCHA_FAILED: timeout" + } + assert ExtensionBridge._is_recaptcha_failure(err_ext) is True + + # Regular 400 (not captcha) + err_other = { + "status": 400, + "data": {"error": {"message": "Invalid aspect ratio"}} + } + assert ExtensionBridge._is_recaptcha_failure(err_other) is False + + # Success + assert ExtensionBridge._is_recaptcha_failure({"status": 200, "data": {}}) is False diff --git a/flow-agent/uv.lock b/flow-agent/uv.lock index 60056b5..cbd0f8a 100644 --- a/flow-agent/uv.lock +++ b/flow-agent/uv.lock @@ -270,7 +270,7 @@ wheels = [ [[package]] name = "flow-agent" -version = "2.0.3" +version = "2.0.5" source = { editable = "." } dependencies = [ { name = "cryptography" }, diff --git a/flow-extension/background.js b/flow-extension/background.js index 3cfc6df..cd1186a 100644 --- a/flow-extension/background.js +++ b/flow-extension/background.js @@ -47,11 +47,12 @@ let metrics = { // ─── URL → Log Type Classifier ───────────────────────────── // Visible log types — only these appear in the request log -const _VISIBLE_TYPES = new Set(['GEN_IMG', 'GEN_VID', 'GEN_VID_REF', 'UPSCALE', 'TRACKING', 'URL_REFRESH']); +const _VISIBLE_TYPES = new Set(['GEN_IMG', 'GEN_VID', 'GEN_VID_REF', 'GEN_MUSIC', 'UPSCALE', 'TRACKING', 'URL_REFRESH']); function _classifyApiUrl(url) { if (url.includes('uploadImage')) return 'UPLOAD'; if (url.includes('batchGenerateImages')) return 'GEN_IMG'; + if (url.includes('sound:generate') || url.includes('soundDemo') || url.includes('batchGenerateMusic')) return 'GEN_MUSIC'; if (url.includes('UpsampleVideo')) return 'UPSCALE'; if (url.includes('ReferenceImages')) return 'GEN_VID_REF'; if (url.includes('batchAsyncGenerateVideo')) return 'GEN_VID'; @@ -139,21 +140,25 @@ chrome.webRequest.onBeforeSendHeaders.addListener( (h) => h.name?.toLowerCase() === 'authorization', ); const value = authHeader?.value || ''; - if (!value.startsWith('Bearer ya29.')) return; + if (!value.startsWith('Bearer ')) return; const token = value.replace(/^Bearer\s+/i, '').trim(); if (!token) return; - // Always update — even if same token string, refresh the timestamp - flowKey = token; - metrics.tokenCapturedAt = Date.now(); - chrome.storage.local.set({ flowKey, metrics }); - console.log('[Flow Agent] Bearer token captured'); - - // Notify whichever transport is active. - sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); + const url = details.url || ''; + if (url.includes('flowmusic.app')) { + flowMusicKey = token; + chrome.storage.local.set({ flowMusicKey }); + console.log('[Flow Agent] FlowMusic token captured'); + } else if (value.startsWith('Bearer ya29.')) { + flowKey = token; + metrics.tokenCapturedAt = Date.now(); + chrome.storage.local.set({ flowKey, metrics }); + console.log('[Flow Agent] Bearer token captured'); + sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); + } }, - { urls: ['https://aisandbox-pa.googleapis.com/*', 'https://labs.google/*'] }, + { urls: ['https://aisandbox-pa.googleapis.com/*', 'https://labs.google/*', 'https://*.flowmusic.app/*', 'https://flowmusic.app/*'] }, ['requestHeaders', 'extraHeaders'], ); @@ -251,7 +256,7 @@ async function _getOrOpenFlowTab() { } scheduleFlowTabClose(); - return retryTabs[0]; + return createdTab; } async function getOrOpenFlowTab() { @@ -272,9 +277,9 @@ function isTokenFresh() { return ageMs < 50 * 60 * 1000; // 50 minutes } -async function captureTokenFromFlowTab() { +async function captureTokenFromFlowTab(forceReload = false) { // Skip if token is still fresh — no need to open/refresh anything - if (isTokenFresh()) { + if (isTokenFresh() && !forceReload) { console.log('[Flow Agent] Token still fresh, skipping tab refresh'); return; } @@ -290,6 +295,16 @@ async function captureTokenFromFlowTab() { console.log('[Flow Agent] Flow tab not ready yet after open'); return; } + if (forceReload && tab.id) { + console.log('[Flow Agent] Force reloading Flow tab (id: ' + tab.id + ')...'); + try { + await chrome.tabs.reload(tab.id, { bypassCache: true }); + await waitForTabComplete(tab.id, 15000); + await sleep(1500); + } catch (err) { + console.warn('[Flow Agent] Tab reload error:', err); + } + } await chrome.scripting.executeScript({ target: { tabId: tab.id }, files: ['content.js'], @@ -305,183 +320,192 @@ async function captureTokenFromFlowTab() { // ─── WebSocket to Agent ───────────────────────────────────── +let _connectingWs = false; + async function connectToAgent() { if (manualDisconnect) return; await connectHttpAgent(); - if (ws?.readyState === WebSocket.CONNECTING) return; - if (ws?.readyState === WebSocket.OPEN) return; - - const data = await chrome.storage.local.get(['clientId']); - const serverIp = CONFIG.DEFAULT_SERVER_HOST; - connectedServerHost = serverIp; - const isLocal = /^(127\.0\.0\.1|localhost|192\.168\.|10\.)/.test(serverIp); - const wsScheme = isLocal ? 'ws' : 'wss'; - const httpScheme = isLocal ? 'http' : 'https'; - const wsUrl = `${wsScheme}://${serverIp}/ws`; - - // Dynamically resolve callbackUrl - callbackUrl = `${httpScheme}://${serverIp}/api/ext/callback`; + if (ws?.readyState === WebSocket.CONNECTING || ws?.readyState === WebSocket.OPEN) return; + if (_connectingWs) return; + _connectingWs = true; try { - ws = new WebSocket(wsUrl); - } catch (e) { - console.error('[Flow Agent] WS connect error:', e); - scheduleReconnect(); - return; - } - - ws.onopen = async () => { - console.log('[Flow Agent] Connected to agent: ' + wsUrl); - chrome.alarms.clear('reconnect'); - setState('idle'); - - const storage = await chrome.storage.local.get(['clientId']); - let clientId = storage.clientId; - if (!clientId) { - const prefix = CONFIG.DEFAULT_CLIENT_ID_PREFIX || 'client'; - clientId = `${prefix}-${Math.random().toString(36).substring(2, 8)}`; - await chrome.storage.local.set({ clientId }); - } - extensionClientId = clientId; + const data = await chrome.storage.local.get(['clientId']); + const serverIp = CONFIG.DEFAULT_SERVER_HOST; + connectedServerHost = serverIp; + const isLocal = /^(127\.0\.0\.1|localhost|192\.168\.|10\.)/.test(serverIp); + const wsScheme = isLocal ? 'ws' : 'wss'; + const httpScheme = isLocal ? 'http' : 'https'; + const wsUrl = `${wsScheme}://${serverIp}/ws`; + + // Dynamically resolve callbackUrl + callbackUrl = `${httpScheme}://${serverIp}/api/ext/callback`; + + const socket = new WebSocket(wsUrl); + ws = socket; + + socket.onopen = async () => { + _connectingWs = false; + if (ws !== socket) return; + console.log('[Flow Agent] Connected to agent: ' + wsUrl); + chrome.alarms.clear('reconnect'); + setState('idle'); - // Send current state + resend token if we have one, along with clientId - ws.send(JSON.stringify({ - type: 'extension_ready', - clientId: clientId, - flowKeyPresent: !!flowKey, - tokenAge: flowKey && metrics.tokenCapturedAt ? Date.now() - metrics.tokenCapturedAt : null, - })); - if (flowKey) { - ws.send(JSON.stringify({ - type: 'token_captured', - clientId: clientId, - flowKey: flowKey - })); - } - // Backend is reachable again — push any responses queued while it was down. - flushOutbox(); - }; + const storage = await chrome.storage.local.get(['clientId']); + let clientId = storage.clientId; + if (!clientId) { + const prefix = CONFIG.DEFAULT_CLIENT_ID_PREFIX || 'client'; + clientId = `${prefix}-${Math.random().toString(36).substring(2, 8)}`; + await chrome.storage.local.set({ clientId }); + } + extensionClientId = clientId; + + // Ensure socket is still open before sending + if (socket.readyState === WebSocket.OPEN) { + socket.send(JSON.stringify({ + type: 'extension_ready', + clientId: clientId, + flowKeyPresent: !!flowKey, + tokenAge: flowKey && metrics.tokenCapturedAt ? Date.now() - metrics.tokenCapturedAt : null, + })); + if (flowKey) { + socket.send(JSON.stringify({ + type: 'token_captured', + clientId: clientId, + flowKey: flowKey + })); + } + } + // Backend is reachable again — push any responses queued while it was down. + flushOutbox(); + }; - ws.onmessage = async ({ data }) => { - try { - const msg = JSON.parse(data); - - if (msg.method === 'api_request') { - await handleApiRequest(msg); - } else if (msg.method === 'get_media_url') { - await handleGetMediaUrl(msg); - } else if (msg.method === 'trpc_request') { - await handleTrpcRequest(msg); - } else if (msg.method === 'upload_video') { - await handleUploadVideo(msg); - } else if (msg.method === 'solve_captcha') { - await handleSolveCaptcha(msg); - } else if (msg.method === 'get_status') { - sendToAgent({ - id: msg.id, - result: { - state, - flowKeyPresent: !!flowKey, - manualDisconnect, - tokenAge: metrics.tokenCapturedAt ? Date.now() - metrics.tokenCapturedAt : null, - metrics, - }, - }); - } else if (msg.method === 'open_flow_tab') { - // Python bridge asks us to open/focus a Flow tab - // If token is still fresh, just send it back — no need to open/reload - if (isTokenFresh()) { - console.log('[Flow Agent] open_flow_tab: token fresh, sending cached token'); - sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); - } else { - console.log('[Flow Agent] open_flow_tab: token missing/expired, opening tab'); - const tabs = await chrome.tabs.query({ - url: ['https://labs.google/fx/tools/flow*', 'https://labs.google/fx/*/tools/flow*'], + socket.onmessage = async ({ data }) => { + try { + const msg = JSON.parse(data); + + if (msg.method === 'api_request') { + await handleApiRequest(msg); + } else if (msg.method === 'get_media_url') { + await handleGetMediaUrl(msg); + } else if (msg.method === 'trpc_request') { + await handleTrpcRequest(msg); + } else if (msg.method === 'upload_video') { + await handleUploadVideo(msg); + } else if (msg.method === 'solve_captcha') { + await handleSolveCaptcha(msg); + } else if (msg.method === 'get_status') { + sendToAgent({ + id: msg.id, + result: { + state, + flowKeyPresent: !!flowKey, + manualDisconnect, + tokenAge: metrics.tokenCapturedAt ? Date.now() - metrics.tokenCapturedAt : null, + metrics, + }, }); - if (tabs.length) { - await chrome.tabs.reload(tabs[0].id); - console.log('[Flow Agent] Refreshed existing Flow tab'); - } else { - await chrome.tabs.create({ url: 'https://labs.google/fx/tools/flow', active: true }); - console.log('[Flow Agent] Opened new Flow tab'); - } - await sleep(5000); - if (flowKey && ws?.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify({ type: 'token_captured', flowKey })); - console.log('[Flow Agent] Sent token after tab open'); + } else if (msg.method === 'open_flow_tab') { + // Python bridge asks us to open/focus a Flow tab + // If token is still fresh, just send it back — no need to open/reload + if (isTokenFresh()) { + console.log('[Flow Agent] open_flow_tab: token fresh, sending cached token'); + sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); } else { - const data = await chrome.storage.local.get(['flowKey']); - if (data.flowKey) { - flowKey = data.flowKey; - if (ws?.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify({ type: 'token_captured', flowKey })); - console.log('[Flow Agent] Sent token from storage after tab open'); + console.log('[Flow Agent] open_flow_tab: token missing/expired, opening tab'); + const tabs = await chrome.tabs.query({ + url: ['https://labs.google/fx/tools/flow*', 'https://labs.google/fx/*/tools/flow*'], + }); + if (tabs.length) { + await chrome.tabs.reload(tabs[0].id); + console.log('[Flow Agent] Refreshed existing Flow tab'); + } else { + await chrome.tabs.create({ url: 'https://labs.google/fx/tools/flow', active: true }); + console.log('[Flow Agent] Opened new Flow tab'); + } + await sleep(5000); + if (flowKey && socket.readyState === WebSocket.OPEN) { + socket.send(JSON.stringify({ type: 'token_captured', flowKey })); + console.log('[Flow Agent] Sent token after tab open'); + } else { + const data = await chrome.storage.local.get(['flowKey']); + if (data.flowKey) { + flowKey = data.flowKey; + if (socket.readyState === WebSocket.OPEN) { + socket.send(JSON.stringify({ type: 'token_captured', flowKey })); + console.log('[Flow Agent] Sent token from storage after tab open'); + } } } } - } - } else if (msg.method === 'refresh_flow_tab' || msg.method === 'force_refresh') { - // Python bridge asks us to refresh token. - // force_refresh (or an explicit msg.force) bypasses the freshness check: - // Google can invalidate a token via inactivity long before its 50-min - // age limit, so a "fresh" token may still be dead (401). In that case we - // must actually reload the tab and re-capture, not resend the cached one. - const force = msg.force === true || msg.method === 'force_refresh'; - if (isTokenFresh() && !force) { - console.log('[Flow Agent] refresh_flow_tab: token fresh, sending cached token'); - if (ws?.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify({ type: 'token_captured', flowKey })); - } - } else { - console.log('[Flow Agent] refresh_flow_tab: forcing tab reload + re-capture'); - // Drop the stale token so captureTokenFromFlowTab can't short-circuit. - if (force) { - flowKey = null; - metrics.tokenCapturedAt = null; - } - await captureTokenFromFlowTab(); - await sleep(3000); - if (flowKey && ws?.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify({ type: 'token_captured', flowKey })); - console.log('[Flow Agent] Sent token after refresh'); + } else if (msg.method === 'refresh_flow_tab' || msg.method === 'force_refresh') { + // Python bridge asks us to refresh token. + // force_refresh (or an explicit msg.force) bypasses the freshness check: + // Google can invalidate a token via inactivity long before its 50-min + // age limit, so a "fresh" token may still be dead (401). In that case we + // must actually reload the tab and re-capture, not resend the cached one. + const force = msg.force === true || msg.method === 'force_refresh'; + if (isTokenFresh() && !force) { + console.log('[Flow Agent] refresh_flow_tab: token fresh, sending cached token'); + if (socket.readyState === WebSocket.OPEN) { + socket.send(JSON.stringify({ type: 'token_captured', flowKey })); + } } else { - const data = await chrome.storage.local.get(['flowKey']); - if (data.flowKey) { - flowKey = data.flowKey; - if (ws?.readyState === WebSocket.OPEN) { - ws.send(JSON.stringify({ type: 'token_captured', flowKey })); + console.log('[Flow Agent] refresh_flow_tab: forcing tab reload + re-capture'); + // Drop the stale token so captureTokenFromFlowTab can't short-circuit. + if (force) { + flowKey = null; + metrics.tokenCapturedAt = null; + } + await captureTokenFromFlowTab(force); + await sleep(3000); + if (flowKey) { + sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); + console.log('[Flow Agent] Sent token after refresh'); + } else { + const data = await chrome.storage.local.get(['flowKey']); + if (data.flowKey) { + flowKey = data.flowKey; + sendToAgent({ type: 'token_captured', flowKey, clientId: extensionClientId }); console.log('[Flow Agent] Sent token from storage after refresh'); } } } + } else if (msg.type === 'callback_config') { + callbackSecret = msg.secret; + callbackUrl = normalizeCallbackUrl(msg.callback_url); + chrome.storage.local.set({ callbackSecret: msg.secret, callbackUrl }); + console.log('[Flow Agent] Received callback config:', callbackUrl); + } else if (msg.type === 'callback_secret') { + callbackSecret = msg.secret; + chrome.storage.local.set({ callbackSecret: msg.secret }); + console.log('[Flow Agent] Received callback secret'); + } else if (msg.type === 'pong') { + // keepalive response } - } else if (msg.type === 'callback_config') { - callbackSecret = msg.secret; - callbackUrl = normalizeCallbackUrl(msg.callback_url); - chrome.storage.local.set({ callbackSecret: msg.secret, callbackUrl }); - console.log('[Flow Agent] Received callback config:', callbackUrl); - } else if (msg.type === 'callback_secret') { - callbackSecret = msg.secret; - chrome.storage.local.set({ callbackSecret: msg.secret }); - console.log('[Flow Agent] Received callback secret'); - } else if (msg.type === 'pong') { - // keepalive response + } catch (e) { + console.error('[Flow Agent] Message error:', e); } - } catch (e) { - console.error('[Flow Agent] Message error:', e); - } - }; + }; - ws.onclose = () => { - setState('off'); - if (!manualDisconnect) scheduleReconnect(); - }; + socket.onclose = () => { + _connectingWs = false; + if (ws === socket) ws = null; + setState('off'); + if (!manualDisconnect) scheduleReconnect(); + }; - ws.onerror = (e) => { - console.error('[Flow Agent] WS error:', e); - metrics.lastError = 'WS_ERROR'; - chrome.storage.local.set({ metrics }); - }; + socket.onerror = (e) => { + _connectingWs = false; + console.error('[Flow Agent] WS error:', e); + metrics.lastError = 'WS_ERROR'; + chrome.storage.local.set({ metrics }); + }; + } catch (e) { + _connectingWs = false; + console.error('[Flow Agent] WS connect error:', e); + scheduleReconnect(); + } } function agentHttpBase() { @@ -837,6 +861,45 @@ async function handleUploadVideo(msg) { } } +let flowMusicKey = null; + +async function getOrCaptureFlowMusicToken() { + if (flowMusicKey) return flowMusicKey; + try { + const tabs = await chrome.tabs.query({ + url: ['https://www.flowmusic.app/*', 'https://flowmusic.app/*'], + }); + if (tabs.length && tabs[0].id) { + const results = await chrome.scripting.executeScript({ + target: { tabId: tabs[0].id }, + func: () => { + for (let i = 0; i < localStorage.length; i++) { + const k = localStorage.key(i); + if (k && (k.includes('auth-token') || k.startsWith('sb-'))) { + try { + const parsed = JSON.parse(localStorage.getItem(k)); + if (parsed?.access_token) return parsed.access_token; + if (parsed?.token) return parsed.token; + } catch {} + } + } + return null; + }, + }); + const token = results?.[0]?.result; + if (token) { + flowMusicKey = token; + chrome.storage.local.set({ flowMusicKey }); + console.log('[Flow Agent] FlowMusic token extracted from tab localStorage'); + return flowMusicKey; + } + } + } catch (e) { + console.warn('[Flow Agent] Error extracting FlowMusic token:', e); + } + return flowMusicKey; +} + async function handleApiRequest(msg) { const { id, params } = msg; const { url, method, headers, body, captchaAction } = params; @@ -846,7 +909,10 @@ async function handleApiRequest(msg) { return; } - if (!url.startsWith('https://aisandbox-pa.googleapis.com/')) { + const isFlowMusic = url.startsWith('https://www.flowmusic.app/') || url.startsWith('https://flowmusic.app/'); + const isGoogleAi = url.startsWith('https://aisandbox-pa.googleapis.com/'); + + if (!isGoogleAi && !isFlowMusic) { sendToAgent({ id, error: 'INVALID_URL' }); return; } @@ -863,6 +929,89 @@ async function handleApiRequest(msg) { } try { + if (isFlowMusic) { + if (!flowMusicKey) { + await getOrCaptureFlowMusicToken(); + } + const fetchHeaders = { + 'content-type': 'application/json', + 'accept': '*/*', + }; + if (flowMusicKey) { + fetchHeaders['authorization'] = `Bearer ${flowMusicKey}`; + } + + // Special handling for FlowMusic SSE event stream + if (url.includes('/__api/messages/') && url.includes('/stream')) { + const streamResp = await fetch(url, { + method: method || 'GET', + headers: fetchHeaders, + credentials: 'include', + }); + if (!streamResp.ok) { + sendToAgent({ id, status: streamResp.status, error: `STREAM_FAILED_${streamResp.status}` }); + setState('idle'); + return; + } + const reader = streamResp.body.getReader(); + const decoder = new TextDecoder(); + let streamBuffer = ''; + let clipIds = []; + const startTime = Date.now(); + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + const chunk = decoder.decode(value, { stream: true }); + streamBuffer += chunk; + + const matches = streamBuffer.matchAll(/"clip_id":\s*"([0-9a-f-]{36})"/g); + for (const m of matches) { + if (!clipIds.includes(m[1])) clipIds.push(m[1]); + } + + if (streamBuffer.includes('event: complete') || streamBuffer.includes('event: final')) { + break; + } + if (Date.now() - startTime > 90000) { + break; + } + } + sendToAgent({ + id, + status: 200, + data: { clip_ids: clipIds } + }); + setState('idle'); + return; + } + + // Standard FlowMusic API calls (conversation, clips, event) + const response = await fetch(url, { + method: method || 'POST', + headers: fetchHeaders, + credentials: 'include', + body: method === 'GET' ? undefined : JSON.stringify(body), + }); + + let responseData; + const responseText = await response.text(); + try { + responseData = JSON.parse(responseText); + } catch { + responseData = responseText; + } + + sendToAgent({ + id, + status: response.status, + data: responseData, + }); + setState('idle'); + return; + } + + // Google AI Sandbox requests // Step 1: Solve captcha if needed let captchaToken = null; if (captchaAction) { @@ -927,10 +1076,7 @@ async function handleApiRequest(msg) { responseData = responseText; } - // Self-heal: a 401 means Google invalidated our cached token (usually via - // inactivity, before our 50-min freshness window). Drop it so the very next - // request / refresh forces a genuine tab reload + re-capture instead of - // resending the same dead token. + // Self-heal: a 401 means Google invalidated our cached token if (response.status === 401) { console.warn('[Flow Agent] 401 UNAUTHENTICATED — invalidating cached token to force refresh'); flowKey = null; @@ -960,10 +1106,9 @@ async function handleApiRequest(msg) { }); if (hasCaptcha) { metrics.failedCount++; metrics.lastError = e.message; } updateRequestLog(logId, { status: 'failed', error: e.message || 'API_REQUEST_FAILED' }); + chrome.storage.local.set({ metrics }); + setState('idle'); } - - chrome.storage.local.set({ metrics }); - setState('idle'); } async function handleGetMediaUrl(msg) { diff --git a/flow-extension/manifest.json b/flow-extension/manifest.json index 3436cac..5d2b00c 100644 --- a/flow-extension/manifest.json +++ b/flow-extension/manifest.json @@ -23,6 +23,8 @@ "https://aisandbox-pa.sandbox.googleapis.com/*", "https://storage.googleapis.com/*", "https://flow-content.google/*", + "https://www.flowmusic.app/*", + "https://flowmusic.app/*", "http://127.0.0.1:8001/*", "http://localhost:8001/*", "https://flow.kodelyx.in/*" @@ -34,7 +36,11 @@ { "matches": [ "https://labs.google/fx/tools/flow*", - "https://labs.google/fx/*/tools/flow*" + "https://labs.google/fx/*/tools/flow*", + "https://labs.google/fx/tools/music-fx*", + "https://labs.google/fx/*/tools/music-fx*", + "https://www.flowmusic.app/*", + "https://flowmusic.app/*" ], "js": [ "content.js" @@ -48,7 +54,9 @@ "injected.js" ], "matches": [ - "https://labs.google/*" + "https://labs.google/*", + "https://www.flowmusic.app/*", + "https://flowmusic.app/*" ] } ],