diff --git a/packages/sprotty-protocol/src/action-handler.ts b/packages/sprotty-protocol/src/action-handler.ts
index 733a349f..2e5180c5 100644
--- a/packages/sprotty-protocol/src/action-handler.ts
+++ b/packages/sprotty-protocol/src/action-handler.ts
@@ -18,7 +18,7 @@ import { Action } from './actions';
 import { DiagramServer } from './diagram-server';
 import { DiagramState } from './diagram-services';
 
-export type ServerActionHandler<A extends Action = Action> = (action: A, state: DiagramState, server: DiagramServer) => Promise<void>;
+export type ServerActionHandler<A extends Action = Action> = (action: A, state: DiagramState, server: DiagramServer) => void | Promise<void>;
 
 /**
  * Use this service to register handlers to specific actions. The `DiagramServer` queries this registry
diff --git a/packages/sprotty-protocol/src/diagram-server.ts b/packages/sprotty-protocol/src/diagram-server.ts
index a6cdcb57..945c7517 100644
--- a/packages/sprotty-protocol/src/diagram-server.ts
+++ b/packages/sprotty-protocol/src/diagram-server.ts
@@ -161,9 +161,9 @@ export class DiagramServer {
         // Find a matching action handler in the registry
         const handlers = this.actionHandlerRegistry?.getHandler(action.kind);
         if (handlers && handlers.length === 1) {
-            return handlers[0](action, this.state, this);
+            return handlers[0](action, this.state, this) ?? Promise.resolve();
         } else if (handlers && handlers.length > 1) {
-            return Promise.all(handlers.map(h => h(action, this.state, this))) as Promise<any>;
+            return Promise.all(handlers.map(h => h(action, this.state, this) ?? Promise.resolve())) as Promise<any>;
         }
         // If no handler is registered, call one of the default handling methods
         switch (action.kind) {
diff --git a/packages/sprotty-protocol/src/diagram-services.ts b/packages/sprotty-protocol/src/diagram-services.ts
index c16db1dc..837bc51e 100644
--- a/packages/sprotty-protocol/src/diagram-services.ts
+++ b/packages/sprotty-protocol/src/diagram-services.ts
@@ -25,9 +25,10 @@ export type DiagramOptions = JsonMap;
  * The current state captured by a `DiagramServer`.
  */
 export interface DiagramState {
-    options?: DiagramOptions;
-    currentRoot: SModelRoot;
-    revision: number;
+    options?: DiagramOptions
+    currentRoot: SModelRoot
+    revision: number
+    index?: SModelIndex
 }
 
 /**
diff --git a/packages/sprotty-protocol/src/index.ts b/packages/sprotty-protocol/src/index.ts
index 00b4b758..fc4e5cce 100644
--- a/packages/sprotty-protocol/src/index.ts
+++ b/packages/sprotty-protocol/src/index.ts
@@ -14,6 +14,7 @@
  * SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
  ********************************************************************************/
 
+export * from './action-handler';
 export * from './actions';
 export * from './diagram-server';
 export * from './diagram-services';
diff --git a/packages/sprotty-protocol/src/utils/model-utils.ts b/packages/sprotty-protocol/src/utils/model-utils.ts
index 271a61bb..401f2cd5 100644
--- a/packages/sprotty-protocol/src/utils/model-utils.ts
+++ b/packages/sprotty-protocol/src/utils/model-utils.ts
@@ -84,13 +84,15 @@ export function getSubType(schema: { type: string }): string {
  * `SModelIndex` might be more effective.
  */
 export function findElement(parent: SModelElement, elementId: string): SModelElement | undefined {
-    if (parent.id === elementId)
+    if (parent.id === elementId) {
         return parent;
-    if (parent.children !== undefined) {
+    }
+    if (parent.children) {
         for (const child of parent.children) {
             const result = findElement(child, elementId);
-            if (result !== undefined)
+            if (result !== undefined) {
                 return result;
+            }
         }
     }
     return undefined;
@@ -105,7 +107,7 @@ export class SModelIndex {
     private readonly id2element: Map<string, SModelElement> = new Map();
     private id2parent: Map<string, SModelElement> = new Map();
 
-    add(element: SModelElement): void {
+    add(element: SModelElement): this {
         if (!element.id) {
             throw new Error("Model element has no ID.");
         } else if (this.contains(element)) {
@@ -118,9 +120,10 @@ export class SModelIndex {
                 this.id2parent.set(child.id, element);
             }
         }
+        return this;
     }
 
-    remove(element: SModelElement): void {
+    remove(element: SModelElement): this {
         this.id2element.delete(element.id);
         if (Array.isArray(element.children)) {
             for (const child of element.children) {
@@ -128,6 +131,7 @@ export class SModelIndex {
                 this.remove(child as any);
             }
         }
+        return this;
     }
 
     contains(element: SModelElement): boolean {