diff --git a/package-lock.json b/package-lock.json
index 51368e807..e4273cd1a 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -179,10 +179,11 @@
       }
     },
     "@azure/cosmos": {
-      "version": "3.16.2",
-      "resolved": "https://registry.npmjs.org/@azure/cosmos/-/cosmos-3.16.2.tgz",
-      "integrity": "sha512-sceY5LWj0BHGj8PSyaVCfDRQLVZyoCfIY78kyIROJVEw0k+p9XFs8fhpykN8JklkCftL0WlaVY+X25SQwnhZsw==",
+      "version": "4.0.0",
+      "resolved": "https://msazure.pkgs.visualstudio.com/_packaging/AzurePortal/npm/registry/@azure/cosmos/-/cosmos-4.0.0.tgz",
+      "integrity": "sha1-X9qLNctiu82lIVm5bEw5gahD1bk=",
       "requires": {
+        "@azure/abort-controller": "^1.0.0",
         "@azure/core-auth": "^1.3.0",
         "@azure/core-rest-pipeline": "^1.2.0",
         "debug": "^4.1.1",
diff --git a/package.json b/package.json
index 302248f9c..8fe8ad2ff 100644
--- a/package.json
+++ b/package.json
@@ -5,7 +5,7 @@
   "main": "index.js",
   "dependencies": {
     "@azure/arm-cosmosdb": "9.1.0",
-    "@azure/cosmos": "3.16.2",
+    "@azure/cosmos": "4.0.0",
     "@azure/cosmos-language-service": "0.0.5",
     "@azure/identity": "1.2.1",
     "@azure/ms-rest-nodeauth": "3.0.7",
diff --git a/src/Common/CosmosClient.test.ts b/src/Common/CosmosClient.test.ts
index 204331d87..c80b810ae 100644
--- a/src/Common/CosmosClient.test.ts
+++ b/src/Common/CosmosClient.test.ts
@@ -125,7 +125,7 @@ describe("requestPlugin", () => {
       const headers = {};
       const endpoint = "https://docs.azure.com";
       const path = "/dbs/foo";
-      requestPlugin({ endpoint, headers, path } as any, next as any);
+      requestPlugin({ endpoint, headers, path } as any, undefined, next as any);
       expect(next.mock.calls[0][0]).toMatchSnapshot();
     });
   });
@@ -137,7 +137,7 @@ describe("requestPlugin", () => {
       const headers = {};
       const endpoint = "";
       const path = "/dbs/foo";
-      requestPlugin({ endpoint, headers, path } as any, next as any);
+      requestPlugin({ endpoint, headers, path } as any, undefined, next as any);
       expect(next.mock.calls[0][0]).toMatchSnapshot();
     });
   });
diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts
index bd4b76d0d..b8a9573a9 100644
--- a/src/Common/CosmosClient.ts
+++ b/src/Common/CosmosClient.ts
@@ -51,7 +51,7 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
   return decodeURIComponent(result.PrimaryReadWriteToken);
 };
 
