Skip to content

Commit 41e6b67

Browse files
jneumsjneums
authored
Merge pull request #7 from prometheus-protocol/8-4-refactor-streaming-callback
refactor: move streaming callback logic into mcp sdk
2 parents 40380da + 3aa995a commit 41e6b67

5 files changed

Lines changed: 51 additions & 103 deletions

File tree

‎examples/private_mcp_server/src/main.mo‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
import Map "mo:map/Map";
22
import { thash } "mo:map/Map";
33
import Result "mo:base/Result";
4-
import Text "mo:base/Text";
5-
import Option "mo:base/Option";
64
import Blob "mo:base/Blob";
7-
import Time "mo:base/Time";
85
import Principal "mo:base/Principal";
96
import Json "../../../src/json";
107
import HttpTypes "mo:http-types";
11-
import BaseX "mo:base-x-encoder";
128

139
// The only SDK import the user needs!
1410
import Mcp "../../../src/mcp/Mcp";
@@ -164,19 +160,9 @@ shared persistent actor class McpServer() = self {
164160
return await HttpHandler.http_request_update(ctx, req);
165161
};
166162

167-
// The streaming callback MUST be a public function of the main actor.
168163
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
169-
let token_key = BaseX.toBase64(token.vals(), #standard({ includePadding = true }));
170-
// It has access to the actor's state.
171-
if (Option.isNull(Map.get(appContext.activeStreams, thash, token_key))) {
172-
return ?{ body = Blob.fromArray([]); token = null };
173-
};
174-
175-
// Update the timestamp to prove the stream is still active.
176-
Map.set(appContext.activeStreams, thash, token_key, Time.now());
177-
178-
let chunk = Text.encodeUtf8("data: {\"type\":\"keep-alive\"}\n\n");
179-
return ?{ body = chunk; token = ?token };
164+
let ctx : HttpHandler.Context = _create_http_context();
165+
return HttpHandler.http_request_streaming_callback(ctx, token);
180166
};
181167

