@armoriq/sdk-dev 0.6.8 → 0.7.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (99) hide show
  1. package/README.md +168 -1
  2. package/dist/_version.d.ts +1 -1
  3. package/dist/_version.js +1 -1
  4. package/dist/cli/commands/auth.d.ts +12 -0
  5. package/dist/cli/commands/auth.d.ts.map +1 -1
  6. package/dist/cli/commands/auth.js +460 -42
  7. package/dist/cli/commands/auth.js.map +1 -1
  8. package/dist/cli/index.js +0 -0
  9. package/dist/client.d.ts +8 -15
  10. package/dist/client.d.ts.map +1 -1
  11. package/dist/client.js +20 -18
  12. package/dist/client.js.map +1 -1
  13. package/dist/config.d.ts +0 -17
  14. package/dist/config.d.ts.map +1 -1
  15. package/dist/config.js +1 -19
  16. package/dist/config.js.map +1 -1
  17. package/dist/index.d.ts +3 -2
  18. package/dist/index.d.ts.map +1 -1
  19. package/dist/index.js +8 -14
  20. package/dist/index.js.map +1 -1
  21. package/dist/integrations/google_adk.d.ts +168 -6
  22. package/dist/integrations/google_adk.d.ts.map +1 -1
  23. package/dist/integrations/google_adk.js +797 -103
  24. package/dist/integrations/google_adk.js.map +1 -1
  25. package/dist/integrations/langchain.d.ts +51 -1
  26. package/dist/integrations/langchain.d.ts.map +1 -1
  27. package/dist/integrations/langchain.js +563 -30
  28. package/dist/integrations/langchain.js.map +1 -1
  29. package/dist/integrations/strands.d.ts +48 -0
  30. package/dist/integrations/strands.d.ts.map +1 -1
  31. package/dist/integrations/strands.js +441 -27
  32. package/dist/integrations/strands.js.map +1 -1
  33. package/dist/models.d.ts +2 -2
  34. package/dist/models.d.ts.map +1 -1
  35. package/dist/observability/content-capture.d.ts +75 -0
  36. package/dist/observability/content-capture.d.ts.map +1 -0
  37. package/dist/observability/content-capture.js +339 -0
  38. package/dist/observability/content-capture.js.map +1 -0
  39. package/dist/observability/index.d.ts +6 -7
  40. package/dist/observability/index.d.ts.map +1 -1
  41. package/dist/observability/index.js +18 -27
  42. package/dist/observability/index.js.map +1 -1
  43. package/dist/observability/otel-config.d.ts +47 -0
  44. package/dist/observability/otel-config.d.ts.map +1 -0
  45. package/dist/observability/otel-config.js +268 -0
  46. package/dist/observability/otel-config.js.map +1 -0
  47. package/dist/observability/otel-export-ceiling.d.ts +96 -0
  48. package/dist/observability/otel-export-ceiling.d.ts.map +1 -0
  49. package/dist/observability/otel-export-ceiling.js +264 -0
  50. package/dist/observability/otel-export-ceiling.js.map +1 -0
  51. package/dist/observability/otel-runtime.d.ts +103 -0
  52. package/dist/observability/otel-runtime.d.ts.map +1 -0
  53. package/dist/observability/otel-runtime.js +668 -0
  54. package/dist/observability/otel-runtime.js.map +1 -0
  55. package/dist/observability/otel-session.d.ts +168 -0
  56. package/dist/observability/otel-session.d.ts.map +1 -0
  57. package/dist/observability/otel-session.js +621 -0
  58. package/dist/observability/otel-session.js.map +1 -0
  59. package/dist/observability/otel-shutdown.d.ts +17 -0
  60. package/dist/observability/otel-shutdown.d.ts.map +1 -0
  61. package/dist/observability/otel-shutdown.js +54 -0
  62. package/dist/observability/otel-shutdown.js.map +1 -0
  63. package/dist/observability/policy-lease.d.ts +22 -0
  64. package/dist/observability/policy-lease.d.ts.map +1 -0
  65. package/dist/observability/policy-lease.js +102 -0
  66. package/dist/observability/policy-lease.js.map +1 -0
  67. package/dist/plan_builder.d.ts +5 -4
  68. package/dist/plan_builder.d.ts.map +1 -1
  69. package/dist/plan_builder.js +14 -15
  70. package/dist/plan_builder.js.map +1 -1
  71. package/dist/session.d.ts +61 -93
  72. package/dist/session.d.ts.map +1 -1
  73. package/dist/session.js +388 -812
  74. package/dist/session.js.map +1 -1
  75. package/dist/token_usage.d.ts +11 -18
  76. package/dist/token_usage.d.ts.map +1 -1
  77. package/dist/token_usage.js +29 -94
  78. package/dist/token_usage.js.map +1 -1
  79. package/dist/tool_name.d.ts +18 -0
  80. package/dist/tool_name.d.ts.map +1 -0
  81. package/dist/tool_name.js +29 -0
  82. package/dist/tool_name.js.map +1 -0
  83. package/dist/tool_push.d.ts +28 -0
  84. package/dist/tool_push.d.ts.map +1 -0
  85. package/dist/tool_push.js +151 -0
  86. package/dist/tool_push.js.map +1 -0
  87. package/dist/tool_registry.d.ts +100 -0
  88. package/dist/tool_registry.d.ts.map +1 -0
  89. package/dist/tool_registry.js +440 -0
  90. package/dist/tool_registry.js.map +1 -0
  91. package/dist/tool_schema.d.ts +22 -0
  92. package/dist/tool_schema.d.ts.map +1 -0
  93. package/dist/tool_schema.js +163 -0
  94. package/dist/tool_schema.js.map +1 -0
  95. package/package.json +13 -7
  96. package/dist/integrations/microsoft_copilot.d.ts +0 -84
  97. package/dist/integrations/microsoft_copilot.d.ts.map +0 -1
  98. package/dist/integrations/microsoft_copilot.js +0 -126
  99. package/dist/integrations/microsoft_copilot.js.map +0 -1
