Update tom_web_app.py

tom_web_app.py

@@ -263,13 +263,13 @@ class Session:
         self.transcript += f"\n{PRESS_ENTER}\n"
 
 def get_session():
-    sid = request.cookies.get("sid")
-    set_cookie = False
+    sid = request.args.get("sid") or request.headers.get("X-Session-Id") or request.cookies.get("sid")
+    new_sid = False
     if not sid or sid not in SESSIONS:
         sid = uuid.uuid4().hex
-        SESSIONS[sid] = Session()
-        set_cookie = True
-    return sid, SESSIONS[sid], set_cookie
+    SESSIONS[sid] = Session()
+    new_sid = True
+    return sid, SESSIONS[sid], new_sid
 
 @app.get("/")
 def index():
@@ -281,46 +281,40 @@ def index():
 
 @app.get("/state")
 def state():
-    sid, sess, set_cookie = get_session()
-    resp = jsonify({
-        "transcript": sess.transcript,
-        "mode": sess.mode,
-        "primary_label": sess.primary_label,
-        "placeholder": sess.placeholder
+    sid, sess, _ = get_session()
+    return jsonify({
+    "sid": sid,
+    "transcript": sess.transcript,
+    "mode": sess.mode,
+    "primary_label": sess.primary_label,
+    "placeholder": sess.placeholder
     })
-    if set_cookie:
-        resp.set_cookie("sid", sid, httponly=True, samesite="Lax")
-    return resp
 
 @app.post("/primary")
 def primary():
-    sid, sess, set_cookie = get_session()
+    sid, sess, _ = get_session()
     sess.start_or_continue()
-    resp = jsonify({
+    return jsonify({
+        "sid": sid,
         "transcript": sess.transcript,
         "mode": sess.mode,
         "primary_label": sess.primary_label,
         "placeholder": sess.placeholder
-    })
-    if set_cookie:
-        resp.set_cookie("sid", sid, httponly=True, samesite="Lax")
-    return resp
+        })
 
 @app.post("/action")
 def action():
-    sid, sess, set_cookie = get_session()
+    sid, sess, _ = get_session()
     data = request.get_json(force=True) or {}
     text = (data.get("text") or "").strip()
     sess.submit_action(text)
-    resp = jsonify({
+    return jsonify({
+        "sid": sid,
         "transcript": sess.transcript,
         "mode": sess.mode,
         "primary_label": sess.primary_label,
         "placeholder": sess.placeholder
-    })
-    if set_cookie:
-        resp.set_cookie("sid", sid, httponly=True, samesite="Lax")
-    return resp
+        })
 
 if __name__ == "__main__":
     import os