snapshot of https://github.com/tldraw/tldraw.git (main, 2026-09-23; blobs > 131072B omitted)

templates/branching-chat/client/nodes/NodeShapeUtil.tsx

@@ -0,0 +1,153 @@
+import {
+	Circle2d,
+	Group2d,
+	HTMLContainer,
+	RecordProps,
+	Rectangle2d,
+	ShapeUtil,
+	TLShape,
+	useEditor,
+	useValue,
+} from 'tldraw'
+import { NODE_WIDTH_PX, PORT_RADIUS_PX } from '../constants'
+import { Port } from '../ports/Port'
+import { getNodePorts } from './nodePorts'
+import {
+	getNodeDefinition,
+	getNodeHeightPx,
+	getNodeTypePorts,
+	getNodeWidthPx,
+	NodeType,
+} from './nodeTypes'
+
+const NODE_TYPE = 'node'
+
+declare module 'tldraw' {
+	export interface TLGlobalShapePropsMap {
+		[NODE_TYPE]: { node: NodeType }
+	}
+}
+
+// Define our custom node shape type that extends tldraw's base shape system
+export type NodeShape = TLShape<typeof NODE_TYPE>
+
+// This class extends tldraw's ShapeUtil to define how our custom node shapes behave
+export class NodeShapeUtil extends ShapeUtil<NodeShape> {
+	static override type = NODE_TYPE
+	static override props: RecordProps<NodeShape> = {
+		node: NodeType,
+	}
+
+	getDefaultProps(): NodeShape['props'] {
+		return {
+			node: getNodeDefinition(this.editor, 'message').getDefault(),
+		}
+	}
+
+	override canEdit(_shape: NodeShape) {
+		return false
+	}
+	override canResize(_shape: NodeShape) {
+		return false
+	}
+	override hideResizeHandles(_shape: NodeShape) {
+		return true
+	}
+	override hideRotateHandle(_shape: NodeShape) {
+		return true
+	}
+	override hideSelectionBoundsBg(_shape: NodeShape) {
+		return true
+	}
+	override hideSelectionBoundsFg(_shape: NodeShape) {
+		return true
+	}
+	override isAspectRatioLocked(_shape: NodeShape) {
+		return false
+	}
+	override getBoundsSnapGeometry(_shape: NodeShape) {
+		return {
+			points: [{ x: 0, y: 0 }],
+		}
+	}
+
+	// Define the geometry of our node shape including ports
+	getGeometry(shape: NodeShape) {
+		const ports = getNodePorts(this.editor, shape)
+
+		const portGeometries = Object.values(ports).map(
+			(port) =>
+				new Circle2d({
+					x: port.x - PORT_RADIUS_PX,
+					y: port.y - PORT_RADIUS_PX,
+					radius: PORT_RADIUS_PX,
+					isFilled: true,
+					// not a label, but this hack excludes them from the selection bounds which is useful
+					isLabel: true,
+				})
+		)
+
+		const bodyGeometry = new Rectangle2d({
+			width: getNodeWidthPx(this.editor, shape),
+			height: getNodeHeightPx(this.editor, shape),
+			isFilled: true,
+		})
+
+		return new Group2d({
+			children: [bodyGeometry, ...portGeometries],
+		})
+	}
+
+	component(shape: NodeShape) {
+		return <NodeShape shape={shape} />
+	}
+
+	getIndicatorPath(shape: NodeShape) {
+		const height = getNodeHeightPx(this.editor, shape)
+		const path = new Path2D()
+		path.rect(0, 0, NODE_WIDTH_PX, height)
+		const ports = getNodePorts(this.editor, shape)
+		for (const port of Object.values(ports)) {
+			path.moveTo(port.x + PORT_RADIUS_PX, port.y)
+			path.arc(port.x, port.y, PORT_RADIUS_PX, 0, Math.PI * 2)
+		}
+		return path
+	}
+}
+
+// Main node component that renders the HTML content
+function NodeShape({ shape }: { shape: NodeShape }) {
+	const editor = useEditor()
+
+	return (
+		<HTMLContainer
+			className="NodeShape"
+			style={{
+				width: getNodeWidthPx(editor, shape),
+				height: getNodeHeightPx(editor, shape),
+			}}
+		>
+			<NodeBody shape={shape} />
+			<NodePorts shape={shape} />
+		</HTMLContainer>
+	)
+}
+
+function NodeBody({ shape }: { shape: NodeShape }) {
+	const node = shape.props.node
+	const editor = useEditor()
+	const { Component } = getNodeDefinition(editor, node)
+	return <Component shape={shape} node={node} />
+}
+
+function NodePorts({ shape }: { shape: NodeShape }) {
+	const editor = useEditor()
+	const ports = useValue('node ports', () => getNodeTypePorts(editor, shape), [shape, editor])
+	return (
+		<>
+			{Object.values(ports).map((port) => (
+				<Port key={port.id} shapeId={shape.id} port={port} />
+			))}
+		</>
+	)
+}

