Add Telegram session picker
Anmol Malik committed
May 28, 2026 at 23:17 UTC
1ff45e71c24e369beaf8c01bbe8b41c8a80e14b0
3 files changed
+281
-1
helpers/integration_commands.py
+26
@@ -35,6 +35,12 @@ COMMAND_REGISTRY: tuple[IntegrationCommandDef, ...] = (
35
"Info",
36
),
37
IntegrationCommandDef("new", "Start a fresh chat context.", "Session"),
38
+ IntegrationCommandDef(
39
+ "sessions",
40
+ "Show or switch recent chat sessions.",
41
+ "Session",
42
+ aliases=("session",),
43
+ ),
44
IntegrationCommandDef("clear", "Reset the current chat context.", "Session", aliases=("reset",)),
45
IntegrationCommandDef(
46
"queue",
@@ -168,6 +174,8 @@ def try_handle_command(context: "AgentContext", text: str) -> str | None:
174
return help_text(full=True)
175
if command == "/status":
176
return _handle_status(context)
177
+ if command == "/sessions":
178
+ return _handle_sessions(context)
179
if command in {"/new", "/clear"}:
180
return _handle_clear(context, new_chat=(command == "/new"))
181
if command == "/send":
@@ -251,6 +259,24 @@ def _handle_status(context: "AgentContext") -> str:
259
)
260
261
262
+def _handle_sessions(context: "AgentContext") -> str:
263
+ from agent import AgentContext
264
+
265
+ contexts = sorted(
266
+ AgentContext.all(),
267
+ key=lambda item: str(item.output().get("last_message") or ""),
268
+ reverse=True,
269
+ )
270
+ lines = ["Recent sessions:"]
271
+ for item in contexts[:4]:
272
+ marker = " (current)" if item.id == context.id else ""
273
+ running = " - running" if item.is_running() else ""
274
+ lines.append(f"- {item.name or item.id}{marker}{running}")
275
+ if len(contexts) > 4:
276
+ lines.append(f"And {len(contexts) - 4} more. Use Telegram buttons to page through them.")
277
+ return "\n".join(lines)
278
+
279
+
280
def _handle_clear(context: "AgentContext", *, new_chat: bool) -> str:
281
context.reset()
282
mq.remove(context)
plugins/_telegram_integration/helpers/command_ui.py
+177
-1
@@ -1,9 +1,10 @@
1
from __future__ import annotations
2
3
+import json
4
from dataclasses import dataclass
5
6
from agent import AgentContext
6
-from helpers import projects, subagents
7
+from helpers import files, projects, subagents
8
from helpers import integration_commands
9
from helpers.persist_chat import save_tmp_chat
10
from helpers.state_monitor_integration import mark_dirty_for_context
@@ -12,9 +13,16 @@ from plugins._telegram_integration.helpers import telegram_client as tc
13
from plugins._telegram_integration.helpers.constants import (
14
CTX_TG_STREAM_ENABLED,
15
CTX_TG_TOOLS_ENABLED,
16
+ CTX_TG_BOT,
17
+ CTX_TG_CHAT_ID,
18
+ CTX_TG_CHAT_TYPE,
19
+ CTX_TG_USER_ID,
20
+ CTX_TG_USERNAME,
21
+ STATE_FILE,
22
)
23
24
PAGE_SIZE = 8
25
+SESSION_PAGE_SIZE = 4
26
CALLBACK_PREFIX = "tg"
27
28
@@ -24,6 +32,14 @@ class PickerItem:
32
label: str
33
34
35
+@dataclass(frozen=True)
36
+class SessionItem:
37
+ context: AgentContext
38
+ label: str
39
+ last_message: str
40
+ running: bool
41
+
42
+
43
async def handle_command(
44
context: AgentContext,
45
token: str,
@@ -44,6 +60,9 @@ async def handle_command(
60
if command in {"/agent", "/profile"} and not args:
61
await send_agent_picker(context, token, chat_id, reply_to_message_id, 0)
62
return True
63
+ if command in {"/sessions", "/session"} and not args:
64
+ await send_session_picker(context, token, chat_id, reply_to_message_id, 0)
65
+ return True
66
if command == "/stream":
67
await send_toggle_picker(
68
context,
@@ -91,6 +110,8 @@ async def handle_callback(
110
await edit_project_picker(context, token, chat_id, message_id, page)
111
elif kind == "agent":
112
await edit_agent_picker(context, token, chat_id, message_id, page)
113
+ elif kind == "session":
114
+ await edit_session_picker(context, token, chat_id, message_id, page)
115
return True
116
if kind == "model" and action in {"set", "clear"}:
117
await _select_model(context, _safe_int(value), clear=(action == "clear"))
@@ -104,6 +125,10 @@ async def handle_callback(
125
await _select_agent(context, _safe_int(value))
126
await edit_agent_picker(context, token, chat_id, message_id, 0, selected=True)
127
return True
128
+ if kind == "session" and action == "set":
129
+ selected_context = await _select_session(context, _safe_int(value))
130
+ await edit_session_picker(selected_context or context, token, chat_id, message_id, 0, selected=bool(selected_context))
131
+ return True
132
if kind in {"stream", "tools"} and action in {"on", "off"}:
133
key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED
134
label = "Response streaming" if kind == "stream" else "Tool progress"
@@ -217,6 +242,30 @@ async def edit_toggle_picker(
242
await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
243
244
245
+async def send_session_picker(
246
+ context: AgentContext,
247
+ token: str,
248
+ chat_id: int,
249
+ reply_to_message_id: int | None,
250
+ page: int,
251
+) -> None:
252
+ text, markup = _session_view(context, page)
253
+ await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup)
254
+
255
+
256
+async def edit_session_picker(
257
+ context: AgentContext,
258
+ token: str,
259
+ chat_id: int,
260
+ message_id: int,
261
+ page: int,
262
+ *,
263
+ selected: bool = False,
264
+) -> None:
265
+ text, markup = _session_view(context, page, selected=selected)
266
+ await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup)
267
+
268
+
269
def _model_view(
270
context: AgentContext,
271
page: int,
@@ -295,6 +344,48 @@ def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict
344
return text, {"inline_keyboard": rows}
345
346
347
+def _session_view(
348
+ context: AgentContext,
349
+ page: int,
350
+ *,
351
+ selected: bool = False,
352
+) -> tuple[str, dict | None]:
353
+ items = _session_items()
354
+ current_label = _session_label(context)
355
+ status = f"Current session: <b>{_html(current_label)}</b>"
356
+ if context.is_running():
357
+ status += "\nSession switching is available after the current run finishes."
358
+ elif selected:
359
+ status = "Session switched.\n" + status
360
+ if not items:
361
+ return status + "\nNo sessions were found.", None
362
+
363
+ total = len(items)
364
+ page = _clamp_page(page, total, SESSION_PAGE_SIZE)
365
+ start = page * SESSION_PAGE_SIZE
366
+ end = min(start + SESSION_PAGE_SIZE, total)
367
+ rows: list[list[dict[str, str]]] = []
368
+ for index, item in enumerate(items[start:end], start=start):
369
+ marker = "• " if item.context.id == context.id else ""
370
+ suffix = " (running)" if item.running else ""
371
+ action = "noop" if context.is_running() or item.running else "set"
372
+ rows.append([{
373
+ "text": f"{marker}{item.label}{suffix}"[:64],
374
+ "callback_data": f"tg:session:{action}:{index}",
375
+ }])
376
+
377
+ nav: list[dict[str, str]] = []
378
+ if page > 0:
379
+ nav.append({"text": "Prev", "callback_data": f"tg:session:page:{page - 1}"})
380
+ if end < total:
381
+ nav.append({"text": "Next", "callback_data": f"tg:session:page:{page + 1}"})
382
+ if nav:
383
+ rows.append(nav)
384
+
385
+ range_text = f"\nShowing {start + 1}-{end} of {total}."
386
+ return status + range_text, {"inline_keyboard": rows}
387
+
388
+
389
async def _select_model(context: AgentContext, index: int, *, clear: bool = False) -> None:
390
if not model_config.is_chat_override_allowed(context.agent0):
391
return
@@ -339,6 +430,86 @@ async def _select_agent(context: AgentContext, index: int) -> None:
430
mark_dirty_for_context(context.id, reason="telegram.agent_select")
431
432
433
+async def _select_session(context: AgentContext, index: int) -> AgentContext | None:
434
+ if context.is_running():
435
+ return None
436
+ items = _session_items()
437
+ if index < 0 or index >= len(items):
438
+ return None
439
+ target = items[index].context
440
+ if target.is_running():
441
+ return None
442
+
443
+ _copy_telegram_binding(context, target)
444
+ _set_session_mapping(target)
445
+ save_tmp_chat(target)
446
+ mark_dirty_for_context(context.id, reason="telegram.session_unselect")
447
+ mark_dirty_for_context(target.id, reason="telegram.session_select")
448
+ return target
449
+
450
+
451
+def _session_items() -> list[SessionItem]:
452
+ contexts = sorted(
453
+ AgentContext.all(),
454
+ key=lambda item: str(item.output().get("last_message") or ""),
455
+ reverse=True,
456
+ )
457
+ return [
458
+ SessionItem(
459
+ context=item,
460
+ label=_session_label(item),
461
+ last_message=str(item.output().get("last_message") or ""),
462
+ running=item.is_running(),
463
+ )
464
+ for item in contexts
465
+ ]
466
+
467
+
468
+def _session_label(context: AgentContext) -> str:
469
+ return str(context.name or context.id or "Session")
470
+
471
+
472
+def _copy_telegram_binding(source: AgentContext, target: AgentContext) -> None:
473
+ for key in (
474
+ CTX_TG_BOT,
475
+ CTX_TG_CHAT_ID,
476
+ CTX_TG_CHAT_TYPE,
477
+ CTX_TG_USER_ID,
478
+ CTX_TG_USERNAME,
479
+ ):
480
+ if key in source.data:
481
+ target.data[key] = source.data[key]
482
+
483
+
484
+def _set_session_mapping(context: AgentContext) -> None:
485
+ bot_name = str(context.data.get(CTX_TG_BOT) or "")
486
+ user_id = context.data.get(CTX_TG_USER_ID)
487
+ chat_id = context.data.get(CTX_TG_CHAT_ID)
488
+ if not bot_name or user_id is None or chat_id is None:
489
+ return
490
+ key = f"{bot_name}:{int(user_id)}:{int(chat_id)}"
491
+ state = _load_telegram_state()
492
+ chats = state.setdefault("chats", {})
493
+ chats[key] = context.id
494
+ _save_telegram_state(state)
495
+
496
+
497
+def _load_telegram_state() -> dict:
498
+ path = files.get_abs_path(STATE_FILE)
499
+ if not files.exists(path):
500
+ return {}
501
+ try:
502
+ return json.loads(files.read_file(path))
503
+ except Exception:
504
+ return {}
505
+
506
+
507
+def _save_telegram_state(state: dict) -> None:
508
+ path = files.get_abs_path(STATE_FILE)
509
+ files.make_dirs(path)
510
+ files.write_file(path, json.dumps(state))
511
+
512
+
513
def _paged_buttons(
514
kind: str,
515
items: list[PickerItem],
@@ -390,6 +561,11 @@ def _safe_int(value: str) -> int:
561
return 0
562
563
564
+def _clamp_page(page: int, total_items: int, page_size: int) -> int:
565
+ total_pages = max(1, (total_items + page_size - 1) // page_size)
566
+ return min(max(0, page), total_pages - 1)
567
+
568
+
569
def _label_for(items: list[PickerItem], key: str) -> str:
570
for item in items:
571
if item.key == key:
tests/test_telegram_sessions_picker.py
new
+78
@@ -0,0 +1,78 @@
1
+import asyncio
2
+from types import SimpleNamespace
3
+
4
+from plugins._telegram_integration.helpers import command_ui
5
+from plugins._telegram_integration.helpers.constants import (
6
+ CTX_TG_BOT,
7
+ CTX_TG_CHAT_ID,
8
+ CTX_TG_USER_ID,
9
+ CTX_TG_USERNAME,
10
+)
11
+
12
+
13
+class FakeContext:
14
+ def __init__(self, id, name, last_message, *, running=False):
15
+ self.id = id
16
+ self.name = name
17
+ self.data = {}
18
+ self._last_message = last_message
19
+ self._running = running
20
+
21
+ def output(self):
22
+ return {"last_message": self._last_message}
23
+
24
+ def is_running(self):
25
+ return self._running
26
+
27
+
28
+def test_session_picker_shows_four_recent_sessions(monkeypatch):
29
+ contexts = [
30
+ FakeContext("ctx-1", "One", "2026-05-01T10:00:00"),
31
+ FakeContext("ctx-2", "Two", "2026-05-02T10:00:00"),
32
+ FakeContext("ctx-3", "Three", "2026-05-03T10:00:00"),
33
+ FakeContext("ctx-4", "Four", "2026-05-04T10:00:00"),
34
+ FakeContext("ctx-5", "Five", "2026-05-05T10:00:00"),
35
+ ]
36
+ current = contexts[4]
37
+ monkeypatch.setattr(command_ui, "AgentContext", SimpleNamespace(all=lambda: contexts))
38
+
39
+ text, markup = command_ui._session_view(current, 0)
40
+
41
+ assert "Current session: <b>Five</b>" in text
42
+ assert "Showing 1-4 of 5." in text
43
+ rows = markup["inline_keyboard"]
44
+ assert [row[0]["text"] for row in rows[:4]] == ["• Five", "Four", "Three", "Two"]
45
+ assert rows[-1] == [{"text": "Next", "callback_data": "tg:session:page:1"}]
46
+
47
+
48
+def test_select_session_remaps_telegram_chat(monkeypatch):
49
+ current = FakeContext("ctx-current", "Current", "2026-05-03T10:00:00")
50
+ target = FakeContext("ctx-target", "Target", "2026-05-04T10:00:00")
51
+ current.data.update(
52
+ {
53
+ CTX_TG_BOT: "main",
54
+ CTX_TG_CHAT_ID: 123,
55
+ CTX_TG_USER_ID: 456,
56
+ CTX_TG_USERNAME: "anmol",
57
+ }
58
+ )
59
+ saved = []
60
+ dirty = []
61
+ state_out = {}
62
+
63
+ monkeypatch.setattr(command_ui, "AgentContext", SimpleNamespace(all=lambda: [current, target]))
64
+ monkeypatch.setattr(command_ui, "save_tmp_chat", lambda context: saved.append(context.id))
65
+ monkeypatch.setattr(command_ui, "mark_dirty_for_context", lambda context_id, *, reason=None: dirty.append((context_id, reason)))
66
+ monkeypatch.setattr(command_ui, "_load_telegram_state", lambda: {"chats": {}})
67
+ monkeypatch.setattr(command_ui, "_save_telegram_state", lambda state: state_out.update(state))
68
+
69
+ selected = asyncio.run(command_ui._select_session(current, 0))
70
+
71
+ assert selected is target
72
+ assert target.data[CTX_TG_BOT] == "main"
73
+ assert target.data[CTX_TG_CHAT_ID] == 123
74
+ assert target.data[CTX_TG_USER_ID] == 456
75
+ assert target.data[CTX_TG_USERNAME] == "anmol"
76
+ assert state_out["chats"]["main:456:123"] == "ctx-target"
77
+ assert saved == ["ctx-target"]
78
+ assert ("ctx-target", "telegram.session_select") in dirty