Skip to content

Commit 393f025

Browse files
yamanahlawatyamanahlawat
committed
Added MCP elicitation support for tool injection in mcp-run-python
1 parent 01c550c commit 393f025

File tree

9 files changed

+1818
-46
lines changed

9 files changed

+1818
-46
lines changed

‎mcp-run-python/build.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,20 @@ if (!import.meta.dirname) {
55
throw new Error('import.meta.dirname is not defined, unable to load prepare_env.py')
66
}
77
const src = path.join(import.meta.dirname, 'src/prepare_env.py')
8+
const toolInjectionSrc = path.join(import.meta.dirname, 'src/tool_injection.py')
89
const dst = path.join(import.meta.dirname, 'src/prepareEnvCode.ts')
910

1011
let pythonCode = await Deno.readTextFile(src)
1112
pythonCode = pythonCode.replace(/\\/g, '\\\\')
13+
14+
// Read tool injection code from separate Python file
15+
let toolInjectionCode = await Deno.readTextFile(toolInjectionSrc)
16+
toolInjectionCode = toolInjectionCode.replace(/\\/g, '\\\\')
17+
1218
const jsCode = `\
1319
// DO NOT EDIT THIS FILE DIRECTLY, INSTEAD RUN "deno run build"
1420
export const preparePythonCode = \`${pythonCode}\`
21+
22+
export const toolInjectionCode = \`${toolInjectionCode}\`
1523
`
1624
await Deno.writeTextFile(dst, jsCode)

‎mcp-run-python/src/main.ts

Lines changed: 176 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,24 @@
11
/// <reference types="npm:@types/[email protected]" />
22

33
import './polyfill.ts'
4-
import http from 'node:http'
4+
import http, { type IncomingMessage, type ServerResponse } from 'node:http'
55
import { parseArgs } from '@std/cli/parse-args'
66
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'
77
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
88
import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/sdk/types.js'
99
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
1010
import { z } from 'zod'
1111

12-
import { asXml, runCode } from './runCode.ts'
12+
import { asXml, runCode, runCodeWithToolInjection, type ToolInjectionConfig } from './runCode.ts'
1313

14-
const VERSION = '0.0.13'
14+
const VERSION = '0.0.14'
1515