templates/branching-chat/client/nodes/nodePorts.tsx

@@ -0,0 +1,96 @@
+/**
+ * This file contains functions for working with ports and connections on nodes.
+ */
+import { createComputedCache, Editor, TLShapeId } from 'tldraw'
+import { getConnectionBindings } from '../connection/ConnectionBindingUtil'
+import { PortId } from '../ports/Port'
+import { NodeShape } from './NodeShapeUtil'
+import { getNodeTypePorts, NodeTypePorts } from './nodeTypes'
+
+/**
+ * Get the ports for a node. This is cached, because we only want to re-evaluate it when the
+ * underlying records change.
+ */
+export function getNodePorts(editor: Editor, shape: NodeShape | TLShapeId): NodeTypePorts {
+	return nodePortsCache.get(editor, typeof shape === 'string' ? shape : shape.id) ?? {}
+}
+const nodePortsCache = createComputedCache(
+	'ports',
+	(editor: Editor, node: NodeShape): NodeTypePorts => getNodeTypePorts(editor, node)
+)
+
+/**
+ * A connection from one node to another.
+ */
+export interface NodePortConnection {
+	connectedShapeId: TLShapeId
+	connectionId: TLShapeId
+	terminal: 'start' | 'end'
+	ownPortId: PortId
+	connectedPortId: PortId
+}
+
+/**
+ * Get the connections for a node. This is cached.
+ */
+export function getNodePortConnections(
+	editor: Editor,
+	shape: NodeShape | TLShapeId
+): NodePortConnection[] {
+	return nodePortConnectionsCache.get(editor, typeof shape === 'string' ? shape : shape.id) ?? []
+}
+const nodePortConnectionsCache = createComputedCache(
+	'port connections',
+	(editor: Editor, node: NodeShape) => {
+		const bindings = editor.getBindingsToShape(node.id, 'connection')
+
+		const connections: NodePortConnection[] = []
+		for (const binding of bindings) {
+			const oppositeTerminal = binding.props.terminal === 'start' ? 'end' : 'start'
+			const oppositeBinding = getConnectionBindings(editor, binding.fromId)[oppositeTerminal]
+			if (!oppositeBinding) continue
+
+			connections.push({
+				connectedShapeId: oppositeBinding.toId,
+				connectionId: binding.fromId,
+				terminal: binding.props.terminal,
+				ownPortId: binding.props.portId,
+				connectedPortId: oppositeBinding.props.portId,
+			})
+		}
+
+		return connections
+	},
+	{
+		areRecordsEqual: (a, b) => a.id === b.id,
+	}
+)
+
+/**
+ * Traverse the graph from a starting node, in a given direction, returning all the nodes that are
+ * connected to it (including the starting node).
+ */
+export function getAllConnectedNodes(
+	editor: Editor,
+	startingNode: TLShapeId | NodeShape,
+	direction?: 'start' | 'end'
+) {
+	const toVisit = [typeof startingNode === 'string' ? startingNode : startingNode.id]
+	const found = new Set<TLShapeId>()
+
+	while (toVisit.length > 0) {
+		const nodeId = toVisit.shift()!
+		const node = editor.getShape(nodeId)
+		if (!node || !editor.isShapeOfType(node, 'node')) continue
+
+		if (found.has(node.id)) continue
+		found.add(node.id)
+
+		for (const connection of getNodePortConnections(editor, node)) {
+			if (direction && connection.terminal !== direction) continue
+			toVisit.push(connection.connectedShapeId)
+		}
+	}
+
+	return found
+}

