Gateway Drain Trigger

Port-less, restart-safe external drain signal for hosted / containerised gateway deployments

from praisonaiagents import Agent

agent = Agent(name="drain-agent", instructions="Drain the gateway queue gracefully.")
agent.start("Drain all pending messages before shutting down the gateway.")

Ask a running gateway to finish in-flight turns and exit — without exposing any inbound port, and without a left-over signal wedging a restarted instance.

The user drops a drain marker file; the gateway finishes in-flight turns and exits without exposing an inbound drain port.

Quick Start

1. Import the policy and epoch function

   from praisonaiagents.gateway import DrainMarkerPolicy, current_epoch

2. Build the policy once at startup

   policy = DrainMarkerPolicy()
   epoch = current_epoch()

3. Run a watcher loop

   import asyncio
   import json
   import os
   from time import monotonic
   from praisonaiagents.gateway import DrainMarkerPolicy, current_epoch
   
   policy = DrainMarkerPolicy()
   last_handled = None
   path = os.path.expanduser("~/.praisonai/gateway/.drain_request.json")
   
   async def watch(gateway):
       global last_handled
       while True:
           if os.path.exists(path):
               with open(path) as f:
                   marker = json.load(f)
           else:
               marker = None
           if policy.drain_requested(
               marker,
               current_epoch(),
               monotonic(),
               last_handled_epoch=last_handled,
           ):
               last_handled = marker.get("epoch")
               await gateway.stop(drain_timeout=30)
               return
           await asyncio.sleep(1)

4. Primary runtime — enable the built-in watcher via gateway.yaml

   The primary WebSocketGateway (praisonai gateway start) ships a built-in marker watcher. Point it at a marker path with a lifecycle: block.

   # gateway.yaml
   lifecycle:
     drain:
       marker_path: "/data/gateway.drain"

   praisonai gateway start --config gateway.yaml

   The gateway polls the marker every 5 seconds and, on an epoch match, drains active sessions bounded by the gateway’s drain_timeout.

5. Primary runtime — enable via CLI flag (no YAML)

   praisonai gateway start --drain-marker /data/gateway.drain

   The flag synthesizes the lifecycle.drain block, so CLI overrides win over gateway.yaml.

6. Trigger a drain from a deploy step

   import json, os
   from praisonaiagents.gateway import current_epoch
   
   marker = {"epoch": current_epoch(), "action": "drain"}
   tmp = "/data/gateway.drain.tmp"
   with open(tmp, "w") as f:
       json.dump(marker, f)
   os.replace(tmp, "/data/gateway.drain")   # atomic write

How It Works

Restart-Safety Story

The entire point of this feature: a marker stamped with current_epoch() is silently ignored by any other instance of the gateway — even if the marker file survives a reboot on a durable volume.

Decision Tree — When Is a Marker Honoured?

The Drain Marker Contract

Write this JSON to the marker path:

{
  "epoch": "<output of current_epoch() at write time>",
  "action": "drain"
}

• Default marker path (convention from PraisonAI #2390, not yet enforced in code): ~/.praisonai/gateway/.drain_request.json
• action defaults to "drain" if absent. Any other value (including a non-string) is ignored.
• Missing / empty / non-string epoch is ignored unless require_epoch=False is passed to the policy.

Configuration Options

DrainMarkerPolicy constructor

from praisonaiagents.gateway import DrainMarkerPolicy

policy = DrainMarkerPolicy(require_epoch=True)

DrainMarkerPolicy.drain_requested() parameters

current_epoch() — what it returns

from praisonaiagents.gateway import current_epoch

epoch = current_epoch()
# Linux: "550e8400-e29b-41d4-a716-446655440000:12345"
# Mac / Windows / no /proc: ""

lifecycle.drain: YAML block (WebSocketGateway)

The primary WebSocketGateway reads the built-in watcher from a lifecycle: block at the top level of gateway.yaml or nested under gateway:.

lifecycle:
  drain:
    marker_path: "/data/gateway.drain"

The watcher polls every 5 seconds. On an epoch match it sets _draining = True and calls _drain_active_sessions(reason="drain-marker", timeout=drain_timeout), so the gateway’s existing drain_timeout bounds the marker-triggered drain.

Observability

When the built-in watcher runs, health() surfaces it under the lifecycle object:

{
  "status": "healthy",
  "lifecycle": {
    "scale_to_zero": false,
    "dormant": false,
    "drain_marker_watch": true
  }
}

drain_marker_watch is true while the watch loop is running. The lifecycle field appears only when a lifecycle feature is configured.

Common Patterns

1. Operator-side: write the marker

import json
import os
from praisonaiagents.gateway import current_epoch

path = os.path.expanduser("~/.praisonai/gateway/.drain_request.json")
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path + ".tmp", "w") as f:
    json.dump({"epoch": current_epoch(), "action": "drain"}, f)
os.replace(path + ".tmp", path)

Writing to <path>.tmp then os.replace() makes the write atomic — a half-written JSON file is treated as malformed and ignored.

2. Gateway-side: minimal watcher loop

import asyncio
import json
import os
from time import monotonic
from praisonaiagents.gateway import DrainMarkerPolicy, current_epoch

policy = DrainMarkerPolicy()
last_handled = None
path = os.path.expanduser("~/.praisonai/gateway/.drain_request.json")

async def watch(gateway):
    global last_handled
    while True:
        if os.path.exists(path):
            with open(path) as f:
                marker = json.load(f)
        else:
            marker = None
        if policy.drain_requested(
            marker,
            current_epoch(),
            monotonic(),
            last_handled_epoch=last_handled,
        ):
            last_handled = marker.get("epoch")
            await gateway.stop(drain_timeout=30)
            return
        await asyncio.sleep(1)

3. Accepting legacy markers without an epoch (advanced)

from praisonaiagents.gateway import DrainMarkerPolicy

policy = DrainMarkerPolicy(require_epoch=False)

Only do this if you own both ends and have a different staleness story.

Best Practices

Related