182168
system func preupgrade() {

‎examples/public_mcp_server/src/main.mo‎

Lines changed: 16 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,9 @@
11
import Map "mo:map/Map";
22
import { thash } "mo:map/Map";
33
import Result "mo:base/Result";
4-
import Text "mo:base/Text";
5-
import Option "mo:base/Option";
6-
import Blob "mo:base/Blob";
7-
import Time "mo:base/Time";
84
import Principal "mo:base/Principal";
95
import Json "../../../src/json";
106
import HttpTypes "mo:http-types";
11-
import BaseX "mo:base-x-encoder";
127

138
// The only SDK import the user needs!
149
import Mcp "../../../src/mcp/Mcp";
@@ -107,44 +102,30 @@ shared persistent actor class McpServer() = self {
107102

108103
// --- PUBLIC ENTRY POINTS ---
109104

110-
// The streaming callback MUST be a public function of the main actor.
111-
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
112-
let token_key = BaseX.toBase64(token.vals(), #standard({ includePadding = true }));
113-
// It has access to the actor's state.
114-
if (Option.isNull(Map.get(appContext.activeStreams, thash, token_key))) {
115-
return ?{ body = Blob.fromArray([]); token = null };
116-
};
117-
118-
// Update the timestamp to prove the stream is still active.
119-
Map.set(appContext.activeStreams, thash, token_key, Time.now());
120-
121-
let chunk = Text.encodeUtf8("data: {\"type\":\"keep-alive\"}\n\n");
122-
return ?{ body = chunk; token = ?token };
123-
};
124-
125-
public query func http_request(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
126-
// Construct the context object on the fly.
127-
let ctx : HttpHandler.Context = {
128-
self = Principal.fromActor(self); // Pass the server principal
105+
// Helper to avoid repeating context creation.
106+
private func _create_http_context() : HttpHandler.Context {
107+
return {
108+
self = Principal.fromActor(self);
129109
active_streams = appContext.activeStreams;
130110
mcp_server = mcpServer;
131111
streaming_callback = http_request_streaming_callback;
132-
auth = null; // No authentication in this example.
133-
http_asset_cache = null; // No HTTP asset cache in this example.
112+
auth = null;
113+
http_asset_cache = null;
134114
};
135-
// Delegate the complex logic to the handler module.
115+
};
116+
117+
public query func http_request(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
118+
let ctx : HttpHandler.Context = _create_http_context();
136119
return HttpHandler.http_request(ctx, req);
137120
};
138121

139122
public func http_request_update(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
140-
let ctx : HttpHandler.Context = {
141-
self = Principal.fromActor(self); // Pass the server principal
142-
active_streams = appContext.activeStreams;
143-
mcp_server = mcpServer;
144-
streaming_callback = http_request_streaming_callback;
145-
auth = null; // No authentication in this example.
146-
http_asset_cache = null; // No HTTP asset cache in this example.
147-
};
123+
let ctx : HttpHandler.Context = _create_http_context();
148124
return await HttpHandler.http_request_update(ctx, req);
149125
};
126+
127+
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
128+
let ctx : HttpHandler.Context = _create_http_context();
129+
return HttpHandler.http_request_streaming_callback(ctx, token);
130+
};
150131
};

‎src/mcp/HttpHandler.mo‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import Blob "mo:base/Blob";
66
import Time "mo:base/Time";
77
import Nat64 "mo:base/Nat64";
88
import Text "mo:base/Text";
9-
import Debug "mo:base/Debug";
9+
import Option "mo:base/Option";
1010
import BaseX "mo:base-x-encoder";
1111
import HttpTypes "mo:http-types";
1212
import Utils "Utils";
@@ -50,6 +50,20 @@ module {
5050
return false;
5151
};
5252

53+
public func http_request_streaming_callback(ctx : Context, token : HttpTypes.StreamingToken) : ?HttpTypes.StreamingCallbackResponse {
54+
let tokenKey = BaseX.toBase64(token.vals(), #standard({ includePadding = true }));
55+
// It has access to the actor's state.
56+
if (Option.isNull(Map.get(ctx.active_streams, thash, tokenKey))) {
57+
return ?{ body = Blob.fromArray([]); token = null };
58+
};
59+
60+
// Update the timestamp to prove the stream is still active.
61+
Map.set(ctx.active_streams, thash, tokenKey, Time.now());
62+
63+
let chunk = Text.encodeUtf8("data: {\"type\":\"keep-alive\"}\n\n");
64+
return ?{ body = chunk; token = ?token };
65+
};
66+
5367
// The public entry point for query calls.
5468
public func http_request(ctx : Context, req : SrvTypes.HttpRequest) : SrvTypes.HttpResponse {
5569
if (req.method == "GET" and Text.contains(req.url, #text "/.well-known/oauth-protected-resource")) {

‎test/e2e/private_mcp_server/main.mo‎

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
import Map "mo:map/Map";
22
import { thash } "mo:map/Map";
33
import Result "mo:base/Result";
4-
import Text "mo:base/Text";
5-
import Option "mo:base/Option";
64
import Blob "mo:base/Blob";
7-
import Time "mo:base/Time";
85
import Principal "mo:base/Principal";
96
import Json "../../../src/json";
107
import HttpTypes "mo:http-types";
11-
import BaseX "mo:base-x-encoder";
128

139
// The only SDK import the user needs!
1410
import Mcp "../../../src/mcp/Mcp";
@@ -164,19 +160,9 @@ shared persistent actor class McpServer() = self {
164160
return await HttpHandler.http_request_update(ctx, req);
165161
};
166162

167-
// The streaming callback MUST be a public function of the main actor.
168163
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
169-
let token_key = BaseX.toBase64(token.vals(), #standard({ includePadding = true }));
170-
// It has access to the actor's state.
171-
if (Option.isNull(Map.get(appContext.activeStreams, thash, token_key))) {
172-
return ?{ body = Blob.fromArray([]); token = null };
173-
};
174-
175-
// Update the timestamp to prove the stream is still active.
176-
Map.set(appContext.activeStreams, thash, token_key, Time.now());
177-
178-
let chunk = Text.encodeUtf8("data: {\"type\":\"keep-alive\"}\n\n");
179-
return ?{ body = chunk; token = ?token };
164+
let ctx : HttpHandler.Context = _create_http_context();
165+
return HttpHandler.http_request_streaming_callback(ctx, token);
180166
};
181167

182168
system func preupgrade() {

‎test/e2e/public_mcp_server/main.mo‎

Lines changed: 16 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,9 @@
11
import Map "mo:map/Map";
22
import { thash } "mo:map/Map";
33
import Result "mo:base/Result";
4-
import Text "mo:base/Text";
5-
import Option "mo:base/Option";
6-
import Blob "mo:base/Blob";
7-
import Time "mo:base/Time";
84
import Principal "mo:base/Principal";
95
import Json "../../../src/json";
106
import HttpTypes "mo:http-types";
11-
import BaseX "mo:base-x-encoder";
127

138
// The only SDK import the user needs!
149
import Mcp "../../../src/mcp/Mcp";
@@ -107,44 +102,30 @@ shared persistent actor class McpServer() = self {
107102

108103
// --- PUBLIC ENTRY POINTS ---
109104

110-
// The streaming callback MUST be a public function of the main actor.
111-
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
112-
let token_key = BaseX.toBase64(token.vals(), #standard({ includePadding = true }));
113-
// It has access to the actor's state.
114-
if (Option.isNull(Map.get(appContext.activeStreams, thash, token_key))) {
115-
return ?{ body = Blob.fromArray([]); token = null };
116-
};
117-
118-
// Update the timestamp to prove the stream is still active.
119-
Map.set(appContext.activeStreams, thash, token_key, Time.now());
120-
121-
let chunk = Text.encodeUtf8("data: {\"type\":\"keep-alive\"}\n\n");
122-
return ?{ body = chunk; token = ?token };
123-
};
124-
125-
public query func http_request(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
126-
// Construct the context object on the fly.
127-
let ctx : HttpHandler.Context = {
128-
self = Principal.fromActor(self); // Pass the server principal
105+
// Helper to avoid repeating context creation.
106+
private func _create_http_context() : HttpHandler.Context {
107+
return {
108+
self = Principal.fromActor(self);
129109
active_streams = appContext.activeStreams;
130110
mcp_server = mcpServer;
131111
streaming_callback = http_request_streaming_callback;
132-
auth = null; // No authentication in this example.
133-
http_asset_cache = null; // No HTTP asset cache in this example.
112+
auth = null;
113+
http_asset_cache = null;
134114
};
135-
// Delegate the complex logic to the handler module.
115+
};
116+
117+
public query func http_request(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
118+
let ctx : HttpHandler.Context = _create_http_context();
136119
return HttpHandler.http_request(ctx, req);
137120
};
138121

139122
public func http_request_update(req : SrvTypes.HttpRequest) : async SrvTypes.HttpResponse {
140-
let ctx : HttpHandler.Context = {
141-
self = Principal.fromActor(self); // Pass the server principal
142-
active_streams = appContext.activeStreams;
143-
mcp_server = mcpServer;
144-
streaming_callback = http_request_streaming_callback;
145-
auth = null; // No authentication in this example.
146-
http_asset_cache = null; // No HTTP asset cache in this example.
147-
};
123+
let ctx : HttpHandler.Context = _create_http_context();
148124
return await HttpHandler.http_request_update(ctx, req);
149125
};
126+
127+
public query func http_request_streaming_callback(token : HttpTypes.StreamingToken) : async ?HttpTypes.StreamingCallbackResponse {
128+
let ctx : HttpHandler.Context = _create_http_context();
129+
return HttpHandler.http_request_streaming_callback(ctx, token);
130+
};
150131
};

0 commit comments

Comments
 (0)