feat: adopt native uv support for Databricks Apps

app.py

@@ -21,6 +21,13 @@
 
 from utils import ensure_https
 
+# 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:
@@ -117,6 +124,11 @@ def _run_step(step_id, command):
         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)
 
@@ -198,14 +210,14 @@ def _setup_git_config():
         f.write('\n')
         f.write('echo "[post-commit] $(date +%H:%M:%S) syncing $REPO_ROOT" >> "$SYNC_LOG"\n')
         f.write('\n')
-        f.write('# Use venv python directly (avoids fragile source activate)\n')
-        f.write('VENV_PYTHON="/app/python/source_code/.venv/bin/python"\n')
-        f.write('SYNC_SCRIPT="/app/python/source_code/sync_to_workspace.py"\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 [ -x "$VENV_PYTHON" ] && [ -f "$SYNC_SCRIPT" ]; then\n')
-        f.write('    nohup "$VENV_PYTHON" "$SYNC_SCRIPT" "$REPO_ROOT" >> "$SYNC_LOG" 2>&1 & disown\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: venv=$VENV_PYTHON script=$SYNC_SCRIPT" >> "$SYNC_LOG"\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}")
@@ -259,12 +271,12 @@ def run_setup():
 
     # --- Parallel agent setup (all independent of each other) ---
     parallel_steps = [
-        ("claude",     ["python", "setup_claude.py"]),
-        ("codex",      ["python", "setup_codex.py"]),
-        ("opencode",   ["python", "setup_opencode.py"]),
-        ("gemini",     ["python", "setup_gemini.py"]),
-        ("databricks", ["python", "setup_databricks.py"]),
-        ("mlflow",     ["python", "setup_mlflow.py"]),
+        ("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"]),
+        ("databricks", ["uv", "run", "python", "setup_databricks.py"]),
+        ("mlflow",     ["uv", "run", "python", "setup_mlflow.py"]),
     ]
 
     with ThreadPoolExecutor(max_workers=len(parallel_steps)) as executor:

app.yaml

@@ -1,4 +1,6 @@
 command:
+  - uv
+  - run
   - gunicorn
   - app:app
 env:
@@ -12,8 +14,5 @@ env:
     value: databricks-gemini-3-1-pro
   - name: CODEX_MODEL
     value: databricks-gpt-5-2
-  #OPTIONAL: Move to the new Databricks Gateway if you have access (recommended), otherwise it will default to the older endpoint
-  - name: DATABRICKS_GATEWAY_HOST
-    valueFrom: DATABRICKS_GATEWAY_HOST
   - name: CLAUDE_CODE_DISABLE_AUTO_MEMORY
     value: 0

app.yaml.template

@@ -10,11 +10,5 @@ env:
     value: databricks-claude-opus-4-6
   - name: GEMINI_MODEL
     value: databricks-gemini-3-1-pro
-  #OPTIONAL: Use the new Databricks AI Gateway if you have access (recommended), otherwise it will default to the older endpoint
-  - name: DATABRICKS_GATEWAY_HOST
-    value: https://<your-gateway-id>.ai-gateway.<env>.cloud.databricks.com
-  # NOTE: CLAUDE_CODE_DISABLE_AUTO_MEMORY=0 enables auto memory, allowing Claude Code to
-  # persist context and history across sessions. This flag is temporary — once Anthropic
-  # completes the rollout and auto memory is on by default, this can be removed entirely.
   - name: CLAUDE_CODE_DISABLE_AUTO_MEMORY
     value: 0

pyproject.toml

@@ -2,9 +2,12 @@
 name = "coda"
 version = "0.16.4"
 description = "CoDA - Coding Agents on Databricks Apps"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 dependencies = [
     "flask>=2.0",
+    "flask-socketio>=5.3",
+    "simple-websocket>=1.0",
+    "gunicorn>=21.0",
     "claude-agent-sdk",
     "databricks-sdk>=0.20.0",
     "mlflow[genai]>=3.4",

tests/test_heartbeat.py

@@ -1,5 +1,6 @@
 """Tests for /api/heartbeat endpoint — lightweight keep-alive."""
 
+import threading
 import time
 from collections import deque
 from unittest import mock
@@ -27,6 +28,7 @@ def _create_fake_session(app_module, session_id="test-session-123", **overrides)
         "output_buffer": deque(maxlen=1000),
         "last_poll_time": time.time() - 60,  # 60s ago
         "created_at": time.time(),
+        "lock": threading.Lock(),
     }
     session.update(overrides)
     with app_module.sessions_lock: