Update tom_web_app.py

tom_web_app.py

@@ -262,58 +262,65 @@ class Session:
         self.placeholder = PRESS_ENTER
         self.transcript += f"\n{PRESS_ENTER}\n"
 
-def get_session(resp=None) -> Tuple[str, Session, Any]:
+def get_session():
     sid = request.cookies.get("sid")
+    set_cookie = False
     if not sid or sid not in SESSIONS:
         sid = uuid.uuid4().hex
-        sess = Session()
-        SESSIONS[sid] = sess
-        if resp is None:
-            resp = make_response()
-        resp.set_cookie("sid", sid, httponly=True, samesite="Lax")
-    else:
-        sess = SESSIONS[sid]
-    return sid, sess, resp
+        SESSIONS[sid] = Session()
+        set_cookie = True
+    return sid, SESSIONS[sid], set_cookie
 
 @app.get("/")
 def index():
+    sid, sess, set_cookie = get_session()
     resp = make_response(render_template("index.html"))
-    _, sess, resp = get_session(resp)
+    if set_cookie:
+        resp.set_cookie("sid", sid, httponly=True, samesite="Lax")
     return resp
 
 @app.get("/state")
 def state():
-    _, sess, _ = get_session()
-    return jsonify({
+    sid, sess, set_cookie = get_session()
+    resp = jsonify({
         "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():
-    _, sess, _ = get_session()
+    sid, sess, set_cookie = get_session()
     sess.start_or_continue()
-    return jsonify({
+    resp = jsonify({
         "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():
-    _, sess, _ = get_session()
+    sid, sess, set_cookie = get_session()
     data = request.get_json(force=True) or {}
     text = (data.get("text") or "").strip()
     sess.submit_action(text)
-    return jsonify({
+    resp = jsonify({
         "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