Skip to content

Commit 2bb2821

Browse files
aaceboaacebo
authored
[mcp/ai/cli] support different transport types (#87)
* Update plugin.ts * add env collection management to cli * add prompt interfaces to `ai` package * Update plugin.ts * add env collection management to cli * add prompt interfaces to `ai` package * support different transport types * update mcp template
1 parent 55442dd commit 2bb2821

File tree

27 files changed

+727
-75
lines changed

27 files changed

+727
-75
lines changed

Diff for: ‎external/mcp/src/plugin.ts

+168-51
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,82 @@
1-
import { ChatPrompt } from '@microsoft/spark.ai';
1+
import { Readable, Writable } from 'stream';
2+
3+
import { IChatPrompt } from '@microsoft/spark.ai';
24
import { App, IActivityContext, IPlugin, IPluginEvents } from '@microsoft/spark.apps';
35
import { ConsoleLogger, EventEmitter, EventHandler, ILogger } from '@microsoft/spark.common';
46

57
import { ServerOptions } from '@modelcontextprotocol/sdk/server/index.js';
68
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
79
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
8-
import { Implementation, CallToolResult } from '@modelcontextprotocol/sdk/types.js';
10+
import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
11+
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
912

1013
import { z } from 'zod';
1114
import { jsonSchemaToZod } from 'json-schema-to-zod';
1215

1316
import { IConnection } from './connection';
1417

15-
export class MCPPlugin implements IPlugin {
18+
/**
19+
* MCP transport options for sse
20+
*/
21+
export type McpSSETransportOptions = {
22+
/**
23+
* the transport type
24+
*/
25+
readonly type: 'sse';
26+
27+
/**
28+
* the url path
29+
* @default /mcp
30+
*/
31+
readonly path?: string;
32+
};
33+
34+
/**
35+
* MCP transport options for stdio
36+
*/
37+
export type McpStdioTransportOptions = {
38+
/**
39+
* the transport type
40+
*/
41+
readonly type: 'stdio';
42+
43+
/**
44+
* stdin to use
45+
*/
46+
readonly stdin?: Readable;
47+
48+
/**
49+
* stdout to use
50+
*/
51+
readonly stdout?: Writable;
52+
};
53+
54+
export type McpPluginOptions = ServerOptions & {
55+
/**
56+
* the MCP server name
57+
* @default mcp
58+
*/
59+
readonly name?: string;
60+
61+
/**
62+
* the MCP server version
63+
* @default 0.0.0
64+
*/
65+
readonly version?: string;
66+
67+
/**
68+
* the transport or transport options
69+
* @default sse
70+
*/
71+
readonly transport?: McpSSETransportOptions | McpStdioTransportOptions;
72+
};
73+
74+
/**
75+
* High-level MCP server that provides a simpler API for working with resources, tools, and prompts.
76+
* For advanced usage (like sending notifications or setting custom request handlers),
77+
* use the underlying Server instance available via the server property.
78+
*/
79+
export class McpPlugin implements IPlugin {
1680
readonly name: string;
1781
readonly version: string;
1882

@@ -24,64 +88,81 @@ export class MCPPlugin implements IPlugin {
2488
protected log: ILogger;
2589
protected id: number = -1;
2690
protected connections: Record<number, IConnection> = {};
27-
protected readonly events = new EventEmitter<IPluginEvents>();
91+
protected events = new EventEmitter<IPluginEvents>();
92+
protected transport: McpSSETransportOptions | McpStdioTransportOptions = { type: 'sse' };
2893

29-
constructor(serverInfo: Implementation | McpServer, options: ServerOptions = {}) {
94+
constructor(options: McpServer | McpPluginOptions = {}) {
3095
this.log = new ConsoleLogger('@spark/mcp');
31-
this.name = serverInfo instanceof McpServer ? 'mcp' : `mcp.${serverInfo.name}`;
32-
this.version = serverInfo instanceof McpServer ? '0.0.0' : serverInfo.version;
33-
this.server = serverInfo instanceof McpServer ? serverInfo : new McpServer(serverInfo, options);
96+
this.name =
97+
options instanceof McpServer ? 'mcp' : `mcp${options.name ? `.${options.name}` : ''}`;
98+
this.version = options instanceof McpServer ? '0.0.0' : options.version || '0.0.0';
99+
this.server =
100+
options instanceof McpServer
101+
? options
102+
: new McpServer(
103+
{
104+
name: this.name,
105+
version: this.version,
106+
},
107+
options
108+
);
109+
110+
if (!(options instanceof McpServer) && options.transport) {
111+
this.transport = options.transport;
112+
}
113+
34114
this.prompt = this.server.prompt.bind(this.server);
35115
this.tool = this.server.tool.bind(this.server);
36116
this.resource = this.server.resource.bind(this.server);
37117
}
38118

39-
use(prompt: ChatPrompt) {
119+
use(prompt: IChatPrompt) {
40120
for (const fn of prompt.functions) {
41121
const schema: z.AnyZodObject = eval(jsonSchemaToZod(fn.parameters, { module: 'cjs' }));
42-
this.server.tool(
43-
fn.name,
44-
fn.description,
45-
schema.shape,
46-
async (args: any): Promise<CallToolResult> => {
47-
try {
48-
const res = await prompt.call(fn.name, args);
49-
50-
return {
51-
content: [
52-
{
53-
type: 'text',
54-
text: typeof res === 'string' ? res : JSON.stringify(res),
55-
},
56-
],
57-
};
58-
} catch (err: any) {
59-
this.log.error(err.toString());
60-
61-
return {
62-
isError: true,
63-
content: [
64-
{
65-
type: 'text',
66-
text: err.toString(),
67-
},
68-
],
69-
};
70-
}
71-
}
72-
);
122+
this.server.tool(fn.name, fn.description, schema.shape, this.onToolCall(fn.name, prompt));
73123
}
74124

75125
return this;
76126
}
77127

78-
onInit(app: App) {
128+
async onInit(app: App) {
79129
this.log = app.log.child(this.name);
80130

81-
app.http.get('/mcp', (_, res) => {
131+
if (this.transport.type === 'sse') {
132+
return this.onInitSSE(app, this.transport);
133+
}
134+
135+
await this.onInitStdio(app, this.transport);
136+
}
137+
138+
async onStart(port: number = 3000) {
139+
this.events.emit('start', this.log);
140+
141+
if (this.transport.type === 'sse') {
142+
this.log.info(`listening at http://localhost:${port}${this.transport.path || '/mcp'}`);
143+
} else {
144+
this.log.info('listening on stdin');
145+
}
146+
}
147+
148+
on<Name extends keyof IPluginEvents>(name: Name, callback: EventHandler<IPluginEvents[Name]>) {
149+
this.events.on(name, callback);
150+
}
151+
152+
onActivity(_: IActivityContext) {}
153+
154+
protected onInitStdio(_: App, options: McpStdioTransportOptions) {
155+
const transport = new StdioServerTransport(options.stdin, options.stdout);
156+
return this.server.connect(transport);
157+
}
158+
159+
protected onInitSSE(app: App, options: McpSSETransportOptions) {
160+
const path = options.path || '/mcp';
161+
162+
app.http.get(path, (_, res) => {
82163
this.id++;
83164
this.log.debug('connecting...');
84-
const transport = new SSEServerTransport(`/mcp/${this.id}/messages`, res);
165+
const transport = new SSEServerTransport(`${path}/${this.id}/messages`, res);
85166
this.connections[this.id] = {
86167
id: this.id,
87168
transport,
@@ -91,7 +172,7 @@ export class MCPPlugin implements IPlugin {
91172
this.server.connect(transport);
92173
});
93174

94-
app.http.post('/mcp/:id/messages', (req, res) => {
175+
app.http.post(`${path}/:id/messages`, (req, res) => {
95176
const id = +req.params.id;
96177
const { transport } = this.connections[id];
97178

@@ -104,14 +185,50 @@ export class MCPPlugin implements IPlugin {
104185
});
105186
}
106187

107-
async onStart(port: number = 3000) {
108-
this.events.emit('start', this.log);
109-
this.log.info(`listening at http://localhost:${port}/mcp 🚀`);
110-
}
188+
protected onToolCall(name: string, prompt: IChatPrompt) {
189+
return async (args: any): Promise<CallToolResult> => {
190+
try {
191+
const res = await prompt.call(name, args);
111192

112-
on<Name extends keyof IPluginEvents>(name: Name, callback: EventHandler<IPluginEvents[Name]>) {
113-
this.events.on(name, callback);
193+
if (this.isCallToolResult(res)) {
194+
return res;
195+
}
196+
197+
return {
198+
content: [
199+
{
200+
type: 'text',
201+
text: typeof res === 'string' ? res : JSON.stringify(res),
202+
},
203+
],
204+
};
205+
} catch (err: any) {
206+
this.log.error(err.toString());
207+
208+
return {
209+
isError: true,
210+
content: [
211+
{
212+
type: 'text',
213+
text: err.toString(),
214+
},
215+
],
216+
};
217+
}
218+
};
114219
}
115220

116-
onActivity(_: IActivityContext) {}
221+
protected isCallToolResult(value: any): value is CallToolResult {
222+
if (!!value || !('content' in value)) return false;
223+
const { content } = value;
224+
225+
return (
226+
Array.isArray(content) &&
227+
content.every(
228+
(item) =>
229+
'type' in item &&
230+
(item.type === 'text' || item.type === 'image' || item.type === 'resource')
231+
)
232+
);
233+
}
117234
}

Diff for: ‎packages/ai/src/prompts/audio.ts

+31-1
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,37 @@ export type AudioPromptOptions = {
1717
readonly model: IAudioModel;
1818
};
1919

20-
export class AudioPrompt {
20+
/**
21+
* a prompt that can interface with
22+
* an audio model
23+
*/
24+
export interface IAudioPrompt {
25+
/**
26+
* the prompt name
27+
*/
28+
readonly name: string;
29+
30+
/**
31+
* the prompt description
32+
*/
33+
readonly description: string;
34+
35+
/**
36+
* convert text to audio
37+
*/
38+
textToAudio?(params: TextToAudioParams): Promise<Buffer>;
39+
40+
/**
41+
* transcribe audio to text
42+
*/
43+
audioToText?(params: AudioToTextParams): Promise<string>;
44+
}
45+
46+
/**
47+
* a prompt that can interface with
48+
* an audio model
49+
*/
50+
export class AudioPrompt implements IAudioPrompt {
2151
get name() {
2252
return this._name;
2353
}

Diff for: ‎packages/ai/src/prompts/chat.ts

+62-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Function, FunctionHandler } from '../function';
22
import { LocalMemory } from '../local-memory';
33
import { IMemory } from '../memory';
4-
import { ContentPart, Message, SystemMessage, UserMessage } from '../message';
4+
import { ContentPart, Message, ModelMessage, SystemMessage, UserMessage } from '../message';
55
import { IChatModel, TextChunkHandler } from '../models';
66
import { Schema } from '../schema';
77
import { ITemplate } from '../template';
@@ -58,7 +58,67 @@ export type ChatPromptSendOptions<TOptions extends Record<string, any> = Record<
5858
readonly onChunk?: TextChunkHandler;
5959
};
6060

61-
export class ChatPrompt<TOptions extends Record<string, any> = Record<string, any>> {
61+
/**
62+
* a prompt that can interface with a
63+
* chat model that provides utility like
64+
* streaming and function calling
65+
*/
66+
export interface IChatPrompt<TOptions extends Record<string, any> = Record<string, any>> {
67+
/**
68+
* the prompt name
69+
*/
70+
readonly name: string;
71+
72+
/**
73+
* the prompt description
74+
*/
75+
readonly description: string;
76+
77+
/**
78+
* the chat history
79+
*/
80+
readonly messages: IMemory;
81+
82+
/**
83+
* the registered functions
84+
*/
85+
readonly functions: Array<Function>;
86+
87+
/**
88+
* add another chat prompt as a
89+
*/
90+
use(prompt: IChatPrompt): this;
91+
use(name: string, prompt: IChatPrompt): this;
92+
93+
/**
94+
* add a function that can be called
95+
* by the model
96+
*/
97+
function(name: string, description: string, handler: FunctionHandler): this;
98+
function(name: string, description: string, parameters: Schema, handler: FunctionHandler): this;
99+
100+
/**
101+
* call a function
102+
*/
103+
call<A extends Record<string, any>, R = any>(name: string, args?: A): Promise<R>;
104+
105+
/**
106+
* send a message to the model and get a response
107+
*/
108+
send(
109+
input: string | ContentPart[],
110+
options?: ChatPromptSendOptions<TOptions>
111+
): Promise<Pick<ModelMessage, 'content'> & Omit<ModelMessage, 'content'>>;
112+
}
113+
114+
/**
115+
* a prompt that can interface with a
116+
* chat model that provides utility like
117+
* streaming and function calling
118+
*/
119+
export class ChatPrompt<TOptions extends Record<string, any> = Record<string, any>>
120+
implements IChatPrompt<TOptions>
121+
{
62122
get name() {
63123
return this._name;
64124
}

Diff for: ‎packages/ai/src/prompts/index.ts

+5
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,7 @@
1+
import { IAudioPrompt } from './audio';
2+
import { IChatPrompt } from './chat';
3+
4+
export type Prompt = IChatPrompt | IAudioPrompt;
5+
16
export * from './chat';
27
export * from './audio';

0 commit comments

Comments
 (0)