feat: Add kernel restart command as this is necessary when updated existing (#16)
jacktday committed
May 20, 2026 at 06:15 UTC
8fdadb94782b3694f091a1926514df06aaa5860d
2 files changed
+38
src/colab_cli/commands/session.py
+32
@@ -27,6 +27,7 @@ from colab_cli.client import (
27
PostAssignmentResponse,
28
Variant,
29
)
30
+from colab_cli.commands.automation import INTERACTIVE_AUTOMATION_TIMEOUT_SEC
31
from colab_cli.utils import get_status_code
32
from colab_cli.state import SessionState
33
from colab_cli.runtime import ColabRuntime
@@ -246,6 +247,36 @@ def new(
247
typer.echo("[colab] Session READY.")
248
249
250
+def restart(
251
+ session: Annotated[
252
+ Optional[str], typer.Option("-s", "--session", help="Session name")
253
+ ] = None,
254
+):
255
+ from colab_cli.common import state
256
+
257
+ name = state.resolve_session(session)
258
+ s = state.store.get(name)
259
+
260
+ def on_started(kid):
261
+ s.kernel_id = kid
262
+ state.store.add(s)
263
+
264
+ def on_sess_started(sid):
265
+ s.session_id = sid
266
+ state.store.add(s)
267
+
268
+ runtime = ColabRuntime(
269
+ s.url,
270
+ s.token,
271
+ kernel_id=s.kernel_id,
272
+ session_id=s.session_id,
273
+ on_kernel_started=on_started,
274
+ on_session_started=on_sess_started,
275
+ )
276
+
277
+ runtime.restart(timeout=INTERACTIVE_AUTOMATION_TIMEOUT_SEC)
278
+
279
+
280
def sessions_command():
281
"""List all active sessions"""
282
from colab_cli.common import state
@@ -471,6 +502,7 @@ def keep_alive(
502
def register(app: typer.Typer):
503
app.command()(new)
504
app.command(name="sessions")(sessions_command)
505
+ app.command()(restart)
506
app.command()(status)
507
app.command()(stop)
508
app.command(hidden=True)(keep_alive)
src/colab_cli/runtime.py
+6
@@ -154,6 +154,12 @@ class ColabRuntime:
154
raise e
155
156
return self._kernel_client
157
+
158
+ def restart(
159
+ self,
160
+ timeout: Optional[float] = None,
161
+ ):
162
+ self.kernel_client.restart(timeout=timeout)
163
164
def execute_code(
165
self,