---
phase: 25-python-tts-server-timestamp-endpoint
plan: 01
type: execute
wave: 1
depends_on: []
files_modified:
  - ~/.local/share/kokoro/kokoro_common.py
  - ~/.local/share/kokoro/tts_server.py
autonomous: true
requirements: [PTS-01, PTS-02, PTS-03]
must_haves:
  truths:
    - "curl localhost:8779/v1/audio/speech-with-timestamps with JSON body returns 200 with audio_b64, words array, audio_duration, and sample_rate"
    - "Each word in the words array has text, onset (float seconds), and duration (float seconds) derived from MToken.start_ts/end_ts"
    - "Punctuation-only tokens are filtered from the words array"
    - "Python server launchd plist has KeepAlive and RunAtLoad for automatic startup"
  artifacts:
    - path: "~/.local/share/kokoro/kokoro_common.py"
      provides: "synthesize_with_timestamps() returning audio + word timing data"
      contains: "synthesize_with_timestamps"
    - path: "~/.local/share/kokoro/tts_server.py"
      provides: "/v1/audio/speech-with-timestamps endpoint"
      contains: "speech-with-timestamps"
  key_links:
    - from: "tts_server.py _handle_speech_with_timestamps"
      to: "kokoro_common.synthesize_with_timestamps"
      via: "function call returning (audio_np, word_timings)"
      pattern: "synthesize_with_timestamps"
    - from: "kokoro_common.synthesize_with_timestamps"
      to: "mlx_audio result.tokens MToken.start_ts/end_ts"
      via: "iterating result.tokens after model.generate()"
      pattern: "start_ts.*end_ts"
---

<objective>
Add `/v1/audio/speech-with-timestamps` endpoint to the existing Python MLX Kokoro TTS server, returning JSON with base64 WAV audio and per-word onset/duration arrays from mlx-audio's native MToken timestamps.

Purpose: Swift companion (Phase 26) will consume this endpoint to drive karaoke subtitle highlighting with zero-drift word timing, replacing the character-weighted fallback that produces visible drift on multi-syllable words.

Output: Modified `kokoro_common.py` with `synthesize_with_timestamps()` and modified `tts_server.py` with the new endpoint. Launchd plist already has KeepAlive (PTS-03 satisfied by existing config).
</objective>

<execution_context>
@$HOME/.claude/get-shit-done/workflows/execute-plan.md
@$HOME/.claude/get-shit-done/templates/summary.md
</execution_context>

<context>
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/STATE.md
@.planning/REQUIREMENTS.md
@.planning/phases/25-python-tts-server-timestamp-endpoint/25-CONTEXT.md

<interfaces>
<!-- Key types and contracts the executor needs. Extracted from mlx-audio. -->

From misaki.en (MToken dataclass):

```python
@dataclass
class MToken:
    text: str           # The word text (e.g., "Hello")
    tag: str            # POS tag
    whitespace: str     # Trailing whitespace
    phonemes: Optional[str] = None
    start_ts: Optional[float] = None  # Onset in seconds (populated by join_timestamps)
    end_ts: Optional[float] = None    # End in seconds (populated by join_timestamps)
```

From mlx_audio KokoroPipeline.Result:

```python
@dataclass
class Result:
    graphemes: str
    phonemes: str
    tokens: Optional[List[en.MToken]] = None  # Word tokens with timestamps
    output: Optional[Any] = None

    @property
    def audio(self) -> Optional[mx.array]: ...
    @property
    def pred_dur(self) -> Optional[mx.array]: ...
```

KokoroPipeline.generate() yields Result objects. When output.pred_dur is not None,
join_timestamps() is called automatically, populating each MToken's start_ts/end_ts.

From kokoro_common.py (current synthesize):

```python
def synthesize(model, text: str, voice: str, lang: str, speed: float) -> np.ndarray:
    """Returns concatenated audio as numpy float32 array."""
    for result in model.generate(text=text, voice=voice, speed=speed, lang_code=lang):
        if result.audio is not None:
            chunks.append(np.array(result.audio))
    return np.concatenate(chunks)
```

