Skip to content

Commit 48675e2

Browse files
yamanahlawatyamanahlawat
committed
feat: Added MCP elicitation support for tool injection in mcp-run-python
1 parent 6e60677 commit 48675e2

File tree

9 files changed

+1581
-41
lines changed

9 files changed

+1581
-41
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 & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
/// <reference types="npm:@types/[email protected]" />
22

33
import './polyfill.ts'
4-
import http from 'node:http'
54
import { randomUUID } from 'node:crypto'
5+
import http, { type IncomingMessage, type ServerResponse } from 'node:http'
66
import { parseArgs } from '@std/cli/parse-args'
77
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'
88
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
@@ -11,14 +11,13 @@ import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'
1111
import { type LoggingLevel, SetLevelRequestSchema } from '@modelcontextprotocol/sdk/types.js'
1212
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
1313
import { z } from 'zod'
14-
15-
import { asXml, runCode } from './runCode.ts'
1614
import { Buffer } from 'node:buffer'
15+
import { asXml, runCode, runCodeWithToolInjection, type ToolInjectionConfig } from './runCode.ts'
1716

18-
const VERSION = '0.0.13'
17+
const VERSION = '0.0.14'
1918

2019
export async function main() {
21-
const { args } = Deno
20+
const args = globalThis.Deno?.args || []
2221
if (args.length === 1 && args[0] === 'stdio') {
2322
await runStdio()
2423
} else if (args.length >= 1 && args[0] === 'streamable_http') {
@@ -29,7 +28,7 @@ export async function main() {
2928
const port = parseInt(flags.port)
3029
runStreamableHttp(port)
3130
} else if (args.length >= 1 && args[0] === 'sse') {
32-
const flags = parseArgs(Deno.args, {
31+
const flags = parseArgs(args, {
3332
string: ['port'],
3433
default: { port: '3001' },
3534
})
@@ -47,7 +46,7 @@ Usage: deno run -N -R=node_modules -W=node_modules --node-modules-dir=auto jsr:@
4746
options:
4847
--port <port> Port to run the SSE server on (default: 3001)`,
4948
)
50-
Deno.exit(1)
49+
globalThis.Deno?.exit(1)
5150
}
5251
}
5352

@@ -57,13 +56,14 @@ options:
5756
function createServer(): McpServer {
5857
const server = new McpServer(
5958
{
60-
name: 'MCP Run Python',
59+
name: 'MCP Run Python with Tool Injection',
6160
version: VERSION,
6261
},
6362
{
6463
instructions: 'Call the "run_python_code" tool with the Python code to run.',
6564
capabilities: {
6665
logging: {},
66+
elicitation: {},
6767
},
6868
},
6969
)
@@ -81,6 +81,13 @@ with a comment of the form:
8181
# dependencies = ['pydantic']
8282
# ///
8383
print('python code here')
84+
85+
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:
86+
87+
result = web_search("search query")
88+
print(result)
89+
90+
The tools are injected into the global namespace automatically - no discovery functions needed.
8491
`
8592

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

170-
const server = http.createServer(async (req, res) => {
271+
const server = http.createServer(async (req: IncomingMessage, res: ServerResponse) => {
171272
const url = httpGetUrl(req)
172273
let pathMatch = false
173274
function match(method: string, path: string): boolean {
@@ -309,21 +410,66 @@ async function warmup() {
309410
console.error(
310411
`Running warmup script for MCP Run Python version ${VERSION}...`,
311412
)
413+
312414
const code = `
313415
import numpy
314416
a = numpy.array([1, 2, 3])
315417
print('numpy array:', a)
316418
a
317419
`
318-
const result = await runCode([{
319-
name: 'warmup.py',
320-
content: code,
321-
active: true,
322-
}], (level, data) =>
323-
// use warn to avoid recursion since console.log is patched in runCode
324-
console.error(`${level}: ${data}`))
325-
console.log('Tool return value:')
420+
const result = await runCode(
421+
[
422+
{
423+
name: 'warmup.py',
424+
content: code,
425+
active: true,
426+
},
427+
],
428+
(level, data) => console.error(`${level}: ${data}`),
429+
)
326430
console.log(asXml(result))
431+
432+
// Test tool injection functionality
433+
console.error('Testing tool injection framework...')
434+
const toolCode = `
435+
# Test tool injection - directly call an injected tool
436+
result = web_search("test query")
437+
print(f"Tool result: {result}")
438+
"tool_test_complete"
439+
`
440+
441+
try {
442+
const toolResult = await runCodeWithToolInjection(
443+
[
444+
{
445+
name: 'tool_test.py',
446+
content: toolCode,
447+
active: true,
448+
},
449+
],
450+
(level, data) => console.error(`${level}: ${data}`),
451+
{
452+
enableToolInjection: true,
453+
availableTools: ['web_search', 'send_email'],
454+
timeoutSeconds: 30,
455+
// deno-lint-ignore no-explicit-any require-await
456+
elicitationCallback: async (_elicitationRequest: any) => {
457+
// Mock callback for warmup test
458+
return {
459+
action: 'accept',
460+
content: {
461+
result: '{"status": "mock success"}',
462+
},
463+
}
464+
},
465+
} as ToolInjectionConfig,
466+
)
467+
console.log('Tool injection result:')
468+
console.log(asXml(toolResult))
469+
} catch (error) {
470+
console.error('Tool injection test failed:', error)
471+
}
472+
327473
console.log('\nwarmup successful 🎉')
328474
}
329475

0 commit comments

Comments
 (0)