weflayr 0.11.0 → 0.13.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "weflayr",
-  "version": "0.11.0",
+  "version": "0.13.0",
   "description": "Weflayr Node.js SDK — instrument any LLM client via JS Proxy",
   "main": "src/index.js",
   "types": "index.d.ts",

package/src/instrument.js

@@ -94,7 +94,7 @@ function _wrapFn(fn, methodName, methodConfig) {
 
       if (methodConfig.streamPath && result && _isAsyncIterable(result[methodConfig.streamPath])) {
         _sendEvent('stream_start', { event_id: eventId, method: methodName, tags });
-        const wrappedStream = _wrapStream(result[methodConfig.streamPath], methodName, tags, requestArgs, settings, startTime, methodConfig, eventId);
+        const wrappedStream = _wrapStream(result[methodConfig.streamPath], methodName, tags, requestArgs, settings, startTime, methodConfig, eventId, result);
         return { ...result, [methodConfig.streamPath]: wrappedStream };
       }
 
@@ -128,9 +128,12 @@ function _wrapFn(fn, methodName, methodConfig) {
   };
 }
 
-async function* _wrapStream(stream, methodName, tags, requestArgs, settings, startTime, methodConfig, eventId) {
+// fullResult is only set for the streamPath case — it's the full SDK result object
+// (e.g. Mastra's StreamTextResult) so streamMiddleware factories can close over it
+// and await properties like result.usage after the stream drains.
+async function* _wrapStream(stream, methodName, tags, requestArgs, settings, startTime, methodConfig, eventId, fullResult) {
   let lastChunk = null;
-  const accumulator = methodConfig.streamMiddleware ? methodConfig.streamMiddleware() : null;
+  const accumulator = methodConfig.streamMiddleware ? methodConfig.streamMiddleware(fullResult) : null;
 
   try {
     for await (const chunk of stream) {
@@ -141,16 +144,17 @@ async function* _wrapStream(stream, methodName, tags, requestArgs, settings, sta
           method: methodName,
           tags,
           elapsed_ms: Date.now() - startTime,
-          ...accumulator.finalize(),
+          ...(await accumulator.finalize()),
         });
       }
       yield chunk;
     }
 
-    const middlewareData = accumulator
-      ? accumulator.finalize()
+    const rawMiddlewareData = accumulator
+      ? (await accumulator.finalize()) || {}
       : (methodConfig.middleware && lastChunk ? methodConfig.middleware(requestArgs, lastChunk) || {}
       : {});
+    const middlewareData = _applyContentPolicy(rawMiddlewareData, settings);
 
     _sendEvent('after', {
       event_id: eventId,

package/tests/index.test.js

@@ -26,6 +26,42 @@
       }
       return OpenAI;
     }
+    if (request === '@mastra/core') {
+      function Agent(opts) {
+        this._name = opts.name;
+      }
+      Agent.prototype.generate = async function(_messages, _options) {
+        return {
+          text:         'Hello from the Mastra fake agent.',
+          modelId:      'gpt-4o-mini',
+          finishReason: 'stop',
+          usage:    { promptTokens: 10, completionTokens: 15, totalTokens: 25 },
+          request:  { body: '{"messages":[{"role":"user","content":"Hello"}]}' },
+          response: {
+            id:       'fake-res',
+            modelId:  'gpt-4o-mini',
+            messages: [{ role: 'assistant', content: [{ type: 'text', text: 'Hello.' }] }],
+            body:     { choices: [{ message: { role: 'assistant', content: 'Hello.' } }] },
+          },
+          steps: [{ text: 'Hello.', request: { body: '...' }, finishReason: 'stop' }],
+        };
+      };
+      // stream() returns a StreamTextResult-like object (NOT directly iterable).
+      // usage and response are DelayedPromise wrappers matching real Mastra behaviour.
+      Agent.prototype.stream = function(_messages, _options) {
+        return {
+          textStream: (async function* () {
+            yield 'Hello'; yield ' Mastra'; yield ' stream.';
+          })(),
+          usage: { status: { type: 'resolved', value: { promptTokens: 10, completionTokens: 15, totalTokens: 25 } } },
+          response: { status: { type: 'resolved', value: { modelId: 'gpt-4o-mini', id: 'fake-stream-res' } } },
+        };
+      };
+      return { Agent };
+    }
+    if (request === '@ai-sdk/openai') {
+      return { openai: (modelId) => ({ _provider: 'openai', modelId }) };
+    }
     if (request === 'weflayr') {
       return _origLoad.call(this, path.resolve(__dirname, '../src/index.js'), parent, isMain);
     }
@@ -580,3 +616,191 @@ test('snippet propagate_runtime: propagated tag in before and after', async () =
   assert.equal(after.tags.customer_id, 'acme-corp');
 });
 
