Skip to content

Commit 107767a

Browse files
KKonstantinovKKonstantinov
committed
save commit
1 parent 5a9de1c commit 107767a

File tree

4 files changed

+462
-42
lines changed

4 files changed

+462
-42
lines changed

‎src/server/context.ts‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
import {
2+
CreateMessageRequest,
3+
CreateMessageResult,
4+
ElicitRequest,
5+
ElicitResult,
6+
ElicitResultSchema,
7+
LoggingMessageNotification,
8+
Notification,
9+
Request,
10+
RequestId,
11+
RequestInfo,
12+
RequestMeta,
13+
Result,
14+
ServerNotification,
15+
ServerRequest,
16+
ServerResult
17+
} from '../types.js';
18+
import { RequestHandlerExtra, RequestOptions, RequestTaskStore } from '../shared/protocol.js';
19+
import { Server } from './index.js';
20+
import { RequestContext } from './requestContext.js';
21+
import { AuthInfo } from './auth/types.js';
22+
import { AnySchema, SchemaOutput } from './zod-compat.js';
23+
24+
export interface ContextInterface<RequestT extends Request = Request, NotificationT extends Notification = Notification> extends RequestHandlerExtra<ServerRequest | RequestT, NotificationT | ServerNotification> {
25+
elicit(params: ElicitRequest['params'], options?: RequestOptions): Promise<ElicitResult>;
26+
requestSampling: (params: CreateMessageRequest['params'], options?: RequestOptions) => Promise<CreateMessageResult>;
27+
log(params: LoggingMessageNotification['params'], sessionId?: string): Promise<void>;
28+
debug(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
29+
info(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
30+
warning(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
31+
error(message: string, extraLogData?: Record<string, unknown>, sessionId?: string): Promise<void>;
32+
}
33+
/**
34+
* A context object that is passed to request handlers.
35+
*
36+
* Implements the RequestHandlerExtra interface for backwards compatibility.
37+
*/
38+
export class Context<
39+
RequestT extends Request = Request,
40+
NotificationT extends Notification = Notification,
41+
ResultT extends Result = Result
42+
> implements ContextInterface<RequestT, NotificationT>
43+
{
44+
private readonly server: Server<RequestT, NotificationT, ResultT>;
45+
46+
/**
47+
* The request context.
48+
* A type-safe context that is passed to request handlers.
49+
*/
50+
public readonly requestCtx: RequestContext<
51+
RequestT | ServerRequest,
52+
NotificationT | ServerNotification,
53+
ResultT | ServerResult
54+
>;
55+
56+
constructor(args: {
57+
server: Server<RequestT, NotificationT, ResultT>;
58+
requestCtx: RequestContext<RequestT | ServerRequest, NotificationT | ServerNotification, ResultT | ServerResult>;
59+
}) {
60+
this.server = args.server;
61+
this.requestCtx = args.requestCtx;
62+
}
63+
64+
public get requestId(): RequestId {
65+
return this.requestCtx.requestId;
66+
}
67+
68+
public get signal(): AbortSignal {
69+
return this.requestCtx.signal;
70+
}
71+
72+
public get authInfo(): AuthInfo | undefined {
73+
return this.requestCtx.authInfo;
74+
}
75+
76+
public get requestInfo(): RequestInfo | undefined {
77+
return this.requestCtx.requestInfo;
78+
}
79+
80+
public get _meta(): RequestMeta | undefined {
81+
return this.requestCtx._meta;
82+
}
83+
84+
public get sessionId(): string | undefined {
85+
return this.requestCtx.sessionId;
86+
}
87+
88+
public get taskId(): string | undefined {
89+
return this.requestCtx.taskId;
90+
}
91+
92+
public get taskStore(): RequestTaskStore | undefined {
93+
return this.requestCtx.taskStore;
94+
}
95+
96+
public get taskRequestedTtl(): number | undefined {
97+
return this.requestCtx.taskRequestedTtl ?? undefined;
98+
}
99+
100+
public closeSSEStream = (): void => {
101+
return this.requestCtx?.closeSSEStream();
102+
}
103+
104+
public closeStandaloneSSEStream = (): void => {
105+
return this.requestCtx?.closeStandaloneSSEStream();
106+
}
107+
108+
/**
109+
* Sends a notification that relates to the current request being handled.
110+
*
111+
* This is used by certain transports to correctly associate related messages.
112+
*/
113+
public sendNotification = (notification: NotificationT | ServerNotification): Promise<void> => {
114+
return this.requestCtx.sendNotification(notification);
115+
};
116+
117+
/**
118+
* Sends a request that relates to the current request being handled.
119+
*
120+
* This is used by certain transports to correctly associate related messages.
121+
*/
122+
public sendRequest = <U extends AnySchema>(
123+
request: RequestT | ServerRequest,
124+
resultSchema: U,
125+
options?: RequestOptions
126+
): Promise<SchemaOutput<U>> => {
127+
return this.requestCtx.sendRequest(request, resultSchema, { ...options, relatedRequestId: this.requestId });
128+
};
129+
130+
/**
131+
* Sends a request to sample an LLM via the client.
132+
*/
133+
public requestSampling(params: CreateMessageRequest['params'], options?: RequestOptions) {
134+
return this.server.createMessage(params, options);
135+
}
136+
137+
/**
138+
* Sends an elicitation request to the client.
139+
*/
140+
public async elicit(params: ElicitRequest['params'], options?: RequestOptions): Promise<ElicitResult> {
141+
const request: ElicitRequest = {
142+
method: 'elicitation/create',
143+
params
144+
};
145+
return await this.server.request(request, ElicitResultSchema, { ...options, relatedRequestId: this.requestId });
146+
}
147+
148+
/**
149+
* Sends a logging message.
150+
*/
151+
public async log(params: LoggingMessageNotification['params'], sessionId?: string) {
152+
await this.server.sendLoggingMessage(params, sessionId);
153+
}
154+
155+
/**
156+
* Sends a debug log message.
157+
*/
158+
public async debug(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
159+
await this.log(
160+
{
161+
level: 'debug',
162+
data: {
163+
...extraLogData,
164+
message
165+
},
166+
logger: 'server'
167+
},
168+
sessionId
169+
);
170+
}
171+
172+
/**
173+
* Sends an info log message.
174+
*/
175+
public async info(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
176+
await this.log(
177+
{
178+
level: 'info',
179+
data: {
180+
...extraLogData,
181+
message
182+
},
183+
logger: 'server'
184+
},
185+
sessionId
186+
);
187+
}
188+
189+
/**
190+
* Sends a warning log message.
191+
*/
192+
public async warning(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
193+
await this.log(
194+
{
195+
level: 'warning',
196+
data: {
197+
...extraLogData,
198+
message
199+
},
200+
logger: 'server'
201+
},
202+
sessionId
203+
);
204+
}
205+
206+
/**
207+
* Sends an error log message.
208+
*/
209+
public async error(message: string, extraLogData?: Record<string, unknown>, sessionId?: string) {
210+
await this.log(
211+
{
212+
level: 'error',
213+
data: {
214+
...extraLogData,
215+
message
216+
},
217+
logger: 'server'
218+
},
219+
sessionId
220+
);
221+
}
222+
}

‎src/server/index.ts‎

Lines changed: 74 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,10 @@ import {
4040
type ToolUseContent,
4141
CallToolRequestSchema,
4242
CallToolResultSchema,
43-
CreateTaskResultSchema
43+
CreateTaskResultSchema,
44+
JSONRPCRequest,
45+
TaskCreationParams,
46+
MessageExtraInfo
4447
} from '../types.js';
4548
import { AjvJsonSchemaValidator } from '../validation/ajv-provider.js';
4649
import type { JsonSchemaType, jsonSchemaValidator } from '../validation/types.js';
@@ -56,6 +59,10 @@ import {
5659
import { RequestHandlerExtra } from '../shared/protocol.js';
5760
import { ExperimentalServerTasks } from '../experimental/tasks/server.js';
5861
import { assertToolsCallTaskCapability, assertClientRequestTaskCapability } from '../experimental/tasks/helpers.js';
62+
import { Context } from './context.js';
63+
import { TaskStore } from '../experimental/index.js';
64+
import { Transport } from '../shared/transport.js';
65+
import { RequestContext } from './requestContext.js';
5966

6067
export type ServerOptions = ProtocolOptions & {
6168
/**
@@ -219,9 +226,23 @@ export class Server<
219226
requestSchema: T,
220227
handler: (
221228
request: SchemaOutput<T>,
222-
extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
229+
extra: Context<RequestT, NotificationT, ResultT>
223230
) => ServerResult | ResultT | Promise<ServerResult | ResultT>
224231
): void {
232+
// Wrap the handler to ensure the extra is a Context and return a decorated handler that can be passed to the base implementation
233+
234+
// Factory function to create a handler decorator that ensures the extra is a Context and returns a decorated handler that can be passed to the base implementation
235+
const handlerDecoratorFactory = (innerHandler: (request: SchemaOutput<T>, extra: Context<RequestT, NotificationT, ResultT>) => ServerResult | ResultT | Promise<ServerResult | ResultT>) => {
236+
const decoratedHandler = (request: SchemaOutput<T>, extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>) => {
237+
if (!this.isContextExtra(extra)) {
238+
throw new Error('Internal error: Expected Context for request handler extra');
239+
}
240+
return innerHandler(request, extra);
241+
}
242+
243+
return decoratedHandler;
244+
}
245+
225246
const shape = getObjectShape(requestSchema);
226247
const methodSchema = shape?.method;
227248
if (!methodSchema) {
@@ -259,7 +280,7 @@ export class Server<
259280

260281
const { params } = validatedRequest.data;
261282

262-
const result = await Promise.resolve(handler(request, extra));
283+
const result = await Promise.resolve(handlerDecoratorFactory(handler)(request, extra));
263284

264285
// When task creation is requested, validate and return CreateTaskResult
265286
if (params.task) {
@@ -286,11 +307,18 @@ export class Server<
286307
};
287308

288309
// Install the wrapped handler
289-
return super.setRequestHandler(requestSchema, wrappedHandler as unknown as typeof handler);
310+
return super.setRequestHandler(requestSchema, handlerDecoratorFactory(wrappedHandler));
290311
}
291312

292313
// Other handlers use default behavior
293-
return super.setRequestHandler(requestSchema, handler);
314+
return super.setRequestHandler(requestSchema, handlerDecoratorFactory(handler));
315+
}
316+
317+
// Runtime type guard: ensure extra is our Context
318+
private isContextExtra(
319+
extra: RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT>
320+
): extra is Context<RequestT, NotificationT, ResultT> {
321+
return extra instanceof Context;
294322
}
295323

296324
protected assertCapabilityForMethod(method: RequestT['method']): void {
@@ -468,6 +496,47 @@ export class Server<
468496
return this._capabilities;
469497
}
470498

499+
protected createRequestExtra(
500+
args: {
501+
request: JSONRPCRequest,
502+
taskStore: TaskStore | undefined,
503+
relatedTaskId: string | undefined,
504+
taskCreationParams: TaskCreationParams | undefined,
505+
abortController: AbortController,
506+
capturedTransport: Transport | undefined,
507+
extra?: MessageExtraInfo
508+
}
509+
): RequestHandlerExtra<ServerRequest | RequestT, ServerNotification | NotificationT> {
510+
const base = super.createRequestExtra(args) as RequestHandlerExtra<
511+
ServerRequest | RequestT,
512+
ServerNotification | NotificationT
513+
>;
514+
515+
// Wrap base in Context to add server utilities while preserving shape
516+
const requestCtx = new RequestContext<
517+
ServerRequest | RequestT,
518+
ServerNotification | NotificationT,
519+
ServerResult | ResultT
520+
>({
521+
signal: base.signal,
522+
authInfo: base.authInfo,
523+
requestInfo: base.requestInfo,
524+
requestId: base.requestId,
525+
_meta: base._meta,
526+
sessionId: base.sessionId,
527+
protocol: this,
528+
closeSSEStream: base.closeSSEStream ?? undefined,
529+
closeStandaloneSSEStream: base.closeStandaloneSSEStream ?? undefined
530+
});
531+
532+
const ctx = new Context<RequestT, NotificationT, ResultT>({
533+
server: this,
534+
requestCtx
535+
});
536+
537+
return ctx;
538+
}
539+
471540
async ping() {
472541
return this.request({ method: 'ping' }, EmptyResultSchema);
473542
}

0 commit comments

Comments
 (0)