templates/branching-chat/client/nodes/nodeTypes.tsx

@@ -0,0 +1,67 @@
+import { Editor, T, WeakCache } from 'tldraw'
+import { PortId, shapePort } from '../ports/Port'
+import { NodeShape } from './NodeShapeUtil'
+import { MessageNodeDefinition } from './types/MessageNode'
+import { NodeDefinition, NodeDefinitionConstructor } from './types/shared'
+
+/** All our node types */
+export const NodeDefinitions = {
+	message: MessageNodeDefinition,
+} satisfies Record<string, NodeDefinitionConstructor<any>>
+
+const nodeDefinitions = new WeakCache<
+	Editor,
+	{ [K in keyof typeof NodeDefinitions]: InstanceType<(typeof NodeDefinitions)[K]> }
+>()
+export function getNodeDefinitions(editor: Editor) {
+	return nodeDefinitions.get(
+		editor,
+		() =>
+			Object.fromEntries(
+				Object.values(NodeDefinitions).map((value) => [value.type, new value(editor)])
+			) as any
+	)
+}
+
+export function getNodeDefinition(
+	editor: Editor,
+	node: NodeType | NodeType['type']
+): NodeDefinition<NodeType> {
+	return getNodeDefinitions(editor)[
+		typeof node === 'string' ? node : node.type
+	] as NodeDefinition<NodeType>
+}
+
+/**
+ * A union type of all our node types.
+ */
+export type NodeType = T.TypeOf<typeof NodeType>
+export const NodeType = T.union(
+	'type',
+	Object.fromEntries(Object.values(NodeDefinitions).map((type) => [type.type, type.validator])) as {
+		[K in keyof typeof NodeDefinitions as (typeof NodeDefinitions)[K]['type']]: (typeof NodeDefinitions)[K]['validator']
+	}
+)
+
+export function getNodeHeightPx(editor: Editor, shape: NodeShape): number {
+	return getNodeDefinition(editor, shape.props.node).getBodyHeightPx(shape, shape.props.node)
+}
+
+export function getNodeWidthPx(editor: Editor, shape: NodeShape): number {
+	return getNodeDefinition(editor, shape.props.node).getBodyWidthPx(shape, shape.props.node)
+}
+
+const _nodeTypePorts = T.dict(T.string, shapePort)
+export type NodeTypePorts = T.TypeOf<typeof _nodeTypePorts>
+
+export function getNodeTypePorts(editor: Editor, shape: NodeShape): NodeTypePorts {
+	return getNodeDefinition(editor, shape.props.node).getPorts(shape, shape.props.node)
+}
+
+export function onNodePortConnect(editor: Editor, shape: NodeShape, port: PortId) {
+	getNodeDefinition(editor, shape.props.node).onPortConnect?.(shape, shape.props.node, port)
+}
+
+export function onNodePortDisconnect(editor: Editor, shape: NodeShape, port: PortId) {
+	getNodeDefinition(editor, shape.props.node).onPortDisconnect?.(shape, shape.props.node, port)
+}

templates/branching-chat/client/nodes/types/MessageNode.tsx

