|
1 | 1 | import { mastra } from '../../../src/mastra' |
2 | | -import { RequestContext } from '@mastra/core/request-context' |
| 2 | +import { |
| 3 | + RequestContext, |
| 4 | + MASTRA_RESOURCE_ID_KEY, |
| 5 | + MASTRA_THREAD_ID_KEY, |
| 6 | +} from '@mastra/core/request-context' |
3 | 7 | import { toAISdkStream } from '@mastra/ai-sdk' |
4 | 8 | import { createUIMessageStream, createUIMessageStreamResponse } from 'ai' |
| 9 | +import type { UIMessage } from 'ai' |
5 | 10 |
|
6 | 11 | export async function POST(req: Request) { |
7 | | - const { messages, data, id } = await req.json() |
| 12 | + const body = await req.json() |
| 13 | + const { messages, data, id } = body |
8 | 14 |
|
9 | | - const agentId: string = |
| 15 | + // Agent selection logic |
| 16 | + const agentId = |
10 | 17 | (typeof data?.agentId === 'string' && data.agentId.length > 0 |
11 | 18 | ? data.agentId |
12 | 19 | : undefined) ?? |
13 | 20 | (typeof id === 'string' && id.length > 0 ? id : undefined) ?? |
14 | 21 | 'weatherAgent' |
15 | 22 |
|
16 | | - const myAgent = mastra.getAgent(agentId) |
| 23 | + const agent = mastra.getAgent(agentId) |
17 | 24 |
|
| 25 | + if (!agent) { |
| 26 | + return Response.json( |
| 27 | + { error: `Agent "${agentId}" not found` }, |
| 28 | + { status: 404 } |
| 29 | + ) |
| 30 | + } |
| 31 | + |
| 32 | + // Extract multi-tenancy IDs |
| 33 | + const threadId = body.threadId ?? data?.threadId |
| 34 | + const resourceId = body.resourceId ?? data?.resourceId |
| 35 | + |
| 36 | + // Create RequestContext for multi-tenancy isolation |
18 | 37 | const requestContext = new RequestContext() |
| 38 | + if (resourceId) requestContext.set(MASTRA_RESOURCE_ID_KEY, resourceId) |
| 39 | + if (threadId) requestContext.set(MASTRA_THREAD_ID_KEY, threadId) |
19 | 40 |
|
20 | | - if (data) { |
| 41 | + // Merge other data into context if present |
| 42 | + if (data && typeof data === 'object') { |
21 | 43 | for (const [key, value] of Object.entries(data)) { |
22 | | - requestContext.set(key, value) |
| 44 | + if ( |
| 45 | + key !== 'agentId' && |
| 46 | + key !== 'threadId' && |
| 47 | + key !== 'resourceId' |
| 48 | + ) { |
| 49 | + requestContext.set(key, value) |
| 50 | + } |
23 | 51 | } |
24 | 52 | } |
25 | 53 |
|
26 | | - const stream = await myAgent.stream(messages, { requestContext }) |
| 54 | + // Prepare stream options |
| 55 | + const streamOptions = { |
| 56 | + threadId, |
| 57 | + resourceId, |
| 58 | + memory: body.memory ?? data?.memory, |
| 59 | + requestContext, |
| 60 | + } |
| 61 | + |
| 62 | + const stream = await agent.stream(messages, streamOptions) |
27 | 63 |
|
28 | 64 | const uiStream = createUIMessageStream({ |
29 | | - originalMessages: messages, |
| 65 | + originalMessages: messages as UIMessage[], |
30 | 66 | execute: async ({ writer }) => { |
31 | | - const aiStream = toAISdkStream(stream, { from: 'agent' }) as any |
| 67 | + const aiStream = toAISdkStream(stream, { |
| 68 | + from: 'agent', |
| 69 | + sendReasoning: true, |
| 70 | + sendSources: true, |
| 71 | + }) |
32 | 72 |
|
33 | | - if (typeof aiStream[Symbol.asyncIterator] === 'function') { |
34 | | - for await (const part of aiStream as AsyncIterable<any>) { |
35 | | - await writer.write(part) |
36 | | - } |
37 | | - } else if (typeof aiStream.getReader === 'function') { |
38 | | - const reader = aiStream.getReader() |
| 73 | + // Support both ReadableStream and AsyncIterable (robust bridge) |
| 74 | + if ( |
| 75 | + aiStream && |
| 76 | + typeof (aiStream as any).getReader === 'function' |
| 77 | + ) { |
| 78 | + const reader = (aiStream as ReadableStream<any>).getReader() |
39 | 79 | try { |
40 | 80 | while (true) { |
41 | 81 | const { value, done } = await reader.read() |
42 | | - if (done) { |
43 | | - break |
44 | | - } |
| 82 | + if (done) break |
45 | 83 | await writer.write(value) |
46 | 84 | } |
47 | 85 | } finally { |
48 | 86 | reader.releaseLock?.() |
49 | 87 | } |
| 88 | + } else if (aiStream && Symbol.asyncIterator in aiStream) { |
| 89 | + for await (const part of aiStream as AsyncIterable<any>) { |
| 90 | + await writer.write(part) |
| 91 | + } |
50 | 92 | } |
51 | 93 | }, |
52 | 94 | }) |
|
0 commit comments