Skip to content

Commit 9331ca6

Browse files
Merge pull request #24 from smallest-inc/feat/chat-with-agent
feat: add chat_with_agent tool (text conversation over realtime WebSocket)
2 parents f0e1441 + 0e0fcf5 commit 9331ca6

4 files changed

Lines changed: 377 additions & 0 deletions

File tree

package.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,12 @@
2828
},
2929
"dependencies": {
3030
"@modelcontextprotocol/sdk": "^1.12.1",
31+
"ws": "^8.18.0",
3132
"zod": "^3.24.1"
3233
},
3334
"devDependencies": {
3435
"@types/node": "^22.0.0",
36+
"@types/ws": "^8.5.0",
3537
"esbuild": "^0.24.0",
3638
"tsx": "^4.19.0",
3739
"typescript": "^5.4.5"

src/chat-client.ts

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
import WebSocket from "ws";
2+
3+
/**
4+
* Minimal text-only client for the Atoms realtime agent — `/connect?mode=chat`.
5+
*
6+
* Speaks the documented WebSocket contract (no audio):
7+
* client → server : {"type":"input_text.send","text":"..."}
8+
* {"type":"session.close"}
9+
* server → client : session.created | transcript {role,text} | session.closed | error
10+
*
11+
* Auth is the raw API key on the `token` query param (the gateway's API-key
12+
* path). The whole exchange is text, so this needs no mic/audio handling — it
13+
* just sends user turns and collects the agent's `transcript` replies.
14+
*/
15+
16+
export interface ChatTurn {
17+
role: "user" | "agent";
18+
text: string;
19+
}
20+
21+
export interface ChatClientOptions {
22+
apiKey: string;
23+
agentId: string;
24+
/** Base URL, e.g. wss://api.smallest.ai/atoms/v1 (no trailing /agent/connect). */
25+
baseWssUrl: string;
26+
variables?: Record<string, string | number | boolean>;
27+
}
28+
29+
interface ReplyWaiter {
30+
resolve: (text: string) => void;
31+
reject: (err: Error) => void;
32+
timer: ReturnType<typeof setTimeout>;
33+
}
34+
35+
export class AtomsChatClient {
36+
private ws: WebSocket | null = null;
37+
private waiters: ReplyWaiter[] = [];
38+
39+
callId = "";
40+
sessionId = "";
41+
sampleRate = 0;
42+
/** Full conversation in arrival order: greeting, user echoes, agent replies. */
43+
readonly transcript: ChatTurn[] = [];
44+
closed = false;
45+
closedReason: string | null = null;
46+
47+
private onSessionCreated: (() => void) | null = null;
48+
/** Reject the in-flight connect() if an error/close arrives before session.created. */
49+
private onConnectFailure: ((err: Error) => void) | null = null;
50+
51+
constructor(private readonly opts: ChatClientOptions) {}
52+
53+
private connectUrl(): string {
54+
const params = new URLSearchParams({
55+
token: this.opts.apiKey,
56+
agent_id: this.opts.agentId,
57+
mode: "chat",
58+
});
59+
if (this.opts.variables && Object.keys(this.opts.variables).length > 0) {
60+
params.set("variables", JSON.stringify(this.opts.variables));
61+
}
62+
return `${this.opts.baseWssUrl.replace(/\/+$/, "")}/agent/connect?${params.toString()}`;
63+
}
64+
65+
/**
66+
* Open the WebSocket and resolve once `session.created` arrives. Then wait up
67+
* to `greetingWaitMs` for the agent's opening message (if the agent greets
68+
* first) so it lands in the transcript before any user turn is sent.
69+
*
70+
* @returns the call id, session id, and the greeting text (or null if the
71+
* agent waits for the user to speak first).
72+
*/
73+
async connect(
74+
greetingWaitMs = 3000,
75+
connectTimeoutMs = 15000
76+
): Promise<{ callId: string; sessionId: string; greeting: string | null }> {
77+
await new Promise<void>((resolve, reject) => {
78+
const ws = new WebSocket(this.connectUrl());
79+
this.ws = ws;
80+
let settled = false;
81+
82+
const timer = setTimeout(() => {
83+
if (settled) return;
84+
settled = true;
85+
ws.close();
86+
reject(new Error("Timed out waiting for session.created"));
87+
}, connectTimeoutMs);
88+
89+
this.onSessionCreated = () => {
90+
if (settled) return;
91+
settled = true;
92+
clearTimeout(timer);
93+
this.onConnectFailure = null;
94+
resolve();
95+
};
96+
this.onConnectFailure = (err: Error) => {
97+
if (settled) return;
98+
settled = true;
99+
clearTimeout(timer);
100+
this.onConnectFailure = null;
101+
reject(err);
102+
};
103+
104+
ws.on("message", (raw: WebSocket.RawData) => this.handleMessage(raw));
105+
ws.on("error", (err: Error) => {
106+
if (!settled) {
107+
settled = true;
108+
clearTimeout(timer);
109+
reject(err);
110+
}
111+
this.failWaiters(err);
112+
});
113+
ws.on("close", () => {
114+
this.closed = true;
115+
this.failWaiters(new Error(`Chat session closed${this.closedReason ? `: ${this.closedReason}` : ""}`));
116+
});
117+
});
118+
119+
// Let the opening message arrive (no-op for wait-for-user-first agents).
120+
// Nothing has been sent yet, so the first agent turn is the greeting.
121+
await this.delay(greetingWaitMs);
122+
const greetingTurn = this.transcript.find((t) => t.role === "agent");
123+
124+
return {
125+
callId: this.callId,
126+
sessionId: this.sessionId,
127+
greeting: greetingTurn ? greetingTurn.text : null,
128+
};
129+
}
130+
131+
/**
132+
* Send one user message and resolve with the agent's next reply.
133+
* Rejects if the session is closed or no reply arrives within `timeoutMs`.
134+
*/
135+
async send(text: string, timeoutMs = 25000): Promise<string> {
136+
if (this.closed || !this.ws || this.ws.readyState !== WebSocket.OPEN) {
137+
throw new Error(`Chat session is not open${this.closedReason ? ` (closed: ${this.closedReason})` : ""}`);
138+
}
139+
const reply = this.nextAgentReply(timeoutMs);
140+
this.ws.send(JSON.stringify({ type: "input_text.send", text }));
141+
return reply;
142+
}
143+
144+
/** Send session.close and shut the socket. */
145+
close(): void {
146+
if (this.ws && this.ws.readyState === WebSocket.OPEN) {
147+
try {
148+
this.ws.send(JSON.stringify({ type: "session.close" }));
149+
} catch {
150+
// ignore — closing anyway
151+
}
152+
}
153+
this.ws?.close();
154+
this.closed = true;
155+
}
156+
157+
// ---------------------------------------------------------------------------
158+
159+
private nextAgentReply(timeoutMs: number): Promise<string> {
160+
return new Promise<string>((resolve, reject) => {
161+
const waiter: ReplyWaiter = {
162+
resolve,
163+
reject,
164+
timer: setTimeout(() => {
165+
const i = this.waiters.indexOf(waiter);
166+
if (i >= 0) this.waiters.splice(i, 1);
167+
reject(new Error(`Timed out after ${timeoutMs}ms waiting for the agent's reply`));
168+
}, timeoutMs),
169+
};
170+
this.waiters.push(waiter);
171+
});
172+
}
173+
174+
private handleMessage(raw: WebSocket.RawData): void {
175+
let event: Record<string, unknown>;
176+
try {
177+
event = JSON.parse(raw.toString());
178+
} catch {
179+
return;
180+
}
181+
182+
switch (event.type) {
183+
case "session.created":
184+
this.callId = String(event.call_id ?? "");
185+
this.sessionId = String(event.session_id ?? "");
186+
this.sampleRate = Number(event.sample_rate ?? 0);
187+
this.onSessionCreated?.();
188+
break;
189+
190+
case "transcript": {
191+
// pipecat emits role "assistant"; the docs/SDK use "agent". Normalize
192+
// everything that isn't the user to "agent".
193+
const role: ChatTurn["role"] = event.role === "user" ? "user" : "agent";
194+
const text = String(event.text ?? "");
195+
this.transcript.push({ role, text });
196+
if (role === "agent") {
197+
const waiter = this.waiters.shift();
198+
if (waiter) {
199+
clearTimeout(waiter.timer);
200+
waiter.resolve(text);
201+
}
202+
}
203+
break;
204+
}
205+
206+
case "session.closed": {
207+
this.closed = true;
208+
this.closedReason = String(event.reason ?? "ended");
209+
const err = new Error(`Chat session closed: ${this.closedReason}`);
210+
if (this.onConnectFailure) this.onConnectFailure(err);
211+
this.failWaiters(err);
212+
break;
213+
}
214+
215+
case "error": {
216+
const message = `[${event.code ?? "error"}] ${event.message ?? "unknown error"}`;
217+
const err = new Error(message);
218+
// Before the session is up, an error means the connect itself failed
219+
// (e.g. the gateway couldn't reach NATS) — surface it to connect().
220+
if (this.onConnectFailure) {
221+
this.onConnectFailure(err);
222+
break;
223+
}
224+
const waiter = this.waiters.shift();
225+
if (waiter) {
226+
clearTimeout(waiter.timer);
227+
waiter.reject(err);
228+
}
229+
break;
230+
}
231+
232+
default:
233+
// transcript.delta / agent.start_talking / etc. — ignored for chat.
234+
break;
235+
}
236+
}
237+
238+
private failWaiters(err: Error): void {
239+
const waiters = this.waiters.splice(0);
240+
for (const w of waiters) {
241+
clearTimeout(w.timer);
242+
w.reject(err);
243+
}
244+
}
245+
246+
private delay(ms: number): Promise<void> {
247+
return new Promise((resolve) => setTimeout(resolve, ms));
248+
}
249+
}

src/tools/chat.ts

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2+
import { z } from "zod";
3+
4+
import { AtomsChatClient, ChatTurn } from "../chat-client.js";
5+
6+
const ATOMS_API_URL = "https://api.smallest.ai/atoms/v1";
7+
8+
/** Derive the realtime WebSocket base from the HTTP API base. */
9+
function wssBaseUrl(): string {
10+
return ATOMS_API_URL.replace(/^http/i, "ws");
11+
}
12+
13+
function renderTranscript(turns: ChatTurn[]): string {
14+
return turns.map((t) => `${t.role === "user" ? "User " : "Agent"}: ${t.text}`).join("\n");
15+
}
16+
17+
export function registerChatWithAgent(server: McpServer) {
18+
server.registerTool(
19+
"chat_with_agent",
20+
{
21+
description:
22+
"Hold a TEXT conversation with a published agent over the realtime chat WebSocket (mode=chat) — " +
23+
"no audio, no phone, just text in / text out. Sends each message in `messages` in order, " +
24+
"waiting for the agent's reply between turns, and returns the full transcript. " +
25+
"Use this to test an agent's prompt/behaviour programmatically (e.g. an automated build → test → " +
26+
"evaluate → refine loop): run a scripted conversation, read the transcript, then adjust the prompt " +
27+
"with update_agent_prompt and run again. This places a real (chargeable) chat session on the agent. " +
28+
"Note: the agent must be published; for unpublished drafts use test_draft (mode=chat) to start one.",
29+
inputSchema: {
30+
agent_id: z.string().describe("The agent ID to chat with (must be a published agent)"),
31+
messages: z
32+
.array(z.string().min(1))
33+
.min(1)
34+
.describe(
35+
"User turns to send, in order. Each is sent only after the previous turn's reply arrives. " +
36+
"For realistic tests, write messages a real caller would send."
37+
),
38+
variables: z
39+
.record(z.string(), z.union([z.string(), z.number(), z.boolean()]))
40+
.optional()
41+
.describe("Per-call values for {{key}} placeholders in the agent prompt"),
42+
reply_timeout_ms: z
43+
.number()
44+
.int()
45+
.min(1000)
46+
.max(120000)
47+
.default(25000)
48+
.describe("Max time to wait for each agent reply before giving up on that turn"),
49+
greeting_wait_ms: z
50+
.number()
51+
.int()
52+
.min(0)
53+
.max(15000)
54+
.default(3000)
55+
.describe("Time to wait after connecting for the agent's opening message (0 if it waits for the user first)"),
56+
},
57+
},
58+
async (params) => {
59+
const apiKey = process.env.ATOMS_API_KEY;
60+
if (!apiKey) {
61+
return {
62+
isError: true,
63+
content: [{ type: "text" as const, text: "ATOMS_API_KEY environment variable is required" }],
64+
};
65+
}
66+
67+
const client = new AtomsChatClient({
68+
apiKey,
69+
agentId: params.agent_id,
70+
baseWssUrl: wssBaseUrl(),
71+
variables: params.variables,
72+
});
73+
74+
let connectError: string | null = null;
75+
let turnError: string | null = null;
76+
let greeting: string | null = null;
77+
78+
try {
79+
const session = await client.connect(params.greeting_wait_ms);
80+
greeting = session.greeting;
81+
82+
for (const message of params.messages) {
83+
try {
84+
await client.send(message, params.reply_timeout_ms);
85+
} catch (err) {
86+
turnError = err instanceof Error ? err.message : String(err);
87+
break; // session likely closed/errored — stop sending
88+
}
89+
}
90+
} catch (err) {
91+
connectError = err instanceof Error ? err.message : String(err);
92+
} finally {
93+
client.close();
94+
}
95+
96+
if (connectError) {
97+
return {
98+
isError: true,
99+
content: [
100+
{
101+
type: "text" as const,
102+
text: `Failed to start chat with agent ${params.agent_id}: ${connectError}`,
103+
},
104+
],
105+
};
106+
}
107+
108+
const result = {
109+
agent_id: params.agent_id,
110+
call_id: client.callId || null,
111+
greeting,
112+
turns_sent: params.messages.length,
113+
ended_reason: client.closedReason,
114+
error: turnError,
115+
transcript: client.transcript,
116+
conversation: renderTranscript(client.transcript),
117+
};
118+
119+
return {
120+
content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
121+
};
122+
}
123+
);
124+
}

0 commit comments

Comments
 (0)