@@ -21,17 +21,31 @@
21
21
  * @langchain/core is an optional peer dep, loaded lazily so importing this module
22
22
  * never requires langchain to be installed.
23
23
  *
24
+ * Compatibility details, tested versions, terminal hooks, and limitations are
25
+ * maintained in docs/integrations/compatibility-matrix.md. Keep callback
26
+ * chaining by adding this handler to the existing `callbacks` array; never
27
+ * replace that array.
28
+ * Root chain/model terminal callbacks close their request automatically. Call
29
+ * `await handler.close()` only for host-aborted runs that emit no terminal
30
+ * callback; it is safe and exact-once.
31
+ *
24
32
  * Usage:
25
33
  * const armoriq = new ArmorIQLangChain({ client, mode: 'sdk' });
26
34
  * const handler = await armoriq.forUser('alice@acme.com', { goal: 'reconcile' });
27
35
  * await agentExecutor.invoke({ input: 'Pay invoice 1234' }, { callbacks: [handler] });
36
+ * await handler.close(); // ends the plan trace and flushes it
28
37
  */
29
38
  Object.defineProperty(exports, "__esModule", { value: true });
30
39
  exports.ArmorIQLangChainEnforcer = exports.ArmorIQLangChain = void 0;
31
40
  exports.toolCallsFromLLMResult = toolCallsFromLLMResult;
41
+ exports.usageFromLLMResult = usageFromLLMResult;
42
+ exports.outputFromLLMResult = outputFromLLMResult;
32
43
  exports.nameFromSerialized = nameFromSerialized;
33
44
  exports.coerceArgs = coerceArgs;
45
+ const tool_registry_1 = require("../tool_registry");
46
+ const tool_push_1 = require("../tool_push");
34
47
  const exceptions_1 = require("../exceptions");