From tts_server.py (existing endpoint pattern):

```python
def _handle_speech(self) -> None:
    body = self._read_body()
    text = _sanitize_text(str(body.get("input", ""))).strip()
    voice = str(body.get("voice", cfg.default_voice))
    lang = str(body.get("language", cfg.default_lang))
    speed = float(body.get("speed", cfg.default_speed))
    audio = synthesize_locked(self.model, text, voice, lang, speed)
    wav_bytes = to_wav_bytes(audio)
    # ... return bytes
```

</interfaces>
</context>

<tasks>

<task type="auto" tdd="true">
  <name>Task 1: Add synthesize_with_timestamps to kokoro_common.py</name>
  <files>/Users/terryli/.local/share/kokoro/kokoro_common.py</files>
  <read_first>
    - /Users/terryli/.local/share/kokoro/kokoro_common.py
    - /Users/terryli/.local/share/kokoro/.venv/lib/python3.13/site-packages/mlx_audio/tts/models/kokoro/pipeline.py (lines 320-355 for join_timestamps, lines 357-383 for Result class, lines 386-429 for generate method)
  </read_first>
  <behavior>
    - Test 1: synthesize_with_timestamps("Hello world") returns (np.ndarray, list[dict]) where each dict has "text", "onset", "duration" keys
    - Test 2: Word timings have onset >= 0.0 and duration > 0.0 for non-punctuation tokens
    - Test 3: Punctuation-only tokens (text matching r'^[^\w]+$') are filtered out of the returned word list
    - Test 4: Audio array is float32 and non-empty (same quality as existing synthesize())
    - Test 5: Multi-chunk synthesis (long text) accumulates timestamps with correct offsets across chunks
  </behavior>
  <action>
Add `synthesize_with_timestamps()` function to `kokoro_common.py` (per D-03 — word timestamps from MToken.start_ts/end_ts).

Implementation:

1. Add new function `synthesize_with_timestamps(model, text, voice, lang, speed) -> tuple[np.ndarray, list[dict]]`
2. Iterate `model.generate(text=text, voice=voice, speed=speed, lang_code=lang)` — same as existing `synthesize()`
3. For each `result`, collect `np.array(result.audio)` into audio chunks (same as now)
4. For each `result`, iterate `result.tokens` (list of MToken). For tokens where `t.start_ts is not None and t.end_ts is not None`:
   - Filter punctuation-only tokens: skip if `re.match(r'^[^\w]+$', t.text)`
   - Compute chunk audio offset: `chunk_offset = sum(len(c) for c in previous_audio_chunks) / SAMPLE_RATE`
   - Build word dict: `{"text": t.text, "onset": chunk_offset + t.start_ts, "duration": t.end_ts - t.start_ts}`
5. Return `(np.concatenate(audio_chunks), word_timings_list)`

The chunk_offset accumulation is critical — model.generate() yields multiple Result objects for long text, and each Result's MToken timestamps are relative to that chunk's audio start. We must offset by the cumulative audio duration of all prior chunks.

Do NOT modify the existing `synthesize()` function — it is used by other callers.

Also create a test file at `/Users/terryli/.local/share/kokoro/test_timestamps.py` that verifies the function works with a real model call (integration test). The test should:

