@@ -903,8 +903,61 @@ def sync_offline_cmd(
903903 _exit (1 )
904904
905905
906- @app .command ("jupyter-setup" )
907- def jupyter_setup (
906+ # ---------------------------------------------------------------------------
907+ # jupyter subcommand group — live session introspection + extension setup
908+ # ---------------------------------------------------------------------------
909+
910+
911+ jupyter_app = typer .Typer (
912+ help = (
913+ "Live Jupyter session introspection (whoami, discover) plus the "
914+ "extension-disable/enable recipe needed before the Jupyter MCP can "
915+ "talk to a cluster JupyterLab."
916+ ),
917+ no_args_is_help = True ,
918+ )
919+ app .add_typer (jupyter_app , name = "jupyter" )
920+
921+
922+ def _do_jupyter_setup (dry_run : bool ) -> None :
923+ """Idempotent Jupyter extension reconfiguration.
924+
925+ Body shared between ``aexp jupyter setup`` (the canonical verb) and the
926+ deprecated flat ``aexp jupyter-setup`` alias.
927+ """
928+ import subprocess
929+ import sys as _sys
930+
931+ cmds : list [list [str ]] = [
932+ [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "disable" , "jupyter_server_documents" ],
933+ [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "enable" , "jupyter_server_ydoc" ],
934+ [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "enable" , "jupyter_server_nbmodel" ],
935+ [_sys .executable , "-m" , "jupyter" , "labextension" , "disable" , "@jupyter-ai-contrib/server-documents" ],
936+ ]
937+
938+ for cmd in cmds :
939+ printable = " " .join (cmd )
940+ if dry_run :
941+ console .print (f"[cyan][dry-run][/cyan] { printable } )
942+ continue
943+ console .print (f"[dim]$[/dim] { printable } )
944+ result = subprocess .run (cmd , capture_output = True , text = True )
945+ if result .stdout :
946+ console .print (result .stdout .rstrip ())
947+ if result .returncode != 0 :
948+ console .print (f"[red]FAILED ({ result .returncode } { result .stderr .rstrip ()} )
949+ _exit (1 )
950+ return
951+
952+ if dry_run :
953+ console .print ("\n [cyan]dry-run complete[/cyan] — no changes applied." )
954+ return
955+ console .print ("\n [green]✓[/green] Jupyter extension state configured." )
956+ console .print (" Restart your JupyterLab process to pick up the changes." )
957+
958+
959+ @jupyter_app .command ("setup" )
960+ def jupyter_setup_cmd (
908961 dry_run : bool = typer .Option (
909962 False ,
910963 "--dry-run" ,
@@ -940,35 +993,129 @@ def jupyter_setup(
940993 See `docs/setup/jupyter-mcp.md` "Investigation log" §1-3 for the full
941994 rationale. Restart your JupyterLab process to pick up the changes.
942995 """
943- import subprocess
944- import sys as _sys
996+ _do_jupyter_setup (dry_run )
945997
946- cmds : list [list [str ]] = [
947- [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "disable" , "jupyter_server_documents" ],
948- [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "enable" , "jupyter_server_ydoc" ],
949- [_sys .executable , "-m" , "jupyter" , "server" , "extension" , "enable" , "jupyter_server_nbmodel" ],
950- [_sys .executable , "-m" , "jupyter" , "labextension" , "disable" , "@jupyter-ai-contrib/server-documents" ],
951- ]
952998
953- for cmd in cmds :
954- printable = " " .join (cmd )
955- if dry_run :
956- console .print (f"[cyan][dry-run][/cyan] { printable } )
957- continue
958- console .print (f"[dim]$[/dim] { printable } )
959- result = subprocess .run (cmd , capture_output = True , text = True )
960- if result .stdout :
961- console .print (result .stdout .rstrip ())
962- if result .returncode != 0 :
963- console .print (f"[red]FAILED ({ result .returncode } { result .stderr .rstrip ()} )
964- _exit (1 )
999+ @app .command ("jupyter-setup" , hidden = True )
1000+ def jupyter_setup_deprecated_alias (
1001+ dry_run : bool = typer .Option (
1002+ False ,
1003+ "--dry-run" ,
1004+ "-n" ,
1005+ help = "Print the commands that would be run without executing them." ,
1006+ ),
1007+ ) -> None :
1008+ """Deprecated alias for ``aexp jupyter setup``. Kept for one release."""
1009+ console .print (
1010+ "[yellow]warning[/yellow] `aexp jupyter-setup` is deprecated; use "
1011+ "`aexp jupyter setup` (subcommand) instead."
1012+ )
1013+ _do_jupyter_setup (dry_run )
1014+
1015+
1016+ def _print_session_info (json_out : bool ) -> None :
1017+ """Shared body for ``aexp jupyter init`` / ``aexp jupyter whoami``."""
1018+ try :
1019+ from aexp .jupyter import _print_info_human , init
1020+ except ImportError as exc :
1021+ console .print (f"[red]{ exc } )
1022+ _exit (1 )
1023+ return
1024+ info = init ()
1025+ if json_out :
1026+ # model_dump_json is the wire format the MCP recipe also produces.
1027+ print (info .model_dump_json (indent = 2 ))
1028+ return
1029+ _print_info_human (info )
1030+
1031+
1032+ @jupyter_app .command ("init" )
1033+ def jupyter_init_cmd (
1034+ json_out : bool = typer .Option (
1035+ False , "--json" , help = "Emit SessionInfo as JSON instead of human-readable text."
1036+ ),
1037+ ) -> None :
1038+ """Introspect the current Jupyter session and print the result.
1039+
1040+ Most useful when run from inside a Jupyter terminal or a notebook cell
1041+ (``!aexp jupyter init``) — populates SLURM context, attached
1042+ notebooks, GPU residents, and sibling Jupyters from live state.
1043+ """
1044+ _print_session_info (json_out )
1045+
1046+
1047+ @jupyter_app .command ("whoami" )
1048+ def jupyter_whoami_cmd (
1049+ json_out : bool = typer .Option (
1050+ False , "--json" , help = "Emit SessionInfo as JSON instead of human-readable text."
1051+ ),
1052+ ) -> None :
1053+ """Alias for ``aexp jupyter init`` — identical output, friendlier verb."""
1054+ _print_session_info (json_out )
1055+
1056+
1057+ @jupyter_app .command ("discover" )
1058+ def jupyter_discover_cmd (
1059+ json_out : bool = typer .Option (
1060+ False , "--json" , help = "Emit the sibling list as JSON instead of a table."
1061+ ),
1062+ describe : bool = typer .Option (
1063+ False ,
1064+ "--describe" ,
1065+ help = (
1066+ "Probe each sibling's /api/sessions and /api/kernels for attached "
1067+ "notebook paths + kernel state (one HTTP roundtrip per server)."
1068+ ),
1069+ ),
1070+ ) -> None :
1071+ """List every Jupyter the user has running, excluding the current one."""
1072+ try :
1073+ from aexp .jupyter import describe_server , discover_other_servers
1074+ except ImportError as exc :
1075+ console .print (f"[red]{ exc } )
1076+ _exit (1 )
1077+ return
1078+ siblings = discover_other_servers ()
1079+ if describe :
1080+ rendered = []
1081+ for s in siblings :
1082+ entry = s .model_dump ()
1083+ entry ["describe" ] = describe_server (s .url , s .token )
1084+ rendered .append (entry )
1085+ if json_out :
1086+ import json as _json
1087+ print (_json .dumps (rendered , indent = 2 , default = str ))
9651088 return
1089+ for entry in rendered :
1090+ console .print (
1091+ f"[cyan]{ entry ['url' ]} { entry ['port' ]}
1092+ f"pid={ entry ['pid' ]} { entry ['hostname' ]}
1093+ )
1094+ d = entry ["describe" ]
1095+ if d .get ("attached_notebooks" ):
1096+ console .print (" notebooks:" )
1097+ for nb in d ["attached_notebooks" ]:
1098+ console .print (f" - { nb } )
1099+ for k in d .get ("kernels" , []):
1100+ console .print (
1101+ f" kernel: id={ k .get ('id' )} { k .get ('execution_state' )}
1102+ f"last={ k .get ('last_activity' )}
1103+ )
1104+ if not rendered :
1105+ console .print ("[dim]no other Jupyter servers visible[/dim]" )
1106+ return
9661107
967- if dry_run :
968- console .print ("\n [cyan]dry-run complete[/cyan] — no changes applied." )
1108+ if json_out :
1109+ import json as _json
1110+ print (_json .dumps ([s .model_dump () for s in siblings ], indent = 2 , default = str ))
9691111 return
970- console .print ("\n [green]✓[/green] Jupyter extension state configured." )
971- console .print (" Restart your JupyterLab process to pick up the changes." )
1112+ if not siblings :
1113+ console .print ("[dim]no other Jupyter servers visible[/dim]" )
1114+ return
1115+ for s in siblings :
1116+ console .print (
1117+ f"[cyan]{ s .url } { s .port } { s .pid } { s .hostname }
1118+ )
9721119
9731120
9741121@app .command ("install-slash-commands" )
0 commit comments