@dxos/edge-client 0.8.4-main.fffef41 → 0.8.4-staging.60fe92afc8

package/src/edge-client.ts

@@ -10,6 +10,7 @@ import {
   scheduleMicroTask,
   scheduleTaskInterval,
 } from '@dxos/async';
+import { Context, TRACE_SPAN_ATTRIBUTE, type TraceContextData } from '@dxos/context';
 import { type Lifecycle, Resource } from '@dxos/context';
 import { log, logInfo } from '@dxos/log';
 import { type Message } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
@@ -35,6 +36,8 @@ export type MessengerConfig = {
   timeout?: number;
   protocol?: Protocol;
   disableAuth?: boolean;
+  /** Sent as `X-DXOS-Client-Tag` on the WebSocket upgrade (Node/`ws` only; ignored in browsers). */
+  clientTag?: string;
 };
 
 export interface EdgeConnection extends Required<Lifecycle> {
@@ -45,7 +48,7 @@ export interface EdgeConnection extends Required<Lifecycle> {
   get isOpen(): boolean;
   get status(): EdgeStatus;
   setIdentity(identity: EdgeIdentity): void;
-  send(message: Message): Promise<void>;
+  send(ctx: Context, message: Message): Promise<void>;
   onMessage(listener: MessageListener): () => void;
   onReconnected(listener: ReconnectListener): () => void;
 }
@@ -126,7 +129,7 @@ export class EdgeClient extends Resource implements EdgeConnection {
    * Send message.
    * NOTE: The message is guaranteed to be delivered but the service must respond with a message to confirm processing.
    */
-  public async send(message: Message) {
+  public async send(ctx: Context, message: Message) {
     if (this._ready.state !== TriggerState.RESOLVED) {
       log('waiting for websocket');
       await this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT });
@@ -143,6 +146,15 @@ export class EdgeClient extends Resource implements EdgeConnection {
       throw new EdgeIdentityChangedError();
     }
 
+    const traceCtx = ctx.getAttribute(TRACE_SPAN_ATTRIBUTE) as TraceContextData | undefined;
+    if (traceCtx) {
+      message.traceContext = {
+        $typeName: 'dxos.edge.messenger.TraceContext',
+        traceparent: traceCtx.traceparent,
+        tracestate: traceCtx.tracestate,
+      };
+    }
+
     this._currentConnection.send(message);
   }
 
@@ -151,7 +163,7 @@ export class EdgeClient extends Resource implements EdgeConnection {
     return () => this._messageListeners.delete(listener);
   }
 
-  public onReconnected(listener: () => void) {
+  public onReconnected(listener: ReconnectListener) {
     this._reconnectListeners.add(listener);
     if (this._ready.state === TriggerState.RESOLVED) {
       // Microtask so that listener is always called asynchronously, no matter the state of the ready trigger
@@ -219,7 +231,11 @@ export class EdgeClient extends Resource implements EdgeConnection {
     log('Opening websocket', { url: url.toString(), protocolHeader });
     const connection = new EdgeWsConnection(
       identity,
-      { url, protocolHeader },
+      {
+        url,
+        protocolHeader,
+        headers: this._config.clientTag ? { 'X-DXOS-Client-Tag': this._config.clientTag } : undefined,
+      },
       {
         onConnected: () => {
           if (this._isActive(connection)) {
@@ -253,9 +269,23 @@ export class EdgeClient extends Resource implements EdgeConnection {
     this._currentConnection = connection;
 
     await connection.open();
-    // Race with restartRequired so that restart is not blocked by _connect execution.
-    // Wait on ready to attempt a reconnect if it times out.
-    await Promise.race([this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT }), restartRequired]);
+
+    // The connection is only a successful start once the socket becomes ready. A socket that
+    // closes or errors before becoming ready (or never connects within the timeout) is a failed
+    // start: throwing lets PersistentLifecycle apply its backoff. Returning here would mark the
+    // attempt as successful and reset the backoff, degenerating reconnects into a hot loop when
+    // the server accepts then immediately drops the socket.
+    const becameReady = await Promise.race([
+      this._ready.wait({ timeout: this._config.timeout ?? DEFAULT_TIMEOUT }).then(
+        () => true,
+        () => false,
+      ),
+      restartRequired.wait().then(() => false),
+    ]);
+    if (!becameReady) {
+      throw new EdgeConnectionClosedError();
+    }
+
     return connection;
   }
 

package/src/edge-http-client.test.ts

@@ -2,7 +2,7 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, it } from 'vitest';
+import { afterEach, describe, it, test, vi } from 'vitest';
 
 import { createEphemeralEdgeIdentity } from './auth';
 import { EdgeHttpClient } from './edge-http-client';
@@ -16,7 +16,41 @@ describe.skipIf(process.env.CI)('EdgeHttpClient', () => {
     const identity = await createEphemeralEdgeIdentity();
     client.setIdentity(identity);
 
-    const result = await client.getStatus();
+    const { Context } = await import('@dxos/context');
+    const result = await client.getStatus(Context.default());
     expect(result).toBeDefined();
   });
 });
+
+describe('EdgeHttpClient.anthropicAiRequest', () => {
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  test('re-bases the request path onto the EDGE /ai/generate/anthropic route', async ({ expect }) => {
+    const fetchMock = vi.fn(async (input: any, _init?: RequestInit) => {
+      const url = String(input instanceof URL ? input : (input.url ?? input));
+      // `/auth` preflight: respond non-401 so no auth header is attached.
+      if (url.endsWith('/auth')) {
+        return new Response(null, { status: 200 });
+      }
+      return new Response(JSON.stringify({ ok: true }), { status: 200 });
+    });
+    vi.stubGlobal('fetch', fetchMock);
+
+    const client = new EdgeHttpClient('https://edge.example.com');
+    const response = await client.anthropicAiRequest(
+      new Request('http://edge/v1/messages?beta=true', {
+        method: 'POST',
+        body: JSON.stringify({ model: 'claude' }),
+      }),
+    );
+
+    expect(response.status).toBe(200);
+
+    const targetCall = fetchMock.mock.calls.find((call) => !String(call[0]).endsWith('/auth'));
+    expect(targetCall).toBeDefined();
+    expect(String(targetCall![0])).toBe('https://edge.example.com/ai/generate/anthropic/v1/messages?beta=true');
+    expect(targetCall![1]?.method).toBe('POST');
+  });
+});