- Import kokoro_common and call synthesize_with_timestamps with "Hello world, this is a test."
- Assert audio is non-empty np.float32
- Assert words list is non-empty
- Assert each word has text (str), onset (float >= 0), duration (float > 0)
- Assert no punctuation-only entries in words list
- Print the word timings for manual inspection
  </action>
  <verify>
  <automated>cd /Users/terryli/.local/share/kokoro && .venv/bin/python -c "
  import kokoro_common as c
  m = c.create_model()
  audio, words = c.synthesize_with_timestamps(m, 'Hello world.', c.DEFAULT_VOICE, c.DEFAULT_LANG, c.DEFAULT_SPEED)
  assert audio.size > 0, 'audio empty'
  assert len(words) > 0, 'no words'
  for w in words:
  assert 'text' in w and 'onset' in w and 'duration' in w, f'bad word: {w}'
  assert w['onset'] >= 0.0, f'negative onset: {w}'
  assert w['duration'] > 0.0, f'non-positive duration: {w}'
  print(f'OK: {len(words)} words, {audio.size/24000:.2f}s audio')
  for w in words: print(f' {w[\"text\"]:15s} onset={w[\"onset\"]:.3f}s dur={w[\"duration\"]:.3f}s')
  "</automated>
  </verify>
  <acceptance_criteria>
  - grep -q "def synthesize_with_timestamps" /Users/terryli/.local/share/kokoro/kokoro_common.py
  - grep -q "start_ts" /Users/terryli/.local/share/kokoro/kokoro_common.py
  - grep -q "end_ts" /Users/terryli/.local/share/kokoro/kokoro_common.py
  - grep -q "chunk_offset" /Users/terryli/.local/share/kokoro/kokoro_common.py
    </acceptance_criteria>
    <done>synthesize_with_timestamps() returns concatenated audio + word timing dicts with onset/duration from MToken.start_ts/end_ts, punctuation filtered, multi-chunk offsets accumulated</done>
    </task>

<task type="auto">
  <name>Task 2: Add /v1/audio/speech-with-timestamps endpoint to tts_server.py</name>
  <files>/Users/terryli/.local/share/kokoro/tts_server.py</files>
  <read_first>
    - /Users/terryli/.local/share/kokoro/tts_server.py (full file — need _handle_speech pattern, do_POST routing, _sanitize_text, synthesize_locked, to_wav_bytes)
    - /Users/terryli/.local/share/kokoro/kokoro_common.py (after Task 1 modifications — need synthesize_with_timestamps signature)
  </read_first>
  <action>
Add `/v1/audio/speech-with-timestamps` endpoint to `tts_server.py` (per D-01 — separate from existing /v1/audio/speech, per D-02 — JSON response format).

Implementation:

1. Add `synthesize_with_timestamps_locked()` wrapper (same pattern as existing `synthesize_locked()`):

   ```python
   def synthesize_with_timestamps_locked(model, text, voice, lang, speed):
       with _synthesis_lock:
           return common.synthesize_with_timestamps(model, text, voice, lang, speed)
   ```

2. Add route in `do_POST()`:

   ```python
   elif self.path == "/v1/audio/speech-with-timestamps":
       self._handle_speech_with_timestamps()
   ```

3. Add `_handle_speech_with_timestamps()` method to KokoroHandler:
   - Parse body same as `_handle_speech()`: input, voice, language, speed
   - Call `synthesize_with_timestamps_locked(self.model, text, voice, lang, speed)`
   - Convert audio to WAV bytes via `to_wav_bytes(audio)`
   - Base64-encode WAV: `base64.b64encode(wav_bytes).decode("ascii")`
   - Build response JSON per D-02:

     ```python
     {
       "audio_b64": "<base64 WAV string>",
       "words": [{"text": "Hello", "onset": 0.0, "duration": 0.45}, ...],
       "audio_duration": len(audio) / common.SAMPLE_RATE,
       "sample_rate": common.SAMPLE_RATE,
     }
     ```

   - Use `json.dumps()` (not orjson) since stdlib is sufficient for this response size
   - Send with `self._send_json(200, response_dict)`

4. Add `import base64` at top of file (not currently imported).

5. Update the module docstring Endpoints list to include the new endpoint.

6. For PTS-03 (launchd auto-start): The existing plist at `~/Library/LaunchAgents/com.terryli.kokoro-tts-server.plist` already has `KeepAlive: true` and `RunAtLoad: true`. No changes needed. Add a verify step to confirm.

After modifying the file, restart the launchd service to pick up the changes:

```bash
launchctl bootout gui/$(id -u) ~/Library/LaunchAgents/com.terryli.kokoro-tts-server.plist 2>/dev/null; sleep 2; launchctl bootstrap gui/$(id -u) ~/Library/LaunchAgents/com.terryli.kokoro-tts-server.plist
```

