@townco/agent 0.1.77 → 0.1.79

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.
@@ -1,5 +1,9 @@
1
1
  import type { ContextEntry } from "../../acp-server/session-storage";
2
2
  import type { HookCallback, HookConfig, HookNotification, ReadonlySession } from "./types";
3
+ /**
4
+ * Callback for streaming hook notifications in real-time
5
+ */
6
+ export type OnHookNotification = (notification: HookNotification) => void;
3
7
  /**
4
8
  * Hook executor manages hook lifecycle
5
9
  */
@@ -7,7 +11,12 @@ export declare class HookExecutor {
7
11
  private hooks;
8
12
  private model;
9
13
  private loadCallback;
10
- constructor(hooks: HookConfig[], model: string, loadCallback: (callbackRef: string) => Promise<HookCallback>);
14
+ private onNotification;
15
+ constructor(hooks: HookConfig[], model: string, loadCallback: (callbackRef: string) => Promise<HookCallback>, onNotification?: OnHookNotification);
16
+ /**
17
+ * Emit a notification - sends immediately if callback provided, otherwise collects for batch return
18
+ */
19
+ private emitNotification;
11
20
  /**
12
21
  * Execute hooks before agent invocation
13
22
  * Returns new context entries to append and any notifications to send
@@ -32,6 +41,7 @@ export declare class HookExecutor {
32
41
  }): Promise<{
33
42
  modifiedOutput?: Record<string, unknown>;
34
43
  truncationWarning?: string;
44
+ notifications: HookNotification[];
35
45
  }>;
36
46
  /**
37
47
  * Execute a single tool_response hook
@@ -1,12 +1,6 @@
1
1
  import { createLogger } from "../../logger.js";
2
- import { DEFAULT_CONTEXT_SIZE, MODEL_CONTEXT_WINDOWS } from "./constants";
2
+ import { getModelContextWindow } from "./constants";
3
3
  const logger = createLogger("hook-executor");
4
- /**
5
- * Get max context size for a model
6
- */
7
- function getModelMaxTokens(model) {
8
- return MODEL_CONTEXT_WINDOWS[model] ?? DEFAULT_CONTEXT_SIZE;
9
- }
10
4
  /**
11
5
  * Hook executor manages hook lifecycle
12
6
  */
@@ -14,10 +8,21 @@ export class HookExecutor {
14
8
  hooks;
15
9
  model;
16
10
  loadCallback;
17
- constructor(hooks, model, loadCallback) {
11
+ onNotification;
12
+ constructor(hooks, model, loadCallback, onNotification) {
18
13
  this.hooks = hooks;
19
14
  this.model = model;
20
15
  this.loadCallback = loadCallback;
16
+ this.onNotification = onNotification;
17
+ }
18
+ /**
19
+ * Emit a notification - sends immediately if callback provided, otherwise collects for batch return
20
+ */
21
+ emitNotification(notification, notifications) {
22
+ notifications.push(notification);
23
+ if (this.onNotification) {
24
+ this.onNotification(notification);
25
+ }
21
26
  }
22
27
  /**
23
28
  * Execute hooks before agent invocation
@@ -49,7 +54,7 @@ export class HookExecutor {
49
54
  * Execute a context_size hook
50
55
  */
51
56
  async executeContextSizeHook(hook, session, actualInputTokens) {
52
- const maxTokens = getModelMaxTokens(this.model);
57
+ const maxTokens = getModelContextWindow(this.model);
53
58
  const percentage = (actualInputTokens / maxTokens) * 100;
54
59
  // Default threshold is 95%
55
60
  const threshold = hook.setting?.threshold ?? 95;
@@ -60,14 +65,16 @@ export class HookExecutor {
60
65
  }
61
66
  logger.info(`Context hook triggered: ${actualInputTokens} tokens (${percentage.toFixed(1)}%) exceeds threshold ${threshold}%`);
62
67
  const notifications = [];
63
- // Notify that hook is triggered
64
- notifications.push({
68
+ const triggeredAt = Date.now();
69
+ // Notify that hook is triggered - sent immediately via callback if available
70
+ this.emitNotification({
65
71
  type: "hook_triggered",
66
72
  hookType: "context_size",
67
73
  threshold,
68
74
  currentPercentage: percentage,
69
75
  callback: hook.callback,
70
- });
76
+ triggeredAt,
77
+ }, notifications);
71
78
  try {
72
79
  // Load and execute callback
73
80
  const callback = await this.loadCallback(hook.callback);
@@ -80,12 +87,13 @@ export class HookExecutor {
80
87
  };
81
88
  const result = await callback(hookContext);
82
89
  // Notify completion
83
- notifications.push({
90
+ this.emitNotification({
84
91
  type: "hook_completed",
85
92
  hookType: "context_size",
86
93
  callback: hook.callback,
87
94
  metadata: result.metadata,
88
- });
95
+ completedAt: Date.now(),
96
+ }, notifications);
89
97
  return {
90
98
  newContextEntry: result.newContextEntry,
91
99
  notifications,
@@ -93,12 +101,13 @@ export class HookExecutor {
93
101
  }
94
102
  catch (error) {
95
103
  // Notify error
96
- notifications.push({
104
+ this.emitNotification({
97
105
  type: "hook_error",
98
106
  hookType: "context_size",
99
107
  callback: hook.callback,
100
108
  error: error instanceof Error ? error.message : String(error),
101
- });
109
+ completedAt: Date.now(),
110
+ }, notifications);
102
111
  // Return no context entry on error
103
112
  return {
104
113
  newContextEntry: null,
@@ -115,21 +124,46 @@ export class HookExecutor {
115
124
  toolName: toolResponse.toolName,
116
125
  outputTokens: toolResponse.outputTokens,
117
126
  });
127
+ const notifications = [];
118
128
  for (const hook of this.hooks) {
119
129
  if (hook.type === "tool_response") {
120
130
  const result = await this.executeToolResponseHook(hook, session, currentContextTokens, toolResponse);
121
131
  if (result) {
122
- return result;
132
+ notifications.push(...result.notifications);
133
+ const response = { notifications };
134
+ if (result.modifiedOutput !== undefined) {
135
+ response.modifiedOutput = result.modifiedOutput;
136
+ }
137
+ if (result.truncationWarning !== undefined) {
138
+ response.truncationWarning = result.truncationWarning;
139
+ }
140
+ return response;
123
141
  }
124
142
  }
125
143
  }
126
- return {}; // No modifications
144
+ return { notifications }; // No modifications
127
145
  }
128
146
  /**
129
147
  * Execute a single tool_response hook
130
148
  */
131
149
  async executeToolResponseHook(hook, session, currentContextTokens, toolResponse) {
132
- const maxTokens = getModelMaxTokens(this.model);
150
+ const maxTokens = getModelContextWindow(this.model);
151
+ const percentage = (currentContextTokens / maxTokens) * 100;
152
+ const notifications = [];
153
+ // Get threshold from settings
154
+ const settings = hook.setting;
155
+ const threshold = settings?.maxContextThreshold ?? 80;
156
+ // Capture start time and emit hook_triggered BEFORE callback runs
157
+ // This allows the UI to show the loading state immediately
158
+ const triggeredAt = Date.now();
159
+ this.emitNotification({
160
+ type: "hook_triggered",
161
+ hookType: "tool_response",
162
+ threshold,
163
+ currentPercentage: percentage,
164
+ callback: hook.callback,
165
+ triggeredAt,
166
+ }, notifications);
133
167
  try {
134
168
  // Load and execute callback
135
169
  const callback = await this.loadCallback(hook.callback);
@@ -145,31 +179,54 @@ export class HookExecutor {
145
179
  session: sessionWithSettings,
146
180
  currentTokens: currentContextTokens,
147
181
  maxTokens,
148
- percentage: (currentContextTokens / maxTokens) * 100,
182
+ percentage,
149
183
  model: this.model,
150
184
  toolResponse,
151
185
  };
152
186
  const result = await callback(hookContext);
153
187
  // Extract modified output and warnings from metadata
154
- if (result.metadata) {
155
- const response = {};
156
- if (result.metadata.modifiedOutput) {
157
- response.modifiedOutput = result.metadata.modifiedOutput;
158
- }
188
+ if (result.metadata && result.metadata.modifiedOutput) {
189
+ // Hook took action - emit completed notification
190
+ const response = { notifications };
191
+ response.modifiedOutput = result.metadata.modifiedOutput;
159
192
  if (result.metadata.truncationWarning) {
160
193
  response.truncationWarning = result.metadata
161
194
  .truncationWarning;
162
195
  }
196
+ // Notify completion
197
+ this.emitNotification({
198
+ type: "hook_completed",
199
+ hookType: "tool_response",
200
+ callback: hook.callback,
201
+ metadata: result.metadata,
202
+ completedAt: Date.now(),
203
+ }, notifications);
163
204
  return response;
164
205
  }
165
- return null;
206
+ // No action was taken - emit completed with no-op metadata
207
+ this.emitNotification({
208
+ type: "hook_completed",
209
+ hookType: "tool_response",
210
+ callback: hook.callback,
211
+ metadata: { action: "no_action_needed" },
212
+ completedAt: Date.now(),
213
+ }, notifications);
214
+ return { notifications };
166
215
  }
167
216
  catch (error) {
217
+ // Notify error
218
+ this.emitNotification({
219
+ type: "hook_error",
220
+ hookType: "tool_response",
221
+ callback: hook.callback,
222
+ error: error instanceof Error ? error.message : String(error),
223
+ completedAt: Date.now(),
224
+ }, notifications);
168
225
  logger.error("Tool response hook execution failed", {
169
226
  callback: hook.callback,
170
227
  error: error instanceof Error ? error.message : String(error),
171
228
  });
172
- return null; // Return original output on error
229
+ return { notifications }; // Return notifications even on error
173
230
  }
174
231
  }
175
232
  }
@@ -167,14 +167,17 @@ export type HookNotification = {
167
167
  threshold: number;
168
168
  currentPercentage: number;
169
169
  callback: string;
170
+ triggeredAt: number;
170
171
  } | {
171
172
  type: "hook_completed";
172
173
  hookType: HookType;
173
174
  callback: string;
174
175
  metadata?: HookResult["metadata"];
176
+ completedAt: number;
175
177
  } | {
176
178
  type: "hook_error";
177
179
  hookType: HookType;
178
180
  callback: string;
179
181
  error: string;
182
+ completedAt: number;
180
183
  };
@@ -12,5 +12,6 @@ export declare class LangchainAgent implements AgentRunner {
12
12
  definition: CreateAgentRunnerParams;
13
13
  constructor(params: CreateAgentRunnerParams);
14
14
  invoke(req: InvokeRequest): AsyncGenerator<ExtendedSessionUpdate, PromptResponse, undefined>;
15
+ private invokeInternal;
15
16
  }
16
17
  export { makeSubagentsTool } from "./tools/subagent.js";
@@ -14,7 +14,7 @@ import { makeFilesystemTools } from "./tools/filesystem";
14
14
  import { makeGenerateImageTool } from "./tools/generate_image";
15
15
  import { SUBAGENT_TOOL_NAME } from "./tools/subagent";
16
16
  import { hashQuery, queryToToolCallId, subagentEvents, } from "./tools/subagent-connections";
17
- import { TODO_WRITE_TOOL_NAME, todoWrite } from "./tools/todo";
17
+ import { makeTodoWriteTool, TODO_WRITE_TOOL_NAME } from "./tools/todo";
18
18
  import { makeTownWebSearchTools, makeWebSearchTools } from "./tools/web_search";
19
19
  const _logger = createLogger("agent-runner");
20
20
  const getWeather = tool(({ city }) => `It's always sunny in ${city}!`, {
@@ -27,8 +27,8 @@ const getWeather = tool(({ city }) => `It's always sunny in ${city}!`, {
27
27
  getWeather.prettyName = "Get Weather";
28
28
  getWeather.icon = "Cloud";
29
29
  export const TOOL_REGISTRY = {
30
- todo_write: todoWrite,
31
- get_weather: getWeather,
30
+ todo_write: () => makeTodoWriteTool(), // Factory function to create fresh instance per invocation
31
+ get_weather: getWeather, // TODO: Convert to factory function for full concurrency safety
32
32
  web_search: () => makeWebSearchTools(),
33
33
  town_web_search: () => makeTownWebSearchTools(),
34
34
  filesystem: () => makeFilesystemTools(process.cwd()),
@@ -64,7 +64,12 @@ export class LangchainAgent {
64
64
  constructor(params) {
65
65
  this.definition = params;
66
66
  }
67
- async *invoke(req) {
67
+ invoke(req) {
68
+ const sessionAttributes = { "agent.session_id": req.sessionId };
69
+ const generator = this.invokeInternal(req);
70
+ return telemetry.bindGeneratorToContext(sessionAttributes, generator);
71
+ }
72
+ async *invokeInternal(req) {
68
73
  // Derive the parent OTEL context for this invocation.
69
74
  // If this is a subagent and the parent process propagated an OTEL trace
70
75
  // context via sessionMeta.otelTraceContext, use that as the parent;
@@ -184,6 +189,7 @@ export class LangchainAgent {
184
189
  "agent.subagent": meta?.[SUBAGENT_MODE_KEY] === true,
185
190
  "agent.message_id": req.messageId,
186
191
  "agent.session_id": req.sessionId,
192
+ "session.id": req.sessionId,
187
193
  }, parentContext);
188
194
  // Create a context with the invocation span as active
189
195
  // This will be used when creating child spans (tool calls)
@@ -242,7 +248,12 @@ export class LangchainAgent {
242
248
  else if (type === "filesystem") {
243
249
  const wd = t.working_directory ??
244
250
  process.cwd();
245
- enabledTools.push(...makeFilesystemTools(wd));
251
+ const fsTools = makeFilesystemTools(wd);
252
+ // Tag filesystem tools with their group name for tool override filtering
253
+ for (const fsTool of fsTools) {
254
+ fsTool.__groupName = "filesystem";
255
+ }
256
+ enabledTools.push(...fsTools);
246
257
  }
247
258
  else if (type === "direct") {
248
259
  // Handle direct tool objects (imported in code)
@@ -264,16 +275,25 @@ export class LangchainAgent {
264
275
  if (!entry) {
265
276
  throw new Error(`Unknown built-in tool "${name}"`);
266
277
  }
278
+ const tagTool = (tool) => {
279
+ // Track which built-in group produced this tool so overrides can
280
+ // filter by the original config name (e.g. "web_search" filters
281
+ // both WebSearch and WebFetch helpers).
282
+ tool.__groupName = name;
283
+ };
267
284
  if (typeof entry === "function") {
268
285
  const result = entry();
269
286
  if (Array.isArray(result)) {
287
+ result.forEach(tagTool);
270
288
  enabledTools.push(...result);
271
289
  }
272
290
  else {
291
+ tagTool(result);
273
292
  enabledTools.push(result);
274
293
  }
275
294
  }
276
295
  else {
296
+ tagTool(entry);
277
297
  enabledTools.push(entry);
278
298
  }
279
299
  }
@@ -422,10 +442,10 @@ export class LangchainAgent {
422
442
  // Apply tool overrides if provided (Town Hall comparison feature)
423
443
  if (req.configOverrides?.tools && req.configOverrides.tools.length > 0) {
424
444
  const allowedToolNames = new Set(req.configOverrides.tools);
425
- finalTools = finalTools.filter((t) => allowedToolNames.has(t.name));
426
- _logger.debug("Applied tool override filter", {
427
- requested: req.configOverrides.tools,
428
- filtered: finalTools.map((t) => t.name),
445
+ finalTools = finalTools.filter((t) => {
446
+ const groupName = t.__groupName;
447
+ return (allowedToolNames.has(groupName ?? "") ||
448
+ allowedToolNames.has(t.name));
429
449
  });
430
450
  }
431
451
  // Create the model instance using the factory
@@ -1110,34 +1130,44 @@ export { makeSubagentsTool } from "./tools/subagent.js";
1110
1130
  function wrapToolWithTracing(originalTool, getIterationContext) {
1111
1131
  const wrappedFunc = async (input) => {
1112
1132
  const toolInputJson = JSON.stringify(input);
1113
- // Get the current iteration context so the tool span is created as a child
1133
+ // CRITICAL: Get the iteration context synchronously when the tool is invoked.
1134
+ // We must capture both the context AND the parent span at this moment.
1114
1135
  const iterationContext = getIterationContext();
1115
- const toolSpan = telemetry.startSpan("agent.tool_call", {
1116
- "tool.name": originalTool.name,
1117
- "tool.input": toolInputJson,
1118
- }, iterationContext);
1119
- // Create a context with the tool span as active
1120
- const spanContext = toolSpan
1121
- ? trace.setSpan(iterationContext, toolSpan)
1122
- : iterationContext;
1123
- try {
1124
- // Execute within the tool span's context
1125
- const result = await context.with(spanContext, () => originalTool.invoke(input));
1126
- const resultStr = typeof result === "string" ? result : JSON.stringify(result);
1127
- if (toolSpan) {
1128
- telemetry.setSpanAttributes(toolSpan, {
1129
- "tool.output": resultStr,
1130
- });
1131
- telemetry.endSpan(toolSpan);
1136
+ const iterationSpan = trace.getSpan(iterationContext);
1137
+ // Execute the ENTIRE tool operation within the iteration context.
1138
+ // Pass the iteration span explicitly as the parent to avoid relying on AsyncLocalStorage,
1139
+ // which LangChain's internals may break.
1140
+ return await context.with(iterationContext, async () => {
1141
+ // Create tool span as direct child of iteration span by explicitly passing parent context
1142
+ const toolSpan = iterationSpan
1143
+ ? telemetry.startSpan("agent.tool_call", {
1144
+ "tool.name": originalTool.name,
1145
+ "tool.input": toolInputJson,
1146
+ }, trace.setSpan(context.active(), iterationSpan))
1147
+ : null;
1148
+ // Create a context with the tool span as active
1149
+ const spanContext = toolSpan
1150
+ ? trace.setSpan(iterationContext, toolSpan)
1151
+ : iterationContext;
1152
+ try {
1153
+ // Execute within the tool span's context
1154
+ const result = await context.with(spanContext, () => originalTool.invoke(input));
1155
+ const resultStr = typeof result === "string" ? result : JSON.stringify(result);
1156
+ if (toolSpan) {
1157
+ telemetry.setSpanAttributes(toolSpan, {
1158
+ "tool.output": resultStr,
1159
+ });
1160
+ telemetry.endSpan(toolSpan);
1161
+ }
1162
+ return result;
1132
1163
  }
1133
- return result;
1134
- }
1135
- catch (error) {
1136
- if (toolSpan) {
1137
- telemetry.endSpan(toolSpan, error);
1164
+ catch (error) {
1165
+ if (toolSpan) {
1166
+ telemetry.endSpan(toolSpan, error);
1167
+ }
1168
+ throw error;
1138
1169
  }
1139
- throw error;
1140
- }
1170
+ });
1141
1171
  };
1142
1172
  // Create new tool with wrapped function
1143
1173
  const wrappedTool = tool(wrappedFunc, {
@@ -1148,5 +1178,6 @@ function wrapToolWithTracing(originalTool, getIterationContext) {
1148
1178
  // Preserve metadata
1149
1179
  wrappedTool.prettyName = originalTool.prettyName;
1150
1180
  wrappedTool.icon = originalTool.icon;
1181
+ wrappedTool.__groupName = originalTool.__groupName;
1151
1182
  return wrappedTool;
1152
1183
  }
@@ -104,7 +104,15 @@ export function makeOtelCallbacks(opts) {
104
104
  // Extract system prompt and serialize messages
105
105
  const systemPrompt = extractSystemPrompt(messages);
106
106
  const serializedMessages = serializeMessages(messages);
107
- // Close previous iteration span if it exists (tools from previous iteration are done)
107
+ // Close previous iteration span before creating new one.
108
+ // By the time this function is called for iteration N+1, all tools from
109
+ // iteration N have been INVOKED (even if still executing asynchronously).
110
+ // Tool spans are created at invocation time using getIterationContext(),
111
+ // so they're already anchored to the correct parent iteration span.
112
+ // It's safe to close iteration N's span now because:
113
+ // 1. Tool spans from iteration N were created with iteration N's context
114
+ // 2. Those tool spans execute within their own context (via context.with())
115
+ // 3. Child operations of tools (like subagents) inherit the tool span's context
108
116
  if (localIterationSpan) {
109
117
  telemetry.endSpan(localIterationSpan);
110
118
  localIterationSpan = null;
@@ -9,6 +9,29 @@ export declare const todoItemSchema: z.ZodObject<{
9
9
  }>;
10
10
  activeForm: z.ZodString;
11
11
  }, z.core.$strip>;
12
+ export declare function makeTodoWriteTool(): import("langchain").DynamicStructuredTool<z.ZodObject<{
13
+ todos: z.ZodArray<z.ZodObject<{
14
+ content: z.ZodString;
15
+ status: z.ZodEnum<{
16
+ pending: "pending";
17
+ in_progress: "in_progress";
18
+ completed: "completed";
19
+ }>;
20
+ activeForm: z.ZodString;
21
+ }, z.core.$strip>>;
22
+ }, z.core.$strip>, {
23
+ todos: {
24
+ content: string;
25
+ status: "pending" | "in_progress" | "completed";
26
+ activeForm: string;
27
+ }[];
28
+ }, {
29
+ todos: {
30
+ content: string;
31
+ status: "pending" | "in_progress" | "completed";
32
+ activeForm: string;
33
+ }[];
34
+ }, string>;
12
35
  export declare const todoWrite: import("langchain").DynamicStructuredTool<z.ZodObject<{
13
36
  todos: z.ZodArray<z.ZodObject<{
14
37
  content: z.ZodString;
@@ -6,12 +6,15 @@ export const todoItemSchema = z.object({
6
6
  status: z.enum(["pending", "in_progress", "completed"]),
7
7
  activeForm: z.string().min(1),
8
8
  });
9
- export const todoWrite = tool(({ todos }) => {
10
- // Simple implementation that confirms the todos were written
11
- return `Successfully updated todo list with ${todos.length} items`;
12
- }, {
13
- name: TODO_WRITE_TOOL_NAME,
14
- description: `Use this tool to create and manage a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.
9
+ // Factory function to create a new todoWrite tool instance for each invocation
10
+ // This prevents sharing tool instances across concurrent invocations
11
+ export function makeTodoWriteTool() {
12
+ const todoWriteTool = tool(({ todos }) => {
13
+ // Simple implementation that confirms the todos were written
14
+ return `Successfully updated todo list with ${todos.length} items`;
15
+ }, {
16
+ name: TODO_WRITE_TOOL_NAME,
17
+ description: `Use this tool to create and manage a structured task list for your current coding session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.
15
18
  It also helps the user understand the progress of the task and overall progress of their requests.
16
19
 
17
20
  ## When to Use This Tool
@@ -72,13 +75,19 @@ NOTE that you should not use this tool if there is only one trivial task to do.
72
75
  - activeForm: "Fixing authentication bug"
73
76
 
74
77
  When in doubt, use this tool. Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully.`,
75
- schema: z.object({
76
- todos: z.array(todoItemSchema),
77
- }),
78
- });
79
- todoWrite.prettyName = "Todo List";
80
- todoWrite.icon = "CheckSquare";
81
- todoWrite.verbiage = {
82
- active: "Updating to-do's",
83
- past: "Updated to-do's",
84
- };
78
+ schema: z.object({
79
+ todos: z.array(todoItemSchema),
80
+ }),
81
+ });
82
+ // Add metadata to the tool instance
83
+ todoWriteTool.prettyName = "Todo List";
84
+ todoWriteTool.icon = "CheckSquare";
85
+ todoWriteTool.verbiage = {
86
+ active: "Updating to-do's",
87
+ past: "Updated to-do's",
88
+ };
89
+ return todoWriteTool;
90
+ }
91
+ // For backwards compatibility, export a default instance
92
+ // NOTE: This instance is shared and should NOT be used in concurrent scenarios
93
+ export const todoWrite = makeTodoWriteTool();
@@ -22,6 +22,24 @@ declare class AgentTelemetry {
22
22
  private enabled;
23
23
  private serviceName;
24
24
  private baseAttributes;
25
+ /**
26
+ * Build a context that carries additional attributes merged with any existing
27
+ * attribute context.
28
+ */
29
+ createContextWithAttributes(attributes: Attributes, parentContext?: Context): Context;
30
+ /**
31
+ * Execute a function within a context that carries the provided attributes.
32
+ */
33
+ withAttributes<T>(attributes: Attributes, fn: () => T): T;
34
+ /**
35
+ * Execute an async function within a context that carries the provided attributes.
36
+ */
37
+ withAttributesAsync<T>(attributes: Attributes, fn: () => Promise<T>): Promise<T>;
38
+ /**
39
+ * Bind an async generator to a context so that every iteration runs with the
40
+ * provided attributes.
41
+ */
42
+ bindGeneratorToContext<T, R, N = unknown>(attributes: Attributes, generator: AsyncGenerator<T, R, N>): AsyncGenerator<T, R, N>;
25
43
  configure(config: TelemetryConfig): void;
26
44
  /**
27
45
  * Update base attributes that will be added to all future spans
@@ -4,12 +4,60 @@
4
4
  */
5
5
  import { context, SpanStatusCode, trace, } from "@opentelemetry/api";
6
6
  import { logs, SeverityNumber, } from "@opentelemetry/api-logs";
7
+ const ATTRIBUTE_CONTEXT_KEY = Symbol("telemetry:attributes");
7
8
  class AgentTelemetry {
8
9
  tracer = null;
9
10
  logger = null;
10
11
  enabled = false;
11
12
  serviceName = "@townco/agent";
12
13
  baseAttributes = {};
14
+ /**
15
+ * Build a context that carries additional attributes merged with any existing
16
+ * attribute context.
17
+ */
18
+ createContextWithAttributes(attributes, parentContext = context.active()) {
19
+ const existing = parentContext.getValue(ATTRIBUTE_CONTEXT_KEY) ?? {};
20
+ return parentContext.setValue(ATTRIBUTE_CONTEXT_KEY, {
21
+ ...existing,
22
+ ...attributes,
23
+ });
24
+ }
25
+ /**
26
+ * Execute a function within a context that carries the provided attributes.
27
+ */
28
+ withAttributes(attributes, fn) {
29
+ const ctx = this.createContextWithAttributes(attributes);
30
+ return context.with(ctx, fn);
31
+ }
32
+ /**
33
+ * Execute an async function within a context that carries the provided attributes.
34
+ */
35
+ async withAttributesAsync(attributes, fn) {
36
+ const ctx = this.createContextWithAttributes(attributes);
37
+ return context.with(ctx, fn);
38
+ }
39
+ /**
40
+ * Bind an async generator to a context so that every iteration runs with the
41
+ * provided attributes.
42
+ */
43
+ bindGeneratorToContext(attributes, generator) {
44
+ const ctx = this.createContextWithAttributes(attributes);
45
+ const boundNext = context.bind(ctx, generator.next.bind(generator));
46
+ const boundReturn = generator.return
47
+ ? context.bind(ctx, generator.return.bind(generator))
48
+ : undefined;
49
+ const boundThrow = generator.throw
50
+ ? context.bind(ctx, generator.throw.bind(generator))
51
+ : undefined;
52
+ return {
53
+ next: (value) => boundNext(value),
54
+ return: boundReturn ? (value) => boundReturn(value) : undefined,
55
+ throw: boundThrow ? (err) => boundThrow(err) : undefined,
56
+ [Symbol.asyncIterator]() {
57
+ return this;
58
+ },
59
+ };
60
+ }
13
61
  configure(config) {
14
62
  this.enabled = config.enabled ?? false;
15
63
  this.serviceName = config.serviceName ?? this.serviceName;
@@ -61,9 +109,11 @@ class AgentTelemetry {
61
109
  }
62
110
  // Use provided context or get the active one
63
111
  const ctx = parentContext ?? context.active();
112
+ const ctxAttributes = ctx.getValue(ATTRIBUTE_CONTEXT_KEY) ?? {};
64
113
  const span = this.tracer.startSpan(name, {
65
114
  attributes: {
66
115
  ...this.baseAttributes,
116
+ ...ctxAttributes,
67
117
  "service.name": this.serviceName,
68
118
  ...attributes,
69
119
  },