Skip to content

Fix streamable hibernation issues #203

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
threepointone merged 4 commits into from
Apr 23, 2025
Merged

Fix streamable hibernation issues #203

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1ded712
Make sure to re-init when checking if the DO is init'd
jmorrell-cloudflare Apr 23, 2025
08041e0
Explicitly track whether the agent has received an initialization req…
jmorrell-cloudflare Apr 23, 2025
3e62e44
Simplify
jmorrell-cloudflare Apr 23, 2025
9c6fb27
Create nasty-ties-jog.md
threepointone Apr 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/nasty-ties-jog.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"agents": patch
---

Jmorrell/fix streamable hibernation issue
16 changes: 11 additions & 5 deletions packages/agents/src/mcp/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ export abstract class McpAgent<
this.#transportType = (await this.ctx.storage.get(
"transportType"
)) as TransportType;
this.init?.();
await this._init(this.props);

// Connect to the MCP server
if (this.#transportType === "sse") {
Expand All @@ -285,16 +285,22 @@ export abstract class McpAgent<

async _init(props: Props) {
await this.ctx.storage.put("props", props ?? {});
await this.ctx.storage.put("transportType", "unset");
if (!this.ctx.storage.get("transportType")) {
await this.ctx.storage.put("transportType", "unset");
}
Comment on lines +288 to +290
Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without this guard this would be overwritten when the DO wakes up. The complexity of our "boot" logic is higher than I would like.

this.props = props;
if (!this.initRun) {
this.initRun = true;
await this.init();
}
}

isInitialized() {
return this.initRun;
async setInitialized() {
await this.ctx.storage.put("initialized", true);
}

async isInitialized() {
return (await this.ctx.storage.get("initialized")) === true;
Comment on lines +298 to +303
Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a little odd because this is only used by the streamable transport, and not the SSE transport.

}

async #initialize(): Promise<void> {
Expand Down Expand Up @@ -898,7 +904,7 @@ export abstract class McpAgent<
const isInitialized = await doStub.isInitialized();

if (isInitializationRequest) {
await doStub._init(ctx.props);
Copy link

@BrandonNoad BrandonNoad May 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a question about this change. ctx is no longer used in static serve. So there doesn't seem to be a way to provide additional context from the original request to init when defining your MCP tools (e.g., header values)?

Also, wasn't the OAuth flow relying on that? https://github.com/cloudflare/workers-oauth-provider/blob/main/src/oauth-provider.ts#L1857

See my comment here for additional details: #194 (comment)

Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare May 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! You're absolutely right. Let me add that back in

Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare May 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that this is only an issue if you need to access the results of auth in your MCP tools. I think the auth flow itself will still work.

Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare May 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#246

await doStub.setInitialized();
} else if (!isInitialized) {
// if we have gotten here, then a session id that was never initialized
// was provided
Expand Down
1 change: 0 additions & 1 deletion packages/agents/src/tests/mcp-streamable-http.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,6 @@ describe("McpAgent Streamable HTTP Transport", () => {

it("should reject invalid session ID", async () => {
const ctx = createExecutionContext();
const sessionId = await initializeServer(ctx);
Copy link
Collaborator Author

@jmorrell-cloudflare jmorrell-cloudflare Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is unused in the rest of the test and can be safely removed without affecting the test logic


// Now try with invalid session ID
const response = await sendPostRequest(
Expand Down