Wait ~15 seconds for model warmup, then test.
</action>
<verify>
<automated>sleep 20 && curl -s -X POST http://127.0.0.1:8779/v1/audio/speech-with-timestamps -H 'Content-Type: application/json' -d '{"input":"Hello world, this is a test of word timestamps."}' | python3 -c "
import sys, json
r = json.load(sys.stdin)
assert 'audio_b64' in r, 'missing audio_b64'
assert 'words' in r, 'missing words'
assert 'audio_duration' in r, 'missing audio_duration'
assert 'sample_rate' in r, 'missing sample_rate'
assert r['sample_rate'] == 24000, f'bad sample_rate: {r[\"sample_rate\"]}'
assert len(r['audio_b64']) > 100, 'audio_b64 too short'
assert len(r['words']) >= 2, f'too few words: {len(r[\"words\"])}'
for w in r['words']:
assert 'text' in w and 'onset' in w and 'duration' in w, f'bad word: {w}'
print(f'OK: {len(r[\"words\"])} words, {r[\"audio_duration\"]:.2f}s audio')
for w in r['words']: print(f' {w[\"text\"]:15s} onset={w[\"onset\"]:.3f}s dur={w[\"duration\"]:.3f}s')
"</automated>
</verify>
<acceptance_criteria> - grep -q "speech-with-timestamps" /Users/terryli/.local/share/kokoro/tts_server.py - grep -q "\_handle_speech_with_timestamps" /Users/terryli/.local/share/kokoro/tts_server.py - grep -q "audio_b64" /Users/terryli/.local/share/kokoro/tts_server.py - grep -q "import base64" /Users/terryli/.local/share/kokoro/tts_server.py - grep -q "synthesize_with_timestamps_locked" /Users/terryli/.local/share/kokoro/tts_server.py - grep -q "KeepAlive" /Users/terryli/Library/LaunchAgents/com.terryli.kokoro-tts-server.plist
</acceptance_criteria>
<done>POST /v1/audio/speech-with-timestamps returns 200 with JSON containing audio_b64 (base64 WAV), words array (text/onset/duration from MToken), audio_duration, and sample_rate. Launchd plist has KeepAlive for automatic startup (PTS-03).</done>
</task>

</tasks>

<verification>
1. Server health check: `curl -s http://127.0.0.1:8779/health | python3 -c "import sys,json; r=json.load(sys.stdin); assert r['status']=='ok', r"`
2. Endpoint returns valid JSON with all required fields: `curl -s -X POST http://127.0.0.1:8779/v1/audio/speech-with-timestamps -H 'Content-Type: application/json' -d '{"input":"Test"}' | python3 -c "import sys,json; r=json.load(sys.stdin); print(f'{len(r[\"words\"])} words, {r[\"audio_duration\"]:.2f}s')" `
3. Word timestamps are from native MToken (not character-weighted): Verify onset values are non-uniform (character-weighted would produce evenly spaced onsets proportional to word length)
4. Existing /v1/audio/speech endpoint still works (no regression): `curl -s -X POST http://127.0.0.1:8779/v1/audio/speech -H 'Content-Type: application/json' -d '{"input":"Hello"}' -o /dev/null -w '%{http_code}'` returns 200
5. PTS-03: `plutil -extract KeepAlive raw ~/Library/LaunchAgents/com.terryli.kokoro-tts-server.plist` returns "true"
</verification>

<success_criteria>

- PTS-01: `curl localhost:8779/v1/audio/speech-with-timestamps -d '{"input":"Hello world"}'` returns JSON with audio_b64, words array, audio_duration, sample_rate
- PTS-02: Words array onset/duration values come from MToken.start_ts/end_ts (verified by non-uniform spacing in output)
- PTS-03: Launchd plist has KeepAlive:true and RunAtLoad:true (existing config, no change needed)
- No regression: existing /v1/audio/speech and /v1/audio/speak endpoints unaffected
  </success_criteria>

<output>
After completion, create `.planning/phases/25-python-tts-server-timestamp-endpoint/25-01-SUMMARY.md`
</output>
