Skip to content

Commit fcd89b9

Browse files
GPSxtremeGPSxtreme
committed
Improve sampling structuring, include more sampling functionality
1 parent 131c656 commit fcd89b9

File tree

7 files changed

+1047
-17
lines changed

7 files changed

+1047
-17
lines changed

‎src/config.ts‎

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
import { config } from "dotenv";
2+
import { z } from "zod";
3+
import { MessageType, TemplateType } from "./sampling/types.js";
4+
5+
config();
6+
7+
const envSchema = z.object({
8+
// Required
9+
TELEGRAM_BOT_TOKEN: z.string(),
10+
11+
// Sampling control
12+
SAMPLING_ENABLED: z.coerce.boolean().default(true),
13+
14+
// Response trigger settings
15+
SAMPLING_MENTION_ONLY: z.coerce.boolean().default(true),
16+
SAMPLING_RESPOND_TO_DMS: z.coerce.boolean().default(true),
17+
18+
// Access control (comma-separated lists)
19+
SAMPLING_ALLOWED_CHATS: z
20+
.string()
21+
.default("")
22+
.transform((val) =>
23+
val
24+
? val
25+
.split(",")
26+
.map((id) => id.trim())
27+
.filter((id) => id.length > 0)
28+
.map((id) => {
29+
// Handle numeric IDs
30+
const numId = Number.parseInt(id);
31+
if (!Number.isNaN(numId)) return numId;
32+
// Handle username strings (normalize to include @)
33+
return id.startsWith("@") ? id : `@${id}`;
34+
})
35+
: [],
36+
),
37+
SAMPLING_BLOCKED_CHATS: z
38+
.string()
39+
.default("")
40+
.transform((val) =>
41+
val
42+
? val
43+
.split(",")
44+
.map((id) => id.trim())
45+
.filter((id) => id.length > 0)
46+
.map((id) => {
47+
// Handle numeric IDs
48+
const numId = Number.parseInt(id);
49+
if (!Number.isNaN(numId)) return numId;
50+
// Handle username strings (normalize to include @)
51+
return id.startsWith("@") ? id : `@${id}`;
52+
})
53+
: [],
54+
),
55+
SAMPLING_ALLOWED_USERS: z
56+
.string()
57+
.default("")
58+
.transform((val) =>
59+
val
60+
? val
61+
.split(",")
62+
.map((id) => Number.parseInt(id.trim()))
63+
.filter((id) => !Number.isNaN(id))
64+
: [],
65+
),
66+
SAMPLING_BLOCKED_USERS: z
67+
.string()
68+
.default("")
69+
.transform((val) =>
70+
val
71+
? val
72+
.split(",")
73+
.map((id) => Number.parseInt(id.trim()))
74+
.filter((id) => !Number.isNaN(id))
75+
: [],
76+
),
77+
SAMPLING_ADMIN_USERS: z
78+
.string()
79+
.default("")
80+
.transform((val) =>
81+
val
82+
? val
83+
.split(",")
84+
.map((id) => Number.parseInt(id.trim()))
85+
.filter((id) => !Number.isNaN(id))
86+
: [],
87+
),
88+
89+
// Message type handlers
90+
SAMPLING_ENABLE_TEXT: z.coerce.boolean().default(true),
91+
SAMPLING_ENABLE_PHOTO: z.coerce.boolean().default(false),
92+
SAMPLING_ENABLE_DOCUMENT: z.coerce.boolean().default(false),
93+
SAMPLING_ENABLE_VOICE: z.coerce.boolean().default(false),
94+
SAMPLING_ENABLE_VIDEO: z.coerce.boolean().default(false),
95+
SAMPLING_ENABLE_STICKER: z.coerce.boolean().default(false),
96+
SAMPLING_ENABLE_LOCATION: z.coerce.boolean().default(false),
97+
SAMPLING_ENABLE_CONTACT: z.coerce.boolean().default(false),
98+
SAMPLING_ENABLE_POLL: z.coerce.boolean().default(false),
99+
100+
// Response behavior
101+
SAMPLING_MAX_TOKENS: z.coerce.number().default(1000),
102+
SAMPLING_SHOW_TYPING: z.coerce.boolean().default(true),
103+
SAMPLING_SILENT_MODE: z.coerce.boolean().default(false),
104+
105+
// Rate limiting
106+
SAMPLING_RATE_LIMIT_USER: z.coerce.number().default(10),
107+
SAMPLING_RATE_LIMIT_CHAT: z.coerce.number().default(20),
108+
109+
// Message filters
110+
SAMPLING_MIN_MESSAGE_LENGTH: z.coerce.number().default(1),
111+
SAMPLING_MAX_MESSAGE_LENGTH: z.coerce.number().default(1000),
112+
SAMPLING_KEYWORD_TRIGGERS: z
113+
.string()
114+
.default("")
115+
.transform((val) =>
116+
val
117+
? val
118+
.split(",")
119+
.map((keyword) => keyword.trim())
120+
.filter((keyword) => keyword.length > 0)
121+
: [],
122+
),
123+
SAMPLING_IGNORE_COMMANDS: z.coerce.boolean().default(true),
124+
});
125+
126+
export const env = envSchema.parse(process.env);
127+
128+
// Create sampling config object from env variables
129+
export const samplingConfig = {
130+
// Sampling control
131+
enabled: env.SAMPLING_ENABLED,
132+
133+
// Response trigger settings
134+
mentionOnly: env.SAMPLING_MENTION_ONLY,
135+
respondToPrivateMessages: env.SAMPLING_RESPOND_TO_DMS,
136+
137+
// Access control
138+
allowedChats: env.SAMPLING_ALLOWED_CHATS,
139+
blockedChats: env.SAMPLING_BLOCKED_CHATS,
140+
allowedUsers: env.SAMPLING_ALLOWED_USERS,
141+
blockedUsers: env.SAMPLING_BLOCKED_USERS,
142+
adminUsers: env.SAMPLING_ADMIN_USERS,
143+
144+
// Message type handlers
145+
enabledListeners: {
146+
[MessageType.TEXT]: env.SAMPLING_ENABLE_TEXT,
147+
[MessageType.PHOTO]: env.SAMPLING_ENABLE_PHOTO,
148+
[MessageType.DOCUMENT]: env.SAMPLING_ENABLE_DOCUMENT,
149+
[MessageType.VOICE]: env.SAMPLING_ENABLE_VOICE,
150+
[MessageType.VIDEO]: env.SAMPLING_ENABLE_VIDEO,
151+
[MessageType.STICKER]: env.SAMPLING_ENABLE_STICKER,
152+
[MessageType.LOCATION]: env.SAMPLING_ENABLE_LOCATION,
153+
[MessageType.CONTACT]: env.SAMPLING_ENABLE_CONTACT,
154+
[MessageType.POLL]: env.SAMPLING_ENABLE_POLL,
155+
},
156+
157+
// Response behavior
158+
maxTokens: env.SAMPLING_MAX_TOKENS,
159+
showTypingIndicator: env.SAMPLING_SHOW_TYPING,
160+
silentMode: env.SAMPLING_SILENT_MODE,
161+
162+
// Rate limiting
163+
rateLimitPerUser: env.SAMPLING_RATE_LIMIT_USER,
164+
rateLimitPerChat: env.SAMPLING_RATE_LIMIT_CHAT,
165+
166+
// Message filters
167+
minMessageLength: env.SAMPLING_MIN_MESSAGE_LENGTH,
168+
maxMessageLength: env.SAMPLING_MAX_MESSAGE_LENGTH,
169+
keywordTriggers: env.SAMPLING_KEYWORD_TRIGGERS,
170+
ignoreCommands: env.SAMPLING_IGNORE_COMMANDS,
171+
172+
// Response templates
173+
templates: {
174+
[TemplateType.TEXT]: `NEW TELEGRAM MESSAGE FROM:
175+
user_id: {userId}
176+
chat_id: {chatId}
177+
isDM: {isDM}
178+
message_id: {messageId}
179+
message_type: ${MessageType.TEXT}
180+
content: {content}`,
181+
182+
[TemplateType.PHOTO]: `NEW PHOTO MESSAGE FROM:
183+
user_id: {userId}
184+
chat_id: {chatId}
185+
isDM: {isDM}
186+
message_id: {messageId}
187+
message_type: ${MessageType.PHOTO}
188+
caption: {caption}
189+
photo_info: {photoInfo}`,
190+
191+
[TemplateType.DOCUMENT]: `NEW DOCUMENT MESSAGE FROM:
192+
user_id: {userId}
193+
chat_id: {chatId}
194+
isDM: {isDM}
195+
message_id: {messageId}
196+
message_type: ${MessageType.DOCUMENT}
197+
filename: {fileName}
198+
mime_type: {mimeType}
199+
caption: {caption}`,
200+
201+
[TemplateType.VOICE]: `NEW VOICE MESSAGE FROM:
202+
user_id: {userId}
203+
chat_id: {chatId}
204+
isDM: {isDM}
205+
message_id: {messageId}
206+
message_type: ${MessageType.VOICE}
207+
duration: {duration}s`,
208+
209+
[TemplateType.VIDEO]: `NEW VIDEO MESSAGE FROM:
210+
user_id: {userId}
211+
chat_id: {chatId}
212+
isDM: {isDM}
213+
message_id: {messageId}
214+
message_type: ${MessageType.VIDEO}
215+
caption: {caption}
216+
duration: {duration}s`,
217+
218+
[TemplateType.STICKER]: `NEW STICKER MESSAGE FROM:
219+
user_id: {userId}
220+
chat_id: {chatId}
221+
isDM: {isDM}
222+
message_id: {messageId}
223+
message_type: ${MessageType.STICKER}
224+
emoji: {stickerEmoji}
225+
set_name: {stickerSetName}`,
226+
227+
[TemplateType.LOCATION]: `NEW LOCATION MESSAGE FROM:
228+
user_id: {userId}
229+
chat_id: {chatId}
230+
isDM: {isDM}
231+
message_id: {messageId}
232+
message_type: ${MessageType.LOCATION}
233+
latitude: {latitude}
234+
longitude: {longitude}`,
235+
236+
[TemplateType.CONTACT]: `NEW CONTACT MESSAGE FROM:
237+
user_id: {userId}
238+
chat_id: {chatId}
239+
isDM: {isDM}
240+
message_id: {messageId}
241+
message_type: ${MessageType.CONTACT}
242+
contact_name: {contactName}
243+
phone_number: {phoneNumber}`,
244+
245+
[TemplateType.POLL]: `NEW POLL MESSAGE FROM:
246+
user_id: {userId}
247+
chat_id: {chatId}
248+
isDM: {isDM}
249+
message_id: {messageId}
250+
message_type: ${MessageType.POLL}
251+
question: {pollQuestion}
252+
options: {pollOptions}`,
253+
254+
[TemplateType.FALLBACK]: `NEW MESSAGE FROM:
255+
user_id: {userId}
256+
chat_id: {chatId}
257+
isDM: {isDM}
258+
message_id: {messageId}
259+
message_type: {messageType}
260+
content: {content}`,
261+
},
262+
} as const;