+// ---------------------------------------------------------------------------
+// Mastra integration — documentation snippets
+// @mastra/core and @ai-sdk/openai are faked by Module._load at top of file.
+// ---------------------------------------------------------------------------
+
+test('snippet mastra_setup: wraps Mastra Agent, sends before+after events on generate()', async () => {
+  freshSDK();
+  setSnippetEnv();
+
+  const sent = await captureHttp(async () => {
+    // SNIPPET_START: mastra_setup
+    const { weflayr_setup, weflayr_instrument } = require('weflayr');
+    const { Agent } = require('@mastra/core');
+    const { openai } = require('@ai-sdk/openai');
+
+    weflayr_setup({
+      intake_url:    process.env.WEFLAYR_INTAKE_URL,
+      client_id:     process.env.WEFLAYR_CLIENT_ID,
+      client_secret: process.env.WEFLAYR_CLIENT_SECRET,
+      methods: [{ call: 'generate' }],
+    });
+
+    const agent = weflayr_instrument(
+      new Agent({
+        name:         'my-agent',
+        model:        openai('gpt-4o-mini'),
+        instructions: 'You are a helpful assistant.',
+      })
+    );
+    // SNIPPET_END: mastra_setup
+
+    await agent.generate([{ role: 'user', content: 'Hello' }]);
+  });
+
+  assert.ok(sent.some(e => e.event_type === 'before'), 'before event sent');
+  assert.ok(sent.some(e => e.event_type === 'after'),  'after event sent');
+  assert.equal(sent[0].method, 'generate');
+});
+
+test('snippet mastra_tags: per-call tags stripped from generate args, appear in events', async () => {
+  const { weflayr_setup, weflayr_instrument } = freshSDK();
+  setSnippetEnv();
+
+  const { Agent } = require('@mastra/core');
+  const { openai } = require('@ai-sdk/openai');
+  weflayr_setup({ ...TEST_CREDS, methods: [{ call: 'generate' }] });
+  const agent = weflayr_instrument(new Agent({ name: 'test-agent', model: openai('gpt-4o-mini'), instructions: '...' }));
+
+  const sent = await captureHttp(async () => {
+    // SNIPPET_START: mastra_tags
+    await agent.generate(
+      [{ role: 'user', content: 'Hello' }],
+      { __weflayr_tags: { feature: 'chat', customer_id: '42', env: 'prod' } },
+    );
+    // SNIPPET_END: mastra_tags
+  });
+
+  const before = sent.find(e => e.event_type === 'before');
+  assert.ok(before, 'before event sent');
+  assert.deepEqual(before.tags, { feature: 'chat', customer_id: '42', env: 'prod' });
+});
+
+test('snippet mastra_content_privacy: text/request/steps/response content stripped from events', async () => {
+  freshSDK();
+  setSnippetEnv();
+
+  const sent = await captureHttp(async () => {
+    // SNIPPET_START: mastra_content_privacy
+    const { weflayr_setup, weflayr_instrument } = require('weflayr');
+    const { Agent } = require('@mastra/core');
+    const { openai } = require('@ai-sdk/openai');
+
+    weflayr_setup({
+      intake_url:    process.env.WEFLAYR_INTAKE_URL,
+      client_id:     process.env.WEFLAYR_CLIENT_ID,
+      client_secret: process.env.WEFLAYR_CLIENT_SECRET,
+      methods: [{ call: 'generate' }],
+      ignore_fields: (data) => {
+        delete data.text;
+        delete data.request;
+        delete data.steps;
+        if (data.response) {
+          delete data.response.body;
+          delete data.response.messages;
+        }
+        return data;
+      },
+    });
+
+    const agent = weflayr_instrument(
+      new Agent({
+        name:         'my-agent',
+        model:        openai('gpt-4o-mini'),
+        instructions: 'You are a helpful assistant.',
+      })
+    );
+    // SNIPPET_END: mastra_content_privacy
+
+    await agent.generate([{ role: 'user', content: 'Hello' }]);
+  });
+
+  const after = sent.find(e => e.event_type === 'after');
+  assert.ok(after, 'after event sent');
+  assert.ok(!after.response || !('text'    in after.response), 'text stripped');
+  assert.ok(!after.response || !('request' in after.response), 'request stripped');
+  assert.ok(!after.response || !('steps'   in after.response), 'steps stripped');
+  assert.ok(!after.response?.response?.body,     'response.body stripped');
+  assert.ok(!after.response?.response?.messages, 'response.messages stripped');
+});
+
+test('snippet mastra_stream: stream_start + after emitted, model/usage in after via streamMiddleware', async () => {
+  freshSDK();
+  setSnippetEnv();
+
+  const sent = await captureHttp(async () => {
+    // SNIPPET_START: mastra_stream
+    const { weflayr_setup, weflayr_instrument } = require('weflayr');
+    const { Agent } = require('@mastra/core');
+    const { openai } = require('@ai-sdk/openai');
+
+    weflayr_setup({
+      intake_url:    process.env.WEFLAYR_INTAKE_URL,
+      client_id:     process.env.WEFLAYR_CLIENT_ID,
+      client_secret: process.env.WEFLAYR_CLIENT_SECRET,
+      methods: [
+        { call: 'generate' },
+        {
+          call: 'stream',
+          streamPath: 'textStream',
+          streamMiddleware: (fullResult) => ({
+            onChunk: () => false,
+            async finalize() {
+              const rawUsage    = await fullResult.usage;
+              const rawResponse = await fullResult.response;
+              const response    = rawResponse?.status?.value ?? rawResponse;
+              return {
+                usage: rawUsage,
+                model: response?.modelId ?? null,
+              };
+            },
+          }),
+        },
+      ],
+      ignore_fields: (data) => {
+        delete data.text;
+        delete data.request;
+        delete data.steps;
+        if (data.response) {
+          delete data.response.body;
+          delete data.response.messages;
+        }
+        return data;
+      },
+    });
+
+    const agent = weflayr_instrument(
+      new Agent({
+        name:         'my-agent',
+        model:        openai('gpt-4o-mini'),
+        instructions: 'You are a helpful assistant.',
+      })
+    );
+
+    const streamResult = await agent.stream(
+      [{ role: 'user', content: 'Hello' }],
+      { __weflayr_tags: { feature: 'chat', customer_id: '42' } },
+    );
+
+    let text = '';
+    for await (const chunk of streamResult.textStream) {
+      text += chunk;
+    }
+    // SNIPPET_END: mastra_stream
+  });
+
+  const types = sent.map(e => e.event_type);
+  assert.ok(types.includes('before'),       'before event sent');
+  assert.ok(types.includes('stream_start'), 'stream_start event sent');
+  assert.ok(types.includes('after'),        'after event sent');
+
+  const after = sent.find(e => e.event_type === 'after');
+  assert.equal(after.method, 'stream');
+  assert.equal(after.model, 'gpt-4o-mini');
+  assert.ok(after.usage, 'usage present in after event');
+  assert.equal(after.tags.feature, 'chat');
+  assert.equal(after.tags.customer_id, '42');
+});
+