@@ -0,0 +1,209 @@
+import { ModelMessage } from 'ai'
+import { useCallback } from 'react'
+import { T, TldrawUiButton, TldrawUiButtonIcon, TldrawUiInput, useEditor } from 'tldraw'
+import { HandleIcon } from '../../components/icons/HandleIcon'
+import { SendIcon } from '../../components/icons/SendIcon'
+import { NODE_HEIGHT_PX, NODE_WIDTH_PX } from '../../constants'
+import { getAllConnectedNodes } from '../nodePorts'
+import { NodeShape } from '../NodeShapeUtil'
+import {
+	NodeComponentProps,
+	NodeDefinition,
+	shapeInputPort,
+	shapeOutputPort,
+	updateNode,
+} from './shared'
+
+/**
+ * This node is a message from the user.
+ */
+export type MessageNode = T.TypeOf<typeof MessageNode>
+export const MessageNode = T.object({
+	type: T.literal('message'),
+	userMessage: T.string,
+	assistantMessage: T.string,
+})
+
+export class MessageNodeDefinition extends NodeDefinition<MessageNode> {
+	static type = 'message'
+	static validator = MessageNode
+	title = 'Message'
+	heading = 'Message'
+	icon = <SendIcon />
+	getDefault(): MessageNode {
+		return {
+			type: 'message',
+			userMessage: 'hello',
+			assistantMessage: '',
+		}
+	}
+	getBodyWidthPx(_shape: NodeShape, _node: MessageNode): number {
+		return NODE_WIDTH_PX
+	}
+	getBodyHeightPx(_shape: NodeShape, node: MessageNode): number {
+		const assistantMessage = node.assistantMessage.trim()
+		if (assistantMessage === '') return NODE_HEIGHT_PX
+		const size = this.editor.textMeasure.measureText(assistantMessage, {
+			fontFamily: 'Inter',
+			fontSize: 12,
+			fontWeight: '500',
+			fontStyle: 'normal',
+			maxWidth: NODE_WIDTH_PX,
+			lineHeight: 1.3,
+			padding: '12px',
+		})
+		return NODE_HEIGHT_PX + size.h
+	}
+	getPorts(shape: NodeShape, node: MessageNode) {
+		return {
+			input: shapeInputPort,
+			output: {
+				...shapeOutputPort,
+				y: this.getBodyHeightPx(shape, node),
+			},
+		}
+	}
+
+	Component = MessageNodeComponent
+}
+
+function MessageNodeComponent({ node, shape }: NodeComponentProps<MessageNode>) {
+	const editor = useEditor()
+
+	const handleSend = useCallback(() => {
+		// 1. gather up parents and create message history
+		// 2. create prompt
+		// 3. send prompt to ai
+		// 4. update node with assistant message
+
+		const messages: ModelMessage[] = []
+		for (const connectedShape of getAllConnectedNodes(editor, shape, 'end')) {
+			const node = editor.getShape(connectedShape)
+			if (!node || !editor.isShapeOfType(node, 'node')) continue
+			if (node.props.node.type !== 'message') continue
+
+			if (node.props.node.assistantMessage && connectedShape !== shape.id) {
+				messages.push({ role: 'assistant', content: node.props.node.assistantMessage })
+			}
+			messages.push({ role: 'user', content: node.props.node.userMessage })
+		}
+		messages.reverse()
+
+		// clear any previous assistant message before starting
+		updateNode<MessageNode>(editor, shape, (node) => ({
+			...node,
+			assistantMessage: '...',
+		}))
+
+		// stream the response and append as chunks arrive
+		;(async () => {
+			try {
+				const response = await fetch('/stream', {
+					method: 'POST',
+					headers: { 'Content-Type': 'application/json' },
+					body: JSON.stringify(messages),
+				})
+				if (!response.body) return
+
+				const reader = response.body.getReader()
+				const decoder = new TextDecoder()
+				let accumulatedText = ''
+
+				while (true) {
+					const { value, done } = await reader.read()
+					if (done) break
+					const chunk = decoder.decode(value, { stream: true })
+					// Some environments may send SSE-style lines; extract data if so, else use raw chunk
+					const maybeSse = chunk
+						.split('\n')
+						.filter((line) => line.startsWith('data:'))
+						.map((line) => line.replace(/^data:\s?/, ''))
+						.join('')
+					accumulatedText += maybeSse || chunk
+					updateNode<MessageNode>(editor, shape, (node) => ({
+						...node,
+						assistantMessage: accumulatedText,
+					}))
+				}
+			} catch (e) {
+				console.error(e)
+			}
+		})()
+	}, [editor, shape])
+
+	const handleMessageChange = useCallback(
+		(value: string) => {
+			updateNode<MessageNode>(editor, shape, (node) => ({
+				...node,
+				userMessage: value,
+			}))
+		},
+		[editor, shape]
+	)
+
+	return (
+		<div
+			style={{
+				pointerEvents: 'auto',
+				display: 'flex',
+				flexDirection: 'column',
+			}}
+		>
+			<div style={{ display: 'flex', flexDirection: 'row', alignItems: 'center' }}>
+				<div
+					style={{
+						height: '100%',
+						width: 32,
+						paddingLeft: 4,
+						display: 'flex',
+						justifyContent: 'center',
+						alignItems: 'center',
+						cursor: 'grab',
+					}}
+				>
+					<TldrawUiButtonIcon icon={<HandleIcon />} />
+				</div>
+				<div
+					style={{ padding: '4px 0px 0px 4px', flexGrow: 2 }}
+					onPointerDown={editor.markEventAsHandled}
+				>
+					<div style={{ padding: '0px 12px', borderRadius: 6, border: '1px solid #e2e8f0' }}>
+						<TldrawUiInput
+							value={node.userMessage}
+							onValueChange={handleMessageChange}
+							onComplete={handleSend}
+						/>
+					</div>
+				</div>
+				<div style={{ display: 'flex', justifyContent: 'flex-end', padding: '0px 0px' }}>
+					<TldrawUiButton
+						type="primary"
+						onClick={handleSend}
+						onPointerDown={editor.markEventAsHandled}
+					>
+						<TldrawUiButtonIcon icon={<SendIcon />} />
+					</TldrawUiButton>
+				</div>
+			</div>
+			{node.assistantMessage && (
+				<div style={{ padding: 4 }}>
+					<div
+						style={{
+							padding: 8,
+							lineHeight: '1.3',
+							fontSize: '12px',
+							borderRadius: 6,
+							border: '1px solid #e2e8f0',
+							fontWeight: '500',
+							fontFamily: 'Inter',
+							overflowWrap: 'normal',
+							whiteSpace: 'pre-wrap',
+						}}
+					>
+						{node.assistantMessage}
+					</div>
+				</div>
+			)}
+		</div>
+	)
+}

