Add high-speed UART baud rate switching

agent/main.c

@@ -288,6 +288,58 @@ static void handle_selfupdate(const uint8_t *data, uint32_t len) {
     /* Never returns */
 }
 
+/*
+ * CMD_SET_BAUD: change UART baud rate.
+ *   Host sends: CMD_SET_BAUD [baud_rate:4LE]
+ *   Agent ACKs at current baud, waits for TX to drain, switches.
+ *   Host should switch immediately after receiving the ACK.
+ */
+static int at_default_baud = 1;
+
+static void handle_set_baud(const uint8_t *data, uint32_t len) {
+    if (len < 4) { proto_send_ack(ACK_CRC_ERROR); return; }
+
+    uint32_t baud = read_le32(&data[0]);
+
+    /* Sanity check: reject absurd baud rates */
+    if (baud < 9600 || baud > 3000000) {
+        proto_send_ack(ACK_FLASH_ERROR);
+        return;
+    }
+
+    /* ACK at current baud rate so host knows we accepted */
+    proto_send_ack(ACK_OK);
+
+    /* Switch to new baud (uart_set_baud waits for TX drain) */
+    uart_set_baud(baud);
+
+    /* Drain any garbage from baud rate transition */
+    while (uart_readable()) uart_getc();
+
+    /* Wait for host to confirm with any valid command within 3 seconds.
+     * If nothing arrives, revert to 115200 — the host may have failed
+     * to switch or the new baud rate doesn't work on this link. */
+    uint8_t pkt[MAX_PAYLOAD + 16];
+    uint32_t pkt_len = 0;
+    uint8_t cmd = proto_recv(pkt, &pkt_len, 3000);
+    if (cmd == 0) {
+        /* No valid command — revert */
+        uart_set_baud(115200);
+        while (uart_readable()) uart_getc();
+        at_default_baud = 1;
+    } else {
+        /* Got a valid command at new baud — confirmed working */
+        at_default_baud = (baud == 115200);
+        switch (cmd) {
+            case CMD_INFO:  handle_info(); break;
+            case CMD_READ:  handle_read(pkt, pkt_len); break;
+            case CMD_WRITE: handle_write(pkt, pkt_len); break;
+            case CMD_CRC32: handle_crc32_cmd(pkt, pkt_len); break;
+            default:        proto_send_ack(ACK_OK); break;
+        }
+    }
+}
+
 int main(void) {
     watchdog_disable();
     uart_init();
@@ -298,22 +350,33 @@ int main(void) {
     proto_send_ready();
 
     uint32_t idle_count = 0;
+    uint32_t baud_idle = 0;
+    at_default_baud = 1;
     while (1) {
         uint32_t data_len = 0;
         uint8_t cmd = proto_recv(cmd_buf, &data_len, 500);
 
         if (cmd == 0) {
             idle_count++;
-            /* Send READY every ~2s (4 x 500ms) so host can detect us
-             * after reconnect. Suppress briefly after a command to
-             * avoid interfering with multi-packet responses. */
             if (idle_count >= 4) {
                 proto_send_ready();
                 idle_count = 0;
             }
+            /* If at non-default baud and idle for ~10s (20 x 500ms),
+             * revert to 115200. Host may have disconnected. */
+            if (!at_default_baud) {
+                baud_idle++;
+                if (baud_idle >= 20) {
+                    uart_set_baud(115200);
+                    while (uart_readable()) uart_getc();
+                    at_default_baud = 1;
+                    baud_idle = 0;
+                }
+            }
             continue;
         }
         idle_count = 0;
+        baud_idle = 0;
 
         switch (cmd) {
             case CMD_INFO:
@@ -331,6 +394,9 @@ int main(void) {
             case CMD_SELFUPDATE:
                 handle_selfupdate(cmd_buf, data_len);
                 break;
+            case CMD_SET_BAUD:
+                handle_set_baud(cmd_buf, data_len);
+                break;
             case CMD_REBOOT:
                 /* Trigger reset via watchdog */
                 WDT_LOCK = WDT_UNLOCK_KEY;

agent/protocol.h

@@ -15,6 +15,7 @@
 #define CMD_CRC32   0x05
 #define CMD_REBOOT      0x06
 #define CMD_SELFUPDATE  0x07
+#define CMD_SET_BAUD    0x08
 
 /* Responses (device → host) */
 #define RSP_INFO    0x81

agent/uart.c

@@ -38,6 +38,29 @@ void uart_init(void) {
     uart_reg(UART_IMSC) = 0;
 }
 
+void uart_set_baud(uint32_t baud) {
+    /* Wait for TX FIFO to drain */
+    while (!(uart_reg(UART_FR) & UART_FR_TXFE)) {}
+    while (uart_reg(UART_FR) & UART_FR_BUSY) {}
+
+    /* Disable UART before changing baud */
+    uart_reg(UART_CR) = 0;
+
+    /* Recompute divisors: IBRD = clock / (16 * baud), FBRD = frac * 64 */
+    uint32_t divisor = UART_CLOCK / (16 * baud);
+    uint32_t remainder = UART_CLOCK % (16 * baud);
+    uint32_t frac = (remainder * 64 + (16 * baud) / 2) / (16 * baud);
+    uart_reg(UART_IBRD) = divisor;
+    uart_reg(UART_FBRD) = frac & 0x3F;
+
+    /* Must write LCR_H after baud rate to latch the divisors */
+    uart_reg(UART_LCR_H) = UART_LCR_WLEN8 | UART_LCR_FEN;
+
+    /* Re-enable */
+    uart_reg(UART_CR) = UART_CR_UARTEN | UART_CR_TXE | UART_CR_RXE;
+    uart_reg(UART_ICR) = 0x7FF;
+}
+
 void uart_putc(uint8_t ch) {
     /* Wait until TX FIFO has space, with timeout */
     volatile uint32_t timeout = 100000;

agent/uart.h

@@ -59,6 +59,7 @@
 #define uart_reg(off)   (*(volatile uint32_t *)(UART_BASE + (off)))
 
 void uart_init(void);
+void uart_set_baud(uint32_t baud);
 void uart_putc(uint8_t ch);
 int uart_putc_safe(uint8_t ch);
 uint8_t uart_getc(void);

src/defib/agent/client.py

@@ -19,6 +19,7 @@
     CMD_READ,
     CMD_REBOOT,
     CMD_SELFUPDATE,
+    CMD_SET_BAUD,
     CMD_WRITE,
     RSP_ACK,
     RSP_CRC32,
@@ -35,11 +36,16 @@
 logger = logging.getLogger(__name__)
 
 # Max bytes per WRITE transfer before PL011 FIFO overflow on uncached DDR.
-# Agent processes COBS+CRC between packets; at 115200 baud with no D-cache,
-# cumulative FIFO overflow occurs after ~30-60KB of continuous streaming.
 WRITE_CHUNK_SIZE = 512
 WRITE_MAX_TRANSFER = 16 * 1024
 
+# Optimal baud rate from hi3516ev300 + FT232R benchmarks.
+# 921600 is the highest rate verified for both READ and WRITE with
+# CRC32 integrity on real hardware. 1152000 works for READ-only but
+# WRITE is marginal. Single rate avoids unnecessary switching.
+DEFAULT_FAST_BAUD = 921600    # ~82 KB/s both directions
+FALLBACK_BAUD = 115200        # Always works
+
 
 def get_agent_binary(chip: str) -> Path | None:
     """Get the path to the pre-compiled agent binary for a chip."""
@@ -87,6 +93,7 @@ def __init__(self, transport: Transport, chip: str = "") -> None:
         self._flash_size = 0
         self._ram_base = 0
         self._sector_size = 0x10000
+        self._current_baud = FALLBACK_BAUD
 
     @property
     def connected(self) -> bool:
@@ -105,6 +112,22 @@ async def connect(self, timeout: float = 10.0) -> bool:
         self._connected = await wait_for_ready(self._transport, timeout)
         return self._connected
 
+    async def _switch_baud(self, target: int) -> bool:
+        """Switch to target baud if not already there. Returns success."""
+        if self._current_baud == target:
+            return True
+        ok = await self.set_baud(target)
+        if ok:
+            self._current_baud = target
+        return ok
+
+    async def _restore_baud(self) -> None:
+        """Return to fallback baud rate."""
+        if self._current_baud != FALLBACK_BAUD:
+            ok = await self.set_baud(FALLBACK_BAUD)
+            if ok:
+                self._current_baud = FALLBACK_BAUD
+
     async def get_info(self) -> dict[str, int | str]:
         """Request device info from the agent."""
         await send_packet(self._transport, CMD_INFO)
@@ -133,8 +156,16 @@ async def read_memory(
         addr: int,
         size: int,
         on_progress: Callable[[int, int], None] | None = None,
+        fast: bool = True,
     ) -> bytes:
-        """Read a memory region from the device. Returns bytes."""
+        """Read a memory region from the device. Returns bytes.
+
+        If fast=True and size > 4KB, switches to DEFAULT_FAST_BAUD,
+        falling back to FALLBACK_BAUD if the switch fails.
+        """
+        if fast and size > 4096:
+            await self._switch_baud(DEFAULT_FAST_BAUD)
+
         payload = struct.pack("<II", addr, size)
         await send_packet(self._transport, CMD_READ, payload)
 
@@ -171,12 +202,17 @@ async def write_memory(
         addr: int,
         data: bytes,
         on_progress: Callable[[int, int], None] | None = None,
+        fast: bool = True,
     ) -> bool:
         """Write data to device RAM in chunked transfers.
 
+        If fast=True and data > 4KB, switches to DEFAULT_FAST_BAUD.
         Splits into WRITE_MAX_TRANSFER-sized blocks to avoid PL011 FIFO
         overflow on uncached DDR. Each block uses WRITE_CHUNK_SIZE packets.
         """
+        if fast and len(data) > 4096:
+            await self._switch_baud(DEFAULT_FAST_BAUD)
+
         total = len(data)
         offset = 0
 
@@ -262,6 +298,49 @@ async def selfupdate(
         cmd, resp = await recv_response(self._transport, timeout=10.0)
         return cmd == RSP_ACK and resp[0] == ACK_OK
 
+    async def set_baud(self, baud: int) -> bool:
+        """Switch UART to a higher baud rate.
+
+        Protocol: send SET_BAUD command, receive ACK at current baud,
+        then both sides switch. Verifies with INFO at new baud.
+        Falls back to original baud on failure.
+        """
+        import asyncio
+
+        port = getattr(self._transport, '_port', None)
+        if port is None:
+            logger.error("set_baud requires serial transport with _port")
+            return False
+
+        old_baud = port.baudrate
+        payload = struct.pack("<I", baud)
+        await send_packet(self._transport, CMD_SET_BAUD, payload)
+
+        cmd, resp = await recv_response(self._transport, timeout=5.0)
+        if cmd != RSP_ACK or resp[0] != ACK_OK:
+            logger.error("Agent rejected baud rate %d", baud)
+            return False
+
+        # Agent has switched — now switch host side
+        await asyncio.sleep(0.05)  # Brief pause for agent to complete switch
+        port.baudrate = baud
+
+        # Verify communication at new baud
+        await asyncio.sleep(0.05)
+        try:
+            await send_packet(self._transport, CMD_INFO)
+            cmd, data = await recv_response(self._transport, timeout=3.0)
+            if cmd == RSP_INFO:
+                logger.info("Baud rate switched to %d", baud)
+                return True
+        except Exception:
+            pass
+
+        # Failed — switch back
+        logger.warning("Verification at %d baud failed, reverting to %d", baud, old_baud)
+        port.baudrate = old_baud
+        return False
+
     async def reboot(self) -> None:
         """Tell the agent to reset the device."""
         await send_packet(self._transport, CMD_REBOOT)

src/defib/agent/protocol.py

@@ -14,6 +14,7 @@
   0x05  CRC32    — CRC32 of flash region: addr(4B LE) + size(4B LE)
   0x06  REBOOT   — Reset device
   0x07  SELFUPDATE — Update agent: addr(4B LE) + size(4B LE) + crc32(4B LE)
+  0x08  SET_BAUD   — Change baud rate: baud(4B LE)
 
 Responses (device → host):
   0x81  INFO_RSP — chip_id(4B) + flash_size(4B) + ram_base(4B)
@@ -39,6 +40,7 @@
 CMD_CRC32 = 0x05
 CMD_REBOOT = 0x06
 CMD_SELFUPDATE = 0x07
+CMD_SET_BAUD = 0x08
 
 # Responses
 RSP_INFO = 0x81