146 lines
4.1 KiB
TypeScript
146 lines
4.1 KiB
TypeScript
// WebSocket Agent with real-time bidirectional communication
|
|
|
|
import { Agent, Connection, ConnectionContext, WSMessage } from "agents";
|
|
|
|
interface Env {
|
|
// Add bindings here
|
|
}
|
|
|
|
interface ChatState {
|
|
messages: Array<{
|
|
id: string;
|
|
text: string;
|
|
sender: string;
|
|
timestamp: number;
|
|
}>;
|
|
participants: string[];
|
|
createdAt: number;
|
|
}
|
|
|
|
export class ChatAgent extends Agent<Env, ChatState> {
|
|
initialState: ChatState = {
|
|
messages: [],
|
|
participants: [],
|
|
createdAt: Date.now()
|
|
};
|
|
|
|
// Called when a client connects via WebSocket
|
|
async onConnect(connection: Connection, ctx: ConnectionContext) {
|
|
// Access original HTTP request for authentication
|
|
const userId = ctx.request.headers.get('X-User-ID') || 'anonymous';
|
|
const authToken = ctx.request.headers.get('Authorization');
|
|
|
|
// Optional: Close connection if unauthorized
|
|
if (!authToken) {
|
|
connection.close(401, "Unauthorized");
|
|
return;
|
|
}
|
|
|
|
console.log(`Client ${connection.id} connected as ${userId}`);
|
|
|
|
// Add to participants
|
|
if (!this.state.participants.includes(userId)) {
|
|
this.setState({
|
|
...this.state,
|
|
participants: [...this.state.participants, userId]
|
|
});
|
|
}
|
|
|
|
// Send welcome message to this connection
|
|
connection.send(JSON.stringify({
|
|
type: 'welcome',
|
|
message: `Welcome, ${userId}!`,
|
|
participants: this.state.participants,
|
|
messageCount: this.state.messages.length
|
|
}));
|
|
|
|
// Broadcast participant joined to all connected clients (via state sync)
|
|
connection.send(JSON.stringify({
|
|
type: 'user_joined',
|
|
userId,
|
|
participants: this.state.participants
|
|
}));
|
|
}
|
|
|
|
// Called for each message received
|
|
async onMessage(connection: Connection, message: WSMessage) {
|
|
// Handle string messages (most common)
|
|
if (typeof message === 'string') {
|
|
try {
|
|
const data = JSON.parse(message);
|
|
|
|
// Handle chat message
|
|
if (data.type === 'chat') {
|
|
const newMessage = {
|
|
id: crypto.randomUUID(),
|
|
text: data.text,
|
|
sender: data.sender || 'anonymous',
|
|
timestamp: Date.now()
|
|
};
|
|
|
|
// Add to state (will sync to all connections)
|
|
this.setState({
|
|
...this.state,
|
|
messages: [...this.state.messages, newMessage]
|
|
});
|
|
|
|
// Send acknowledgement to sender
|
|
connection.send(JSON.stringify({
|
|
type: 'message_sent',
|
|
messageId: newMessage.id
|
|
}));
|
|
}
|
|
|
|
// Handle typing indicator
|
|
if (data.type === 'typing') {
|
|
// Broadcast to all except sender
|
|
connection.send(JSON.stringify({
|
|
type: 'user_typing',
|
|
userId: data.sender
|
|
}));
|
|
}
|
|
} catch (e) {
|
|
// Handle parse error
|
|
connection.send(JSON.stringify({
|
|
type: 'error',
|
|
message: 'Invalid message format'
|
|
}));
|
|
}
|
|
}
|
|
|
|
// Handle binary messages (ArrayBuffer or ArrayBufferView)
|
|
if (message instanceof ArrayBuffer) {
|
|
console.log('Received binary message:', message.byteLength, 'bytes');
|
|
// Process binary data...
|
|
}
|
|
}
|
|
|
|
// Called when connection has an error
|
|
async onError(connection: Connection, error: unknown): Promise<void> {
|
|
console.error('WebSocket error:', error);
|
|
|
|
// Send error message to client
|
|
connection.send(JSON.stringify({
|
|
type: 'error',
|
|
message: 'Connection error occurred'
|
|
}));
|
|
}
|
|
|
|
// Called when connection closes
|
|
async onClose(connection: Connection, code: number, reason: string, wasClean: boolean): Promise<void> {
|
|
console.log(`Connection ${connection.id} closed:`, code, reason, wasClean);
|
|
|
|
// Clean up connection-specific state if needed
|
|
// Note: Agent state persists even after all connections close
|
|
}
|
|
|
|
// React to state updates (from any source)
|
|
onStateUpdate(state: ChatState, source: "server" | Connection) {
|
|
console.log('Chat state updated');
|
|
console.log('Message count:', state.messages.length);
|
|
console.log('Participants:', state.participants);
|
|
}
|
|
}
|
|
|
|
export default ChatAgent;
|