From 4b2f8122e37bd53f8bb28716987609f2d9601e7d Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 8 Oct 2026 07:31:09 -1000
Subject: [PATCH] [udp] Keep constant write payloads in shared flash tables
 (#20268)

---
 esphome/components/udp/__init__.py            | 15 ++------
 esphome/components/udp/automation.h           | 34 ++++---------------
 .../component_tests/udp/config/udp_write.yaml | 22 ++++++++++++
 tests/component_tests/udp/test_udp_write.py   | 17 ++++++++++
 4 files changed, 49 insertions(+), 39 deletions(-)
 create mode 100644 tests/component_tests/udp/config/udp_write.yaml
 create mode 100644 tests/component_tests/udp/test_udp_write.py

diff --git a/esphome/components/udp/__init__.py b/esphome/components/udp/__init__.py
index 316d4fc311..cd3cb44ae1 100644
--- a/esphome/components/udp/__init__.py
+++ b/esphome/components/udp/__init__.py
@@ -184,16 +184,7 @@ async def udp_write_to_code(
     udp_var = await cg.get_variable(config[CONF_ID])
     await cg.register_parented(var, udp_var)
     cg.add(udp_var.set_should_broadcast())
-    data = config[CONF_DATA]
-    if isinstance(data, bytes):
-        data = list(data)
-
-    if cg.is_template(data):
-        templ = await cg.templatable(data, args, cg.std_vector.template(cg.uint8))
-        cg.add(var.set_data_template(templ))
-    else:
-        # Generate static array in flash to avoid RAM copy
-        arr_id = ID(f"{action_id}_data", is_declaration=True, type=cg.uint8)
-        arr = cg.static_const_array(arr_id, cg.ArrayInitializer(*data))
-        cg.add(var.set_data_static(arr, len(data)))
+    await automation.templatable_bytes(
+        config[CONF_DATA], args, var.set_data_template, var.set_data_static, "udp_data"
+    )
     return var
diff --git a/esphome/components/udp/automation.h b/esphome/components/udp/automation.h
index 9b3ba7e0dd..a6a40829b1 100644
--- a/esphome/components/udp/automation.h
+++ b/esphome/components/udp/automation.h
@@ -4,39 +4,19 @@
 #ifdef USE_NETWORK
 #include "esphome/core/automation.h"
 
-#include <vector>
-
 namespace esphome::udp {
 
-template<typename... Ts> class UDPWriteAction final : public Action<Ts...>, public Parented<UDPComponent> {
- public:
-  void set_data_template(std::vector<uint8_t> (*func)(Ts...)) {
-    this->data_.func = func;
-    this->len_ = -1;  // Sentinel value indicates template mode
-  }
+/// Constant payloads up to this size are sent from the stack on ESP8266.
+static constexpr size_t STACK_PAYLOAD_SIZE = 128;
 
-  void set_data_static(const uint8_t *data, size_t len) {
-    this->data_.data = data;
-    this->len_ = len;  // Length >= 0 indicates static mode
-  }
+template<typename... Ts> class UDPWriteAction final : public Action<Ts...>, public Parented<UDPComponent> {
+  TEMPLATABLE_BYTES(data)
 
   void play(const Ts &...x) override {
-    if (this->len_ >= 0) {
-      // Static mode: pass pointer directly to send_packet(const uint8_t *, size_t)
-      this->parent_->send_packet(this->data_.data, static_cast<size_t>(this->len_));
-    } else {
-      // Template mode: call function and pass vector to send_packet(const std::vector<uint8_t> &)
-      auto val = this->data_.func(x...);
-      this->parent_->send_packet(val);
-    }
+    // One write per packet: WiFiUDP appends all or nothing, so a failed allocation never sends a partial packet
+    this->data_.template visit<STACK_PAYLOAD_SIZE>(
+        [this](const uint8_t *data, size_t len) { this->parent_->send_packet(data, len); }, x...);
   }
-
- protected:
-  ssize_t len_{-1};  // -1 = template mode, >=0 = static mode with length
-  union Data {
-    std::vector<uint8_t> (*func)(Ts...);  // Function pointer (stateless lambdas)
-    const uint8_t *data;                  // Pointer to static data in flash
-  } data_;
 };
 
 }  // namespace esphome::udp
diff --git a/tests/component_tests/udp/config/udp_write.yaml b/tests/component_tests/udp/config/udp_write.yaml
new file mode 100644
index 0000000000..645879bfcd
--- /dev/null
+++ b/tests/component_tests/udp/config/udp_write.yaml
@@ -0,0 +1,22 @@
+esphome:
+  name: test
+
+esp8266:
+  board: d1_mini
+
+wifi:
+  ssid: test
+  password: password1
+
+udp:
+  addresses:
+    - 10.0.0.1
+
+button:
+  - platform: template
+    name: send
+    on_press:
+      - udp.write: [0xAA, 0x55, 0x01, 0x02]
+      - udp.write: [0xAA, 0x55, 0x01, 0x02]
+      - udp.write: ""
+      - udp.write: !lambda return {0x10, 0x20};
diff --git a/tests/component_tests/udp/test_udp_write.py b/tests/component_tests/udp/test_udp_write.py
new file mode 100644
index 0000000000..d999891d95
--- /dev/null
+++ b/tests/component_tests/udp/test_udp_write.py
@@ -0,0 +1,17 @@
+"""Tests for the udp.write action codegen."""
+
+from collections.abc import Callable
+from pathlib import Path
+
+
+def test_write_payloads_share_a_progmem_table(
+    generate_main: Callable[[str | Path], str],
+    component_config_path: Callable[[str], Path],
+) -> None:
+    """Identical constant payloads share one flash table; an empty one needs none."""
+    main_cpp = generate_main(component_config_path("udp_write.yaml"))
+
+    assert main_cpp.count("udp_data[] PROGMEM = {0xAA, 0x55, 0x01, 0x02};") == 1
+    assert main_cpp.count("->set_data_static(udp_data, 4);") == 2
+    assert "->set_data_static(nullptr, 0);" in main_cpp
+    assert "->set_data_template(" in main_cpp