From a35e98023dad5e9732e4f20993849f5a8e0669d2 Mon Sep 17 00:00:00 2001
From: toby <toby@tobydeMacBook-Air.local>
Date: Sun, 15 Mar 2026 14:39:30 +0800
Subject: [PATCH] Add org Start URL support and fix multi-worker auth session
 issue

- auth_flow.py: make start_url configurable (default personal Q URL)
- app.py: add start_url field to AuthStartBody, pass to device_authorize
- frontend/index.html: add Start URL input field in login section
- docker-compose.yml: set workers=1 to fix in-memory session loss across workers

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 app.py              | 3 ++-
 auth_flow.py        | 6 +++---
 docker-compose.yml  | 2 +-
 frontend/index.html | 4 +++-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/app.py b/app.py
index 1f632d9..8443317 100644
--- a/app.py
+++ b/app.py
@@ -1271,6 +1271,7 @@ def _load_auth_flow_module():
 class AuthStartBody(BaseModel):
     label: Optional[str] = None
     enabled: Optional[bool] = True
+    start_url: Optional[str] = None
 
 class AdminLoginRequest(BaseModel):
     password: str
@@ -1353,7 +1354,7 @@ async def auth_start(body: AuthStartBody, _: bool = Depends(verify_admin_passwor
         """
         try:
             cid, csec = await register_client_min()
-            dev = await device_authorize(cid, csec)
+            dev = await device_authorize(cid, csec, start_url=body.start_url)
         except httpx.HTTPError as e:
             raise HTTPException(status_code=502, detail=f"OIDC error: {str(e)}")
 
diff --git a/auth_flow.py b/auth_flow.py
index 45cc226..e5412a3 100644
--- a/auth_flow.py
+++ b/auth_flow.py
@@ -18,7 +18,7 @@ def _get_proxies() -> Optional[Dict[str, str]]:
 REGISTER_URL = f"{OIDC_BASE}/client/register"
 DEVICE_AUTH_URL = f"{OIDC_BASE}/device_authorization"
 TOKEN_URL = f"{OIDC_BASE}/token"
-START_URL = "https://view.awsapps.com/start"
+DEFAULT_START_URL = "https://view.awsapps.com/start"
 
 USER_AGENT = "aws-sdk-rust/1.3.9 os/windows lang/rust/1.87.0"
 X_AMZ_USER_AGENT = "aws-sdk-rust/1.3.9 ua/2.1 api/ssooidc/1.88.0 os/windows lang/rust/1.87.0 m/E app/AmazonQ-For-CLI"
@@ -72,7 +72,7 @@ async def register_client_min() -> Tuple[str, str]:
         return data["clientId"], data["clientSecret"]
 
 
-async def device_authorize(client_id: str, client_secret: str) -> Dict:
+async def device_authorize(client_id: str, client_secret: str, start_url: Optional[str] = None) -> Dict:
     """
     Start device authorization. Returns dict that includes:
     - deviceCode
@@ -84,7 +84,7 @@ async def device_authorize(client_id: str, client_secret: str) -> Dict:
     payload = {
         "clientId": client_id,
         "clientSecret": client_secret,
-        "startUrl": START_URL,
+        "startUrl": start_url or DEFAULT_START_URL,
     }
     proxies = _get_proxies()
     mounts = None
diff --git a/docker-compose.yml b/docker-compose.yml
index 52e1ddb..6294aff 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -13,7 +13,7 @@ services:
     volumes:
       - ./:/app
     restart: unless-stopped
-    command: uvicorn app:app --host 0.0.0.0 --port 8000 --workers 4
+    command: uvicorn app:app --host 0.0.0.0 --port 8000 --workers 1
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/healthz"]
       interval: 30s
diff --git a/frontend/index.html b/frontend/index.html
index 1d52d94..ad472be 100644
--- a/frontend/index.html
+++ b/frontend/index.html
@@ -266,6 +266,7 @@ <h2>创建账号</h2>
         <h2>URL 登录（5分钟超时）</h2>
         <div class="row">
           <div class="field"><label>label（可选）</label><input id="auth_label" /></div>
+          <div class="field"><label>Start URL（留空使用个人版；企业版填组织 URL，如 https://your-company.awsapps.com/start）</label><input id="auth_start_url" placeholder="https://view.awsapps.com/start" /></div>
           <div class="field" style="max-width:220px">
             <label>启用（登录成功后新账号是否启用）</label>
             <div>
@@ -757,7 +758,8 @@ <h2>Chat 测试（/v2/chat/test）</h2>
  async function startAuth(){
    const body = {
      label: (document.getElementById('auth_label').value || '').trim() || null,
-     enabled: document.getElementById('auth_enabled').checked
+     enabled: document.getElementById('auth_enabled').checked,
+     start_url: (document.getElementById('auth_start_url').value || '').trim() || null,
    };
    try {
      const r = await authFetch(api('/v2/auth/start'), {