-export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, next) => {
+export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, diagnosticNode, next) => {
   requestContext.endpoint = new URL(configContext.PROXY_PATH, window.location.href).href;
   requestContext.headers["x-ms-proxy-target"] = endpoint();
   return next(requestContext);
diff --git a/src/Explorer/Tabs/DocumentsTab.ts b/src/Explorer/Tabs/DocumentsTab.ts
index 064220511..8d95bbf18 100644
--- a/src/Explorer/Tabs/DocumentsTab.ts
+++ b/src/Explorer/Tabs/DocumentsTab.ts
@@ -1,4 +1,4 @@
-import { extractPartitionKey, ItemDefinition, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos";
+import { ItemDefinition, PartitionKey, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import { querySampleDocuments, readSampleDocument } from "Explorer/QueryCopilot/QueryCopilotUtilities";
 import * as ko from "knockout";
 import Q from "q";
@@ -6,10 +6,10 @@ import { format } from "react-string-format";
 import { QueryConstants } from "Shared/Constants";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import DeleteDocumentIcon from "../../../images/DeleteDocument.svg";
-import DiscardIcon from "../../../images/discard.svg";
 import NewDocumentIcon from "../../../images/NewDocument.svg";
-import SaveIcon from "../../../images/save-cosmos.svg";
 import UploadIcon from "../../../images/Upload_16x16.svg";
+import DiscardIcon from "../../../images/discard.svg";
+import SaveIcon from "../../../images/save-cosmos.svg";
 import * as Constants from "../../Common/Constants";
 import {
   DocumentsGridMetrics,
@@ -17,15 +17,15 @@ import {
   QueryCopilotSampleContainerId,
   QueryCopilotSampleDatabaseId,
 } from "../../Common/Constants";
+import editable from "../../Common/EditableUtility";
+import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
+import * as HeadersUtility from "../../Common/HeadersUtility";
+import { Splitter, SplitterBounds, SplitterDirection } from "../../Common/Splitter";
 import { createDocument } from "../../Common/dataAccess/createDocument";
 import { deleteDocument } from "../../Common/dataAccess/deleteDocument";
 import { queryDocuments } from "../../Common/dataAccess/queryDocuments";
 import { readDocument } from "../../Common/dataAccess/readDocument";
 import { updateDocument } from "../../Common/dataAccess/updateDocument";
-import editable from "../../Common/EditableUtility";
-import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
-import * as HeadersUtility from "../../Common/HeadersUtility";
-import { Splitter, SplitterBounds, SplitterDirection } from "../../Common/Splitter";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { Action } from "../../Shared/Telemetry/TelemetryConstants";
@@ -33,6 +33,7 @@ import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
 import { userContext } from "../../UserContext";
 import { logConsoleError } from "../../Utils/NotificationConsoleUtils";
 import * as QueryUtils from "../../Utils/QueryUtils";
+import { extractPartitionKeyValues } from "../../Utils/QueryUtils";
 import { CommandButtonComponentProps } from "../Controls/CommandButton/CommandButtonComponent";
 import { useDialog } from "../Controls/Dialog";
 import Explorer from "../Explorer";
@@ -479,7 +480,7 @@ export default class DocumentsTab extends TabsBase {
           const value: string = this.renderObjectForEditor(savedDocument || {}, null, 4);
           this.selectedDocumentContent.setBaseline(value);
           this.initialDocumentContent(value);
-          const partitionKeyValueArray = extractPartitionKey(
+          const partitionKeyValueArray: PartitionKey[] = extractPartitionKeyValues(
             savedDocument,
             this.partitionKey as PartitionKeyDefinition,
           );
@@ -530,7 +531,10 @@ export default class DocumentsTab extends TabsBase {
     const selectedDocumentId = this.selectedDocumentId();
     const documentContent = JSON.parse(this.selectedDocumentContent());
 
-    const partitionKeyValueArray = extractPartitionKey(documentContent, this.partitionKey as PartitionKeyDefinition);
+    const partitionKeyValueArray: PartitionKey[] = extractPartitionKeyValues(
+      documentContent,
+      this.partitionKey as PartitionKeyDefinition,
+    );
     selectedDocumentId.partitionKeyValue = partitionKeyValueArray;
 
     this.isExecutionError(false);
diff --git a/src/Explorer/Tabs/MongoDocumentsTab.ts b/src/Explorer/Tabs/MongoDocumentsTab.ts
index ba3a145a7..d3bf31a54 100644
--- a/src/Explorer/Tabs/MongoDocumentsTab.ts
+++ b/src/Explorer/Tabs/MongoDocumentsTab.ts
@@ -1,4 +1,5 @@
-import { extractPartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
+import { PartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
+import { extractPartitionKeyValues } from "Utils/QueryUtils";
 import * as ko from "knockout";
 import Q from "q";
 import * as Constants from "../../Common/Constants";
@@ -88,7 +89,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
     )
       .then(
         (savedDocument: any) => {
-          let partitionKeyArray = extractPartitionKey(
+          const partitionKeyArray: PartitionKey[] = extractPartitionKeyValues(
             savedDocument,
             this._getPartitionKeyDefinition() as PartitionKeyDefinition,
           );
@@ -150,7 +151,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
 
           this.documentIds().forEach((documentId: DocumentId) => {
             if (documentId.rid === updatedDocument._rid) {
-              const partitionKeyArray = extractPartitionKey(
+              const partitionKeyArray: PartitionKey[] = extractPartitionKeyValues(
                 updatedDocument,
                 this._getPartitionKeyDefinition() as PartitionKeyDefinition,
               );
@@ -289,7 +290,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
   }
 
   private _hasShardKeySpecified(document: any): boolean {
-    return Boolean(extractPartitionKey(document, this._getPartitionKeyDefinition() as PartitionKeyDefinition));
+    return Boolean(extractPartitionKeyValues(document, this._getPartitionKeyDefinition() as PartitionKeyDefinition));
   }
 
   private _getPartitionKeyDefinition(): DataModels.PartitionKey {
diff --git a/src/Explorer/Tree/ConflictId.ts b/src/Explorer/Tree/ConflictId.ts
index 4225d80be..0e54b6a37 100644
--- a/src/Explorer/Tree/ConflictId.ts
+++ b/src/Explorer/Tree/ConflictId.ts
@@ -1,4 +1,4 @@
-import { extractPartitionKey } from "@azure/cosmos";
+import { extractPartitionKeyValues } from "Utils/QueryUtils";
 import * as ko from "knockout";
 import * as Constants from "../../Common/Constants";
 import { readDocument } from "../../Common/dataAccess/readDocument";
@@ -42,7 +42,7 @@ export default class ConflictId {
     }
     this.partitionKeyProperty = container && container.partitionKeyProperty;
     this.partitionKey = container && container.partitionKey;
-    this.partitionKeyValue = extractPartitionKey(this.parsedContent, this.partitionKey as any);
+    this.partitionKeyValue = extractPartitionKeyValues(this.parsedContent, this.partitionKey as any);
     this.stringPartitionKeyValue = this.getPartitionKeyValueAsString();
     this.id = ko.observable(data.id);
     this.isDirty = ko.observable(false);
diff --git a/src/Utils/PriorityBasedExecutionUtils.ts b/src/Utils/PriorityBasedExecutionUtils.ts
index 1295f4e61..efb80386d 100644
--- a/src/Utils/PriorityBasedExecutionUtils.ts
+++ b/src/Utils/PriorityBasedExecutionUtils.ts
@@ -29,7 +29,7 @@ export function getPriorityLevel(): PriorityLevel {
   }
 }
 
-export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, next) => {
+export const requestPlugin: Cosmos.Plugin<any> = async (requestContext, undefined, next) => {
   if (isRelevantRequest(requestContext)) {
     const priorityLevel: PriorityLevel = getPriorityLevel();
     requestContext.headers["x-ms-cosmos-priority-level"] = priorityLevel as string;
diff --git a/src/Utils/QueryUtils.test.ts b/src/Utils/QueryUtils.test.ts
index 72a568a3a..f59ee4189 100644
--- a/src/Utils/QueryUtils.test.ts
+++ b/src/Utils/QueryUtils.test.ts
@@ -1,8 +1,10 @@
+import { PartitionKey, PartitionKeyDefinition, PartitionKeyKind } from "@azure/cosmos";
 import * as Q from "q";
 import * as sinon from "sinon";
 import * as DataModels from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 import * as QueryUtils from "./QueryUtils";
+import { extractPartitionKeyValues } from "./QueryUtils";
 
 describe("Query Utils", () => {
   const generatePartitionKeyForPath = (path: string): DataModels.PartitionKey => {
@@ -94,4 +96,69 @@ describe("Query Utils", () => {
       expect(queryStub.getCall(0).args[0]).toBe(0);
     });
   });
+
+  describe("extractPartitionKey", () => {
+    const documentContent = {
+      "Volcano Name": "Adams",
+      Country: "United States",
+      Region: "US-Washington",
+      Location: {
+        type: "Point",
+        coordinates: [-121.49, 46.206],
+      },
+      Elevation: 3742,
+      Type: "Stratovolcano",
+      Status: "Tephrochronology",
+      "Last Known Eruption": "Last known eruption from A.D. 1-1499, inclusive",
+      id: "9e3c494e-8367-3f50-1f56-8c6fcb961363",
+      _rid: "xzo0AJRYUxUFAAAAAAAAAA==",
+      _self: "dbs/xzo0AA==/colls/xzo0AJRYUxU=/docs/xzo0AJRYUxUFAAAAAAAAAA==/",
+      _etag: '"ce00fa43-0000-0100-0000-652840440000"',
+      _attachments: "attachments/",
+      _ts: 1697136708,
+    };
+
+    it("should extract single partition key value", () => {
+      const singlePartitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.Hash,
+        paths: ["/Elevation"],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
+        documentContent,
+        singlePartitionKeyDefinition,
+      );
+      expect(partitionKeyValues.length).toBe(1);
+      expect(partitionKeyValues[0]).toEqual(3742);
+    });
+
+    it("should extract two partition key values", () => {
+      const multiPartitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.MultiHash,
+        paths: ["/Type", "/Status"],
+      };
+      const expectedPartitionKeyValues: string[] = ["Stratovolcano", "Tephrochronology"];
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
+        documentContent,
+        multiPartitionKeyDefinition,
+      );
+      expect(partitionKeyValues.length).toBe(2);
+      expect(expectedPartitionKeyValues).toContain(documentContent["Type"]);
+      expect(expectedPartitionKeyValues).toContain(documentContent["Status"]);
+    });
+
+    it("should extract no partition key values", () => {
+      const singlePartitionKeyDefinition: PartitionKeyDefinition = {
+        kind: PartitionKeyKind.Hash,
+        paths: ["/InvalidPartitionKeyPath"],
+      };
+
+      const partitionKeyValues: PartitionKey[] = extractPartitionKeyValues(
+        documentContent,
+        singlePartitionKeyDefinition,
+      );
+
+      expect(partitionKeyValues.length).toBe(0);
+    });
+  });
 });
diff --git a/src/Utils/QueryUtils.ts b/src/Utils/QueryUtils.ts
index d56ec73eb..5440c2dda 100644
--- a/src/Utils/QueryUtils.ts
+++ b/src/Utils/QueryUtils.ts
@@ -1,3 +1,4 @@
+import { PartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
 import * as DataModels from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 
@@ -82,3 +83,22 @@ export const queryPagesUntilContentPresent = async (
 
   return await doRequest(firstItemIndex);
 };
+
+/* eslint-disable  @typescript-eslint/no-explicit-any */
+export const extractPartitionKeyValues = (
+  documentContent: any,
+  partitionKeyDefinition: PartitionKeyDefinition,
+): PartitionKey[] => {
+  if (!partitionKeyDefinition.paths || partitionKeyDefinition.paths.length === 0) {
+    return undefined;
+  }
+
+  const partitionKeyValues: PartitionKey[] = [];
+  partitionKeyDefinition.paths.forEach((partitionKeyPath: string) => {
+    const partitionKeyPathWithoutSlash: string = partitionKeyPath.substring(1);
+    if (documentContent[partitionKeyPathWithoutSlash]) {
+      partitionKeyValues.push(documentContent[partitionKeyPathWithoutSlash]);
+    }
+  });
+  return partitionKeyValues;
+};