@dxos/edge-client 0.8.4-main.72ec0f3 → 0.8.4-main.74a063c4e0

package/src/edge-http-client.ts

@@ -4,11 +4,13 @@
 
 import * as FetchHttpClient from '@effect/platform/FetchHttpClient';
 import * as HttpClient from '@effect/platform/HttpClient';
+import { type Context as OtelContext, propagation } from '@opentelemetry/api';
 import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
 
 import { sleep } from '@dxos/async';
 import { Context } from '@dxos/context';
+import { runAndForwardErrors } from '@dxos/effect';
 import { invariant } from '@dxos/invariant';
 import { type PublicKey, type SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
@@ -17,12 +19,15 @@ import {
   type CreateAgentResponseBody,
   type CreateSpaceRequest,
   type CreateSpaceResponseBody,
+  EDGE_CLIENT_TAG_HEADER,
   EdgeAuthChallengeError,
   EdgeCallFailedError,
+  type EdgeFailure,
   type EdgeStatus,
   type ExecuteWorkflowResponseBody,
   type ExportBundleRequest,
   type ExportBundleResponse,
+  type FeedProtocol,
   type GetAgentStatusResponseBody,
   type GetNotarizationResponseBody,
   type ImportBundleRequest,
@@ -32,13 +37,16 @@ import {
   type JoinSpaceResponseBody,
   type ObjectId,
   type PostNotarizationRequestBody,
-  type QueryResult,
-  type QueueQuery,
   type RecoverIdentityRequest,
   type RecoverIdentityResponseBody,
   type UploadFunctionRequest,
   type UploadFunctionResponseBody,
 } from '@dxos/protocols';
+import {
+  type QueryRequest as QueryRequestProto,
+  type QueryResponse as QueryResponseProto,
+} from '@dxos/protocols/proto/dxos/echo/query';
+import { TRACE_PROCESSOR, TRACE_SPAN_ATTRIBUTE } from '@dxos/tracing';
 import { createUrl } from '@dxos/util';
 
 import { type EdgeIdentity, handleAuthChallenge } from './edge-identity';
@@ -67,7 +75,6 @@ export type RetryConfig = {
 
 type EdgeHttpRequestArgs = {
   method: string;
-  context?: Context;
   retry?: RetryConfig;
   body?: any;
   /**
@@ -75,12 +82,6 @@ type EdgeHttpRequestArgs = {
    */
   json?: boolean;
 
-  /**
-   * Do not expect a standard EDGE JSON response with a `success` field.
-   * @deprecated Use only for debugging.
-   */
-  rawResponse?: boolean;
-
   /**
    * Force authentication.
    * This should be used for requests with large bodies to avoid sending the body twice.
@@ -89,11 +90,23 @@ type EdgeHttpRequestArgs = {
   auth?: boolean;
 };
 
-export type EdgeHttpGetArgs = Pick<EdgeHttpRequestArgs, 'context' | 'retry' | 'auth'>;
-export type EdgeHttpPostArgs = Pick<EdgeHttpRequestArgs, 'context' | 'retry' | 'body' | 'auth'>;
+export type EdgeHttpCallArgs = Pick<EdgeHttpRequestArgs, 'retry' | 'auth'>;
+
+export type GetCronTriggersResponse = {
+  cronIds: string[];
+};
+
+export type EdgeHttpClientOptions = {
+  /**
+   * Tag included in the {@link EDGE_CLIENT_TAG_HEADER} header on every request.
+   * Used on Edge to classify traffic for metering (e.g. `ci-e2e`).
+   */
+  clientTag?: string;
+};
 
 export class EdgeHttpClient {
   private readonly _baseUrl: string;
+  private readonly _clientTag: string | undefined;
 
   private _edgeIdentity: EdgeIdentity | undefined;
 
@@ -102,8 +115,9 @@ export class EdgeHttpClient {
    */
   private _authHeader: string | undefined;
 
-  constructor(baseUrl: string) {
+  constructor(baseUrl: string, options?: EdgeHttpClientOptions) {
     this._baseUrl = getEdgeUrlWithProtocol(baseUrl, 'http');
+    this._clientTag = options?.clientTag;
     log('created', { url: this._baseUrl });
   }
 
@@ -122,23 +136,28 @@ export class EdgeHttpClient {
   // Status
   //
 
-  public async getStatus(args?: EdgeHttpGetArgs): Promise<EdgeStatus> {
-    return this._call(new URL('/status', this.baseUrl), { ...args, method: 'GET' });
+  public async getStatus(ctx: Context, args?: EdgeHttpCallArgs): Promise<EdgeStatus> {
+    return this._call(ctx, new URL('/status', this.baseUrl), { ...args, method: 'GET', auth: true });
   }
 
   //
   // Agents
   //
 
-  public createAgent(body: CreateAgentRequestBody, args?: EdgeHttpGetArgs): Promise<CreateAgentResponseBody> {
-    return this._call(new URL('/agents/create', this.baseUrl), { ...args, method: 'POST', body });
+  public createAgent(
+    ctx: Context,
+    body: CreateAgentRequestBody,
+    args?: EdgeHttpCallArgs,
+  ): Promise<CreateAgentResponseBody> {
+    return this._call(ctx, new URL('/agents/create', this.baseUrl), { ...args, method: 'POST', body });
   }
 
   public getAgentStatus(
+    ctx: Context,
     request: { ownerIdentityKey: PublicKey },
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<GetAgentStatusResponseBody> {
-    return this._call(new URL(`/users/${request.ownerIdentityKey.toHex()}/agent/status`, this.baseUrl), {
+    return this._call(ctx, new URL(`/users/${request.ownerIdentityKey.toHex()}/agent/status`, this.baseUrl), {
       ...args,
       method: 'GET',
     });
@@ -148,16 +167,21 @@ export class EdgeHttpClient {
   // Credentials
   //
 
-  public getCredentialsForNotarization(spaceId: SpaceId, args?: EdgeHttpGetArgs): Promise<GetNotarizationResponseBody> {
-    return this._call(new URL(`/spaces/${spaceId}/notarization`, this.baseUrl), { ...args, method: 'GET' });
+  public getCredentialsForNotarization(
+    ctx: Context,
+    spaceId: SpaceId,
+    args?: EdgeHttpCallArgs,
+  ): Promise<GetNotarizationResponseBody> {
+    return this._call(ctx, new URL(`/spaces/${spaceId}/notarization`, this.baseUrl), { ...args, method: 'GET' });
   }
 
   public async notarizeCredentials(
+    ctx: Context,
     spaceId: SpaceId,
     body: PostNotarizationRequestBody,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<void> {
-    await this._call(new URL(`/spaces/${spaceId}/notarization`, this.baseUrl), { ...args, body, method: 'POST' });
+    await this._call(ctx, new URL(`/spaces/${spaceId}/notarization`, this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
@@ -165,10 +189,11 @@ export class EdgeHttpClient {
   //
 
   public async recoverIdentity(
+    ctx: Context,
     body: RecoverIdentityRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<RecoverIdentityResponseBody> {
-    return this._call(new URL('/identity/recover', this.baseUrl), { ...args, body, method: 'POST' });
+    return this._call(ctx, new URL('/identity/recover', this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
@@ -176,11 +201,12 @@ export class EdgeHttpClient {
   //
 
   public async joinSpaceByInvitation(
+    ctx: Context,
     spaceId: SpaceId,
     body: JoinSpaceRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<JoinSpaceResponseBody> {
-    return this._call(new URL(`/spaces/${spaceId}/join`, this.baseUrl), { ...args, body, method: 'POST' });
+    return this._call(ctx, new URL(`/spaces/${spaceId}/join`, this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
@@ -188,18 +214,19 @@ export class EdgeHttpClient {
   //
 
   public async initiateOAuthFlow(
+    ctx: Context,
     body: InitiateOAuthFlowRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<InitiateOAuthFlowResponse> {
-    return this._call(new URL('/oauth/initiate', this.baseUrl), { ...args, body, method: 'POST' });
+    return this._call(ctx, new URL('/oauth/initiate', this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
   // Spaces
   //
 
-  async createSpace(body: CreateSpaceRequest, args?: EdgeHttpGetArgs): Promise<CreateSpaceResponseBody> {
-    return this._call(new URL('/spaces/create', this.baseUrl), { ...args, body, method: 'POST' });
+  async createSpace(ctx: Context, body: CreateSpaceRequest, args?: EdgeHttpCallArgs): Promise<CreateSpaceResponseBody> {
+    return this._call(ctx, new URL('/spaces/create', this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
@@ -207,13 +234,16 @@ export class EdgeHttpClient {
   //
 
   public async queryQueue(
+    ctx: Context,
     subspaceTag: string,
     spaceId: SpaceId,
-    query: QueueQuery,
-    args?: EdgeHttpGetArgs,
-  ): Promise<QueryResult> {
-    const { queueId } = query;
+    query: FeedProtocol.QueueQuery,
+    args?: EdgeHttpCallArgs,
+  ): Promise<FeedProtocol.QueryResult> {
+    const queueId = query.queueIds?.[0];
+    invariant(queueId, 'queueId required');
     return this._call(
+      ctx,
       createUrl(new URL(`/spaces/${subspaceTag}/${spaceId}/queue/${queueId}/query`, this.baseUrl), {
         after: query.after,
         before: query.before,
@@ -229,13 +259,14 @@ export class EdgeHttpClient {
   }
 
   public async insertIntoQueue(
+    ctx: Context,
     subspaceTag: string,
     spaceId: SpaceId,
     queueId: ObjectId,
     objects: unknown[],
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<void> {
-    return this._call(new URL(`/spaces/${subspaceTag}/${spaceId}/queue/${queueId}`, this.baseUrl), {
+    return this._call(ctx, new URL(`/spaces/${subspaceTag}/${spaceId}/queue/${queueId}`, this.baseUrl), {
       ...args,
       body: { objects },
       method: 'POST',
@@ -243,13 +274,15 @@ export class EdgeHttpClient {
   }
 
   public async deleteFromQueue(
+    ctx: Context,
     subspaceTag: string,
     spaceId: SpaceId,
     queueId: ObjectId,
     objectIds: ObjectId[],
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<void> {
     return this._call(
+      ctx,
       createUrl(new URL(`/spaces/${subspaceTag}/${spaceId}/queue/${queueId}`, this.baseUrl), {
         ids: objectIds.join(','),
       }),
@@ -265,9 +298,10 @@ export class EdgeHttpClient {
   //
 
   public async uploadFunction(
+    ctx: Context,
     pathParts: { functionId?: string },
     body: UploadFunctionRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<UploadFunctionResponseBody> {
     const formData = new FormData();
     formData.append('name', body.name ?? '');
@@ -284,7 +318,7 @@ export class EdgeHttpClient {
     }
 
     const path = ['functions', ...(pathParts.functionId ? [pathParts.functionId] : [])].join('/');
-    return this._call(new URL(path, this.baseUrl), {
+    return this._call(ctx, new URL(path, this.baseUrl), {
       ...args,
       body: formData,
       method: 'PUT',
@@ -292,11 +326,12 @@ export class EdgeHttpClient {
     });
   }
 
-  public async listFunctions(args?: EdgeHttpGetArgs): Promise<any> {
-    return this._call(new URL('/functions', this.baseUrl), { ...args, method: 'GET' });
+  public async listFunctions(ctx: Context, args?: EdgeHttpCallArgs): Promise<any> {
+    return this._call(ctx, new URL('/functions', this.baseUrl), { ...args, method: 'GET' });
   }
 
   public async invokeFunction(
+    ctx: Context,
     params: {
       functionId: string;
       version?: string;
@@ -305,7 +340,7 @@ export class EdgeHttpClient {
       subrequestsLimit?: number;
     },
     input: unknown,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<any> {
     const url = new URL(`/functions/${params.functionId}`, this.baseUrl);
     if (params.version) {
@@ -321,11 +356,10 @@ export class EdgeHttpClient {
       url.searchParams.set('subrequestsLimit', params.subrequestsLimit.toString());
     }
 
-    return this._call(url, {
+    return this._call(ctx, url, {
       ...args,
       body: input,
       method: 'POST',
-      rawResponse: true,
     });
   }
 
@@ -334,12 +368,13 @@ export class EdgeHttpClient {
   //
 
   public async executeWorkflow(
+    ctx: Context,
     spaceId: SpaceId,
     graphId: ObjectId,
     input: any,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<ExecuteWorkflowResponseBody> {
-    return this._call(new URL(`/workflows/${spaceId}/${graphId}`, this.baseUrl), {
+    return this._call(ctx, new URL(`/workflows/${spaceId}/${graphId}`, this.baseUrl), {
       ...args,
       body: input,
       method: 'POST',
@@ -350,8 +385,32 @@ export class EdgeHttpClient {
   // Triggers
   //
 
-  public async getCronTriggers(spaceId: SpaceId) {
-    return this._call(new URL(`/test/functions/${spaceId}/triggers/crons`, this.baseUrl), { method: 'GET' });
+  public async getCronTriggers(ctx: Context, spaceId: SpaceId): Promise<GetCronTriggersResponse> {
+    return this._call<GetCronTriggersResponse>(ctx, new URL(`/functions/${spaceId}/triggers/crons`, this.baseUrl), {
+      method: 'GET',
+    });
+  }
+
+  public async forceRunCronTrigger(ctx: Context, spaceId: SpaceId, triggerId: ObjectId) {
+    return this._call(ctx, new URL(`/functions/${spaceId}/triggers/crons/${triggerId}/run`, this.baseUrl), {
+      method: 'POST',
+    });
+  }
+
+  //
+  // Query
+  //
+
+  /**
+   * Execute a QueryAST query against a space.
+   */
+  public async execQuery(
+    ctx: Context,
+    spaceId: SpaceId,
+    body: QueryRequestProto,
+    args?: EdgeHttpCallArgs,
+  ): Promise<QueryResponseProto> {
+    return this._call(ctx, new URL(`/spaces/${spaceId}/exec-query`, this.baseUrl), { ...args, body, method: 'POST' });
   }
 
   //
@@ -359,19 +418,21 @@ export class EdgeHttpClient {
   //
 
   public async importBundle(
-    spaceId: SpaceId, //
+    ctx: Context,
+    spaceId: SpaceId,
     body: ImportBundleRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<void> {
-    return this._call(new URL(`/spaces/${spaceId}/import`, this.baseUrl), { ...args, body, method: 'PUT' });
+    return this._call(ctx, new URL(`/spaces/${spaceId}/import`, this.baseUrl), { ...args, body, method: 'PUT' });
   }
 
   public async exportBundle(
+    ctx: Context,
     spaceId: SpaceId,
     body: ExportBundleRequest,
-    args?: EdgeHttpGetArgs,
+    args?: EdgeHttpCallArgs,
   ): Promise<ExportBundleResponse> {
-    return this._call(new URL(`/spaces/${spaceId}/export`, this.baseUrl), {
+    return this._call(ctx, new URL(`/spaces/${spaceId}/export`, this.baseUrl), {
       ...args,
       body,
       method: 'POST',
@@ -390,16 +451,17 @@ export class EdgeHttpClient {
       Effect.provide(FetchHttpClient.layer),
       Effect.provide(HttpConfig.default),
       Effect.withSpan('EdgeHttpClient'),
-      Effect.runPromise,
+      runAndForwardErrors,
     ) as T;
   }
 
   // TODO(burdon): Refactor with effect (see edge-http-client.test.ts).
-  private async _call<T>(url: URL, args: EdgeHttpRequestArgs): Promise<T> {
+  private async _call<T>(ctx: Context, url: URL, args: EdgeHttpRequestArgs): Promise<T> {
     const shouldRetry = createRetryHandler(args);
-    const requestContext = args.context ?? Context.default();
     log('fetch', { url, request: args.body });
 
+    const traceHeaders = getTraceHeaders(ctx);
+
     let handledAuth = false;
     const tryCount = 1;
     while (true) {
@@ -412,15 +474,12 @@ export class EdgeHttpClient {
           }
         }
 
-        const request = createRequest(args, this._authHeader);
+        const request = createRequest(args, this._authHeader, traceHeaders, this._clientTag);
         log('call edge', { url, tryCount, authHeader: !!this._authHeader });
         const response = await fetch(url, request);
 
         if (response.ok) {
           const body = await response.clone().json();
-          if (args.rawResponse) {
-            return body as any;
-          }
           invariant(body, 'Expected body to be present');
           if (!('success' in body)) {
             return body;
@@ -434,13 +493,13 @@ export class EdgeHttpClient {
           continue;
         }
 
-        const body =
+        const body: EdgeFailure =
           response.headers.get('Content-Type') === 'application/json' ? await response.clone().json() : undefined;
 
         invariant(!body?.success, 'Expected body to not be a failure response or undefined.');
 
-        if (body?.errorData?.type === 'auth_challenge' && typeof body?.errorData?.challenge === 'string') {
-          processingError = new EdgeAuthChallengeError(body.errorData.challenge, body.errorData);
+        if (body?.data?.type === 'auth_challenge' && typeof body?.data?.challenge === 'string') {
+          processingError = new EdgeAuthChallengeError(body.data.challenge, body.data);
         } else if (body?.success === false) {
           processingError = EdgeCallFailedError.fromUnsuccessfulResponse(response, body);
         } else {
@@ -451,7 +510,7 @@ export class EdgeHttpClient {
         processingError = EdgeCallFailedError.fromProcessingFailureCause(error);
       }
 
-      if (processingError?.isRetryable && (await shouldRetry(requestContext, processingError.retryAfterMs))) {
+      if (processingError?.isRetryable && (await shouldRetry(ctx, processingError.retryAfterMs))) {
         log.verbose('retrying edge request', { url, processingError });
       } else {
         throw processingError!;
@@ -473,6 +532,8 @@ export class EdgeHttpClient {
 const createRequest = (
   { method, body, json = true }: EdgeHttpRequestArgs,
   authHeader: string | undefined,
+  traceHeaders?: Record<string, string>,
+  clientTag?: string,
 ): RequestInit => {
   let requestBody: BodyInit | undefined;
   const headers: HeadersInit = {};
@@ -492,6 +553,14 @@ const createRequest = (
     headers['Authorization'] = authHeader;
   }
 
+  if (traceHeaders) {
+    Object.assign(headers, traceHeaders);
+  }
+
+  if (clientTag) {
+    headers[EDGE_CLIENT_TAG_HEADER] = clientTag;
+  }
+
   return {
     method,
     body: requestBody,
@@ -499,6 +568,25 @@ const createRequest = (
   };
 };
 
+/**
+ * Extract W3C Trace Context headers (traceparent/tracestate) from a DXOS Context.
+ */
+const getTraceHeaders = (ctx: Context): Record<string, string> | undefined => {
+  const spanId = ctx.getAttribute(TRACE_SPAN_ATTRIBUTE);
+  const otlpContext =
+    typeof spanId === 'number'
+      ? (TRACE_PROCESSOR.remoteTracing.getSpanContext(spanId) as OtelContext | undefined)
+      : undefined;
+
+  if (!otlpContext) {
+    return undefined;
+  }
+
+  const headers: Record<string, string> = {};
+  propagation.inject(otlpContext, headers);
+  return Object.keys(headers).length > 0 ? headers : undefined;
+};
+
 /**
  * @deprecated
  */

package/src/edge-ws-connection.ts

@@ -50,7 +50,7 @@ export class EdgeWsConnection extends Resource {
 
   constructor(
     private readonly _identity: EdgeIdentity,
-    private readonly _connectionInfo: { url: URL; protocolHeader?: string },
+    private readonly _connectionInfo: { url: URL; protocolHeader?: string; headers?: Record<string, string> },
     private readonly _callbacks: EdgeWsConnectionCallbacks,
   ) {
     super();
@@ -121,6 +121,7 @@ export class EdgeWsConnection extends Resource {
       this._connectionInfo.protocolHeader
         ? [...baseProtocols, this._connectionInfo.protocolHeader]
         : [...baseProtocols],
+      this._connectionInfo.headers ? { headers: this._connectionInfo.headers } : undefined,
     );
     const muxer = new WebSocketMuxer(this._ws);
     this._wsMuxer = muxer;

package/src/http-client.test.ts

@@ -8,6 +8,7 @@ import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
 import { afterEach, beforeEach, describe, it } from 'vitest';
 
+import { runAndForwardErrors } from '@dxos/effect';
 import { invariant } from '@dxos/invariant';
 
 import { HttpConfig, withLogging, withRetry, withRetryConfig } from './http-client';
@@ -36,7 +37,7 @@ describe('HttpClient', () => {
         withRetry(HttpClient.get(server.url)),
         Effect.provide(FetchHttpClient.layer),
         Effect.withSpan('EdgeHttpClient'),
-        Effect.runPromise,
+        runAndForwardErrors,
       );
       expect(result).toMatchObject({ success: true, data: { value: 100 } });
     }
@@ -49,7 +50,7 @@ describe('HttpClient', () => {
         Effect.provide(FetchHttpClient.layer),
         Effect.provide(HttpConfig.default), // TODO(burdon): Swap out to mock.
         Effect.withSpan('EdgeHttpClient'), // TODO(burdon): OTEL.
-        Effect.runPromise,
+        runAndForwardErrors,
       );
       expect(result).toMatchObject({ success: true, data: { value: 100 } });
     }

package/src/testing/test-utils.ts

@@ -16,13 +16,13 @@ import { toUint8Array } from '../protocol';
 
 export const DEFAULT_PORT = 8080;
 
-type TestEdgeWsServerParams = {
+type TestEdgeWsServerProps = {
   admitConnection?: Trigger;
   payloadDecoder?: (payload: Uint8Array) => any;
   messageHandler?: (payload: any) => Promise<Uint8Array | undefined>;
 };
 
-export const createTestEdgeWsServer = async (port = DEFAULT_PORT, params?: TestEdgeWsServerParams) => {
+export const createTestEdgeWsServer = async (port = DEFAULT_PORT, params?: TestEdgeWsServerProps) => {
   const wsServer = new WebSocket.Server({
     port,
     verifyClient: createConnectionDelayHandler(params),
@@ -86,7 +86,7 @@ export const createTestEdgeWsServer = async (port = DEFAULT_PORT, params?: TestE
   };
 };
 
-const createConnectionDelayHandler = (params: TestEdgeWsServerParams | undefined) => {
+const createConnectionDelayHandler = (params: TestEdgeWsServerProps | undefined) => {
   return (_: any, callback: (admit: boolean) => void) => {
     if (params?.admitConnection) {
       log('delaying edge connection admission');
@@ -116,7 +116,7 @@ const createResponseSender = (connection: () => WebSocketMuxer) => {
   };
 };
 
-const decodePayload = async (request: Message, params: TestEdgeWsServerParams | undefined) => {
+const decodePayload = async (request: Message, params: TestEdgeWsServerProps | undefined) => {
   const requestPayload = params?.payloadDecoder
     ? params.payloadDecoder(request.payload!.value!)
     : protocol.getPayload(request, TextMessageSchema);