From c286fbecd3ae65ab243c738fefba2a39a2be1c80 Mon Sep 17 00:00:00 2001
From: Austin Mroz <AustinMroz@utexas.edu>
Date: Fri, 19 Jul 2024 14:21:47 -0500
Subject: [PATCH] Support auth token in orchestration connection

---
 orchestration.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/orchestration.py b/orchestration.py
index c5f0a20..1f08971 100644
--- a/orchestration.py
+++ b/orchestration.py
@@ -7,6 +7,8 @@ import aiohttp
 import asyncio
 from .workflowcheckpointing import post_prompt_remote
 
+STATIC_AUTH_TOKEN = os.environ.get("STATIC_AUTH_TOKEN", None)
+
 web = server.web
 ps = server.PromptServer.instance
 
@@ -46,7 +48,12 @@ async def websocket_loop():
         if 'ORCHESTRATION_SERVER' not in os.environ:
             while True:
                 await asyncio.sleep(60)
-        async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"]) as ws:
+        if STATIC_AUTH_TOKEN:
+            headers = {"Authorization": f"Bearer {STATIC_AUTH_TOKEN}"}
+        else:
+            headers = None
+        async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"],
+                                      headers=headers) as ws:
             print("connected to server")
             async for msg in ws:
                 try: