feat: Runtime.path API (#1014)

.changeset/quick-masks-jump.md

@@ -0,0 +1,6 @@
+---
+"@assistant-ui/react-playground": patch
+"@assistant-ui/react": patch
+---
+
+feat: Runtime.path API

packages/react-playground/src/lib/playground-runtime.ts

@@ -522,6 +522,15 @@ export class PlaygroundThreadRuntimeCore implements INTERNAL.ThreadRuntimeCore {
   public export(): never {
     throw new Error("Playground does not support exporting messages.");
   }
+
+  public getMessageById(messageId: string) {
+    const idx = this.messages.findIndex((m) => m.id === messageId);
+    if (idx === -1) return undefined;
+    return {
+      message: this.messages[idx]!,
+      parentId: this.messages[idx - 1]?.id ?? null,
+    };
+  }
 }
 
 type PlaygroundThreadRuntime = ThreadRuntime & {
@@ -558,7 +567,7 @@ class PlaygroundRuntimeImpl
   public static override create(_core: PlaygroundRuntimeCore) {
     return new PlaygroundRuntimeImpl(
       _core,
-      AssistantRuntimeImpl.createThreadRuntime(
+      AssistantRuntimeImpl.createMainThreadRuntime(
         _core,
         PlaygroundThreadRuntimeImpl,
       ),

packages/react/src/api/AssistantRuntime.ts

@@ -28,7 +28,7 @@ export type AssistantRuntime = {
 };
 
 export class AssistantRuntimeImpl
-  implements AssistantRuntimeCore, AssistantRuntime
+  implements Omit<AssistantRuntimeCore, "thread">, AssistantRuntime
 {
   protected constructor(
     private readonly _core: AssistantRuntimeCore,
@@ -64,14 +64,18 @@ export class AssistantRuntimeImpl
     return this._core.subscribe(callback);
   }
 
-  protected static createThreadRuntime(
+  protected static createMainThreadRuntime(
     _core: AssistantRuntimeCore,
     CustomThreadRuntime: new (
       binding: ThreadRuntimeCoreBinding,
     ) => ThreadRuntime = ThreadRuntimeImpl,
   ) {
     return new CustomThreadRuntime(
       new NestedSubscriptionSubject({
+        path: {
+          ref: "threads.main",
+          threadSelector: { type: "main" },
+        },
         getState: () => _core.thread,
         subscribe: (callback) => _core.subscribe(callback),
       }),
@@ -86,7 +90,7 @@ export class AssistantRuntimeImpl
   ) {
     return new AssistantRuntimeImpl(
       _core,
-      AssistantRuntimeImpl.createThreadRuntime(_core, CustomThreadRuntime),
+      AssistantRuntimeImpl.createMainThreadRuntime(_core, CustomThreadRuntime),
     ) as AssistantRuntime;
   }
 }

packages/react/src/api/AttachmentRuntime.ts

@@ -7,6 +7,7 @@ import {
   PendingAttachment,
   Unsubscribe,
 } from "../types";
+import { AttachmentRuntimePath } from "./PathTypes";
 
 type MessageAttachmentState = CompleteAttachment & {
   source: "message";
@@ -38,13 +39,17 @@ export type AttachmentState =
   | MessageAttachmentState;
 
 type AttachmentSnapshotBinding<Source extends AttachmentRuntimeSource> =
-  SubscribableWithState<AttachmentState & { source: Source }>;
+  SubscribableWithState<
+    AttachmentState & { source: Source },
+    AttachmentRuntimePath & { attachmentSource: Source }
+  >;
 
 type AttachmentRuntimeSource = AttachmentState["source"];
 
 export type AttachmentRuntime<
   TSource extends AttachmentRuntimeSource = AttachmentRuntimeSource,
 > = {
+  path: AttachmentRuntimePath & { attachmentSource: TSource };
   readonly source: TSource;
   getState(): AttachmentState & { source: TSource };
   remove(): Promise<void>;
@@ -55,6 +60,10 @@ export abstract class AttachmentRuntimeImpl<
   Source extends AttachmentRuntimeSource = AttachmentRuntimeSource,
 > implements AttachmentRuntime
 {
+  public get path() {
+    return this._core.path;
+  }
+
   public abstract get source(): Source;
 
   constructor(private _core: AttachmentSnapshotBinding<Source>) {}

packages/react/src/api/ComposerRuntime.ts

@@ -14,13 +14,21 @@ import {
 } from "./AttachmentRuntime";
 import { ShallowMemoizeSubject } from "./subscribable/ShallowMemoizeSubject";
 import { SKIP_UPDATE } from "./subscribable/SKIP_UPDATE";
+import { ComposerRuntimePath } from "./PathTypes";
 
 export type ThreadComposerRuntimeCoreBinding = SubscribableWithState<
-  ThreadComposerRuntimeCore | undefined
+  ThreadComposerRuntimeCore | undefined,
+  ComposerRuntimePath & { composerSource: "thread" }
+>;
+
+export type EditComposerRuntimeCoreBinding = SubscribableWithState<
+  ComposerRuntimeCore | undefined,
+  ComposerRuntimePath & { composerSource: "edit" }
 >;
 
 export type ComposerRuntimeCoreBinding = SubscribableWithState<
-  ComposerRuntimeCore | undefined
+  ComposerRuntimeCore | undefined,
+  ComposerRuntimePath
 >;
 
 type LegacyEditComposerState = Readonly<{
@@ -165,6 +173,7 @@ const getEditComposerState = (
 };
 
 export type ComposerRuntime = {
+  path: ComposerRuntimePath;
   readonly type: "edit" | "thread";
   getState(): ComposerState;
 
@@ -207,6 +216,10 @@ export type ComposerRuntime = {
 export abstract class ComposerRuntimeImpl
   implements ComposerRuntimeCore, ComposerRuntime
 {
+  public get path() {
+    return this._core.path;
+  }
+
   public abstract get type(): "edit" | "thread";
 
   constructor(protected _core: ComposerRuntimeCoreBinding) {}
@@ -318,6 +331,7 @@ export type ThreadComposerRuntime = Omit<
   ComposerRuntime,
   "getState" | "getAttachmentByIndex"
 > & {
+  readonly path: ComposerRuntimePath & { composerSource: "thread" };
   readonly type: "thread";
   getState(): ThreadComposerState;
 
@@ -335,6 +349,12 @@ export class ThreadComposerRuntimeImpl
   extends ComposerRuntimeImpl
   implements ThreadComposerRuntime, ThreadComposerState
 {
+  public override get path() {
+    return this._core.path as ComposerRuntimePath & {
+      composerSource: "thread";
+    };
+  }
+
   public get type() {
     return "thread" as const;
   }
@@ -343,10 +363,12 @@ export class ThreadComposerRuntimeImpl
 
   constructor(core: ThreadComposerRuntimeCoreBinding) {
     const stateBinding = new LazyMemoizeSubject({
+      path: core.path,
       getState: () => getThreadComposerState(core.getState()),
       subscribe: (callback) => core.subscribe(callback),
     });
     super({
+      path: core.path,
       getState: () => core.getState(),
       subscribe: (callback) => stateBinding.subscribe(callback),
     });
@@ -364,6 +386,12 @@ export class ThreadComposerRuntimeImpl
   public getAttachmentByIndex(idx: number) {
     return new ThreadComposerAttachmentRuntimeImpl(
       new ShallowMemoizeSubject({
+        path: {
+          ...this.path,
+          attachmentSource: "thread-composer",
+          attachmentSelector: { type: "index", index: idx },
+          ref: this.path.ref + `${this.path.ref}.attachments[${idx}]`,
+        },
         getState: () => {
           const attachments = this.getState().attachments;
           const attachment = attachments[idx];
@@ -386,6 +414,7 @@ export type EditComposerRuntime = Omit<
   ComposerRuntime,
   "getState" | "getAttachmentByIndex"
 > & {
+  readonly path: ComposerRuntimePath & { composerSource: "edit" };
   readonly type: "edit";
 
   getState(): EditComposerState;
@@ -405,21 +434,27 @@ export class EditComposerRuntimeImpl
   extends ComposerRuntimeImpl
   implements EditComposerRuntime, EditComposerState
 {
+  public override get path() {
+    return this._core.path as ComposerRuntimePath & { composerSource: "edit" };
+  }
+
   public get type() {
     return "edit" as const;
   }
 
   private _getState;
   constructor(
-    core: ComposerRuntimeCoreBinding,
+    core: EditComposerRuntimeCoreBinding,
     private _beginEdit: () => void,
   ) {
     const stateBinding = new LazyMemoizeSubject({
+      path: core.path,
       getState: () => getEditComposerState(core.getState(), this._beginEdit),
       subscribe: (callback) => core.subscribe(callback),
     });
 
     super({
+      path: core.path,
       getState: () => core.getState(),
       subscribe: (callback) => stateBinding.subscribe(callback),
     });
@@ -445,6 +480,12 @@ export class EditComposerRuntimeImpl
   public getAttachmentByIndex(idx: number) {
     return new EditComposerAttachmentRuntimeImpl(
       new ShallowMemoizeSubject({
+        path: {
+          ...this.path,
+          attachmentSource: "edit-composer",
+          attachmentSelector: { type: "index", index: idx },
+          ref: this.path.ref + `${this.path.ref}.attachments[${idx}]`,
+        },
         getState: () => {
           const attachments = this.getState().attachments;
           const attachment = attachments[idx];

packages/react/src/api/ContentPartRuntime.ts

@@ -8,6 +8,7 @@ import { ThreadRuntimeCoreBinding } from "./ThreadRuntime";
 import { MessageStateBinding } from "./MessageRuntime";
 import { SubscribableWithState } from "./subscribable/Subscribable";
 import { Unsubscribe } from "../types";
+import { ContentPartRuntimePath } from "./PathTypes";
 
 export type ContentPartState = (
   | ThreadUserContentPart
@@ -20,15 +21,24 @@ export type ContentPartState = (
   status: ContentPartStatus | ToolCallContentPartStatus;
 };
 
-type ContentPartSnapshotBinding = SubscribableWithState<ContentPartState>;
+type ContentPartSnapshotBinding = SubscribableWithState<
+  ContentPartState,
+  ContentPartRuntimePath
+>;
 
 export type ContentPartRuntime = {
+  path: ContentPartRuntimePath;
+
   getState(): ContentPartState;
   addToolResult(result: any): void;
   subscribe(callback: () => void): Unsubscribe;
 };
 
 export class ContentPartRuntimeImpl implements ContentPartRuntime {
+  public get path() {
+    return this.contentBinding.path;
+  }
+
   constructor(
     private contentBinding: ContentPartSnapshotBinding,
     private messageApi: MessageStateBinding,

packages/react/src/api/MessageRuntime.ts

@@ -149,6 +149,10 @@ export type MessageRuntime = {
 };
 
 export class MessageRuntimeImpl implements MessageRuntime {
+  public get path() {
+    return this._core.path;
+  }
+
   constructor(
     private _core: MessageStateBinding,
     private _threadBinding: ThreadRuntimeCoreBinding,

packages/react/src/api/PathTypes.ts

@@ -0,0 +1,47 @@
+export type ThreadRuntimePath = {
+  ref: string;
+  threadSelector: { type: "main" };
+};
+
+export type MessageRuntimePath = ThreadRuntimePath & {
+  messageSelector:
+    | { type: "messageId"; messageId: string }
+    | { type: "index"; index: number };
+};
+
+export type ContentPartRuntimePath = MessageRuntimePath & {
+  contentPartSelector:
+    | { type: "index"; index: number }
+    | { type: "toolCallId"; toolCallId: string };
+};
+
+export type AttachmentRuntimePath = (
+  | (MessageRuntimePath & {
+      attachmentSource: "message" | "edit-composer";
+    })
+  | (ThreadRuntimePath & {
+      attachmentSource: "thread-composer";
+    })
+) & {
+  attachmentSelector:
+    | {
+        type: "index";
+        index: number;
+      }
+    | {
+        type: "index";
+        index: number;
+      }
+    | {
+        type: "index";
+        index: number;
+      };
+};
+
+export type ComposerRuntimePath =
+  | (ThreadRuntimePath & {
+      composerSource: "thread";
+    })
+  | (MessageRuntimePath & {
+      composerSource: "edit";
+    });