48
+ const tool_name_1 = require("../tool_name");
35
49
  class ArmorIQLangChain {
36
50
  client;
37
51
  mode;
@@ -48,8 +62,7 @@ class ArmorIQLangChain {
48
62
  this.defaultMcpName = opts.defaultMcpName;
49
63
  this.customParser = opts.toolNameParser;
50
64
  this.approvalWaitSeconds =
51
- opts.approvalWaitSeconds ??
52
- Number(process.env.ARMORIQ_APPROVAL_WAIT_SECONDS ?? '300');
65
+ opts.approvalWaitSeconds ?? Number(process.env.ARMORIQ_APPROVAL_WAIT_SECONDS ?? '300');
53
66
  this.approvalPollInterval = opts.approvalPollInterval ?? 5;
54
67
  }
55
68
  async bootstrap() {
@@ -60,16 +73,24 @@ class ArmorIQLangChain {
60
73
  async toolNameParser() {
61
74
  if (this.customParser)
62
75
  return this.customParser;
63
- const toolMap = (await this.bootstrap()).toolMap ?? {};
64
76
  const defaultMcp = this.defaultMcpName ?? 'unknown';
77
+ let toolMap = {};
78
+ try {
79
+ toolMap = (await this.bootstrap()).toolMap ?? {};
80
+ }
81
+ catch {
82
+ // Bootstrap enriches tool-name routing only. A model-only LangChain
83
+ // request must still complete when that optional control-plane call is
84
+ // unavailable; a later tool enforcement check remains fail-closed.
85
+ console.warn('[armoriq] langchain bootstrap unavailable; using safe tool-name fallback');
86
+ }
65
87
  return (toolName) => {
66
88
  const mcp = toolMap[toolName];
67
89
  if (mcp)
68
90
  return { mcp, action: toolName };
69
- if (toolName.includes('__')) {
70
- const [prefix, ...rest] = toolName.split('__');
71
- return { mcp: prefix, action: rest.join('__') };
72
- }
91
+ const split = (0, tool_name_1.splitPrefixedToolName)(toolName);
92
+ if (split)
93
+ return split;
73
94
  return { mcp: defaultMcp, action: toolName };
74
95
  };
75
96
  }
@@ -87,6 +108,7 @@ class ArmorIQLangChain {
87
108
  goal: opts?.goal,
88
109
  onEvent: opts?.onEvent,
89
110
  parser: await this.toolNameParser(),
111
+ session: opts?.session,
90
112
  });
91
113
  return buildCallback(enforcer);
92
114
  }
@@ -95,13 +117,38 @@ exports.ArmorIQLangChain = ArmorIQLangChain;
95
117
  /** Framework-agnostic enforcement core (no langchain import), so the allow/hold/
96
118
  * block logic is unit-testable without the langchain package. */
97
119
  class ArmorIQLangChainEnforcer {
120
+ /** A tool ran; record its name for the inventory (names-only until a model call declares it). */
121
+ observeTool(name) {
122
+ (0, tool_registry_1.observeForClient)(this.factory.client, name);
123
+ }
124
+ get client() {
125
+ return this.factory.client;
126
+ }
127
+ /**
128
+ * Fill the inventory from the tool definitions a chat-model call carries.
129
+ * LangChain never hands the callback the tool objects, but every model call
130
+ * carries the bound definitions, so dynamic tool sets are seen as they change.
131
+ */
132
+ async registerModelTools(tools) {
133
+ try {
134
+ const client = this.factory.client;
135
+ const d = (0, tool_registry_1.declarationsFromOpenAiTools)(tools, await (0, tool_registry_1.knownServerNames)(client), await (0, tool_registry_1.knownToolMap)(client));
136
+ (0, tool_registry_1.registerForClient)(this.factory.client, d.tools, d.servers);
137
+ }
138
+ catch {
139
+ /* inventory problems never stop an agent */
140
+ }
141
+ (0, tool_push_1.noteModelCall)(this.factory.client);
142
+ }
98
143
  factory;
99
144
  scope;
100
145
  userEmail;
101
146
  goal;
102
147
  onEvent;
103
148
  parser;
149
+ preboundSession;
104
150
  toolCallNames = new Map();
151
+ pendingPlanCapture;
105
152
  session;
106
153
  planStarted = false;
107
154
  constructor(args) {
@@ -111,6 +158,8 @@ class ArmorIQLangChainEnforcer {
111
158
  this.goal = args.goal;
112
159
  this.onEvent = args.onEvent;
113
160
  this.parser = args.parser;
161
+ this.session = args.session;
162
+ this.preboundSession = args.session;
114
163
  }
115
164
  emit(kind, payload) {
116
165
  if (!this.onEvent)
@@ -133,11 +182,71 @@ class ArmorIQLangChainEnforcer {
133
182
  }
134
183
  return this.session;
135
184
  }
185
+ /** Lazily creates the request session so native OTel and ArmorIQ spans use
186
+ * one stable session id, including a direct model-only completion. */
187
+ telemetrySession() {
188
+ return this.ensureSession();
189
+ }
190
+ hasPreboundSession() {
191
+ return this.preboundSession !== undefined;
192
+ }
193
+ /** Each framework root run owns an independent session and mutable callback
194
+ * state. Sharing one enforcer would cross-contaminate concurrent requests. */
195
+ fork() {
196
+ return new ArmorIQLangChainEnforcer({
197
+ factory: this.factory,
198
+ scope: this.scope,
199
+ userEmail: this.userEmail,
200
+ goal: this.goal,
201
+ onEvent: this.onEvent,
202
+ parser: this.parser,
203
+ session: this.session,
204
+ });
205
+ }
136
206
  async capturePlan(toolCalls) {
137
207
  if (!toolCalls.length)
138
208
  return;
139
- await this.ensureSession().startPlan(toolCalls.map((c) => ({ name: c.name, args: c.args })), this.goal);
140
- this.planStarted = true;
209
+ const session = this.ensureSession();
210
+ const pendingPlanCapture = {
211
+ session,
212
+ promise: session.startPlan(toolCalls.map((c) => ({ name: c.name, args: c.args })), this.goal),
213
+ };
214
+ this.pendingPlanCapture = pendingPlanCapture;
215
+ try {
216
+ await pendingPlanCapture.promise;
217
+ if (this.session === session)
218
+ this.planStarted = true;
219
+ }
220
+ finally {
221
+ if (this.pendingPlanCapture === pendingPlanCapture)
222
+ this.pendingPlanCapture = undefined;
223
+ }
224
+ }
225
+ /** End and ship the request-owned plan session without closing the shared client. */
226
+ async close(status = 'ok', taskOutcome, content) {
227
+ const session = this.session;
228
+ const pendingPlanCapture = this.pendingPlanCapture;
229
+ this.session = undefined;
230
+ this.planStarted = false;
231
+ this.toolCallNames.clear();
232
+ if (session && pendingPlanCapture && pendingPlanCapture.session === session) {
233
+ try {
234
+ await pendingPlanCapture.promise;
235
+ }
236
+ catch {
237
+ // Plan capture failures are handled by the callback; teardown remains best-effort.
238
+ }
239
+ }
240
+ if (session) {
241
+ // Content is authorized, bounded, and redacted by the request-owned core
242
+ // session. Keep the legacy call shape when no content/outcome exists.
243
+ if (content !== undefined)
244
+ await session.close(status, taskOutcome, content);
245
+ else if (taskOutcome !== undefined)
246
+ await session.close(status, taskOutcome);
247
+ else
248
+ await session.close(status);
249
+ }
141
250
  }
142
251
  async captureFromLlm(toolCalls) {
143
252
  this.noteToolCallIds(toolCalls);
@@ -168,6 +277,17 @@ class ArmorIQLangChainEnforcer {
168
277
  }
169
278
  return runName || undefined;
170
279
  }
280
+ frameworkMcpOperation(toolName, itemOrdinal, toolCallId) {
281
+ return {
282
+ category: 'mcp',
283
+ name: 'mcp.execute',
284
+ toolType: 'mcp',
285
+ toolName,
286
+ callId: toolCallId,
287
+ mcpServer: this.parser(toolName).mcp,
288
+ planItemOrdinal: itemOrdinal,
289
+ };
290
+ }
171
291
  /** A tool we cannot name cannot be enforced: a policy keyed on the real name
172
292
  * would silently never match. Refuse it instead of guessing. */
173
293
  refuseUnnamedTool(serializedHint) {
@@ -177,6 +297,11 @@ class ArmorIQLangChainEnforcer {
177
297
  }
178
298
  /** Allow -> return; block/hold-denied -> throw (LangChain stops the tool). */
179
299
  async enforceTool(toolName, args) {
300
+ await this.enforceToolWithDecision(toolName, args);
301
+ }
302
+ /** Internal callback path that retains safe policy metadata without changing
303
+ * the public enforceTool() compatibility contract. */
304
+ async enforceToolWithDecision(toolName, args) {
180
305
  args = args || {};
181
306
  try {
182
307
  if (!this.planStarted) {
@@ -185,9 +310,9 @@ class ArmorIQLangChainEnforcer {
185
310
  await this.capturePlan([{ name: toolName, args }]);
186
311
  }
187
312
  const session = this.ensureSession();
188
- const decision = await session.check(toolName, args, this.userEmail);
313
+ const decision = await session.check(toolName, args, this.userEmail, { emitOtel: false });
189
314
  if (decision.allowed)
190
- return; // explicit allow -> tool runs
315
+ return decision; // explicit allow -> tool runs
191
316
  if (decision.action === 'hold') {
192
317
  this.emit('hold', {
193
318
  tool: toolName,
@@ -201,7 +326,7 @@ class ArmorIQLangChainEnforcer {
201
326
  });
202
327
  if (outcome === 'approved') {
203
328
  this.emit('approved', { tool: toolName, delegationId: decision.delegationId });
204
- return; // approved -> tool runs
329
+ return { ...decision, allowed: true, action: 'allow', approvalOutcome: outcome }; // approved -> tool runs
205
330
  }
206
331
  this.emit(outcome, {
207
332
  tool: toolName,
@@ -228,13 +353,15 @@ class ArmorIQLangChainEnforcer {
228
353
  * 'completed' once all declared steps have run). Only called after a tool actually
229
354
  * ran — enforcement-cancelled tools never reach here. Best-effort: audit failures
230
355
  * are logged, never allowed to break the run (mirrors the strands AfterToolCall). */
231
- async reportExecution(toolName, args, result, error) {
356
+ async reportExecution(toolName, args, result, error, frameworkCall) {
232
357
  if (!this.session || !this.planStarted)
233
358
  return;
234
359
  try {
235
360
  await this.session.report(toolName, args ?? {}, result, {
236
361
  status: error ? 'error' : 'success',
237
362
  errorMessage: error,
363
+ emitOtel: false,
364
+ operation: this.frameworkMcpOperation(toolName, frameworkCall?.itemOrdinal, frameworkCall?.toolCallId),
238
365
  });
239
366
  }
240
367
  catch (exc) {
@@ -257,6 +384,35 @@ function toolCallsFromLLMResult(output) {
257
384
  }
258
385
  return calls;
259
386
  }
387
+ /** Extract only provider-reported aggregate usage; prompt/content is never
388
+ * copied into telemetry attributes. LangChain provider adapters use several
389
+ * casing variants, so accept their stable aliases. */
390
+ function usageFromLLMResult(output) {
391
+ const usage = output?.llmOutput?.tokenUsage ??
392
+ output?.llmOutput?.usage ??
393
+ output?.generations?.[0]?.[0]?.message?.usage_metadata ??
394
+ {};
395
+ const input = usage?.promptTokens ?? usage?.prompt_tokens ?? usage?.inputTokens ?? usage?.input_tokens;
396
+ const outputTokens = usage?.completionTokens ??
397
+ usage?.completion_tokens ??
398
+ usage?.outputTokens ??
399
+ usage?.output_tokens;
400
+ return {
401
+ ...(typeof input === 'number' && Number.isFinite(input)
402
+ ? { 'gen_ai.usage.input_tokens': input }
403
+ : {}),
404
+ ...(typeof outputTokens === 'number' && Number.isFinite(outputTokens)
405
+ ? { 'gen_ai.usage.output_tokens': outputTokens }
406
+ : {}),
407
+ };
408
+ }
409
+ /** Preserve framework payload shape. The core owns redaction, limits, and
410
+ * capture-mode authorization; adapters must not pre-redact or stringify away
411
+ * structured content before the policy gate sees it. */
412
+ function outputFromLLMResult(output) {
413
+ const first = output?.generations?.[0]?.[0];
414
+ return first?.message?.content ?? first?.text ?? output?.output;
415
+ }
260
416
  /** Best-effort name from LangChain's Serialized payload. `id` is deliberately
261
417
  * ignored: its last element is the tool CLASS name, which looks like a tool
262
418
  * name but never is one. */
@@ -268,6 +424,29 @@ function serializedHint(tool) {
268
424
  return tool.id[tool.id.length - 1] ?? 'unknown';
269
425
  return String(tool?.id ?? 'unknown');
270
426
  }
427
+ /**
428
+ * Real model identifier (e.g. "gpt-4o-mini") from LangChain's chat-model-start
429
+ * or LLM-start callback params, falling back to the run name. `runName` is
430
+ * the model class's display name (e.g. "ChatOpenAI"), not a model id, so
431
+ * passing it straight to `OtelSession.beginModel` as the model would leave
432
+ * `gen_ai.system` undetected -- `modelSystem()` only recognizes known
433
+ * model-id prefixes. Both callbacks carry the same `invocation_params` shape
434
+ * in their `extra` argument.
435
+ *
436
+ * Some providers (e.g. `ChatGoogleGenerativeAI`) never put a model id in
437
+ * `invocation_params` at all -- they only report it through `getLsParams()`,
438
+ * which @langchain/core's `CallbackManager` folds into its inheritable
439
+ * metadata and passes as the callback's `metadata` argument (LangSmith's
440
+ * `ls_model_name` field). Fall back to that before giving up on the run name.
441
+ */
442
+ function modelNameFromInvocationParams(extra, metadata, fallback) {
443
+ const invocationParams = extra
444
+ ?.invocation_params;
445
+ const model = invocationParams?.model ??
446
+ invocationParams?.model_name ??
447
+ metadata?.ls_model_name;
448
+ return typeof model === 'string' && model ? model : fallback;
449
+ }
271
450
  function coerceArgs(input) {
272
451
  if (input && typeof input === 'object')
273
452
  return input;
@@ -282,6 +461,151 @@ function coerceArgs(input) {
282
461
  }
283
462
  return {};
284
463
  }
464
+ /**
465
+ * Adapter state over the request-owned OTel session. The session owns the
466
+ * provider, root span, bounded flush, and terminal shutdown. This class owns
467
+ * only LangChain's run-id correlation, so separate handlers never share state.
468
+ */
469
+ class LangChainTelemetry {
470
+ session;
471
+ runs = new Map();
472
+ toolOrdinals = new Map();
473
+ nextToolOrdinal = 0;
474
+ constructor(session) {
475
+ this.session = session;
476
+ }
477
+ noteToolCalls(calls) {
478
+ for (const call of calls) {
479
+ const ordinal = this.nextToolOrdinal++;
480
+ if (call.id)
481
+ this.toolOrdinals.set(call.id, ordinal);
482
+ }
483
+ }
484
+ toolOrdinal(toolCallId) {
485
+ return toolCallId ? this.toolOrdinals.get(toolCallId) : undefined;
486
+ }
487
+ otel() {
488
+ return this.session().otelSession;
489
+ }
490
+ async beginRoot(input) {
491
+ try {
492
+ await this.otel()?.beginRoot(input === undefined ? undefined : { input });
493
+ }
494
+ catch {
495
+ // Native telemetry must never change a LangChain run result.
496
+ }
497
+ }
498
+ async beginModel(runId, model = 'langchain', input) {
499
+ if (!runId || this.runs.has(runId))
500
+ return;
501
+ try {
502
+ // Direct LLM callbacks have no chain start. Ensure their root receives
503
+ // the same prompt instead of creating a content-less root in beginModel.
504
+ await this.beginRoot(input);
505
+ const handle = await this.otel()?.beginModel(model, input === undefined ? undefined : { input });
506
+ if (handle)
507
+ this.runs.set(runId, { kind: 'model', handle });
508
+ }
509
+ catch {
510
+ // Native telemetry must never change a LangChain run result.
511
+ }
512
+ }
513
+ async beginPolicy(runId, toolName, itemOrdinal, toolCallId, args) {
514
+ if (!runId || this.runs.has(runId))
515
+ return;
516
+ try {
517
+ const handle = await this.otel()?.beginPolicy({
518
+ toolName,
519
+ itemOrdinal,
520
+ toolCallId,
521
+ ...(args === undefined ? {} : { arguments: args }),
522
+ });
523
+ if (handle)
524
+ this.runs.set(runId, { kind: 'policy', handle });
525
+ }
526
+ catch {
527
+ // Native telemetry must never change a LangChain run result.
528
+ }
529
+ }
530
+ async beginTool(runId, toolName, itemOrdinal, toolCallId, args, operation) {
531
+ if (!runId || this.runs.has(runId))
532
+ return;
533
+ try {
534
+ const handle = await this.otel()?.beginTool({
535
+ toolName,
536
+ itemOrdinal,
537
+ toolCallId,
538
+ operation,
539
+ ...(args === undefined ? {} : { arguments: args }),
540
+ });
541
+ if (handle)
542
+ this.runs.set(runId, { kind: 'tool', handle });
543
+ }
544
+ catch {
545
+ // Native telemetry must never change a LangChain run result.
546
+ }
547
+ }
548
+ async end(runId, outcome = 'ok', modelValues = {}, payload = {}) {
549
+ if (!runId)
550
+ return;
551
+ const run = this.runs.get(runId);
552
+ if (!run)
553
+ return;
554
+ this.runs.delete(runId);
555
+ try {
556
+ if (run.kind === 'model') {
557
+ await this.otel()?.endModel(run.handle, modelValues, outcome === 'ok' ? undefined : new Error('redacted'), outcome, payload.output === undefined ? undefined : { output: payload.output });
558
+ }
559
+ else if (run.kind === 'policy') {
560
+ await this.otel()?.endPolicy(run.handle, {
561
+ ...(outcome === 'ok'
562
+ ? { decision: 'allow' }
563
+ : outcome === 'denied'
564
+ ? { decision: 'block' }
565
+ : {}),
566
+ ...(outcome === 'ok' ? {} : { error: new Error('redacted') }),
567
+ terminalStatus: outcome,
568
+ ...payload.policy,
569
+ });
570
+ }
571
+ else {
572
+ await this.otel()?.endTool(run.handle, {
573
+ outcome: outcome === 'ok'
574
+ ? 'success'
575
+ : outcome === 'timeout' || outcome === 'disconnected' || outcome === 'cancelled'
576
+ ? outcome
577
+ : 'error',
578
+ ...(outcome === 'ok' ? {} : { error: new Error('redacted') }),
579
+ ...(payload.toolResult === undefined ? {} : { result: payload.toolResult }),
580
+ });
581
+ }
582
+ }
583
+ catch {
584
+ // Native telemetry must never change a LangChain run result.
585
+ }
586
+ }
587
+ async endAll(outcome = 'ok') {
588
+ await Promise.all([...this.runs.keys()].map((runId) => this.end(runId, outcome)));
589
+ }
590
+ }
591
+ /** Translate the framework's terminal error vocabulary without treating an
592
+ * expected abort/disconnect as an ordinary model failure. */
593
+ function terminalStatusFromError(error) {
594
+ const value = error;
595
+ const name = String(value?.name ?? 'Error').toLowerCase();
596
+ const code = String(value?.code ?? '').toLowerCase();
597
+ const message = String(value?.message ?? '').toLowerCase();
598
+ if (name === 'aborterror' || code === 'abort_err' || code === 'err_abort')
599
+ return 'cancelled';
600
+ if (name.includes('timeout') || code === 'etimedout' || message.includes('timed out'))
601
+ return 'timeout';
602
+ if (code === 'econnreset' ||
603
+ code === 'err_stream_premature_close' ||
604
+ message.includes('disconnect') ||
605
+ message.includes('connection reset'))
606
+ return 'disconnected';
607
+ return 'error';
608
+ }
285
609
  function buildCallback(enforcer) {
286
610
  let BaseCallbackHandler;
287
611
  try {
@@ -296,39 +620,248 @@ function buildCallback(enforcer) {
296
620
  // Propagate our block/hold throws so LangChain actually stops the tool.
297
621
  raiseError = true;
298
622
  awaitHandlers = true;
299
- // runId -> {name, args} for tools that passed enforcement and ran, so
300
- // handleToolEnd/handleToolError can report the execution (close the plan).
301
- pending = {};
302
- async handleLLMEnd(output) {
623
+ requests = new Map();
624
+ runRoots = new Map();
625
+ toolCallRoots = new Map();
626
+ retiredRuns = new Set();
627
+ preboundRootRunId;
628
+ handlerClosed = false;
629
+ createRequest(rootRunId) {
630
+ if (enforcer.hasPreboundSession() &&
631
+ this.preboundRootRunId !== undefined &&
632
+ this.preboundRootRunId !== rootRunId) {
633
+ throw new exceptions_1.PolicyBlockedException('ArmorIQ enforcement error (fail-closed): a prebound session supports one root run only');
634
+ }
635
+ if (enforcer.hasPreboundSession())
636
+ this.preboundRootRunId = rootRunId;
637
+ const requestEnforcer = enforcer.fork();
638
+ const request = {
639
+ enforcer: requestEnforcer,
640
+ telemetry: new LangChainTelemetry(() => requestEnforcer.telemetrySession()),
641
+ runs: new Set([rootRunId]),
642
+ pending: new Map(),
643
+ };
644
+ this.requests.set(rootRunId, request);
645
+ this.runRoots.set(rootRunId, rootRunId);
646
+ return request;
647
+ }
648
+ toolCallKey(rootRunId, toolCallId) {
649
+ return `${rootRunId}\u0000${toolCallId}`;
650
+ }
651
+ requestFor(runId, parentRunId) {
652
+ if (!runId ||
653
+ this.handlerClosed ||
654
+ this.retiredRuns.has(runId) ||
655
+ (parentRunId !== undefined && this.retiredRuns.has(parentRunId)))
656
+ return undefined;
657
+ const rootRunId = this.runRoots.get(runId) ?? (parentRunId ? this.runRoots.get(parentRunId) : undefined);
658
+ const request = rootRunId ? this.requests.get(rootRunId) : this.createRequest(runId);
659
+ if (!request)
660
+ return undefined;
661
+ request.runs.add(runId);
662
+ this.runRoots.set(runId, rootRunId ?? runId);
663
+ return request;
664
+ }
665
+ requestForTool(runId, parentRunId, toolCallId) {
666
+ if (toolCallId) {
667
+ const rootRunId = (parentRunId ? this.runRoots.get(parentRunId) : undefined) ??
668
+ (runId ? this.runRoots.get(runId) : undefined);
669
+ let mappedRoot = rootRunId
670
+ ? this.toolCallRoots.get(this.toolCallKey(rootRunId, toolCallId))
671
+ : undefined;
672
+ // Older/direct tool invocations may omit parentRunId. Reuse a model
673
+ // plan only when that id is unambiguous across live roots; otherwise
674
+ // fail closed rather than cross-routing a duplicate provider id.
675
+ if (!mappedRoot && !rootRunId) {
676
+ const matches = [...this.toolCallRoots]
677
+ .filter(([key]) => key.endsWith(`\u0000${toolCallId}`))
678
+ .map(([, root]) => root);
679
+ if (matches.length === 1)
680
+ mappedRoot = matches[0];
681
+ }
682
+ const request = mappedRoot ? this.requests.get(mappedRoot) : undefined;
683
+ if (request && runId && !this.retiredRuns.has(runId)) {
684
+ request.runs.add(runId);
685
+ this.runRoots.set(runId, mappedRoot);
686
+ return request;
687
+ }
688
+ }
689
+ return this.requestFor(runId, parentRunId);
690
+ }
691
+ async closeRequest(rootRunId, status = 'ok', output) {
692
+ const request = this.requests.get(rootRunId);
693
+ if (!request)
694
+ return;
695
+ // Retire before awaiting teardown. Late callbacks must never create a
696
+ // fresh session while the old request drains its final spans.
697
+ this.requests.delete(rootRunId);
698
+ for (const runId of request.runs) {
699
+ this.runRoots.delete(runId);
700
+ this.retiredRuns.add(runId);
701
+ }
702
+ for (const [toolCallId, mappedRoot] of this.toolCallRoots) {
703
+ if (mappedRoot === rootRunId)
704
+ this.toolCallRoots.delete(toolCallId);
705
+ }
706
+ request.pending.clear();
707
+ await request.telemetry.endAll(status);
708
+ await request.enforcer.close(status, undefined, output === undefined ? undefined : { output });
709
+ }
710
+ async closeRoot(runId, status = 'ok', output) {
711
+ if (!runId)
712
+ return;
713
+ const rootRunId = this.runRoots.get(runId);
714
+ if (rootRunId)
715
+ await this.closeRequest(rootRunId, status, output);
716
+ }
717
+ async handleChainStart(_chain, _inputs, runId, parentRunId, _tags, _metadata, _runType, runName) {
718
+ void _chain;
719
+ void runName;
720
+ const request = this.requestFor(runId, parentRunId);
721
+ if (request && !parentRunId)
722
+ await request.telemetry.beginRoot(_inputs);
723
+ }
724
+ async handleChainEnd(_output, runId) {
725
+ if (runId && this.runRoots.get(runId) === runId)
726
+ await this.closeRoot(runId, 'ok', _output);
727
+ }
728
+ async handleChainError(_error, runId) {
729
+ if (runId && this.runRoots.get(runId) === runId) {
730
+ await this.closeRoot(runId, terminalStatusFromError(_error));
731
+ }
732
+ }
733
+ async handleChatModelStart(_model, _messages, runId, parentRunId, extra, _tags, metadata, runName) {
734
+ const request = this.requestFor(runId, parentRunId);
735
+ if (!request)
736
+ return;
737
+ const tools = extra?.invocation_params?.tools;
738
+ if (Array.isArray(tools) && tools.length)
739
+ await request.enforcer.registerModelTools(tools);
740
+ else
741
+ (0, tool_push_1.noteModelCall)(request.enforcer.client);
742
+ const model = modelNameFromInvocationParams(extra, metadata, runName ?? 'langchain');
743
+ await request.telemetry.beginModel(runId, model, _messages);
744
+ }
745
+ /**
746
+ * Non-chat LLMs report an array of prompts instead of message objects.
747
+ * `extra` must stay in this exact slot -- @langchain/core's callback
748
+ * manager invokes `handleLLMStart` with 8 positional args (llm, prompts,
749
+ * runId, parentRunId, extraParams, tags, metadata, runName); omitting
750
+ * `extra` here previously shifted every argument after it by one, so
751
+ * `runName` was silently receiving the `metadata` object instead of the
752
+ * run name string.
753
+ */
754
+ async handleLLMStart(_llm, prompts, runId, parentRunId, extra, _tags, metadata, runName) {
755
+ const request = this.requestFor(runId, parentRunId);
756
+ if (request) {
757
+ const model = modelNameFromInvocationParams(extra, metadata, runName ?? 'langchain');
758
+ await request.telemetry.beginModel(runId, model, prompts);
759
+ }
760
+ }
761
+ async handleLLMEnd(output, runId, parentRunId) {
303
762
  const calls = toolCallsFromLLMResult(output);
763
+ const request = this.requestFor(runId, parentRunId);
764
+ if (!request)
765
+ return;
766
+ await request.telemetry.beginModel(runId);
767
+ request.telemetry.noteToolCalls(calls);
768
+ const rootRunId = runId ? this.runRoots.get(runId) : undefined;
769
+ for (const call of calls) {
770
+ if (call.id && rootRunId) {
771
+ this.toolCallRoots.set(this.toolCallKey(rootRunId, call.id), rootRunId);
772
+ }
773
+ }
304
774
  if (calls.length)
305
- await enforcer.captureFromLlm(calls);
775
+ await request.enforcer.captureFromLlm(calls);
776
+ // Direct model invocations do not always have a chain callback, but the
777
+ // model completion itself is a complete OTel lifecycle.
778
+ await request.telemetry.end(runId, 'ok', usageFromLLMResult(output), {
779
+ output: outputFromLLMResult(output),
780
+ });
781
+ if (!parentRunId && calls.length === 0) {
782
+ await this.closeRoot(runId, 'ok', outputFromLLMResult(output));
783
+ }
784
+ }
785
+ async handleLLMError(_error, runId, parentRunId) {
786
+ const request = this.requestFor(runId, parentRunId);
787
+ if (!request)
788
+ return;
789
+ const status = terminalStatusFromError(_error);
790
+ await request.telemetry.beginModel(runId);
791
+ await request.telemetry.end(runId, status);
792
+ if (!parentRunId)
793
+ await this.closeRoot(runId, status);
306
794
  }
307
- async handleToolStart(tool, input, runId, _parentRunId, _tags, _metadata, runName, toolCallId) {
308
- const toolName = enforcer.resolveToolName(toolCallId, runName) ??
795
+ async handleToolStart(tool, input, runId, parentRunId, _tags, _metadata, runName, toolCallId) {
796
+ const request = this.requestForTool(runId, parentRunId, toolCallId);
797
+ if (!request) {
798
+ throw new exceptions_1.PolicyBlockedException('ArmorIQ enforcement error (fail-closed): callback request is closed');
799
+ }
800
+ const toolName = request.enforcer.resolveToolName(toolCallId, runName) ??
309
801
  nameFromSerialized(tool) ??
310
- enforcer.refuseUnnamedTool(serializedHint(tool));
802
+ request.enforcer.refuseUnnamedTool(serializedHint(tool));
803
+ request.enforcer.observeTool(toolName);
311
804
  const args = coerceArgs(input);
312
- await enforcer.enforceTool(toolName, args); // throws to block/hold-deny
805
+ await request.telemetry.beginPolicy(runId, toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId, args);
806
+ try {
807
+ const decision = await request.enforcer.enforceToolWithDecision(toolName, args); // throws to block/hold-deny
808
+ await request.telemetry.end(runId, 'ok', {}, {
809
+ policy: {
810
+ policyName: decision?.matchedPolicy,
811
+ ...(decision?.policyId ? { policyId: decision.policyId } : {}),
812
+ policyVersion: decision?.policyVersion,
813
+ policySource: decision?.policySource,
814
+ policyReasonCode: decision?.reason,
815
+ defaultAction: decision?.defaultAction,
816
+ matchedRuleId: decision?.matchedRuleId,
817
+ },
818
+ });
819
+ }
820
+ catch (error) {
821
+ await request.telemetry.end(runId, error instanceof exceptions_1.PolicyBlockedException || error instanceof exceptions_1.PolicyHoldException
822
+ ? 'denied'
823
+ : 'error', {}, { policy: { policyReasonCode: 'enforcement_error' } });
824
+ throw error;
825
+ }
826
+ await request.telemetry.beginTool(runId, toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId, args, request.enforcer.frameworkMcpOperation(toolName, request.telemetry.toolOrdinal(toolCallId), toolCallId));
313
827
  // Only reached when allowed -> the tool will run; remember it so
314
828
  // handleToolEnd can report the execution and close the plan lifecycle.
315
- if (runId)
316
- this.pending[runId] = { name: toolName, args };
829
+ if (runId) {
830
+ request.pending.set(runId, {
831
+ name: toolName,
832
+ args,
833
+ itemOrdinal: request.telemetry.toolOrdinal(toolCallId),
834
+ toolCallId,
835
+ });
836
+ }
317
837
  }
318
838
  async handleToolEnd(output, runId) {
319
- const info = runId ? this.pending[runId] : undefined;
839
+ const rootRunId = runId ? this.runRoots.get(runId) : undefined;
840
+ const request = rootRunId ? this.requests.get(rootRunId) : undefined;
841
+ const info = runId ? request?.pending.get(runId) : undefined;
320
842
  if (info) {
321
- delete this.pending[runId];
322
- await enforcer.reportExecution(info.name, info.args, output);
843
+ request.pending.delete(runId);
844
+ await request.enforcer.reportExecution(info.name, info.args, output, undefined, info);
323
845
  }
846
+ await request?.telemetry.end(runId, 'ok', {}, { toolResult: output });
847
+ }
848
+ /** LangChain has no end-of-request hook, so the host finalizes the session here. */
849
+ async close(status = 'ok') {
850
+ if (this.handlerClosed)
851
+ return;
852
+ this.handlerClosed = true;
853
+ await Promise.all([...this.requests.keys()].map((rootRunId) => this.closeRequest(rootRunId, status)));
324
854
  }
325
855
  async handleToolError(err, runId) {
326
- const info = runId ? this.pending[runId] : undefined;
856
+ const rootRunId = runId ? this.runRoots.get(runId) : undefined;
857
+ const request = rootRunId ? this.requests.get(rootRunId) : undefined;
858
+ const info = runId ? request?.pending.get(runId) : undefined;
327
859
  if (info) {
328
- delete this.pending[runId];
860
+ request.pending.delete(runId);
329
861
  const msg = err?.message ?? String(err);
330
- await enforcer.reportExecution(info.name, info.args, msg, msg);
862
+ await request.enforcer.reportExecution(info.name, info.args, msg, msg, info);
331
863
  }
864
+ await request?.telemetry.end(runId, terminalStatusFromError(err));
332
865
  }
333
866
  }
334
867
  return new ArmorIQLangChainCallback();