Assistant/backend/DeviceHub/DeviceHubEndpoint.cs

using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using backend.Conversations;
using backend.Data;
using backend.Devices;
using backend.Realtime;
using backend.Tools;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;

namespace backend.DeviceHub;

public static class DeviceHubEndpoint
{
    public static IEndpointRouteBuilder MapDeviceHub(this IEndpointRouteBuilder app)
    {
        app.Map("/device", async (
            HttpContext ctx,
            DeviceAuth auth,
            DeviceRegistry reg,
            AppDbContext db,
            RealtimeSessionFactory sessions,
            IServiceScopeFactory scopes) =>
        {
            if (!ctx.WebSockets.IsWebSocketRequest)
            {
                ctx.Response.StatusCode = StatusCodes.Status400BadRequest;
                return;
            }

            var device = await auth.AuthenticateAsync(
                ctx.Request.Headers.Authorization.ToString(), ctx.RequestAborted);
            if (device is null)
            {
                ctx.Response.StatusCode = StatusCodes.Status401Unauthorized;
                return;
            }

            using var ws = await ctx.WebSockets.AcceptWebSocketAsync();
            await using var active = new ActiveDevice(device.Id, ws);
            if (!reg.TryRegister(active))
            {
                await ws.CloseAsync(WebSocketCloseStatus.PolicyViolation,
                    "device already connected", ctx.RequestAborted);
                return;
            }

            try
            {
                await HandleAsync(active, device, db, sessions, scopes, ctx.RequestAborted);
            }
            finally
            {
                reg.Unregister(device.Id);
            }
        });
        return app;
    }

    private static async Task HandleAsync(
        ActiveDevice active, Device device, AppDbContext db,
        RealtimeSessionFactory sessions, IServiceScopeFactory scopes,
        CancellationToken ct)
    {
        var buffer = new byte[8192];
        var ws = active.Ws;

        while (ws.State == WebSocketState.Open && !ct.IsCancellationRequested)
        {
            using var msg = new MemoryStream();
            WebSocketReceiveResult res;
            while (true)
            {
                res = await ws.ReceiveAsync(buffer, ct);
                if (res.MessageType == WebSocketMessageType.Close) return;
                msg.Write(buffer, 0, res.Count);
                if (res.EndOfMessage) break;
            }

            if (res.MessageType == WebSocketMessageType.Binary)
            {
                if (msg.Length == 3840 && active.CurrentSession is not null)
                {
                    await active.CurrentSession.WriteUplinkAsync(msg.ToArray(), ct);
                }
                continue;
            }

            var text = Encoding.UTF8.GetString(msg.ToArray());
            using var doc = JsonDocument.Parse(text);
            var type = doc.RootElement.GetProperty("type").GetString();
            switch (type)
            {
                case "hello":
                    await SendHelloAckAsync(active, device, ct);
                    break;

                case "ping":
                    device.LastSeenAt = DateTimeOffset.UtcNow;
                    await db.SaveChangesAsync(ct);
                    await active.SendEnvelopeAsync(new PongEnvelope("pong"), ct);
                    break;

                case "wake":
                    if (active.CurrentSession is null)
                        await StartSessionAsync(active, device, sessions, scopes, ct);
                    break;

                case "cancel":
                    active.SessionCts?.Cancel();
                    break;

                case "playback_done":
                    break; // v1: client coordination only

                case "tool_result":
                    var env = JsonSerializer.Deserialize<ToolResultEnvelope>(text)!;
                    active.CompleteToolResult(env);
                    break;
            }
        }
    }

    private static async Task StartSessionAsync(
        ActiveDevice active, Device device,
        RealtimeSessionFactory sessions, IServiceScopeFactory scopes,
        CancellationToken parentCt)
    {
        var enabledTools = JsonSerializer.Deserialize<string[]>(device.Config?.EnabledToolsJson ?? "[]")
            ?? Array.Empty<string>();
        var cfg = new RealtimeSettings(
            Model: device.Config?.Model ?? "gpt-4o-realtime-preview",
            Voice: device.Config?.Voice ?? "alloy",
            SystemPrompt: device.Config?.SystemPrompt ?? "",
            IdleTimeoutSeconds: device.Config?.IdleTimeoutSeconds ?? 30,
            EnabledTools: new HashSet<string>(enabledTools));

        active.SessionCts = CancellationTokenSource.CreateLinkedTokenSource(parentCt);
        var ct = active.SessionCts.Token;

        var sessionScope = scopes.CreateAsyncScope();
        var scopedSessions = sessionScope.ServiceProvider.GetRequiredService<RealtimeSessionFactory>();

        RealtimeSession session;
        try
        {
            session = await scopedSessions.CreateAsync(
                device.Id, cfg, active, active, ct);
        }
        catch (Exception ex)
        {
            await active.SendEnvelopeAsync(
                new ErrorEnvelope("error", "upstream", ex.Message, true), parentCt);
            await sessionScope.DisposeAsync();
            return;
        }

        active.CurrentSession = session;
        _ = Task.Run(async () =>
        {
            try { await session.RunAsync(ct); }
            finally
            {
                active.CurrentSession = null;
                active.SessionCts?.Dispose();
                active.SessionCts = null;
                await sessionScope.DisposeAsync();
            }
        }, ct);
    }

    private static async Task SendHelloAckAsync(ActiveDevice active, Device device, CancellationToken ct)
    {
        var enabled = device.Config?.EnabledToolsJson ?? "[]";
        var config = JsonNode.Parse($$"""
        {
          "voice": "{{Escape(device.Config?.Voice)}}",
          "model": "{{Escape(device.Config?.Model)}}",
          "system_prompt": "{{Escape(device.Config?.SystemPrompt)}}",
          "idle_timeout_seconds": {{device.Config?.IdleTimeoutSeconds ?? 30}},
          "enabled_tools": {{enabled}}
        }
        """)!.AsObject();

        await active.WriteEnvelopeAsync(new JsonObject
        {
            ["type"] = "hello_ack",
            ["config"] = config,
        }, ct);
    }

    private static string Escape(string? s) =>
        (s ?? "").Replace("\\", "\\\\").Replace("\"", "\\\"");
}