File size: 5,013 Bytes
a8052e2
 
 
9a55333
0bd57b6
a8052e2
9a55333
8c553b3
9a55333
a8052e2
 
 
 
 
 
0bd57b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a8052e2
 
 
f6fd017
0bd57b6
a8052e2
 
0bd57b6
a8052e2
 
 
0bd57b6
a8052e2
 
 
 
 
 
 
 
 
0bd57b6
a8052e2
 
0bd57b6
a8052e2
0bd57b6
9a55333
 
a8052e2
 
 
 
 
8c553b3
0bd57b6
a8052e2
 
 
8c553b3
 
a8052e2
 
 
 
0bd57b6
 
 
a8052e2
 
 
 
 
 
 
 
0bd57b6
 
 
 
 
 
 
 
 
 
 
 
 
a8052e2
 
 
 
 
 
8c553b3
a8052e2
8c553b3
0bd57b6
 
 
 
 
a8052e2
0bd57b6
8c553b3
 
a8052e2
8c553b3
0bd57b6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import asyncio
import websockets
import json
import numpy as np
import traceback # Import the traceback module
from music_generator import MusicGenerator

# --- WebSocket Server Setup ---

# In-memory storage for connected clients
clients = {
    "webapp": set()
}
audio_source = None

# --- Main Application Logic ---

def initialize_dependencies():
    """Loads all necessary files and initializes objects."""
    global notes, generator
    print("Initializing dependencies...")
    try:
        # It's better to load files relative to the script's location
        import os
        dir_path = os.path.dirname(os.path.realpath(__file__))
        
        with open(os.path.join(dir_path, 'consonance_matrix.json')) as f:
            consonance_matrix = np.array(json.load(f))

        notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        generator = MusicGenerator(len(notes))
        print("Dependencies initialized successfully.")
        return True
    except Exception as e:
        print("--- CRITICAL: FAILED TO INITIALIZE DEPENDENCIES ---")
        print(traceback.format_exc())
        print("----------------------------------------------------")
        return False

async def broadcast_to_webapps(message):
    """Sends a message to all connected webapp clients."""
    if clients["webapp"]:
        tasks = [client.send(message) for client in clients["webapp"]]
        await asyncio.gather(*tasks, return_exceptions=True)

async def handle_audio_data(data):
    """Mocks the audio analysis."""
    import random
    detected_chord = random.choice(notes)
    predicted_chord = random.choice(notes)
    key = "C Major"

    analysis_result = {
        "type": "analysis_update",
        "current_chord": detected_chord,
        "predicted_chord": predicted_chord,
        "musical_key": key
    }
    await broadcast_to_webapps(json.dumps(analysis_result))

# --- WebSocket Connection Management ---

async def connection_handler(websocket, path):
    """Handles incoming WebSocket connections with robust error logging."""
    global audio_source
    print(f"New client connected: {websocket.remote_address}")

    try:
        initial_message = await websocket.recv()
        message_data = json.loads(initial_message)
        client_type = message_data.get("type")

        if client_type == "extension_hello":
            if audio_source is not None:
                await audio_source.close(code=1012, reason="New extension connected.")
            audio_source = websocket
            print("Audio capture extension connected.")
            await websocket.send(json.dumps({"status": "connected", "role": "audio_source"}))
            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source connected."}))

        elif client_type == "webapp_hello":
            clients["webapp"].add(websocket)
            print("Web app client connected.")
            await websocket.send(json.dumps({"status": "connected", "role": "viewer"}))
            status_msg = "Audio source connected." if audio_source else "Waiting for audio source..."
            await websocket.send(json.dumps({"type": "status_update", "message": status_msg}))
        
        else:
            print(f"Unknown client type: {client_type}. Disconnecting.")
            return

        async for message in websocket:
            if websocket == audio_source:
                await handle_audio_data(message)

    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed normally for {websocket.remote_address}. Code: {e.code}, Reason: {e.reason}")
    except Exception as e:
        # THIS IS THE CRITICAL PART
        print(f"--- UNEXPECTED SERVER ERROR in connection_handler ---")
        print(f"Error Type: {type(e).__name__}")
        print(f"Error Message: {e}")
        print("Traceback:")
        print(traceback.format_exc())
        print("----------------------------------------------------")
        # Close the connection with a specific error code if it's not already closed
        if not websocket.closed:
            await websocket.close(code=1011, reason="Internal Server Error")
    finally:
        if websocket in clients["webapp"]:
            clients["webapp"].remove(websocket)
        if websocket == audio_source:
            audio_source = None
            print("Audio capture extension disconnected.")
            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source disconnected."}))

async def main():
    """Initializes dependencies and starts the WebSocket server."""
    if not initialize_dependencies():
        print("Server cannot start due to initialization failure.")
        return

    websocket_port = 7860
    print(f"Starting WebSocket server on port {websocket_port}...")
    async with websockets.serve(connection_handler, "0.0.0.0", websocket_port):
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    asyncio.run(main())