From 530922d6e0ad23556da1cf2341b1f23f01f37861 Mon Sep 17 00:00:00 2001
From: Pigbibi <20649888+Pigbibi@users.noreply.github.com>
Date: Mon, 8 Jun 2026 17:47:36 +0800
Subject: [PATCH] fix: dismiss small IBKR gateway dialogs

---
 2fa_bot.py                         | 33 ++++++++++++++++++++++++++----
 tests/test_docker_compose_ports.sh | 22 ++++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/2fa_bot.py b/2fa_bot.py
index 9fe05ab..f07928b 100644
--- a/2fa_bot.py
+++ b/2fa_bot.py
@@ -37,6 +37,15 @@
     "no",
     "off",
 }
+DISMISS_SMALL_GATEWAY_DIALOGS = os.environ.get(
+    "IBKR_DISMISS_SMALL_GATEWAY_DIALOGS",
+    "yes",
+).strip().lower() not in {
+    "0",
+    "false",
+    "no",
+    "off",
+}
 
 # Window titles to search for 2FA prompts. Live IBKR accounts can show mobile
 # push / IB Key wording instead of the shorter TOTP-oriented prompts.
@@ -71,10 +80,14 @@
 )
 DISMISSIBLE_DIALOG_SEARCH_PATTERNS = [
     "Login Messages",
+    "IBKR Gateway",
 ]
 DISMISSIBLE_DIALOG_TITLE_KEYWORDS = (
     "login messages",
 )
+SMALL_GATEWAY_DIALOG_TITLE = "ibkr gateway"
+SMALL_GATEWAY_DIALOG_MAX_WIDTH = 650
+SMALL_GATEWAY_DIALOG_MAX_HEIGHT = 220
 # Current IBKR Gateway TOTP prompts place the code field in the upper half of
 # the compact dialog. Keep the click centered on the text field instead of the
 # button/link area below it.
@@ -205,9 +218,21 @@ def is_auth_candidate(title):
     return any(keyword in normalized_title for keyword in AUTH_TITLE_KEYWORDS)
 
 
-def is_dismissible_dialog_candidate(title):
+def is_small_gateway_dialog(title, width, height):
+    if not DISMISS_SMALL_GATEWAY_DIALOGS:
+        return False
+    if title.lower() != SMALL_GATEWAY_DIALOG_TITLE:
+        return False
+    if not width or not height:
+        return False
+    return width <= SMALL_GATEWAY_DIALOG_MAX_WIDTH and height <= SMALL_GATEWAY_DIALOG_MAX_HEIGHT
+
+
+def is_dismissible_dialog_candidate(title, width=None, height=None):
     normalized_title = title.lower()
-    return any(keyword in normalized_title for keyword in DISMISSIBLE_DIALOG_TITLE_KEYWORDS)
+    if any(keyword in normalized_title for keyword in DISMISSIBLE_DIALOG_TITLE_KEYWORDS):
+        return True
+    return is_small_gateway_dialog(title, width, height)
 
 
 def find_windows_by_patterns(patterns):
@@ -243,9 +268,9 @@ def find_dismissible_dialogs():
     candidates = []
     for window_id in reversed(find_windows_by_patterns(DISMISSIBLE_DIALOG_SEARCH_PATTERNS)):
         title = get_window_title(window_id)
-        if not is_dismissible_dialog_candidate(title):
-            continue
         width, height = get_window_geometry(window_id)
+        if not is_dismissible_dialog_candidate(title, width, height):
+            continue
         candidates.append(WindowCandidate(window_id, title, width, height))
     return candidates
 
diff --git a/tests/test_docker_compose_ports.sh b/tests/test_docker_compose_ports.sh
index fae582a..572f5ac 100644
--- a/tests/test_docker_compose_ports.sh
+++ b/tests/test_docker_compose_ports.sh
@@ -51,4 +51,26 @@ grep -Fq '      - ACCEPT_API_FROM_IP=${ACCEPT_API_FROM_IP:?Set ACCEPT_API_FROM_I
 
 grep -Fq 'INPUT_CLICK_POSITION = (0.50, 0.40)' "$repo_dir/2fa_bot.py"
 grep -Fq 'MIN_TOTP_SECONDS_REMAINING = 15' "$repo_dir/2fa_bot.py"
+grep -Fq 'SMALL_GATEWAY_DIALOG_MAX_WIDTH = 650' "$repo_dir/2fa_bot.py"
+grep -Fq 'SMALL_GATEWAY_DIALOG_MAX_HEIGHT = 220' "$repo_dir/2fa_bot.py"
+grep -Fq 'def is_small_gateway_dialog(title, width, height):' "$repo_dir/2fa_bot.py"
+grep -Fq 'is_dismissible_dialog_candidate(title, width, height)' "$repo_dir/2fa_bot.py"
 grep -Fq 'def type_totp_into_active_window(code):' "$repo_dir/2fa_bot.py"
+
+REPO_DIR="$repo_dir" python3 <<'PY'
+import importlib.util
+import os
+import sys
+import types
+
+sys.modules["pyotp"] = types.SimpleNamespace()
+module_path = os.path.join(os.environ["REPO_DIR"], "2fa_bot.py")
+spec = importlib.util.spec_from_file_location("ibkr_2fa_bot_test", module_path)
+module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(module)
+
+assert module.is_dismissible_dialog_candidate("Login Messages")
+assert module.is_dismissible_dialog_candidate("IBKR Gateway", 509, 131)
+assert not module.is_dismissible_dialog_candidate("IBKR Gateway", 700, 550)
+assert not module.is_dismissible_dialog_candidate("IBKR Gateway", 790, 610)
+PY