‎src/index.ts‎

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
#!/usr/bin/env node
22
import { FastMCP, type FastMCPSession } from "fastmcp";
3-
import { SamplingHandler } from "./sampling.js";
43
import { forwardMessageTool } from "./tools/forward-message.js";
54
import { getChannelInfoTool } from "./tools/get-channel-info.js";
65
import { getChannelMembersTool } from "./tools/get-channel-members.js";
76
import { pinMessageTool } from "./tools/pin-message.js";
87
import { sendMessageTool } from "./tools/send-message.js";
8+
import { SamplingHandler } from "./sampling/handler.js";
9+
import { samplingConfig } from "./config.js";
910

1011
// =============================================================================
1112
// CONSTANTS
@@ -43,18 +44,22 @@ function setupSessionEventHandlers(server: FastMCP): SamplingHandler | null {
4344
server.on("connect", (event) => {
4445
console.log("🔌 Client connected:", event.session);
4546

46-
if (!samplingHandler) {
47-
initializeSamplingHandler(event.session)
48-
.then((handler) => {
49-
samplingHandler = handler;
50-
console.log("✅ Telegram sampling handler initialized");
51-
})
52-
.catch((error) => {
53-
console.error("❌ Failed to initialize sampling handler:", error);
54-
});
47+
if (samplingConfig.enabled) {
48+
if (!samplingHandler) {
49+
initializeSamplingHandler(event.session)
50+
.then((handler) => {
51+
samplingHandler = handler;
52+
console.log("✅ Telegram sampling handler initialized");
53+
})
54+
.catch((error) => {
55+
console.error("❌ Failed to initialize sampling handler:", error);
56+
});
57+
} else {
58+
samplingHandler.updateSession(event.session);
59+
console.log("🔄 Session updated for existing sampling handler");
60+
}
5561
} else {
56-
samplingHandler.updateSession(event.session);
57-
console.log("🔄 Session updated for existing sampling handler");
62+
console.log("ℹ️ Sampling is disabled via SAMPLING_ENABLED=false");
5863
}
5964
});
6065

@@ -82,13 +87,17 @@ function setupGracefulShutdown(samplingHandler: SamplingHandler | null): void {
8287
`\n🛑 Received ${signal}, shutting down Telegram MCP Server...`,
8388
);
8489

85-
if (samplingHandler) {
90+
if (samplingHandler && samplingConfig.enabled) {
8691
try {
8792
await samplingHandler.stop();
8893
console.log("✅ Telegram bot stopped gracefully");
8994
} catch (error) {
9095
console.error("❌ Error stopping Telegram bot:", error);
9196
}
97+
} else if (samplingConfig.enabled) {
98+
console.log("ℹ️ No Telegram bot to stop (not initialized yet)");
99+
} else {
100+
console.log("ℹ️ No Telegram bot to stop (sampling was disabled)");
92101
}
93102

94103
process.exit(0);
@@ -102,10 +111,16 @@ function logStartupInfo(): void {
102111
console.log("✅ Telegram MCP Server started successfully over stdio");
103112
console.log("📡 Ready to accept MCP client connections");
104113
console.log(`🛠️ Available tools: ${AVAILABLE_TOOLS.join(", ")}`);
105-
console.log(
106-
"🤖 Telegram bot will start when first client connects for AI sampling",
107-
);
108-
console.log("💡 Make sure TELEGRAM_BOT_TOKEN environment variable is set");
114+
115+
if (samplingConfig.enabled) {
116+
console.log(
117+
"🤖 Telegram bot will start when first client connects for AI sampling",
118+
);
119+
console.log("💡 Make sure TELEGRAM_BOT_TOKEN environment variable is set");
120+
} else {
121+
console.log("⚠️ AI sampling is disabled (SAMPLING_ENABLED=false)");
122+
console.log("💡 Only core MCP tools will be available");
123+
}
109124
}
110125

111126
// =============================================================================

0 commit comments

Comments
 (0)