1616
export async function main() {
17-
const { args } = Deno
17+
const args = globalThis.Deno?.args || []
1818
if (args.length === 1 && args[0] === 'stdio') {
1919
await runStdio()
2020
} else if (args.length >= 1 && args[0] === 'sse') {
21-
const flags = parseArgs(Deno.args, {
21+
const flags = parseArgs(args, {
2222
string: ['port'],
2323
default: { port: '3001' },
2424
})
@@ -36,7 +36,7 @@ Usage: deno run -N -R=node_modules -W=node_modules --node-modules-dir=auto jsr:@
3636
options:
3737
--port <port> Port to run the SSE server on (default: 3001)`,
3838
)
39-
Deno.exit(1)
39+
globalThis.Deno?.exit(1)
4040
}
4141
}
4242

@@ -46,13 +46,14 @@ options:
4646
function createServer(): McpServer {
4747
const server = new McpServer(
4848
{
49-
name: 'MCP Run Python',
49+
name: 'MCP Run Python with Tool Injection',
5050
version: VERSION,
5151
},
5252
{
5353
instructions: 'Call the "run_python_code" tool with the Python code to run.',
5454
capabilities: {
5555
logging: {},
56+
elicitation: {},
5657
},
5758
},
5859
)
@@ -70,6 +71,13 @@ with a comment of the form:
7071
# dependencies = ['pydantic']
7172
# ///
7273
print('python code here')
74+
75+
TOOL INJECTION: When 'tools' parameter is provided, the specified tool functions become available directly in Python's global namespace. You can call them directly like any other function. For example, if 'web_search' is provided as a tool, you can call it directly:
76+
77+
result = web_search("search query")
78+
print(result)
79+
80+
The tools are injected into the global namespace automatically - no discovery functions needed.
7381
`
7482

7583
let setLogLevel: LoggingLevel = 'emergency'
@@ -82,21 +90,115 @@ print('python code here')
8290
server.tool(
8391
'run_python_code',
8492
toolDescription,
85-
{ python_code: z.string().describe('Python code to run') },
86-
async ({ python_code }: { python_code: string }) => {
93+
{
94+
python_code: z.string().describe('Python code to run'),
95+
tools: z
96+
.array(z.string())
97+
.optional()
98+
.describe('List of available tools for injection (enables tool injection when provided)'),
99+
},
100+
async ({
101+
python_code,
102+
tools = [],
103+
}: {
104+
python_code: string
105+
tools?: string[]
106+
}) => {
87107
const logPromises: Promise<void>[] = []
88-
const result = await runCode([{
89-
name: 'main.py',
90-
content: python_code,
91-
active: true,
92-
}], (level, data) => {
93-
if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
94-
logPromises.push(server.server.sendLoggingMessage({ level, data }))
108+
109+
// Check if tools are provided
110+
if (tools.length > 0) {
111+
// Create elicitation callback
112+
// deno-lint-ignore no-explicit-any
113+
const elicitationCallback = async (elicitationRequest: any) => {
114+
// Convert Python dict to JavaScript object if needed
115+
let jsRequest
116+
if (elicitationRequest && typeof elicitationRequest === 'object' && elicitationRequest.toJs) {
117+
jsRequest = elicitationRequest.toJs()
118+
} else if (elicitationRequest && typeof elicitationRequest === 'object') {
119+
// Handle Python dict-like objects
120+
jsRequest = {
121+
message: elicitationRequest.message || elicitationRequest.get?.('message'),
122+
requestedSchema: elicitationRequest.requestedSchema || elicitationRequest.get?.('requestedSchema'),
123+
}
124+
} else {
125+
jsRequest = elicitationRequest
126+
}
127+
128+
try {
129+
const elicitationResult = await server.server.request(
130+
{
131+
method: 'elicitation/create',
132+
params: {
133+
message: jsRequest.message,
134+
requestedSchema: jsRequest.requestedSchema,
135+
},
136+
},
137+
z.object({
138+
action: z.enum(['accept', 'decline', 'cancel']),
139+
content: z.optional(z.record(z.string(), z.unknown())),
140+
}),
141+
)
142+
143+
return elicitationResult
144+
} catch (error) {
145+
logPromises.push(
146+
server.server.sendLoggingMessage({
147+
level: 'error',
148+
data: `Elicitation error: ${error}`,
149+
}),
150+
)
151+
throw error
152+
}
153+
}
154+
155+
// Use tool injection mode
156+
const result = await runCodeWithToolInjection(
157+
[
158+
{
159+
name: 'main.py',
160+
content: python_code,
161+
active: true,
162+
},
163+
],
164+
(level, data) => {
165+
if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
166+
logPromises.push(server.server.sendLoggingMessage({ level, data }))
167+
}
168+
},
169+
{
170+
enableToolInjection: true,
171+
availableTools: tools,
172+
timeoutSeconds: 30,
173+
elicitationCallback,
174+
} as ToolInjectionConfig,
175+
)
176+
177+
await Promise.all(logPromises)
178+
179+
return {
180+
content: [{ type: 'text', text: asXml(result) }],
181+
}
182+
} else {
183+
// Use basic mode without tool injection
184+
const result = await runCode(
185+
[
186+
{
187+
name: 'main.py',
188+
content: python_code,
189+
active: true,
190+
},
191+
],
192+
(level, data) => {
193+
if (LogLevels.indexOf(level) >= LogLevels.indexOf(setLogLevel)) {
194+
logPromises.push(server.server.sendLoggingMessage({ level, data }))
195+
}
196+
},
197+
)
198+
await Promise.all(logPromises)
199+
return {
200+
content: [{ type: 'text', text: asXml(result) }],
95201
}
96-
})
97-
await Promise.all(logPromises)
98-
return {
99-
content: [{ type: 'text', text: asXml(result) }],
100202
}
101203
},
102204
)
@@ -110,7 +212,7 @@ function runSse(port: number) {
110212
const mcpServer = createServer()
111213
const transports: { [sessionId: string]: SSEServerTransport } = {}
112214

113-
const server = http.createServer(async (req, res) => {
215+
const server = http.createServer(async (req: IncomingMessage, res: ServerResponse) => {
114216
const url = new URL(
115217
req.url ?? '',
116218
`http://${req.headers.host ?? 'unknown'}`,
@@ -175,21 +277,66 @@ async function warmup() {
175277
console.error(
176278
`Running warmup script for MCP Run Python version ${VERSION}...`,
177279
)
280+
178281
const code = `
179282
import numpy
180283
a = numpy.array([1, 2, 3])
181284
print('numpy array:', a)
182285
a
183286
`
184-
const result = await runCode([{
185-
name: 'warmup.py',
186-
content: code,
187-
active: true,
188-
}], (level, data) =>
189-
// use warn to avoid recursion since console.log is patched in runCode
190-
console.error(`${level}: ${data}`))
191-
console.log('Tool return value:')
287+
const result = await runCode(
288+
[
289+
{
290+
name: 'warmup.py',
291+
content: code,
292+
active: true,
293+
},
294+
],
295+
(level, data) => console.error(`${level}: ${data}`),
296+
)
192297
console.log(asXml(result))
298+
299+
// Test tool injection functionality
300+
console.error('Testing tool injection framework...')
301+
const toolCode = `
302+
# Test tool injection - directly call an injected tool
303+
result = web_search("test query")
304+
print(f"Tool result: {result}")
305+
"tool_test_complete"
306+
`
307+
308+
try {
309+
const toolResult = await runCodeWithToolInjection(
310+
[
311+
{
312+
name: 'tool_test.py',
313+
content: toolCode,
314+
active: true,
315+
},
316+
],
317+
(level, data) => console.error(`${level}: ${data}`),
318+
{
319+
enableToolInjection: true,
320+
availableTools: ['web_search', 'send_email'],
321+
timeoutSeconds: 30,
322+
// deno-lint-ignore no-explicit-any require-await
323+
elicitationCallback: async (_elicitationRequest: any) => {
324+
// Mock callback for warmup test
325+
return {
326+
action: 'accept',
327+
content: {
328+
result: '{"status": "mock success"}',
329+
},
330+
}
331+
},
332+
} as ToolInjectionConfig,
333+
)
334+
console.log('Tool injection result:')
335+
console.log(asXml(toolResult))
336+
} catch (error) {
337+
console.error('Tool injection test failed:', error)
338+
}
339+
193340
console.log('\nwarmup successful 🎉')
194341
}
195342

0 commit comments

Comments
 (0)