Skip to content

feat(mastra): support shared state for local agents #123

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

Closed
wants to merge 1 commit into from
Closed

feat(mastra): support shared state for local agents #123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion typescript-sdk/integrations/mastra/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
},
"peerDependencies": {
"@copilotkit/runtime": "^1.8.13",
"@mastra/core": "^0.10.1",
"@mastra/core": "^0.10.6",
"zod": "^3.0.0"
},
"devDependencies": {
Expand Down
305 changes: 196 additions & 109 deletions typescript-sdk/integrations/mastra/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import type {
RunAgentInput,
RunFinishedEvent,
RunStartedEvent,
StateSnapshotEvent,
TextMessageChunkEvent,
ToolCall,
ToolCallArgsEvent,
Expand All @@ -22,9 +23,9 @@ import {
ExperimentalEmptyAdapter,
} from "@copilotkit/runtime";
import { processDataStream } from "@ai-sdk/ui-utils";
import type { CoreMessage, Mastra } from "@mastra/core";
import type { CoreMessage, Mastra, StorageThreadType } from "@mastra/core";
import { registerApiRoute } from "@mastra/core/server";
import type { Agent as LocalMastraAgent } from "@mastra/core/agent";
import { Agent as LocalMastraAgent } from "@mastra/core/agent";
import type { Context } from "hono";
import { RuntimeContext } from "@mastra/core/runtime-context";
import { randomUUID } from "crypto";
Expand Down Expand Up @@ -185,88 +186,148 @@ export class MastraAgent extends AbstractAgent {
finalMessages.push(assistantMessage);

return new Observable<BaseEvent>((subscriber) => {
subscriber.next({
type: EventType.RUN_STARTED,
threadId: input.threadId,
runId: input.runId,
} as RunStartedEvent);

this.streamMastraAgent(input, {
onTextPart: (text) => {
assistantMessage.content += text;
const event: TextMessageChunkEvent = {
type: EventType.TEXT_MESSAGE_CHUNK,
role: "assistant",
messageId,
delta: text,
};
subscriber.next(event);
},
onToolCallPart: (streamPart) => {
let toolCall: ToolCall = {
id: streamPart.toolCallId,
type: "function",
function: {
name: streamPart.toolName,
arguments: JSON.stringify(streamPart.args),
const run = async () => {
subscriber.next({
type: EventType.RUN_STARTED,
threadId: input.threadId,
runId: input.runId,
} as RunStartedEvent);

// Handle local agent memory management (from Mastra implementation)
if ('metrics' in this.agent) {
const memory = this.agent.getMemory();

if (memory && input.state && Object.keys(input.state || {}).length > 0) {
let thread: StorageThreadType | null = await memory.getThreadById({ threadId: input.threadId });

if (!thread) {
thread = {
id: input.threadId,
title: '',
metadata: {},
resourceId: this.resourceId as string,
createdAt: new Date(),
updatedAt: new Date(),
};
}

if (thread.resourceId && thread.resourceId !== this.resourceId) {
throw new Error(
`Thread with id ${input.threadId} resourceId does not match the current resourceId ${this.resourceId}`,
);
}

const { messages, ...rest } = input.state;
const workingMemory = JSON.stringify(rest);

// Update thread metadata with new working memory
await memory.saveThread({
thread: {
...thread,
metadata: {
...thread.metadata,
workingMemory,
},
},
});
}
}

try {
await this.streamMastraAgent(input, {
onTextPart: (text) => {
assistantMessage.content += text;
const event: TextMessageChunkEvent = {
type: EventType.TEXT_MESSAGE_CHUNK,
role: "assistant",
messageId,
delta: text,
};
subscriber.next(event);
},
onToolCallPart: (streamPart) => {
let toolCall: ToolCall = {
id: streamPart.toolCallId,
type: "function",
function: {
name: streamPart.toolName,
arguments: JSON.stringify(streamPart.args),
},
};
assistantMessage.toolCalls!.push(toolCall);

const startEvent: ToolCallStartEvent = {
type: EventType.TOOL_CALL_START,
parentMessageId: messageId,
toolCallId: streamPart.toolCallId,
toolCallName: streamPart.toolName,
};
subscriber.next(startEvent);

const argsEvent: ToolCallArgsEvent = {
type: EventType.TOOL_CALL_ARGS,
toolCallId: streamPart.toolCallId,
delta: JSON.stringify(streamPart.args),
};
subscriber.next(argsEvent);

const endEvent: ToolCallEndEvent = {
type: EventType.TOOL_CALL_END,
toolCallId: streamPart.toolCallId,
};
subscriber.next(endEvent);
},
};
assistantMessage.toolCalls!.push(toolCall);

const startEvent: ToolCallStartEvent = {
type: EventType.TOOL_CALL_START,
parentMessageId: messageId,
toolCallId: streamPart.toolCallId,
toolCallName: streamPart.toolName,
};
subscriber.next(startEvent);

const argsEvent: ToolCallArgsEvent = {
type: EventType.TOOL_CALL_ARGS,
toolCallId: streamPart.toolCallId,
delta: JSON.stringify(streamPart.args),
};
subscriber.next(argsEvent);

const endEvent: ToolCallEndEvent = {
type: EventType.TOOL_CALL_END,
toolCallId: streamPart.toolCallId,
};
subscriber.next(endEvent);
},
onToolResultPart(streamPart) {
const toolMessage: ToolMessage = {
role: "tool",
id: randomUUID(),
toolCallId: streamPart.toolCallId,
content: JSON.stringify(streamPart.result),
};
finalMessages.push(toolMessage);
},
onFinishMessagePart: () => {
// Emit message snapshot
const event: MessagesSnapshotEvent = {
type: EventType.MESSAGES_SNAPSHOT,
messages: finalMessages,
};
subscriber.next(event);

// Emit run finished event
subscriber.next({
type: EventType.RUN_FINISHED,
threadId: input.threadId,
runId: input.runId,
} as RunFinishedEvent);

// Complete the observable
subscriber.complete();
},
onError: (error) => {
console.error("error", error);
// Handle error
onToolResultPart(streamPart) {
const toolMessage: ToolMessage = {
role: "tool",
id: randomUUID(),
toolCallId: streamPart.toolCallId,
content: JSON.stringify(streamPart.result),
};
finalMessages.push(toolMessage);
},
onFinishMessagePart: async () => {
// Emit message snapshot
const event: MessagesSnapshotEvent = {
type: EventType.MESSAGES_SNAPSHOT,
messages: finalMessages,
};
subscriber.next(event);

if ('metrics' in this.agent){
const memory = this.agent.getMemory();
if (memory) {
const workingMemory = await memory.getWorkingMemory({ threadId: input.threadId, format: 'json' });
subscriber.next({
type: EventType.STATE_SNAPSHOT,
snapshot: workingMemory,
} as StateSnapshotEvent);
}
}

// Emit run finished event
subscriber.next({
type: EventType.RUN_FINISHED,
threadId: input.threadId,
runId: input.runId,
} as RunFinishedEvent);

// Complete the observable
subscriber.complete();
},
onError: (error) => {
console.error("error", error);
// Handle error
subscriber.error(error);
},
});
} catch (error) {
console.error("Stream error:", error);
subscriber.error(error);
},
});
}
};

run();

return () => {};
});
Expand All @@ -278,7 +339,7 @@ export class MastraAgent extends AbstractAgent {
* @param options - The options for the mastra agent.
* @returns The stream of the mastra agent.
*/
private streamMastraAgent(
private async streamMastraAgent(
{ threadId, runId, messages, tools }: RunAgentInput,
{
onTextPart,
Expand All @@ -287,7 +348,7 @@ export class MastraAgent extends AbstractAgent {
onToolResultPart,
onError,
}: MastraAgentStreamOptions,
) {
): Promise<void> {
const clientTools = tools.reduce(
(acc, tool) => {
acc[tool.name as string] = {
Expand All @@ -310,48 +371,74 @@ export class MastraAgent extends AbstractAgent {
}

if (isLocalMastraAgent(this.agent)) {
// in process agent
return this.agent
.stream(convertedMessages, {
// Local agent - use the agent's stream method directly
try {
const response = await this.agent.stream(convertedMessages, {
threadId,
resourceId,
runId,
clientTools,
runtimeContext,
})
.then((response) => {
return processDataStream({
stream: (response as any).toDataStreamResponse().body!,
onTextPart,
onToolCallPart,
onToolResultPart,
onFinishMessagePart,
});
})
.catch((error) => {
onError?.(error);
});

// For local agents, the response should already be a stream
// Process it using the agent's built-in streaming mechanism
if (response && typeof response === 'object') {
// If the response has a toDataStreamResponse method, use it
if ('toDataStreamResponse' in response && typeof response.toDataStreamResponse === 'function') {
const dataStreamResponse = response.toDataStreamResponse();
if (dataStreamResponse && dataStreamResponse.body) {
await processDataStream({
stream: dataStreamResponse.body,
onTextPart,
onToolCallPart,
onToolResultPart,
onFinishMessagePart,
});
} else {
throw new Error('Invalid data stream response from local agent');
}
} else {
// If it's already a readable stream, process it directly
await processDataStream({
stream: response as any,
onTextPart,
onToolCallPart,
onToolResultPart,
onFinishMessagePart,
});
}
} else {
throw new Error('Invalid response from local agent');
}
} catch (error) {
onError?.(error as Error);
}
} else {
// remote agent
return this.agent
.stream({
// Remote agent - use the remote agent's stream method
try {
const response = await this.agent.stream({
threadId,
resourceId,
runId,
messages: convertedMessages,
clientTools,
})
.then((response) => {
return response.processDataStream({
});

// Remote agents should have a processDataStream method
if (response && typeof response.processDataStream === 'function') {
await response.processDataStream({
onTextPart,
onToolCallPart,
onToolResultPart,
onFinishMessagePart,
});
})
.catch((error) => {
onError?.(error);
});
} else {
throw new Error('Invalid response from remote agent');
}
} catch (error) {
onError?.(error as Error);
}
}
}
}
Expand Down
Loading