templates/branching-chat/client/nodes/types/shared.tsx

@@ -0,0 +1,73 @@
+import { Editor, T, TLUiIconJsx } from 'tldraw'
+import { NODE_HEIGHT_PX, NODE_WIDTH_PX } from '../../constants'
+import { PortId, ShapePort } from '../../ports/Port'
+import { NodeShape } from '../NodeShapeUtil'
+import { NodeType, NodeTypePorts } from '../nodeTypes'
+
+export interface NodeComponentProps<Node extends { type: string }> {
+	shape: NodeShape
+	node: Node
+}
+
+export abstract class NodeDefinition<Node extends { type: string }> {
+	constructor(public readonly editor: Editor) {
+		const ctor = this.constructor as NodeDefinitionConstructor<Node>
+		this.type = ctor.type
+		this.validator = ctor.validator
+	}
+
+	readonly type: Node['type']
+	readonly validator: T.Validator<Node>
+	abstract readonly title: string
+	abstract readonly heading?: string
+	abstract readonly icon: TLUiIconJsx
+
+	abstract getDefault(): Node
+	abstract getBodyWidthPx(shape: NodeShape, node: Node): number
+	abstract getBodyHeightPx(shape: NodeShape, node: Node): number
+	abstract getPorts(shape: NodeShape, node: Node): NodeTypePorts
+	onPortConnect(_shape: NodeShape, _node: Node, _port: PortId): void {}
+	onPortDisconnect(_shape: NodeShape, _node: Node, _port: PortId): void {}
+	abstract Component: React.ComponentType<NodeComponentProps<Node>>
+}
+
+export interface NodeDefinitionConstructor<Node extends { type: string }> {
+	new (editor: Editor): NodeDefinition<Node>
+	readonly type: Node['type']
+	readonly validator: T.Validator<Node>
+}
+
+/**
+ * The standard input port for a node.
+ */
+export const shapeInputPort: ShapePort = {
+	id: 'input',
+	terminal: 'end',
+	x: NODE_WIDTH_PX / 2,
+	y: 0,
+}
+
+/**
+ * The standard output port for a node.
+ */
+export const shapeOutputPort: ShapePort = {
+	id: 'output',
+	terminal: 'start',
+	x: NODE_WIDTH_PX / 2,
+	y: NODE_HEIGHT_PX,
+}
+
+/**
+ * Update the `node` prop within a node shape.
+ */
+export function updateNode<T extends NodeType>(
+	editor: Editor,
+	shape: NodeShape,
+	update: (node: T) => T
+) {
+	editor.updateShape({
+		id: shape.id,
+		type: shape.type,
+		props: { node: update(shape.props.node as T) },
+	})
+}