-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathapp.py
More file actions
1305 lines (1063 loc) Β· 49.7 KB
/
app.py
File metadata and controls
1305 lines (1063 loc) Β· 49.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
import os
import pty
import fcntl
import struct
import termios
import select
import subprocess
import uuid
import threading
import signal
import time
import copy
import logging
from concurrent.futures import ThreadPoolExecutor, wait
from flask import Flask, send_from_directory, request, jsonify, session
from flask_socketio import SocketIO, emit, join_room, leave_room, disconnect
from werkzeug.utils import secure_filename
from collections import deque
import tomllib
import requests
import app_state
from utils import ensure_https, get_gateway_host
from pat_rotator import PATRotator
from telemetry import log_telemetry, set_product_info
# Sanitize DATABRICKS_TOKEN early β the platform sometimes injects trailing
# newlines / whitespace which causes auth failures. Cleaning it here prevents
# the agent from "fixing" it in the terminal and leaking the raw token.
_raw_token = os.environ.get("DATABRICKS_TOKEN", "")
if _raw_token != _raw_token.strip():
os.environ["DATABRICKS_TOKEN"] = _raw_token.strip()
# App version (single source of truth: pyproject.toml)
_pyproject_file = os.path.join(os.path.dirname(__file__), 'pyproject.toml')
try:
with open(_pyproject_file, 'rb') as _f:
APP_VERSION = tomllib.load(_f)['project']['version']
except Exception:
APP_VERSION = '0.0.0'
# Session timeout configuration
SESSION_TIMEOUT_SECONDS = 86400 # No poll for 24 hours = dead session
CLEANUP_INTERVAL_SECONDS = 900 # Check for stale sessions every 15 min
GRACEFUL_SHUTDOWN_WAIT = 3 # Seconds to wait after SIGHUP before SIGKILL
MAX_CONCURRENT_SESSIONS = int(os.environ.get("MAX_CONCURRENT_SESSIONS", "5"))
# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# PAT auto-rotation β initialized after sessions dict is defined (see below)
app = Flask(__name__, static_folder='static', static_url_path='/static')
app.secret_key = os.urandom(24)
app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024 # 32 MB β aligned with Claude Code's 30 MB file limit
# WebSocket support via Flask-SocketIO (simple-websocket transport, threading mode)
socketio = SocketIO(app, async_mode='threading', cors_allowed_origins=[], logger=False, engineio_logger=False)
# Store sessions: {session_id: {"master_fd": fd, "pid": pid, "output_buffer": deque, "lock": Lock, ...}}
# sessions_lock guards dict-level ops (add/remove/iterate); each session["lock"] guards per-session state
sessions = {}
sessions_lock = threading.Lock()
# PAT auto-rotation (short-lived tokens, background refresh)
# Only rotates while active sessions exist β stops when all sessions are reaped
pat_rotator = PATRotator(
session_count_fn=lambda: len(sessions),
)
# SIGTERM graceful shutdown: notify clients before gunicorn stops the worker
shutting_down = False
_start_time = time.time()
def handle_sigterm(signum, frame):
"""Notify clients that app is shutting down, then let gunicorn handle the rest."""
global shutting_down
# Ignore SIGTERMs in the first 10s β likely stale signals from a prior process kill
if time.time() - _start_time < 10:
logger.info("SIGTERM received during startup β ignoring (likely stale signal)")
return
shutting_down = True
logger.info("SIGTERM received β setting shutting_down flag for clients")
# Notify WS clients immediately (HTTP poll clients will see shutting_down on next poll)
try:
socketio.emit('shutting_down', {})
except Exception:
pass
# NOTE: Do not register SIGTERM handler at module level.
# It is installed in initialize_app() for gunicorn only.
# For local dev (__main__), we keep SIG_DFL so the process just exits.
# Setup state tracking
setup_lock = threading.Lock()
setup_state = {
"status": "pending",
"started_at": None,
"completed_at": None,
"error": None,
"steps": [
{"id": "git", "label": "Configuring git identity", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "micro", "label": "Installing micro editor", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "gh", "label": "Installing GitHub CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "dbcli", "label": "Upgrading Databricks CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "proxy", "label": "Starting content-filter proxy", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "claude", "label": "Configuring Claude CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "codex", "label": "Configuring Codex CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "opencode", "label": "Configuring OpenCode CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "gemini", "label": "Configuring Gemini CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "hermes", "label": "Configuring Hermes Agent", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "databricks", "label": "Setting up Databricks CLI", "status": "pending", "started_at": None, "completed_at": None, "error": None},
{"id": "mlflow", "label": "Enabling MLflow tracing", "status": "pending", "started_at": None, "completed_at": None, "error": None},
]
}
def _update_step(step_id, **kwargs):
with setup_lock:
for step in setup_state["steps"]:
if step["id"] == step_id:
step.update(kwargs)
break
def _get_setup_state_snapshot():
with setup_lock:
return copy.deepcopy(setup_state)
# Single-user security: only the token owner can access the terminal
app_owner = None
def _run_step(step_id, command):
_update_step(step_id, status="running", started_at=time.time())
try:
env = os.environ.copy()
if not env.get("HOME") or env["HOME"] == "/":
env["HOME"] = "/app/python/source_code"
home = env.get("HOME", "/app/python/source_code")
# Ensure uv and other tools in ~/.local/bin are on PATH
local_bin = os.path.join(home, ".local", "bin")
if local_bin not in env.get("PATH", ""):
env["PATH"] = f"{local_bin}:{env.get('PATH', '')}"
env.pop("DATABRICKS_CLIENT_ID", None)
env.pop("DATABRICKS_CLIENT_SECRET", None)
result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=300)
if result.returncode == 0:
_update_step(step_id, status="complete", completed_at=time.time())
else:
err = result.stderr.strip() or result.stdout.strip() or "Unknown error"
_update_step(step_id, status="error", completed_at=time.time(), error=err[:500])
except subprocess.TimeoutExpired:
_update_step(step_id, status="error", completed_at=time.time(), error="Timed out after 300s")
except Exception as e:
_update_step(step_id, status="error", completed_at=time.time(), error=str(e))
def _setup_git_config():
"""Configure git identity and hooks by writing files directly (no subprocess)."""
home = os.environ.get("HOME", "/app/python/source_code")
if not home or home == "/":
home = "/app/python/source_code"
# Get user identity from Databricks token
user_email = None
display_name = None
try:
from databricks.sdk import WorkspaceClient
db_host = ensure_https(os.environ.get("DATABRICKS_HOST", ""))
db_token = os.environ.get("DATABRICKS_TOKEN")
if db_host and db_token:
w = WorkspaceClient(host=db_host, token=db_token, auth_type="pat")
set_product_info(w)
me = w.current_user.me()
user_email = me.user_name
display_name = me.display_name or user_email.split("@")[0]
except Exception as e:
logger.warning(f"Could not get user identity from token: {e}")
# Write ~/.gitconfig directly (more reliable than subprocess git config)
gitconfig_path = os.path.join(home, ".gitconfig")
hooks_dir = os.path.join(home, ".githooks")
os.makedirs(hooks_dir, exist_ok=True)
lines = []
if user_email and display_name:
lines.append("[user]")
lines.append(f"\temail = {user_email}")
lines.append(f"\tname = {display_name}")
lines.append("[core]")
lines.append(f"\thooksPath = {hooks_dir}")
with open(gitconfig_path, "w") as f:
f.write("\n".join(lines) + "\n")
logger.info(f"Git config written to {gitconfig_path}")
# Write post-commit hook for workspace sync (works from any CLI: Claude, Gemini, OpenCode, etc.)
# Only syncs repos inside ~/projects/ β skips the app source and any other repos
post_commit = os.path.join(hooks_dir, "post-commit")
with open(post_commit, "w") as f:
f.write('#!/bin/bash\n')
f.write('# Auto-sync to Databricks Workspace on commit (works from any CLI)\n')
f.write('SYNC_LOG="$HOME/.sync.log"\n')
f.write('\n')
f.write('# Resolve git repo root (handles commits from subdirectories)\n')
f.write('REPO_ROOT="$(git rev-parse --show-toplevel 2>/dev/null)"\n')
f.write('if [ -z "$REPO_ROOT" ]; then\n')
f.write(' echo "[post-commit] $(date +%H:%M:%S) SKIP: not inside a git repo" >> "$SYNC_LOG"\n')
f.write(' exit 0\n')
f.write('fi\n')
f.write('\n')
f.write('# Only sync repos inside ~/projects/\n')
f.write('PROJECTS_DIR="$HOME/projects"\n')
f.write('case "$REPO_ROOT" in\n')
f.write(' "$PROJECTS_DIR"/*)\n')
f.write(' ;; # allowed - continue\n')
f.write(' *)\n')
f.write(' echo "[post-commit] $(date +%H:%M:%S) SKIP: $REPO_ROOT is outside $PROJECTS_DIR" >> "$SYNC_LOG"\n')
f.write(' exit 0\n')
f.write(' ;;\n')
f.write('esac\n')
f.write('\n')
f.write('echo "[post-commit] $(date +%H:%M:%S) syncing $REPO_ROOT" >> "$SYNC_LOG"\n')
f.write('\n')
f.write('# Use uv run so sync script gets the correct Python + deps\n')
f.write('APP_DIR="/app/python/source_code"\n')
f.write('SYNC_SCRIPT="$APP_DIR/sync_to_workspace.py"\n')
f.write('\n')
f.write('if [ -f "$SYNC_SCRIPT" ]; then\n')
f.write(' nohup uv run --project "$APP_DIR" python "$SYNC_SCRIPT" "$REPO_ROOT" >> "$SYNC_LOG" 2>&1 & disown\n')
f.write('else\n')
f.write(' echo "[post-commit] $(date +%H:%M:%S) SKIP: sync script not found" >> "$SYNC_LOG"\n')
f.write('fi\n')
os.chmod(post_commit, 0o755)
logger.info(f"Post-commit hook written to {post_commit}")
# Reinit app source git to remove template origin (Databricks Apps only)
_reinit_app_git()
def _reinit_app_git():
"""On Databricks Apps, reinit git to remove template origin remote."""
app_dir = os.path.dirname(os.path.abspath(__file__))
if app_dir != "/app/python/source_code":
return # Local dev β leave git intact
git_dir = os.path.join(app_dir, ".git")
if not os.path.isdir(git_dir):
return # Already clean
import shutil
shutil.rmtree(git_dir)
subprocess.run(["git", "init"], cwd=app_dir, capture_output=True)
subprocess.run(["git", "add", "."], cwd=app_dir, capture_output=True)
subprocess.run(
["git", "commit", "-m", "Initial commit from coding-agents template"],
cwd=app_dir, capture_output=True,
)
logger.info("Reinitialized app source git (template origin removed)")
def _configure_all_cli_auth(token):
"""Configure auth for ALL coding-agent CLIs after a PAT is provided.
Called from /api/configure-pat when a user supplies a PAT interactively.
Handles: Claude CLI (inline), Databricks CLI (via pat_rotator), and
Codex/OpenCode/Gemini CLIs (by re-running their setup scripts with token in env).
"""
import json
from utils import resolve_and_cache_gateway
resolve_and_cache_gateway()
home = os.environ.get("HOME", "/app/python/source_code")
if not home or home == "/":
home = "/app/python/source_code"
# 1. Configure Claude CLI (~/.claude/settings.json)
claude_dir = os.path.join(home, ".claude")
os.makedirs(claude_dir, exist_ok=True)
gateway_host = get_gateway_host()
databricks_host = ensure_https(os.environ.get("DATABRICKS_HOST", "").rstrip("/"))
if gateway_host:
anthropic_base_url = f"{gateway_host}/anthropic"
else:
anthropic_base_url = f"{databricks_host}/serving-endpoints/anthropic"
# Read-merge-write to preserve env vars from other setup scripts (e.g. setup_mlflow.py)
settings_path = os.path.join(claude_dir, "settings.json")
try:
with open(settings_path) as f:
settings = json.load(f)
except (FileNotFoundError, json.JSONDecodeError):
settings = {}
settings.setdefault("env", {})
settings["env"]["ANTHROPIC_MODEL"] = os.environ.get("ANTHROPIC_MODEL", "databricks-claude-opus-4-7")
settings["env"]["ANTHROPIC_BASE_URL"] = anthropic_base_url
settings["env"]["ANTHROPIC_AUTH_TOKEN"] = token
settings["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] = "databricks-claude-opus-4-7"
settings["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] = "databricks-claude-sonnet-4-6"
settings["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = "databricks-claude-haiku-4-5"
settings["env"]["ANTHROPIC_CUSTOM_HEADERS"] = "x-databricks-use-coding-agent-mode: true"
settings["env"]["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] = "1"
with open(settings_path, "w") as f:
json.dump(settings, f, indent=2)
logger.info(f"Claude CLI auth configured: {settings_path}")
# 2. Configure Databricks CLI (~/.databrickscfg) β already called by
# configure_pat() via pat_rotator, but explicit for clarity
pat_rotator._write_databrickscfg(token)
logger.info("Databricks CLI auth configured: ~/.databrickscfg")
# 3. Re-run Codex, OpenCode, Gemini setup scripts with token in env
# They are idempotent: detect CLI already installed, just write config files
env = {**os.environ, "DATABRICKS_TOKEN": token}
for script in ["setup_codex.py", "setup_opencode.py", "setup_gemini.py", "setup_hermes.py"]:
try:
result = subprocess.run(
["uv", "run", "python", script],
env=env, capture_output=True, text=True, timeout=60
)
if result.returncode == 0:
logger.info(f"CLI config updated: {script}")
else:
logger.warning(f"CLI config failed: {script}: {result.stderr[:200]}")
except Exception as e:
logger.warning(f"CLI config error: {script}: {e}")
def run_setup():
with setup_lock:
setup_state["status"] = "running"
setup_state["started_at"] = time.time()
# Probe AI Gateway once; result is cached in _GATEWAY_RESOLVED for subprocesses
from utils import resolve_and_cache_gateway
resolve_and_cache_gateway()
# --- Sequential prerequisites (git identity + editor) ---
# Git config β done directly in Python, not as a subprocess
_update_step("git", status="running", started_at=time.time())
try:
_setup_git_config()
_update_step("git", status="complete", completed_at=time.time())
except Exception as e:
_update_step("git", status="error", completed_at=time.time(), error=str(e))
_run_step("micro", ["bash", "-c",
"mkdir -p ~/.local/bin && bash install_micro.sh && mv micro ~/.local/bin/ 2>/dev/null || true"])
_run_step("gh", ["bash", "install_gh.sh"])
# --- Upgrade Databricks CLI (runtime image ships an older version) ---
_run_step("dbcli", ["bash", "install_databricks_cli.sh"])
# --- Content-filter proxy (must be running before OpenCode starts) ---
# Sanitizes requests/responses between OpenCode and Databricks
# (see OpenCode #5028, docs/plans/2026-03-11-litellm-empty-content-blocks-design.md)
_run_step("proxy", ["uv", "run", "python", "setup_proxy.py"])
# --- Parallel agent setup (all independent of each other) ---
parallel_steps = [
("claude", ["uv", "run", "python", "setup_claude.py"]),
("codex", ["uv", "run", "python", "setup_codex.py"]),
("opencode", ["uv", "run", "python", "setup_opencode.py"]),
("gemini", ["uv", "run", "python", "setup_gemini.py"]),
("hermes", ["uv", "run", "python", "setup_hermes.py"]),
("databricks", ["uv", "run", "python", "setup_databricks.py"]),
]
with ThreadPoolExecutor(max_workers=len(parallel_steps)) as executor:
futures = [
executor.submit(_run_step, step_id, command)
for step_id, command in parallel_steps
]
wait(futures)
# --- MLflow setup runs AFTER claude setup to avoid settings.json race ---
# setup_mlflow.py merges env vars into ~/.claude/settings.json which
# setup_claude.py also writes; running sequentially prevents clobbering.
_run_step("mlflow", ["uv", "run", "python", "setup_mlflow.py"])
# Sync latest token into all CLI configs β covers the race where PAT
# rotation happened while a setup script was still installing (the
# rotation's update_cli_tokens() call silently skips missing config files).
current_token = os.environ.get("DATABRICKS_TOKEN", "")
if current_token:
try:
from cli_auth import update_cli_tokens
update_cli_tokens(current_token)
logger.info("Post-setup token sync: all CLI configs updated with current token")
except Exception as e:
logger.warning(f"Post-setup token sync failed: {e}")
with setup_lock:
any_error = any(s["status"] == "error" for s in setup_state["steps"])
setup_state["status"] = "error" if any_error else "complete"
setup_state["completed_at"] = time.time()
def get_token_owner():
"""Get the owner email. Priority: Apps API (app.creator) > PAT (current_user.me).
Uses the auto-provisioned SP to call the Apps API β no PAT needed for
owner resolution. Falls back to PAT-based lookup for backward compat.
"""
from databricks.sdk import WorkspaceClient
# 1. Try Apps API via SP credentials (no PAT needed)
app_name = os.environ.get("DATABRICKS_APP_NAME")
if app_name:
try:
w = WorkspaceClient() # auto-detects SP credentials
set_product_info(w)
app = w.apps.get(name=app_name)
owner = (app.creator or "").lower()
logger.info(f"Owner resolved from app.creator: {owner}")
return owner
except Exception as e:
logger.warning(f"Could not resolve owner via Apps API: {e}")
# 2. Fallback: PAT-based resolution
try:
host = ensure_https(os.environ.get("DATABRICKS_HOST", ""))
token = os.environ.get("DATABRICKS_TOKEN")
if not host or not token:
return None
w = WorkspaceClient(host=host, token=token, auth_type="pat")
set_product_info(w)
username = w.current_user.me().user_name
return username.lower() if username else username
except Exception as e:
logger.warning(f"Could not determine token owner: {e}")
return None
def get_request_user():
"""Extract user email from Databricks Apps request headers.
Returns lowercase email to ensure case-insensitive matching against app_owner.
"""
email = (
request.headers.get("X-Forwarded-Email")
or request.headers.get("X-Forwarded-User")
or request.headers.get("X-Databricks-User-Email")
)
return email.lower() if email else email
def _is_databricks_apps():
"""Detect if we're running on Databricks Apps (not local dev)."""
return os.environ.get("DATABRICKS_APP_PORT") or os.path.isdir("/app/python/source_code")
def check_authorization():
"""Check if the current user is authorized to access the app.
Fails CLOSED on Databricks Apps: if we can't determine the owner,
deny all access rather than allowing unauthenticated terminal access.
Fails open only for local development.
Fixes: https://github.com/datasciencemonkey/coding-agents-databricks-apps/issues/57
"""
# Fail closed on Databricks Apps if owner couldn't be resolved
if not app_owner:
if _is_databricks_apps():
logger.error("SECURITY: app_owner not resolved β denying all access (fail-closed)")
return False, "unknown"
return True, None # Local dev only
current_user = get_request_user()
# If no user identity in request (local dev), allow access
if not current_user:
if _is_databricks_apps():
logger.warning("No user identity in request on Databricks Apps β denying access")
return False, "unknown"
return True, None
# Check if current user is the owner
if current_user != app_owner:
logger.warning(f"Unauthorized access attempt by {current_user} (owner: {app_owner})")
return False, current_user
return True, None
def _check_ws_authorization():
"""Check authorization for WebSocket connections β mirrors HTTP check_authorization().
Fails CLOSED on Databricks Apps: if app_owner is unresolved or no user identity
in headers, deny WebSocket access. Matches the HTTP handler's behavior exactly.
"""
if not app_owner:
if _is_databricks_apps():
logger.error("SECURITY: app_owner not resolved β denying WebSocket (fail-closed)")
return False
return True # Local dev only
# Socket.IO passes HTTP headers from the initial handshake via request context
raw_user = (
request.headers.get("X-Forwarded-Email")
or request.headers.get("X-Forwarded-User")
or request.headers.get("X-Databricks-User-Email")
)
current_user = raw_user.lower() if raw_user else raw_user
if not current_user:
if _is_databricks_apps():
logger.warning("No user identity in WebSocket request on Databricks Apps β denying")
return False
return True # Local dev only
if current_user != app_owner:
logger.warning(f"WebSocket unauthorized: {current_user} (owner: {app_owner})")
return False
return True
# ββ WebSocket Event Handlers ββββββββββββββββββββββββββββββββββββββββββββββ
@socketio.on('connect')
def handle_ws_connect():
"""Authenticate WebSocket connections (AC-3)."""
if not _check_ws_authorization():
disconnect()
return False
logger.info("WebSocket client connected")
@socketio.on('join_session')
def handle_join_session(data):
"""Client joins a session room to receive output (AC-4)."""
session_id = data.get('session_id')
if not session_id:
return {'status': 'error', 'message': 'session_id required'}
session = _get_session(session_id)
if not session:
return {'status': 'error', 'message': 'Session not found'}
with session["lock"]:
session["last_poll_time"] = time.time()
session["output_buffer"].clear() # Prevent duplicate output on WSβHTTP switch
join_room(session_id)
logger.info(f"WebSocket client joined session room {session_id}")
return {'status': 'ok'}
@socketio.on('leave_session')
def handle_leave_session(data):
"""Client leaves a session room (AC-5)."""
session_id = data.get('session_id')
if session_id:
leave_room(session_id)
logger.info(f"WebSocket client left session room {session_id}")
@socketio.on('terminal_input')
def handle_terminal_input(data):
"""Receive keystrokes from client, write to PTY (AC-6)."""
session_id = data.get('session_id')
input_data = data.get('input', '')
session = _get_session(session_id)
if not session:
return
with session["lock"]:
session["last_poll_time"] = time.time()
fd = session["master_fd"]
try:
os.write(fd, input_data.encode())
except OSError as e:
logger.warning(f"WebSocket input write error for {session_id}: {e}")
@socketio.on('terminal_resize')
def handle_terminal_resize(data):
"""Receive resize events from client (AC-7)."""
session_id = data.get('session_id')
cols = data.get('cols', 80)
rows = data.get('rows', 24)
session = _get_session(session_id)
if not session:
return
with session["lock"]:
session["last_poll_time"] = time.time()
fd = session["master_fd"]
try:
winsize = struct.pack("HHHH", rows, cols, 0, 0)
fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
except OSError as e:
logger.warning(f"WebSocket resize error for {session_id}: {e}")
@socketio.on('heartbeat')
def handle_ws_heartbeat(data):
"""Periodic keepalive from WS client β prevents idle session reaping (AC-17)."""
session_ids = data.get('session_ids', [])
now = time.time()
for sid in session_ids:
session = _get_session(sid)
if session:
with session["lock"]:
session["last_poll_time"] = now
@socketio.on('disconnect')
def handle_ws_disconnect():
"""Log WebSocket disconnections. Do NOT auto-close PTY β client may reconnect."""
logger.info("WebSocket client disconnected")
def _get_session(session_id):
"""Get a session dict reference under the global lock. Returns None if not found."""
with sessions_lock:
return sessions.get(session_id)
def read_pty_output(session_id, fd):
"""Background thread to read PTY output into buffer and push via WebSocket."""
session = _get_session(session_id)
if not session:
return
pid = session["pid"]
session_lock = session["lock"]
while True:
with sessions_lock:
if session_id not in sessions:
break
try:
readable, _, errors = select.select([fd], [], [fd], 0.05)
if readable or errors:
output = os.read(fd, 65536)
if not output:
# EOF β process exited
break
decoded = output.decode(errors="replace")
with session_lock:
# Buffer for HTTP polling fallback (AC-15)
session["output_buffer"].append(decoded)
session["last_poll_time"] = time.time() # Keep session alive during WS output
# Push via WebSocket to the session room (AC-8)
try:
socketio.emit('terminal_output',
{'session_id': session_id, 'output': decoded},
room=session_id)
except Exception:
pass # No WebSocket clients β HTTP polling handles it
else:
# select timed out β check if process is still alive
try:
pid_result, _ = os.waitpid(pid, os.WNOHANG)
if pid_result != 0:
# Process exited
break
except ChildProcessError:
# Process already reaped
break
except OSError:
break
# Process exited or fd closed β notify WebSocket clients (AC-9)
try:
socketio.emit('session_exited', {'session_id': session_id}, room=session_id)
except Exception:
pass
logger.info(f"Session {session_id} process exited")
# Clean up immediately β no zombie sessions in the picker
if session:
terminate_session(session_id, session["pid"], session["master_fd"])
def terminate_session(session_id, pid, master_fd):
"""Gracefully terminate a session: SIGHUP -> wait -> SIGKILL -> cleanup."""
logger.info(f"Terminating stale session {session_id} (pid={pid})")
# Notify WebSocket clients that the session is closed
try:
socketio.emit('session_closed', {'session_id': session_id}, room=session_id)
except Exception:
pass
try:
os.kill(pid, signal.SIGHUP)
time.sleep(GRACEFUL_SHUTDOWN_WAIT)
# Check if still alive, force kill if needed
try:
os.kill(pid, 0) # Check if process exists
os.kill(pid, signal.SIGKILL)
logger.info(f"Force killed session {session_id} (pid={pid})")
except OSError:
pass # Already dead
os.close(master_fd)
except OSError:
pass # Process or fd already gone
with sessions_lock:
sessions.pop(session_id, None)
def _get_session_process(pid):
"""Return the name of the foreground child process for *pid*.
Uses ``pgrep -P`` to find children (works on both macOS and Linux),
then ``ps -o comm=`` to resolve the process name.
Returns:
str: process name, or ``"unknown"`` on any error / dead PID.
"""
if not isinstance(pid, int) or pid <= 0:
return "unknown"
try:
# Step 1 β find child PIDs via pgrep (cross-platform)
child_result = subprocess.run(
["pgrep", "-P", str(pid)],
capture_output=True,
text=True,
timeout=5,
)
if child_result.returncode == 0 and child_result.stdout.strip():
child_pids = child_result.stdout.strip().splitlines()
last_child_pid = child_pids[-1].strip()
# Step 2 β resolve child name
name_result = subprocess.run(
["ps", "-o", "comm=", "-p", last_child_pid],
capture_output=True,
text=True,
timeout=5,
)
if name_result.returncode == 0 and name_result.stdout.strip():
name = name_result.stdout.strip().splitlines()[0].strip()
# ps may return the full path; take basename
return os.path.basename(name)
# Step 3 β no children: fall back to the process itself
self_result = subprocess.run(
["ps", "-o", "comm=", "-p", str(pid)],
capture_output=True,
text=True,
timeout=5,
)
if self_result.returncode == 0 and self_result.stdout.strip():
name = self_result.stdout.strip().splitlines()[0].strip()
return os.path.basename(name)
return "unknown"
except Exception:
return "unknown"
def cleanup_stale_sessions():
"""Background thread that removes sessions with no recent polling."""
while True:
time.sleep(CLEANUP_INTERVAL_SECONDS)
now = time.time()
stale_sessions = []
warning_threshold = SESSION_TIMEOUT_SECONDS * 0.8
with sessions_lock:
session_snapshot = list(sessions.items())
for session_id, session in session_snapshot:
with session["lock"]:
idle = now - session["last_poll_time"]
if idle > SESSION_TIMEOUT_SECONDS:
stale_sessions.append((session_id, session["pid"], session["master_fd"]))
elif idle > warning_threshold:
session["timeout_warning"] = True
if stale_sessions:
logger.info(f"Found {len(stale_sessions)} stale session(s) to clean up")
# Terminate each stale session (outside the lock)
for session_id, pid, master_fd in stale_sessions:
terminate_session(session_id, pid, master_fd)
@app.before_request
def authorize_request():
"""Check authorization before processing any request."""
# Skip auth for health check, setup status, and Socket.IO (has own auth via connect event)
if request.path in ("/health", "/api/setup-status", "/api/pat-status", "/api/configure-pat", "/api/app-state") or request.path.startswith("/socket.io"):
return None
authorized, user = check_authorization()
if not authorized:
return jsonify({
"error": "Unauthorized",
"message": f"This app belongs to {app_owner}. You are logged in as {user}."
}), 403
return None
@app.after_request
def set_security_headers(response):
response.headers["X-Content-Type-Options"] = "nosniff"
response.headers["X-Frame-Options"] = "DENY"
response.headers["X-XSS-Protection"] = "1; mode=block"
response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
# CSP: restrict scripts to self + inline (needed for embedded <script> block),
# styles to self + inline, block all other sources. Prevents external script injection.
# connect-src allows WebSocket + API calls to self.
# Fixes: https://github.com/datasciencemonkey/coding-agents-databricks-apps/issues/58
response.headers["Content-Security-Policy"] = (
"default-src 'none'; "
"script-src 'self' 'unsafe-inline' 'wasm-unsafe-eval'; "
"style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; "
"img-src 'self' data:; "
"font-src 'self' https://fonts.gstatic.com; "
"connect-src 'self' ws: wss:; "
"frame-ancestors 'none'; "
"base-uri 'self'; "
"form-action 'self'"
)
return response
@app.route("/")
def index():
return send_from_directory("static", "index.html")
@app.route("/api/setup-status")
def get_setup_status():
return jsonify(_get_setup_state_snapshot())
@app.route("/api/app-state")
def get_app_state():
"""Admin endpoint: persisted app state (owner, last rotation)."""
return jsonify(app_state.get_state())
@app.route("/api/sessions")
def list_sessions():
"""Return a JSON array of active (non-exited) sessions with metadata."""
now = time.time()
with sessions_lock:
snapshot = list(sessions.items())
result = []
for session_id, sess in snapshot:
if sess.get("exited"):
continue
result.append({
"session_id": session_id,
"label": sess.get("label", ""),
"created_at": sess.get("created_at"),
"last_poll_time": sess.get("last_poll_time"),
"exited": False,
"process": _get_session_process(sess["pid"]),
"idle_seconds": round(now - sess.get("last_poll_time", now), 1),
})
return jsonify(result)
@app.route("/api/session/attach", methods=["POST"])
def attach_session():
"""Reattach to an existing session β returns buffered output for replay."""
data = request.get_json(silent=True) or {}
session_id = data.get("session_id", "")
sess = _get_session(session_id)
if not sess or sess.get("exited"):
return jsonify({"error": "Session not found or exited"}), 404
# Reset idle clock so the 24h reaper starts fresh
sess["last_poll_time"] = time.time()
return jsonify({
"session_id": session_id,
"label": sess.get("label", ""),
"output": list(sess["output_buffer"]),
"process": _get_session_process(sess["pid"]),
"created_at": sess.get("created_at"),
})
@app.route("/health")
def health():
with sessions_lock:
session_count = len(sessions)
with setup_lock:
current_setup_status = setup_state["status"]
return jsonify({
"status": "healthy",
"version": APP_VERSION,
"setup_status": current_setup_status,
"active_sessions": session_count,
"session_timeout_seconds": SESSION_TIMEOUT_SECONDS
})
@app.route("/api/version")
def get_version():
return jsonify({"version": APP_VERSION})
@app.route("/api/pat-status")
def pat_status():
"""Check if a valid, usable PAT is configured."""
host = ensure_https(os.environ.get("DATABRICKS_HOST", ""))
token = os.environ.get("DATABRICKS_TOKEN", "").strip()
if not token or pat_rotator.is_token_expired:
# No token, or token lifetime exceeded (rotation stopped while no sessions)
return jsonify({"configured": False, "valid": False,
"workspace_host": host})
# Validate with direct HTTP β avoids SDK auth fallback to SP
try:
resp = requests.get(f"{host}/api/2.0/preview/scim/v2/Me",
headers={"Authorization": f"Bearer {token}"}, timeout=10)
if resp.status_code == 200:
user = resp.json().get("userName", "unknown")
return jsonify({"configured": True, "valid": True, "user": user})
return jsonify({"configured": True, "valid": False,
"workspace_host": host})
except Exception:
return jsonify({"configured": True, "valid": False,
"workspace_host": host})
@app.route("/api/configure-pat", methods=["POST"])
def configure_pat():
"""Accept a user-provided PAT, validate it, and start rotation."""
data = request.json
token = data.get("token", "").strip()
if not token:
return jsonify({"error": "Token required"}), 400
# Validate the token β direct HTTP, no SDK fallback
host = ensure_https(os.environ.get("DATABRICKS_HOST", ""))
try:
resp = requests.get(f"{host}/api/2.0/preview/scim/v2/Me",
headers={"Authorization": f"Bearer {token}"}, timeout=10)
if resp.status_code != 200:
return jsonify({"error": "Invalid token"}), 400
user = resp.json().get("userName", "unknown")
except Exception as e:
return jsonify({"error": f"Token validation failed: {e}"}), 400
# Immediately mint a controlled short-lived token from the user-pasted PAT.
# This gives us a token ID we own β all future rotations can revoke the old one.
os.environ["DATABRICKS_TOKEN"] = token
pat_rotator._current_token = token
pat_rotator._current_token_id = None
rotated = pat_rotator._rotate_once()
if rotated:
token = pat_rotator.token # use the newly minted token from here on
# Revoke only the bootstrap PAT β leave other user PATs intact (#98)
pat_rotator.revoke_bootstrap_token()
else:
# Rotation failed β fall back to user-pasted token (still valid)
pat_rotator._write_databrickscfg(token)
pat_rotator.start()
# Configure all CLI tools (Claude, Codex, OpenCode, Gemini, Databricks)
_configure_all_cli_auth(pat_rotator.token or token)
# Run setup now that we have a valid token (installs CLIs, configures agents)
# Only run if setup hasn't completed yet
with setup_lock:
if setup_state["status"] != "complete":
setup_thread = threading.Thread(target=run_setup, daemon=True, name="setup-thread")
setup_thread.start()
logger.info("Setup triggered after PAT configuration")
logger.info(f"PAT configured interactively by {user} β rotation started")
return jsonify({"status": "ok", "user": user, "message": "Token configured. Auto-rotation started."})
@app.route("/api/session", methods=["POST"])