mirror of
https://github.com/n8n-io/n8n.git
synced 2026-08-29 01:39:24 +08:00
feat(core): Add local agent token counting (#35836)
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Robin Braumann
+97
-101
File diff suppressed because one or more lines are too long
+52
-53
File diff suppressed because one or more lines are too long
+148
-155
File diff suppressed because one or more lines are too long
+166
-169
File diff suppressed because one or more lines are too long
+97
-98
File diff suppressed because one or more lines are too long
+20
-19
@@ -25,30 +25,31 @@
|
||||
]
|
||||
},
|
||||
"status": 200,
|
||||
"response": "{\"model\":\"claude-haiku-4-5-20251001\",\"id\":\"msg_011CdF25YmR1SgVB1EfdRoVg\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Acknowledged and stored.\\n\\n**Durable Context:**\\n- Customer: Orion Basin\\n- Marker: OBSERVATION_MARKER_ORION\\n- Routing: Tier 3 Support\\n\\nReady for future interactions with this customer.\"}],\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"stop_details\":null,\"usage\":{\"input_tokens\":83,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0,\"cache_creation\":{\"ephemeral_5m_input_tokens\":0,\"ephemeral_1h_input_tokens\":0},\"output_tokens\":57,\"service_tier\":\"standard\",\"inference_geo\":\"not_available\"}}",
|
||||
"response": "{\"model\":\"claude-haiku-4-5-20251001\",\"id\":\"msg_011CdoJaZiFzLkxHGKHXgguD\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Acknowledged. Durable context stored:\\n\\n- **Customer:** Orion Basin\\n- **Marker:** OBSERVATION_MARKER_ORION\\n- **Routing rule:** Orion Basin escalations → Tier 3 Support\\n\\nReady for future turns.\"}],\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"stop_details\":null,\"usage\":{\"input_tokens\":83,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0,\"cache_creation\":{\"ephemeral_5m_input_tokens\":0,\"ephemeral_1h_input_tokens\":0},\"output_tokens\":62,\"service_tier\":\"standard\",\"inference_geo\":\"not_available\"}}",
|
||||
"rawHeaders": {
|
||||
"anthropic-ratelimit-input-tokens-limit": "10000000",
|
||||
"anthropic-ratelimit-input-tokens-remaining": "10000000",
|
||||
"anthropic-ratelimit-input-tokens-reset": "2026-07-21T10:08:58Z",
|
||||
"anthropic-ratelimit-input-tokens-reset": "2026-08-07T11:24:21Z",
|
||||
"anthropic-ratelimit-output-tokens-limit": "2000000",
|
||||
"anthropic-ratelimit-output-tokens-remaining": "2000000",
|
||||
"anthropic-ratelimit-output-tokens-reset": "2026-07-21T10:08:58Z",
|
||||
"anthropic-ratelimit-output-tokens-reset": "2026-08-07T11:24:22Z",
|
||||
"anthropic-ratelimit-requests-limit": "20000",
|
||||
"anthropic-ratelimit-requests-remaining": "19999",
|
||||
"anthropic-ratelimit-requests-reset": "2026-07-21T10:08:57Z",
|
||||
"anthropic-ratelimit-requests-reset": "2026-08-07T11:24:21Z",
|
||||
"anthropic-ratelimit-tokens-limit": "12000000",
|
||||
"anthropic-ratelimit-tokens-remaining": "12000000",
|
||||
"anthropic-ratelimit-tokens-reset": "2026-07-21T10:08:58Z",
|
||||
"anthropic-ratelimit-tokens-reset": "2026-08-07T11:24:21Z",
|
||||
"cf-cache-status": "DYNAMIC",
|
||||
"cf-ray": "a1e9684748914f59-VIE",
|
||||
"cf-ray": "a275ea0fcebba2c1-VIE",
|
||||
"connection": "keep-alive",
|
||||
"content-security-policy": "default-src 'none'; frame-ancestors 'none'",
|
||||
"content-type": "application/json",
|
||||
"date": "Tue, 21 Jul 2026 10:08:58 GMT",
|
||||
"request-id": "req_011CdF25YGQV5EQxRsfPxzsr",
|
||||
"date": "Fri, 07 Aug 2026 11:24:22 GMT",
|
||||
"request-id": "req_011CdoJaTukxmdj5cSgqDbPN",
|
||||
"server": "cloudflare",
|
||||
"server-timing": "x-originResponse;dur=2278",
|
||||
"strict-transport-security": "max-age=31536000; includeSubDomains; preload",
|
||||
"traceresponse": "00-5af887e1e2734e9dac4d8c0ce678a227-457831853787f4e5-01",
|
||||
"traceresponse": "00-9f93f02a02d9bffb75b6fbb089930db0-3cf18bebffb634d6-01",
|
||||
"vary": "Accept-Encoding",
|
||||
"x-robots-tag": "none"
|
||||
},
|
||||
@@ -73,37 +74,37 @@
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Current timestamp: 2026-06-01T12:00:00.000Z\n\nUnobserved transcript tokens: 120\n\nCurrent observation log tail:\n(empty)\n\nNew transcript delta since the last observation:\n[2026-06-01T12:00:00.000Z] user:\nIMPORTANT durable context for future turns.\nCustomer: Orion Basin.\nDurable marker exactly: OBSERVATION_MARKER_ORION.\nRouting rule exactly: Orion Basin escalations go to Tier 3 Support.\nRemember this for future conversations.\n[2026-06-01T12:00:00.001Z] assistant:\nAcknowledged and stored.\n\n**Durable Context:**\n- Customer: Orion Basin\n- Marker: OBSERVATION_MARKER_ORION\n- Routing: Tier 3 Support\n\nReady for future interactions with this customer."
|
||||
"text": "Current timestamp: 2026-06-01T12:00:00.000Z\n\nUnobserved transcript tokens: 129\n\nCurrent observation log tail:\n(empty)\n\nNew transcript delta since the last observation:\n[2026-06-01T12:00:00.000Z] user:\nIMPORTANT durable context for future turns.\nCustomer: Orion Basin.\nDurable marker exactly: OBSERVATION_MARKER_ORION.\nRouting rule exactly: Orion Basin escalations go to Tier 3 Support.\nRemember this for future conversations.\n[2026-06-01T12:00:00.001Z] assistant:\nAcknowledged. Durable context stored:\n\n- **Customer:** Orion Basin\n- **Marker:** OBSERVATION_MARKER_ORION\n- **Routing rule:** Orion Basin escalations → Tier 3 Support\n\nReady for future turns."
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"status": 200,
|
||||
"response": "{\"model\":\"claude-haiku-4-5-20251001\",\"id\":\"msg_011CdF25gr5vgAoAMWaubxu7\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"* CRITICAL (12:00) Customer is Orion Basin; durable marker is OBSERVATION_MARKER_ORION; routing rule is Orion Basin escalations go to Tier 3 Support.\"}],\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"stop_details\":null,\"usage\":{\"input_tokens\":3367,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0,\"cache_creation\":{\"ephemeral_5m_input_tokens\":0,\"ephemeral_1h_input_tokens\":0},\"output_tokens\":50,\"service_tier\":\"standard\",\"inference_geo\":\"not_available\"}}",
|
||||
"response": "{\"model\":\"claude-haiku-4-5-20251001\",\"id\":\"msg_011CdoJagcFJS4BogxnBJpS9\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"* CRITICAL (12:00) Customer: Orion Basin; durable marker: OBSERVATION_MARKER_ORION; routing rule: Orion Basin escalations go to Tier 3 Support.\"}],\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"stop_details\":null,\"usage\":{\"input_tokens\":3372,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":0,\"cache_creation\":{\"ephemeral_5m_input_tokens\":0,\"ephemeral_1h_input_tokens\":0},\"output_tokens\":50,\"service_tier\":\"standard\",\"inference_geo\":\"not_available\"}}",
|
||||
"rawHeaders": {
|
||||
"anthropic-ratelimit-input-tokens-limit": "10000000",
|
||||
"anthropic-ratelimit-input-tokens-remaining": "9998000",
|
||||
"anthropic-ratelimit-input-tokens-reset": "2026-07-21T10:08:59Z",
|
||||
"anthropic-ratelimit-input-tokens-reset": "2026-08-07T11:24:23Z",
|
||||
"anthropic-ratelimit-output-tokens-limit": "2000000",
|
||||
"anthropic-ratelimit-output-tokens-remaining": "2000000",
|
||||
"anthropic-ratelimit-output-tokens-reset": "2026-07-21T10:09:00Z",
|
||||
"anthropic-ratelimit-output-tokens-reset": "2026-08-07T11:24:23Z",
|
||||
"anthropic-ratelimit-requests-limit": "20000",
|
||||
"anthropic-ratelimit-requests-remaining": "19999",
|
||||
"anthropic-ratelimit-requests-reset": "2026-07-21T10:08:59Z",
|
||||
"anthropic-ratelimit-requests-reset": "2026-08-07T11:24:22Z",
|
||||
"anthropic-ratelimit-tokens-limit": "12000000",
|
||||
"anthropic-ratelimit-tokens-remaining": "11998000",
|
||||
"anthropic-ratelimit-tokens-reset": "2026-07-21T10:08:59Z",
|
||||
"anthropic-ratelimit-tokens-reset": "2026-08-07T11:24:23Z",
|
||||
"cf-cache-status": "DYNAMIC",
|
||||
"cf-ray": "a1e9685028b34f59-VIE",
|
||||
"cf-ray": "a275ea21b9aca2c1-VIE",
|
||||
"connection": "keep-alive",
|
||||
"content-security-policy": "default-src 'none'; frame-ancestors 'none'",
|
||||
"content-type": "application/json",
|
||||
"date": "Tue, 21 Jul 2026 10:09:00 GMT",
|
||||
"request-id": "req_011CdF25gYjBAsRKrSM4WR5h",
|
||||
"date": "Fri, 07 Aug 2026 11:24:23 GMT",
|
||||
"request-id": "req_011CdoJagFQnXXC9rmiDmwKz",
|
||||
"server": "cloudflare",
|
||||
"strict-transport-security": "max-age=31536000; includeSubDomains; preload",
|
||||
"traceresponse": "00-67ad6e180d0b8ee787122886f64349c4-4d77d51a1e151ef1-01",
|
||||
"traceresponse": "00-fe010d76ff9f95ad6595c20edc43889d-42f3d13dcec435bc-01",
|
||||
"vary": "Accept-Encoding",
|
||||
"x-robots-tag": "none"
|
||||
},
|
||||
|
||||
@@ -98,16 +98,18 @@ export type {
|
||||
ObservationLogStatus,
|
||||
ObservationLogTaskKind,
|
||||
ObservationLogTaskLockHandle,
|
||||
TokenCounter,
|
||||
} from './types';
|
||||
export type { ProviderOptions } from '@ai-sdk/provider-utils';
|
||||
export { AgentEvent } from './types';
|
||||
export type { AgentEventData, AgentEventHandler } from './types';
|
||||
export {
|
||||
estimateObservationTokens,
|
||||
OBSERVATION_LOG_MARKERS,
|
||||
OBSERVATION_LOG_STATUSES,
|
||||
} from './types';
|
||||
export {
|
||||
estimateObservationTokens,
|
||||
type TokenCounter,
|
||||
} from './runtime/model/model-token-counter';
|
||||
|
||||
export { createCancellation, isCancellation, CANCELLATION_TYPE } from './sdk/cancellation';
|
||||
export type { Cancellation } from './sdk/cancellation';
|
||||
|
||||
@@ -24,7 +24,7 @@ import {
|
||||
type DelegateSubAgentRunner,
|
||||
} from '../tools/delegate-sub-agent-tool';
|
||||
import { toAiSdkTools } from '../tools/tool-adapter';
|
||||
import { MAX_MODEL_TOOL_RESULT_CHARS } from '../tools/tool-result-guard';
|
||||
import { MAX_MODEL_TOOL_RESULT_TOKENS } from '../tools/tool-result-guard';
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Module mocks
|
||||
@@ -7383,6 +7383,14 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
}
|
||||
}
|
||||
|
||||
async function expectWithinTokenLimit(value: unknown): Promise<void> {
|
||||
const { getEncoding } = await import('@n8n/ai-utilities/tokenizer');
|
||||
const encoder = await getEncoding('o200k_base');
|
||||
expect(encoder.encode(JSON.stringify(value)).length).toBeLessThanOrEqual(
|
||||
MAX_MODEL_TOOL_RESULT_TOKENS,
|
||||
);
|
||||
}
|
||||
|
||||
function getModelToolResult(callIndex = 1): unknown {
|
||||
const call = generateText.mock.calls[callIndex][0] as {
|
||||
messages: Array<{
|
||||
@@ -7416,9 +7424,9 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
|
||||
it('bounds an oversized transformed result while preserving raw output', async () => {
|
||||
const rawOutput = {
|
||||
value: `RAW_HEAD${'r'.repeat(MAX_MODEL_TOOL_RESULT_CHARS)}RAW_TAIL`,
|
||||
value: `RAW_HEAD${'r '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}RAW_TAIL`,
|
||||
};
|
||||
const transformedOutput = `TRANSFORM_HEAD${'t'.repeat(MAX_MODEL_TOOL_RESULT_CHARS)}TRANSFORM_TAIL`;
|
||||
const transformedOutput = `TRANSFORM_HEAD${'t '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}TRANSFORM_TAIL`;
|
||||
const tool: BuiltTool = {
|
||||
name: 'large_result',
|
||||
description: 'Return a large result',
|
||||
@@ -7436,7 +7444,7 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
const result = await runtime.generate('run');
|
||||
const envelope = getModelToolResult() as TruncationEnvelope;
|
||||
|
||||
expect(JSON.stringify(envelope).length).toBeLessThanOrEqual(MAX_MODEL_TOOL_RESULT_CHARS);
|
||||
await expectWithinTokenLimit(envelope);
|
||||
expect(envelope).toMatchObject({
|
||||
_truncated: true,
|
||||
originalCharCount: JSON.stringify(transformedOutput).length,
|
||||
@@ -7487,7 +7495,9 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
description: 'Throw a large error',
|
||||
inputSchema: z.object({}),
|
||||
handler: () => {
|
||||
throw new Error(`ERROR_HEAD${'e'.repeat(MAX_MODEL_TOOL_RESULT_CHARS)}ERROR_TAIL`);
|
||||
throw new Error(
|
||||
`ERROR_HEAD${'e '.repeat(MAX_MODEL_TOOL_RESULT_TOKENS + 10_000)}ERROR_TAIL`,
|
||||
);
|
||||
},
|
||||
};
|
||||
const { runtime } = createRuntimeWithTools([tool], 1);
|
||||
@@ -7508,14 +7518,14 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
|
||||
expect(result.finishReason).toBe('stop');
|
||||
expect(toolCall?.state).toBe('rejected');
|
||||
expect(JSON.stringify(envelope).length).toBeLessThanOrEqual(MAX_MODEL_TOOL_RESULT_CHARS);
|
||||
await expectWithinTokenLimit(envelope);
|
||||
expect(envelope.head).toContain('ERROR_HEAD');
|
||||
expect(envelope.tail).toContain('ERROR_TAIL');
|
||||
});
|
||||
|
||||
it('bounds aggregate toMessage text while preserving file content', async () => {
|
||||
const fileData = Buffer.from('file').toString('base64');
|
||||
const textBlockLength = Math.floor(MAX_MODEL_TOOL_RESULT_CHARS / 2);
|
||||
const textBlockTokenCount = Math.floor(MAX_MODEL_TOOL_RESULT_TOKENS * 0.6);
|
||||
const tool: BuiltTool = {
|
||||
name: 'large_message',
|
||||
description: 'Return a large message',
|
||||
@@ -7524,9 +7534,9 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
toMessage: () => ({
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: 'text', text: `MESSAGE_HEAD${'m'.repeat(textBlockLength)}` },
|
||||
{ type: 'text', text: `MESSAGE_HEAD${'m '.repeat(textBlockTokenCount)}` },
|
||||
{ type: 'file', mediaType: 'text/plain', data: fileData },
|
||||
{ type: 'text', text: `${'n'.repeat(textBlockLength)}MESSAGE_TAIL` },
|
||||
{ type: 'text', text: `${'n '.repeat(textBlockTokenCount)}MESSAGE_TAIL` },
|
||||
],
|
||||
}),
|
||||
};
|
||||
@@ -7554,7 +7564,7 @@ describe('AgentRuntime — oversized tool results', () => {
|
||||
const envelope = parseEnvelope(textBlocks[0]?.text ?? '');
|
||||
|
||||
expect(textBlocks).toHaveLength(1);
|
||||
expect(JSON.stringify(envelope).length).toBeLessThanOrEqual(MAX_MODEL_TOOL_RESULT_CHARS);
|
||||
await expectWithinTokenLimit(envelope);
|
||||
expect(envelope.head).toContain('MESSAGE_HEAD');
|
||||
expect(envelope.tail).toContain('MESSAGE_TAIL');
|
||||
expect(fileBlock).toMatchObject({ type: 'file', mediaType: 'text/plain', data: fileData });
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { createModelTokenCounter } from '../model/model-token-counter';
|
||||
|
||||
describe('createModelTokenCounter', () => {
|
||||
it('selects the encoding for the model family', async () => {
|
||||
const text =
|
||||
'こんにちは世界。ユーザーは毎週金曜日にレポートを送信します。مرحبا بالعالم — يجب إرسال التقرير يوم الجمعة. 🚀🔐🧪';
|
||||
|
||||
expect(await createModelTokenCounter('openai/gpt-5')(text)).toBe(37);
|
||||
expect(await createModelTokenCounter('anthropic/claude-haiku-4-5')(text)).toBe(68);
|
||||
});
|
||||
|
||||
it('counts tokenizer special markers as ordinary text', async () => {
|
||||
expect(await createModelTokenCounter('openai/gpt-5')('<|endoftext|>')).toBe(7);
|
||||
});
|
||||
});
|
||||
@@ -400,6 +400,10 @@ describe('runObservationLogObserver', () => {
|
||||
|
||||
it('writes parsed observations and advances the cursor after observing', async () => {
|
||||
const store = new InMemoryMemory();
|
||||
const parentText = 'User needs the current request remembered.';
|
||||
const childText = 'Observer pipeline parsed the child row.';
|
||||
const tokenCounter = async (text: string) =>
|
||||
await Promise.resolve(text === parentText ? 7 : text === childText ? 9 : 10);
|
||||
await store.saveThread({ id: 'thread-1', resourceId: 'user-1' });
|
||||
await store.saveMessages({
|
||||
threadId: 'thread-1',
|
||||
@@ -407,19 +411,17 @@ describe('runObservationLogObserver', () => {
|
||||
messages: [message('m1', 'user', 'I need this remembered.', new Date(2026, 4, 12, 14, 30))],
|
||||
});
|
||||
|
||||
const now = new Date(2026, 4, 12, 14, 31);
|
||||
const result = await runObservationLogObserver({
|
||||
memory: store,
|
||||
observationScopeId: 'thread-1',
|
||||
observerThresholdTokens: 1,
|
||||
observationLogTailLimit: 20,
|
||||
tokenCounter: () => 10,
|
||||
now: new Date(2026, 4, 12, 14, 31),
|
||||
tokenCounter,
|
||||
now,
|
||||
observe: async () =>
|
||||
await Promise.resolve(
|
||||
[
|
||||
'* CRITICAL (14:31) User needs the current request remembered.',
|
||||
' * COMPLETION (14:31) Observer pipeline parsed the child row.',
|
||||
].join('\n'),
|
||||
[`* CRITICAL (14:31) ${parentText}`, ` * COMPLETION (14:31) ${childText}`].join('\n'),
|
||||
),
|
||||
});
|
||||
|
||||
@@ -430,13 +432,17 @@ describe('runObservationLogObserver', () => {
|
||||
expect(observations).toMatchObject([
|
||||
{
|
||||
marker: 'critical',
|
||||
text: 'User needs the current request remembered.',
|
||||
text: parentText,
|
||||
parentId: null,
|
||||
tokenCount: 7,
|
||||
createdAt: now,
|
||||
},
|
||||
{
|
||||
marker: 'completion',
|
||||
text: 'Observer pipeline parsed the child row.',
|
||||
text: childText,
|
||||
parentId: observations[0]?.id,
|
||||
tokenCount: 9,
|
||||
createdAt: new Date(now.getTime() + 1),
|
||||
},
|
||||
]);
|
||||
expect(await store.getCursor('thread-1')).toMatchObject({
|
||||
|
||||
@@ -376,6 +376,10 @@ describe('runObservationLogReflector', () => {
|
||||
observationScopeId: 'thread-1',
|
||||
reflectorThresholdTokens: 10,
|
||||
now: new Date('2026-05-12T15:00:00.000Z'),
|
||||
tokenCounter: async (text) => {
|
||||
expect(text).toBe('User compared old plan A and old plan B.');
|
||||
return await Promise.resolve(6);
|
||||
},
|
||||
reflect: async (input) => {
|
||||
expect(input.renderedObservationLog).toContain(`[${stale.id}] INFO`);
|
||||
return await Promise.resolve(
|
||||
@@ -404,6 +408,7 @@ describe('runObservationLogReflector', () => {
|
||||
marker: 'important',
|
||||
text: 'User compared old plan A and old plan B.',
|
||||
createdAt: new Date('2026-05-12T15:00:00.000Z'),
|
||||
tokenCount: 6,
|
||||
}),
|
||||
],
|
||||
},
|
||||
@@ -420,6 +425,9 @@ describe('runObservationLogReflector', () => {
|
||||
expect.objectContaining({ id: oldB.id, status: 'superseded' }),
|
||||
]),
|
||||
);
|
||||
await expect(
|
||||
store.getObservationLog({ observationScopeId: 'thread-1', status: 'active' }),
|
||||
).resolves.toMatchObject([{ tokenCount: 6 }]);
|
||||
});
|
||||
|
||||
it('warns but still applies reflection output that remains over budget', async () => {
|
||||
|
||||
@@ -76,6 +76,14 @@ describe('renderObservationLog', () => {
|
||||
expect(renderObservationLog([entry({ tokenCount: 2 })], { renderTokenBudget: 1 })).toBeNull();
|
||||
});
|
||||
|
||||
it('does not treat an invalid persisted token count as free', () => {
|
||||
expect(
|
||||
renderObservationLog([entry({ text: 'This entry does not fit.', tokenCount: 0 })], {
|
||||
renderTokenBudget: 1,
|
||||
}),
|
||||
).toBeNull();
|
||||
});
|
||||
|
||||
it('does not render a child as a root when its parent is outside the budget', () => {
|
||||
const parent = entry({ id: 'parent', tokenCount: 3 });
|
||||
const child = entry({
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
import { InMemoryMemory } from '../memory/memory-store';
|
||||
|
||||
describe('observation log store', () => {
|
||||
@@ -22,7 +21,7 @@ describe('observation log store', () => {
|
||||
parentId: null,
|
||||
status: 'active',
|
||||
supersededBy: null,
|
||||
tokenCount: estimateObservationTokens('User chose the observation log model.'),
|
||||
tokenCount: 7,
|
||||
createdAt,
|
||||
});
|
||||
|
||||
@@ -150,4 +149,29 @@ describe('observation log store', () => {
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it('persists concurrent appends to the same initially empty scope', async () => {
|
||||
const store = new InMemoryMemory();
|
||||
|
||||
const [[first], [second]] = await Promise.all([
|
||||
store.appendObservationLogEntries([
|
||||
{
|
||||
observationScopeId: 'thread-1',
|
||||
marker: 'info',
|
||||
text: 'First concurrent observation',
|
||||
},
|
||||
]),
|
||||
store.appendObservationLogEntries([
|
||||
{
|
||||
observationScopeId: 'thread-1',
|
||||
marker: 'info',
|
||||
text: 'Second concurrent observation',
|
||||
},
|
||||
]),
|
||||
]);
|
||||
|
||||
const persisted = await store.getObservationLog({ observationScopeId: 'thread-1' });
|
||||
expect(persisted).toEqual(expect.arrayContaining([first, second]));
|
||||
expect(persisted).toHaveLength(2);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import type { ModelConfig } from '../../../types/sdk/agent';
|
||||
import { getModelIdString } from '../runtime-context';
|
||||
import { getModelIdString } from '../../../utils/model';
|
||||
|
||||
describe('getModelIdString', () => {
|
||||
it('returns a plain string model id unchanged', () => {
|
||||
|
||||
@@ -7,7 +7,7 @@ import { incrementMessageCount, incrementTokenCountFromUsage } from './execution
|
||||
import { GenerateSink } from './generate-sink';
|
||||
import { hydrateFileParts } from './hydrate-file-parts';
|
||||
import type { RunOutputSink, RunServices } from './run-output-sink';
|
||||
import { RuntimeContextBuilder, getModelIdString } from './runtime-context';
|
||||
import { RuntimeContextBuilder } from './runtime-context';
|
||||
import {
|
||||
extractSettledToolCalls,
|
||||
formatMcpConnectionNote,
|
||||
@@ -51,12 +51,14 @@ import type {
|
||||
} from '../../types/sdk/agent';
|
||||
import type { AgentMessage, ContentToolCall } from '../../types/sdk/message';
|
||||
import type { JSONValue } from '../../types/utils/json';
|
||||
import { getModelIdString } from '../../utils/model';
|
||||
import { parseWithSchema } from '../../utils/parse';
|
||||
import { MemoryOrchestrator } from '../memory/memory-orchestrator';
|
||||
import type { ScopedMemoryTaskEvent } from '../memory/scoped-memory-task-runner';
|
||||
import { generateThreadTitle } from '../memory/title-generation';
|
||||
import { AgentMessageList, type SerializedMessageList } from '../model/message-list';
|
||||
import type { FetchFn } from '../model/model-factory';
|
||||
import { createModelTokenCounter } from '../model/model-token-counter';
|
||||
import {
|
||||
applyRuntimeCacheBreakpoints,
|
||||
buildInstructionPromptCacheOptions,
|
||||
@@ -193,6 +195,7 @@ export class AgentRuntime {
|
||||
|
||||
constructor(config: AgentRuntimeConfig) {
|
||||
this.config = config;
|
||||
const tokenCounter = createModelTokenCounter(config.model);
|
||||
this.telemetry = new RuntimeTelemetry(config);
|
||||
this.runId = config.runId ?? generateRunId();
|
||||
if (config.deferredTools && config.deferredTools.length > 0) {
|
||||
@@ -206,12 +209,14 @@ export class AgentRuntime {
|
||||
this.backgroundTasks,
|
||||
this.eventBus,
|
||||
this.telemetry,
|
||||
tokenCounter,
|
||||
);
|
||||
this.toolExecutor = new ToolCallExecutor({
|
||||
telemetry: this.telemetry,
|
||||
eventBus: this.eventBus,
|
||||
concurrency: config.toolCallConcurrency ?? 1,
|
||||
onCancelled: () => this.updateState({ status: 'cancelled' }),
|
||||
tokenCounter,
|
||||
});
|
||||
this.modelCost = config.modelCost;
|
||||
this.currentState = {
|
||||
|
||||
@@ -3,8 +3,9 @@ import type { LanguageModel, Output } from 'ai';
|
||||
|
||||
import type { AgentRuntimeConfig } from './agent-runtime';
|
||||
import type { AgentExecutionCounter, BuiltTool, JSONObject } from '../../types';
|
||||
import type { AgentPersistenceOptions, ExecutionOptions, ModelConfig } from '../../types/sdk/agent';
|
||||
import type { AgentPersistenceOptions, ExecutionOptions } from '../../types/sdk/agent';
|
||||
import { lockAdditionalProperties } from '../../utils/json-schema';
|
||||
import { getModelIdString } from '../../utils/model';
|
||||
import { isZodSchema } from '../../utils/zod';
|
||||
import {
|
||||
createRecallMemoryTool,
|
||||
@@ -31,21 +32,6 @@ function wrapBuiltInRules(fragments: string[]): string | undefined {
|
||||
return `<built_in_rules>\n${fragments.map((f) => `- ${f}`).join('\n')}\n</built_in_rules>`;
|
||||
}
|
||||
|
||||
/** Resolve a model config to its canonical `provider/model` id string. */
|
||||
export function getModelIdString(model: ModelConfig): string {
|
||||
if (typeof model === 'string') return model;
|
||||
if ('id' in model && typeof model.id === 'string') return model.id;
|
||||
if ('modelId' in model && typeof model.modelId === 'string') {
|
||||
const rawProvider = 'provider' in model ? String(model.provider) : 'unknown';
|
||||
// AI SDK providers stamp a dotted sub-namespace (e.g. 'anthropic.messages',
|
||||
// 'openai.chat'); strip it so the id matches the canonical 'provider/model'
|
||||
// the billing rate table is keyed on.
|
||||
const provider = rawProvider.split('.')[0];
|
||||
return `${provider}/${model.modelId}`;
|
||||
}
|
||||
return 'unknown';
|
||||
}
|
||||
|
||||
export interface StaticLoopContext {
|
||||
model: LanguageModel;
|
||||
aiProviderTools: ReturnType<typeof toAiSdkProviderTools>;
|
||||
|
||||
@@ -28,6 +28,7 @@ import type { AgentDbMessage } from '../../types/sdk/message';
|
||||
import type { ObservationLogScope, ObservationLogTaskKind } from '../../types/sdk/observation-log';
|
||||
import type { AgentRuntimeConfig } from '../loop/agent-runtime';
|
||||
import type { AgentMessageList } from '../model/message-list';
|
||||
import { estimateObservationTokens, type TokenCounter } from '../model/model-token-counter';
|
||||
import type { BackgroundTaskTracker } from '../state/background-task-tracker';
|
||||
import type { AgentEventBus } from '../state/event-bus';
|
||||
import {
|
||||
@@ -75,6 +76,7 @@ export class MemoryOrchestrator {
|
||||
private readonly backgroundTasks: BackgroundTaskTracker,
|
||||
private readonly eventBus: AgentEventBus,
|
||||
private readonly runtimeTelemetry: RuntimeTelemetry,
|
||||
private readonly tokenCounter: TokenCounter = estimateObservationTokens,
|
||||
) {}
|
||||
|
||||
async loadHistoryMessages(
|
||||
@@ -361,6 +363,7 @@ export class MemoryOrchestrator {
|
||||
observerThresholdTokens,
|
||||
observationLogTailLimit: observationalMemory.observationLogTailLimit ?? 0,
|
||||
observe,
|
||||
tokenCounter: this.tokenCounter,
|
||||
executionCounter,
|
||||
telemetry,
|
||||
}),
|
||||
@@ -382,6 +385,7 @@ export class MemoryOrchestrator {
|
||||
...scope,
|
||||
reflectorThresholdTokens,
|
||||
reflect,
|
||||
tokenCounter: this.tokenCounter,
|
||||
executionCounter,
|
||||
telemetry,
|
||||
}),
|
||||
|
||||
@@ -30,7 +30,6 @@ import type {
|
||||
} from '../../types';
|
||||
import type { AgentDbMessage } from '../../types/sdk/message';
|
||||
import type { ObservationCursor } from '../../types/sdk/observation';
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
import type {
|
||||
BuiltObservationLogStore,
|
||||
BuiltObservationLogTaskLockStore,
|
||||
@@ -43,6 +42,7 @@ import type {
|
||||
ObservationLogTaskKind,
|
||||
ObservationLogTaskLockHandle,
|
||||
} from '../../types/sdk/observation-log';
|
||||
import { estimateObservationTokens } from '../model/model-token-counter';
|
||||
|
||||
interface StoredMessage {
|
||||
message: AgentDbMessage;
|
||||
@@ -247,12 +247,25 @@ export class InMemoryMemory
|
||||
|
||||
// ── Observational memory ─────────────────────────────────────────────
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/require-await
|
||||
async appendObservationLogEntries(
|
||||
rows: NewObservationLogEntry[],
|
||||
): Promise<ObservationLogEntry[]> {
|
||||
const prepared = await Promise.all(
|
||||
rows.map(async (row) => {
|
||||
const createdAt = row.createdAt ?? new Date();
|
||||
return {
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? (await estimateObservationTokens(row.text)),
|
||||
createdAt,
|
||||
};
|
||||
}),
|
||||
);
|
||||
|
||||
const persisted: ObservationLogEntry[] = [];
|
||||
for (const row of rows) {
|
||||
for (const row of prepared) {
|
||||
const key = row.observationScopeId;
|
||||
const bucket = this.observationLogByScope.get(key) ?? [];
|
||||
const entry: ObservationLogEntry = {
|
||||
@@ -260,10 +273,10 @@ export class InMemoryMemory
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? estimateObservationTokens(row.text),
|
||||
parentId: row.parentId,
|
||||
tokenCount: row.tokenCount,
|
||||
...activeLifecycleState(),
|
||||
createdAt: row.createdAt ?? new Date(),
|
||||
createdAt: row.createdAt,
|
||||
};
|
||||
bucket.push(entry);
|
||||
this.observationLogByScope.set(key, bucket);
|
||||
|
||||
@@ -8,8 +8,8 @@ import type {
|
||||
} from './observation-log-reflector';
|
||||
import type { ModelConfig } from '../../types/sdk/agent';
|
||||
import type { MemoryTaskUsageReport } from '../../types/sdk/observation-log';
|
||||
import { getModelIdString } from '../../utils/model';
|
||||
import { incrementTokenCountFromUsage } from '../loop/execution-counter';
|
||||
import { getModelIdString } from '../loop/runtime-context';
|
||||
import { loadAi } from '../model/lazy-ai';
|
||||
import { createModel } from '../model/model-factory';
|
||||
import { toTokenUsage } from '../streaming/stream';
|
||||
|
||||
@@ -12,10 +12,9 @@ import type {
|
||||
ObservationLogMarker,
|
||||
ObservationLogObserveFn,
|
||||
ObservationLogObserverInput,
|
||||
TokenCounter,
|
||||
} from '../../types/sdk/observation-log';
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
import type { BuiltTelemetry } from '../../types/telemetry';
|
||||
import { estimateObservationTokens, type TokenCounter } from '../model/model-token-counter';
|
||||
|
||||
export type { ObservationLogObserveFn, ObservationLogObserverInput };
|
||||
|
||||
@@ -182,7 +181,7 @@ export async function runObservationLogObserver(
|
||||
|
||||
const tokenCounter = opts.tokenCounter ?? estimateObservationTokens;
|
||||
const transcript = renderObserverTranscript(deltaMessages);
|
||||
const tokenCount = tokenCounter(transcript);
|
||||
const tokenCount = await tokenCounter(transcript);
|
||||
if (tokenCount < opts.observerThresholdTokens) {
|
||||
return { status: 'skipped', reason: 'below-threshold', tokenCount };
|
||||
}
|
||||
@@ -213,15 +212,28 @@ export async function runObservationLogObserver(
|
||||
opts.onMalformedLine?.(line);
|
||||
}
|
||||
|
||||
const prepared = await Promise.all(
|
||||
parsed.entries.map(async (entry) => {
|
||||
const text = redactText(entry.text).text;
|
||||
return {
|
||||
marker: entry.marker,
|
||||
parentIndex: entry.parentIndex,
|
||||
text,
|
||||
tokenCount: await tokenCounter(text),
|
||||
};
|
||||
}),
|
||||
);
|
||||
|
||||
const inserted: ObservationLogEntry[] = [];
|
||||
for (const entry of parsed.entries) {
|
||||
for (const entry of prepared) {
|
||||
const parentId = entry.parentIndex === null ? null : (inserted[entry.parentIndex]?.id ?? null);
|
||||
const [row] = await memory.appendObservationLogEntries([
|
||||
{
|
||||
observationScopeId,
|
||||
marker: entry.marker,
|
||||
text: redactText(entry.text).text,
|
||||
text: entry.text,
|
||||
parentId,
|
||||
tokenCount: entry.tokenCount,
|
||||
createdAt: new Date(now.getTime() + inserted.length),
|
||||
},
|
||||
]);
|
||||
|
||||
@@ -12,10 +12,10 @@ import type {
|
||||
ObservationLogMerge,
|
||||
ObservationLogReflection,
|
||||
ObservationLogReflectionResult,
|
||||
TokenCounter,
|
||||
} from '../../types/sdk/observation-log';
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
import { getStoredObservationTokenCount } from '../../types/sdk/observation-log';
|
||||
import type { BuiltTelemetry } from '../../types/telemetry';
|
||||
import { estimateObservationTokens, type TokenCounter } from '../model/model-token-counter';
|
||||
|
||||
export type { ObservationLogReflectFn, ObservationLogReflectorInput };
|
||||
|
||||
@@ -172,7 +172,7 @@ export async function runObservationLogReflector(
|
||||
observationScopeId,
|
||||
order: 'asc',
|
||||
});
|
||||
const tokenCount = countObservationTokens(activeObservationLog, tokenCounter);
|
||||
const tokenCount = countObservationTokens(activeObservationLog);
|
||||
if (tokenCount <= reflectorThresholdTokens) {
|
||||
return { status: 'skipped', reason: 'below-threshold', tokenCount };
|
||||
}
|
||||
@@ -193,15 +193,19 @@ export async function runObservationLogReflector(
|
||||
activeObservationLog,
|
||||
withCreatedAt(parseObservationLogReflectionJson(output), now),
|
||||
);
|
||||
const reflection = {
|
||||
const reflection: ObservationLogReflection = {
|
||||
...normalized,
|
||||
merge: normalized.merge.map((merge) => ({ ...merge, text: redactText(merge.text).text })),
|
||||
merge: await Promise.all(
|
||||
normalized.merge.map(async (merge) => {
|
||||
const text = redactText(merge.text).text;
|
||||
return { ...merge, text, tokenCount: await tokenCounter(text) };
|
||||
}),
|
||||
),
|
||||
};
|
||||
const result = await memory.applyObservationLogReflection({ observationScopeId }, reflection);
|
||||
|
||||
const remainingTokenCount = countObservationTokens(
|
||||
await memory.getActiveObservationLog({ observationScopeId }),
|
||||
tokenCounter,
|
||||
);
|
||||
const overBudgetAfterReflection = remainingTokenCount > reflectorThresholdTokens;
|
||||
if (overBudgetAfterReflection) {
|
||||
@@ -349,16 +353,8 @@ function withCreatedAt(reflection: ObservationLogReflection, now: Date): Observa
|
||||
};
|
||||
}
|
||||
|
||||
function countObservationTokens(
|
||||
entries: ObservationLogEntry[],
|
||||
tokenCounter: TokenCounter,
|
||||
): number {
|
||||
return entries.reduce((total, entry) => total + observationTokenCount(entry, tokenCounter), 0);
|
||||
}
|
||||
|
||||
function observationTokenCount(entry: ObservationLogEntry, tokenCounter: TokenCounter): number {
|
||||
if (Number.isFinite(entry.tokenCount) && entry.tokenCount > 0) return entry.tokenCount;
|
||||
return tokenCounter(entry.text);
|
||||
function countObservationTokens(entries: ObservationLogEntry[]): number {
|
||||
return entries.reduce((total, entry) => total + getStoredObservationTokenCount(entry), 0);
|
||||
}
|
||||
|
||||
function compareEntries(a: ObservationLogEntry, b: ObservationLogEntry): number {
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
import type {
|
||||
ObservationLogEntry,
|
||||
ObservationLogMarker,
|
||||
TokenCounter,
|
||||
import {
|
||||
getStoredObservationTokenCount,
|
||||
type ObservationLogEntry,
|
||||
type ObservationLogMarker,
|
||||
} from '../../types/sdk/observation-log';
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
|
||||
const MARKER_LABELS: Record<ObservationLogMarker, string> = {
|
||||
critical: 'CRITICAL',
|
||||
@@ -19,7 +18,6 @@ const MARKER_LEGEND =
|
||||
|
||||
export interface RenderObservationLogOptions {
|
||||
renderTokenBudget?: number;
|
||||
tokenCounter?: TokenCounter;
|
||||
}
|
||||
|
||||
function compareEntries(a: ObservationLogEntry, b: ObservationLogEntry): number {
|
||||
@@ -34,11 +32,6 @@ function formatObservationTime(date: Date): string {
|
||||
return `${hours}:${minutes}`;
|
||||
}
|
||||
|
||||
function observationTokenCount(entry: ObservationLogEntry, tokenCounter: TokenCounter): number {
|
||||
if (Number.isFinite(entry.tokenCount) && entry.tokenCount > 0) return entry.tokenCount;
|
||||
return tokenCounter(entry.text);
|
||||
}
|
||||
|
||||
function renderBullet(entry: ObservationLogEntry, indent = ''): string {
|
||||
return `${indent}* ${MARKER_LABELS[entry.marker]} (${formatObservationTime(entry.createdAt)}) ${entry.text}`;
|
||||
}
|
||||
@@ -48,13 +41,12 @@ export function renderObservationLog(
|
||||
options: RenderObservationLogOptions = {},
|
||||
): string | null {
|
||||
const activeEntries = entries.filter((entry) => entry.status === 'active').sort(compareEntries);
|
||||
const tokenCounter = options.tokenCounter ?? estimateObservationTokens;
|
||||
const renderTokenBudget = options.renderTokenBudget;
|
||||
let remainingTokens = renderTokenBudget ?? Number.POSITIVE_INFINITY;
|
||||
|
||||
const included = new Set<string>();
|
||||
for (const entry of activeEntries) {
|
||||
const tokenCount = observationTokenCount(entry, tokenCounter);
|
||||
const tokenCount = getStoredObservationTokenCount(entry);
|
||||
if (tokenCount > remainingTokens) continue;
|
||||
included.add(entry.id);
|
||||
remainingTokens -= tokenCount;
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import type { ModelConfig } from '../../types/sdk/agent';
|
||||
import { getModelIdString } from '../../utils/model';
|
||||
|
||||
export type TokenCounter = (text: string) => number | Promise<number>;
|
||||
|
||||
function createTokenCounter(encoding: 'cl100k_base' | 'o200k_base'): TokenCounter {
|
||||
return async (text) => {
|
||||
if (text.length === 0) return 0;
|
||||
const { getEncoding } = await import('@n8n/ai-utilities/tokenizer');
|
||||
const encoder = await getEncoding(encoding);
|
||||
return encoder.encode(text, [], []).length;
|
||||
};
|
||||
}
|
||||
|
||||
export const estimateObservationTokens = createTokenCounter('cl100k_base');
|
||||
|
||||
const estimateOpenAiTokens = createTokenCounter('o200k_base');
|
||||
|
||||
export function createModelTokenCounter(model: ModelConfig): TokenCounter {
|
||||
const modelId = getModelIdString(model);
|
||||
return modelId.startsWith('openai/') ? estimateOpenAiTokens : estimateObservationTokens;
|
||||
}
|
||||
@@ -31,6 +31,7 @@ import { incrementToolCallCount } from '../loop/execution-counter';
|
||||
import { stringifyError } from '../loop/runtime-helpers';
|
||||
import type { AgentMessageList } from '../model/message-list';
|
||||
import { normalizeToolInputForModel } from '../model/messages';
|
||||
import type { TokenCounter } from '../model/model-token-counter';
|
||||
import type { AgentEventBus } from '../state/event-bus';
|
||||
import type { RuntimeTelemetry } from '../telemetry/runtime-telemetry';
|
||||
|
||||
@@ -181,6 +182,7 @@ export interface ToolCallExecutorDeps {
|
||||
concurrency: number;
|
||||
/** Invoked when a run is aborted mid-batch so the runtime can set cancelled state. */
|
||||
onCancelled: () => void;
|
||||
tokenCounter: TokenCounter;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -663,7 +665,7 @@ export class ToolCallExecutor {
|
||||
const builtTool = toolMap.get(toolName);
|
||||
|
||||
if (!builtTool) {
|
||||
return this.toolError(params, new Error(`Tool ${toolName} not found`));
|
||||
return await this.toolError(params, new Error(`Tool ${toolName} not found`));
|
||||
}
|
||||
|
||||
// Already settled by the AI SDK (e.g. provider-executed tools): emit the
|
||||
@@ -743,14 +745,14 @@ export class ToolCallExecutor {
|
||||
this.deps.onCancelled();
|
||||
return this.buildCancelledOutcome(params, 'Run aborted');
|
||||
}
|
||||
return this.toolError(params, error as Error);
|
||||
return await this.toolError(params, error as Error);
|
||||
}
|
||||
|
||||
if (isSuspendedToolResult(toolResult)) {
|
||||
return await this.buildSuspendedOutcome(params, builtTool, toolResult);
|
||||
}
|
||||
|
||||
return this.buildSuccessOutcome(params, builtTool, input, toolResult);
|
||||
return await this.buildSuccessOutcome(params, builtTool, input, toolResult);
|
||||
}
|
||||
|
||||
private async runCancellationCleanup(
|
||||
@@ -847,7 +849,7 @@ export class ToolCallExecutor {
|
||||
}
|
||||
|
||||
/** Emit a failed ToolExecutionEnd, record the error on the list, return an error outcome. */
|
||||
private toolError(params: ProcessToolCallParams, error: unknown): ToolCallOutcome {
|
||||
private async toolError(params: ProcessToolCallParams, error: unknown): Promise<ToolCallOutcome> {
|
||||
this.eventBus.emit({
|
||||
type: AgentEvent.ToolExecutionEnd,
|
||||
toolCallId: params.toolCallId,
|
||||
@@ -855,7 +857,10 @@ export class ToolCallExecutor {
|
||||
result: error,
|
||||
isError: true,
|
||||
});
|
||||
params.list.setToolCallError(params.toolCallId, guardToolErrorForModel(stringifyError(error)));
|
||||
params.list.setToolCallError(
|
||||
params.toolCallId,
|
||||
await guardToolErrorForModel(stringifyError(error), this.deps.tokenCounter),
|
||||
);
|
||||
return { outcome: 'error', error };
|
||||
}
|
||||
|
||||
@@ -924,7 +929,7 @@ export class ToolCallExecutor {
|
||||
if (!result.success) {
|
||||
return {
|
||||
ok: false,
|
||||
outcome: this.toolError(params, new Error(`Invalid tool input: ${result.error}`)),
|
||||
outcome: await this.toolError(params, new Error(`Invalid tool input: ${result.error}`)),
|
||||
};
|
||||
}
|
||||
return { ok: true, input: result.data as JSONValue };
|
||||
@@ -983,13 +988,19 @@ export class ToolCallExecutor {
|
||||
if (builtTool.suspendSchema) {
|
||||
const parseResult = await parseWithSchema(builtTool.suspendSchema, toolResult.payload);
|
||||
if (!parseResult.success) {
|
||||
return this.toolError(params, new Error(`Invalid suspend payload: ${parseResult.error}`));
|
||||
return await this.toolError(
|
||||
params,
|
||||
new Error(`Invalid suspend payload: ${parseResult.error}`),
|
||||
);
|
||||
}
|
||||
toolResult.payload = parseResult.data as JSONValue;
|
||||
}
|
||||
const resumeSchema = getToolResumeJsonSchema(builtTool, toolResult.resumeSchema);
|
||||
if (!resumeSchema) {
|
||||
return this.toolError(params, new Error(`Tool ${params.toolName} has no resume schema`));
|
||||
return await this.toolError(
|
||||
params,
|
||||
new Error(`Tool ${params.toolName} has no resume schema`),
|
||||
);
|
||||
}
|
||||
return {
|
||||
outcome: 'suspended',
|
||||
@@ -1000,12 +1011,12 @@ export class ToolCallExecutor {
|
||||
}
|
||||
|
||||
/** Apply toModelOutput, emit ToolExecutionEnd, build the success outcome. */
|
||||
private buildSuccessOutcome(
|
||||
private async buildSuccessOutcome(
|
||||
params: ProcessToolCallParams,
|
||||
builtTool: BuiltTool,
|
||||
input: JSONValue,
|
||||
toolResult: unknown,
|
||||
): ToolCallOutcome {
|
||||
): Promise<ToolCallOutcome> {
|
||||
const { toolCallId, toolName, list } = params;
|
||||
|
||||
// Apply toModelOutput transform before emitting the success event.
|
||||
@@ -1015,9 +1026,9 @@ export class ToolCallExecutor {
|
||||
try {
|
||||
modelResult = builtTool.toModelOutput ? builtTool.toModelOutput(toolResult) : toolResult;
|
||||
} catch (error) {
|
||||
return this.toolError(params, error);
|
||||
return await this.toolError(params, error);
|
||||
}
|
||||
const guardedResult = guardToolResultForModel(modelResult);
|
||||
const guardedResult = await guardToolResultForModel(modelResult, this.deps.tokenCounter);
|
||||
|
||||
this.eventBus.emit({
|
||||
type: AgentEvent.ToolExecutionEnd,
|
||||
@@ -1031,7 +1042,7 @@ export class ToolCallExecutor {
|
||||
|
||||
const customMessage = builtTool.toMessage?.(toolResult);
|
||||
const guardedCustomMessage = customMessage
|
||||
? guardToolMessageForModel(customMessage)
|
||||
? await guardToolMessageForModel(customMessage, this.deps.tokenCounter)
|
||||
: undefined;
|
||||
if (guardedCustomMessage) {
|
||||
list.addResponse([guardedCustomMessage]);
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
import { toJsonValue } from '@n8n/utils/json/to-json-value';
|
||||
|
||||
import type { AgentMessage, MessageContent } from '../../types/sdk/message';
|
||||
import { estimateObservationTokens } from '../../types/sdk/observation-log';
|
||||
import type { JSONObject, JSONValue } from '../../types/utils/json';
|
||||
import { estimateObservationTokens, type TokenCounter } from '../model/model-token-counter';
|
||||
|
||||
export const MAX_MODEL_TOOL_RESULT_TOKENS = 50_000;
|
||||
export const MAX_MODEL_TOOL_RESULT_CHARS = MAX_MODEL_TOOL_RESULT_TOKENS * 4;
|
||||
// BPE tokens cannot exceed UTF-8 bytes; reserve room for provider message framing.
|
||||
const SMALL_PAYLOAD_OVERHEAD_BYTES = 1_024;
|
||||
|
||||
interface TruncatedToolResult extends JSONObject {
|
||||
_truncated: true;
|
||||
@@ -21,33 +22,52 @@ export interface GuardedToolResult {
|
||||
truncated: boolean;
|
||||
}
|
||||
|
||||
export function guardToolResultForModel(output: unknown): GuardedToolResult {
|
||||
export async function guardToolResultForModel(
|
||||
output: unknown,
|
||||
tokenCounter: TokenCounter = estimateObservationTokens,
|
||||
): Promise<GuardedToolResult> {
|
||||
const historyOutput = toJsonValue(output);
|
||||
const serialized = JSON.stringify(historyOutput);
|
||||
|
||||
if (estimateObservationTokens(serialized) <= MAX_MODEL_TOOL_RESULT_TOKENS) {
|
||||
if (isClearlyWithinTokenLimit(serialized)) {
|
||||
return { historyOutput, wireOutput: output, truncated: false };
|
||||
}
|
||||
|
||||
const truncated = buildTruncationEnvelope(serialized);
|
||||
const tokenCount = await tokenCounter(serialized);
|
||||
if (tokenCount <= MAX_MODEL_TOOL_RESULT_TOKENS) {
|
||||
return { historyOutput, wireOutput: output, truncated: false };
|
||||
}
|
||||
|
||||
const truncated = await buildTruncationEnvelope(serialized, tokenCount, tokenCounter);
|
||||
return { historyOutput: truncated, wireOutput: truncated, truncated: true };
|
||||
}
|
||||
|
||||
export function guardToolErrorForModel(errorText: string): string {
|
||||
const guarded = guardToolResultForModel(errorText);
|
||||
export async function guardToolErrorForModel(
|
||||
errorText: string,
|
||||
tokenCounter: TokenCounter = estimateObservationTokens,
|
||||
): Promise<string> {
|
||||
const guarded = await guardToolResultForModel(errorText, tokenCounter);
|
||||
return guarded.truncated ? JSON.stringify(guarded.historyOutput) : errorText;
|
||||
}
|
||||
|
||||
export function guardToolMessageForModel(message: AgentMessage): AgentMessage {
|
||||
export async function guardToolMessageForModel(
|
||||
message: AgentMessage,
|
||||
tokenCounter: TokenCounter = estimateObservationTokens,
|
||||
): Promise<AgentMessage> {
|
||||
if (!('content' in message)) return message;
|
||||
|
||||
const textBlocks = message.content.filter((block) => block.type === 'text');
|
||||
if (textBlocks.length === 0) return message;
|
||||
|
||||
const serialized = JSON.stringify(textBlocks.map(({ text }) => text));
|
||||
if (estimateObservationTokens(serialized) <= MAX_MODEL_TOOL_RESULT_TOKENS) return message;
|
||||
if (isClearlyWithinTokenLimit(serialized)) return message;
|
||||
|
||||
const replacement = JSON.stringify(buildTruncationEnvelope(serialized));
|
||||
const tokenCount = await tokenCounter(serialized);
|
||||
if (tokenCount <= MAX_MODEL_TOOL_RESULT_TOKENS) return message;
|
||||
|
||||
const replacement = JSON.stringify(
|
||||
await buildTruncationEnvelope(serialized, tokenCount, tokenCounter),
|
||||
);
|
||||
let replacedText = false;
|
||||
const content = message.content.flatMap((block): MessageContent[] => {
|
||||
if (block.type !== 'text') return [block];
|
||||
@@ -60,31 +80,42 @@ export function guardToolMessageForModel(message: AgentMessage): AgentMessage {
|
||||
return { ...message, content };
|
||||
}
|
||||
|
||||
function buildTruncationEnvelope(serialized: string): TruncatedToolResult {
|
||||
function isClearlyWithinTokenLimit(serialized: string): boolean {
|
||||
return (
|
||||
Buffer.byteLength(serialized, 'utf8') + SMALL_PAYLOAD_OVERHEAD_BYTES <=
|
||||
MAX_MODEL_TOOL_RESULT_TOKENS
|
||||
);
|
||||
}
|
||||
|
||||
async function buildTruncationEnvelope(
|
||||
serialized: string,
|
||||
originalTokenCount: number,
|
||||
tokenCounter: TokenCounter,
|
||||
): Promise<TruncatedToolResult> {
|
||||
const base = {
|
||||
_truncated: true,
|
||||
originalCharCount: serialized.length,
|
||||
estimatedTokenCount: estimateObservationTokens(serialized),
|
||||
estimatedTokenCount: originalTokenCount,
|
||||
} as const;
|
||||
let low = 0;
|
||||
let high = Math.min(serialized.length, MAX_MODEL_TOOL_RESULT_CHARS);
|
||||
let best: TruncatedToolResult = { ...base, head: '', tail: '' };
|
||||
let excerptLength = Math.floor(
|
||||
(serialized.length * MAX_MODEL_TOOL_RESULT_TOKENS * 0.9) / originalTokenCount,
|
||||
);
|
||||
|
||||
while (low <= high) {
|
||||
const excerptLength = Math.floor((low + high) / 2);
|
||||
while (true) {
|
||||
const { head, tail } = splitHeadAndTail(serialized, excerptLength);
|
||||
const candidate: TruncatedToolResult = { ...base, head, tail };
|
||||
|
||||
if (JSON.stringify(candidate).length <= MAX_MODEL_TOOL_RESULT_CHARS) {
|
||||
best = candidate;
|
||||
low = excerptLength + 1;
|
||||
} else {
|
||||
high = excerptLength - 1;
|
||||
const candidateTokenCount = await tokenCounter(JSON.stringify(candidate));
|
||||
if (candidateTokenCount <= MAX_MODEL_TOOL_RESULT_TOKENS || excerptLength === 0) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
return best;
|
||||
excerptLength = Math.min(
|
||||
excerptLength - 1,
|
||||
Math.floor((excerptLength * MAX_MODEL_TOOL_RESULT_TOKENS * 0.9) / candidateTokenCount),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function splitHeadAndTail(value: string, excerptLength: number): { head: string; tail: string } {
|
||||
const headLength = Math.ceil(excerptLength / 2);
|
||||
const tailLength = Math.floor(excerptLength / 2);
|
||||
|
||||
@@ -145,10 +145,8 @@ export type {
|
||||
ObservationLogStatus,
|
||||
ObservationLogTaskKind,
|
||||
ObservationLogTaskLockHandle,
|
||||
TokenCounter,
|
||||
} from './sdk/observation-log';
|
||||
export {
|
||||
estimateObservationTokens,
|
||||
OBSERVATION_LOG_MARKERS,
|
||||
OBSERVATION_LOG_STATUSES,
|
||||
} from './sdk/observation-log';
|
||||
|
||||
@@ -68,9 +68,12 @@ export interface ObservationLogReflectionResult {
|
||||
inserted: ObservationLogEntry[];
|
||||
}
|
||||
|
||||
export type TokenCounter = (text: string) => number;
|
||||
|
||||
export const estimateObservationTokens: TokenCounter = (text) => Math.ceil(text.length / 4);
|
||||
export function getStoredObservationTokenCount(
|
||||
entry: Pick<ObservationLogEntry, 'text' | 'tokenCount'>,
|
||||
): number {
|
||||
if (Number.isFinite(entry.tokenCount) && entry.tokenCount > 0) return entry.tokenCount;
|
||||
return Buffer.byteLength(entry.text, 'utf8');
|
||||
}
|
||||
|
||||
export interface ObservationLogObserverInput {
|
||||
observationScopeId: string;
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
import type { ModelConfig } from '../types/sdk/agent';
|
||||
|
||||
/** Resolve a model config to its canonical `provider/model` id string. */
|
||||
export function getModelIdString(model: ModelConfig): string {
|
||||
if (typeof model === 'string') return model;
|
||||
if ('id' in model && typeof model.id === 'string') return model.id;
|
||||
if ('modelId' in model && typeof model.modelId === 'string') {
|
||||
const rawProvider = 'provider' in model ? String(model.provider) : 'unknown';
|
||||
const provider = rawProvider.split('.')[0];
|
||||
return `${provider}/${model.modelId}`;
|
||||
}
|
||||
return 'unknown';
|
||||
}
|
||||
@@ -124,7 +124,11 @@ const MODELS_DEV_FIXTURE = {
|
||||
},
|
||||
};
|
||||
|
||||
const SENSITIVE_RESPONSE_HEADERS = ['anthropic-organization-id'];
|
||||
const SENSITIVE_RESPONSE_HEADERS = [
|
||||
'anthropic-organization-id',
|
||||
'anthropic-workspace-id',
|
||||
'set-cookie',
|
||||
];
|
||||
const ENCODING_RESPONSE_HEADERS = ['content-encoding', 'content-length', 'transfer-encoding'];
|
||||
|
||||
function sanitizeCassette(defs: nock.Definition[]): nock.Definition[] {
|
||||
|
||||
@@ -25,6 +25,9 @@
|
||||
"http-proxy-agent": [
|
||||
"dist/esm/utils/http-proxy-agent.d.ts"
|
||||
],
|
||||
"tokenizer": [
|
||||
"dist/esm/utils/tokenizer/tiktoken.d.ts"
|
||||
],
|
||||
"web-search": [
|
||||
"dist/esm/web-search/index.d.ts"
|
||||
],
|
||||
@@ -69,6 +72,11 @@
|
||||
"import": "./dist/esm/utils/http-proxy-agent.js",
|
||||
"require": "./dist/cjs/utils/http-proxy-agent.js"
|
||||
},
|
||||
"./tokenizer": {
|
||||
"types": "./dist/esm/utils/tokenizer/tiktoken.d.ts",
|
||||
"import": "./dist/esm/utils/tokenizer/tiktoken.js",
|
||||
"require": "./dist/cjs/utils/tokenizer/tiktoken.js"
|
||||
},
|
||||
"./web-search": {
|
||||
"types": "./dist/esm/web-search/index.d.ts",
|
||||
"import": "./dist/esm/web-search/index.js",
|
||||
|
||||
@@ -24,7 +24,6 @@ import type { AgentResourceRepository } from '../../repositories/agent-resource.
|
||||
import type { AgentThreadRepository } from '../../repositories/agent-thread.repository';
|
||||
import { N8nMemory } from '../n8n-memory';
|
||||
|
||||
const estimateObservationTokens = (text: string) => Math.ceil(text.length / 4);
|
||||
type N8nMemoryImplementation = ReturnType<N8nMemory['getImplementation']>;
|
||||
|
||||
describe('N8nMemory', () => {
|
||||
@@ -795,7 +794,7 @@ describe('N8nMemory', () => {
|
||||
marker: 'important',
|
||||
text: 'hello',
|
||||
parentId: null,
|
||||
tokenCount: estimateObservationTokens('hello'),
|
||||
tokenCount: 1,
|
||||
status: 'active',
|
||||
supersededBy: null,
|
||||
}),
|
||||
@@ -929,7 +928,7 @@ describe('N8nMemory', () => {
|
||||
marker: 'important',
|
||||
text: 'Merged observation',
|
||||
parentId: null,
|
||||
tokenCount: estimateObservationTokens('Merged observation'),
|
||||
tokenCount: 3,
|
||||
status: 'active',
|
||||
supersededBy: null,
|
||||
}),
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import {
|
||||
activeLifecycleState,
|
||||
droppedLifecycleState,
|
||||
estimateObservationTokens,
|
||||
normalizeObservationLogReflection,
|
||||
hashEpisodicMemoryContent,
|
||||
hashEpisodicMemoryEvidence,
|
||||
@@ -68,8 +69,6 @@ import { AgentObservationRepository } from '../repositories/agent-observation.re
|
||||
import { AgentResourceRepository } from '../repositories/agent-resource.repository';
|
||||
import { AgentThreadRepository } from '../repositories/agent-thread.repository';
|
||||
|
||||
const estimateObservationTokens = (text: string) => Math.ceil(text.length / 4);
|
||||
|
||||
@Service()
|
||||
export class N8nMemory {
|
||||
constructor(
|
||||
@@ -325,17 +324,19 @@ export class N8nMemoryImpl
|
||||
): Promise<ObservationLogEntry[]> {
|
||||
if (rows.length === 0) return [];
|
||||
|
||||
const entities: AgentObservationEntity[] = rows.map((row) =>
|
||||
this.observationRepository.create({
|
||||
agentId: this.agentId,
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? estimateObservationTokens(row.text),
|
||||
...activeLifecycleState(),
|
||||
createdAt: row.createdAt,
|
||||
}),
|
||||
const entities: AgentObservationEntity[] = await Promise.all(
|
||||
rows.map(async (row) =>
|
||||
this.observationRepository.create({
|
||||
agentId: this.agentId,
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? (await estimateObservationTokens(row.text)),
|
||||
...activeLifecycleState(),
|
||||
createdAt: row.createdAt,
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
const saved = await this.observationRepository.save(entities);
|
||||
@@ -425,17 +426,19 @@ export class N8nMemoryImpl
|
||||
);
|
||||
const inserted = normalized.merge.length
|
||||
? await repo.save(
|
||||
normalized.merge.map((entry) =>
|
||||
repo.create({
|
||||
agentId: this.agentId,
|
||||
observationScopeId: scope.observationScopeId,
|
||||
marker: entry.marker,
|
||||
text: entry.text,
|
||||
parentId: entry.parentId ?? null,
|
||||
tokenCount: entry.tokenCount ?? estimateObservationTokens(entry.text),
|
||||
...activeLifecycleState(),
|
||||
createdAt: entry.createdAt,
|
||||
}),
|
||||
await Promise.all(
|
||||
normalized.merge.map(async (entry) =>
|
||||
repo.create({
|
||||
agentId: this.agentId,
|
||||
observationScopeId: scope.observationScopeId,
|
||||
marker: entry.marker,
|
||||
text: entry.text,
|
||||
parentId: entry.parentId ?? null,
|
||||
tokenCount: entry.tokenCount ?? (await estimateObservationTokens(entry.text)),
|
||||
...activeLifecycleState(),
|
||||
createdAt: entry.createdAt,
|
||||
}),
|
||||
),
|
||||
),
|
||||
)
|
||||
: [];
|
||||
|
||||
+3
-2
@@ -53,7 +53,7 @@ describe('TypeORMObservationLogStore', () => {
|
||||
marker: 'important',
|
||||
text: 'User prefers Slack alerts',
|
||||
parentId: null,
|
||||
tokenCount: 8,
|
||||
tokenCount: 4,
|
||||
status: 'active',
|
||||
supersededBy: null,
|
||||
createdAt,
|
||||
@@ -66,7 +66,7 @@ describe('TypeORMObservationLogStore', () => {
|
||||
marker: 'important',
|
||||
text: 'User prefers Slack alerts',
|
||||
parentId: null,
|
||||
tokenCount: 8,
|
||||
tokenCount: 4,
|
||||
status: 'active',
|
||||
supersededBy: null,
|
||||
createdAt,
|
||||
@@ -93,6 +93,7 @@ describe('TypeORMObservationLogStore', () => {
|
||||
}),
|
||||
]);
|
||||
expect(active).toEqual(inserted);
|
||||
expect(observationRepo.create).toHaveBeenCalledWith(expect.objectContaining({ tokenCount: 4 }));
|
||||
});
|
||||
|
||||
it('stores and reads observation cursors', async () => {
|
||||
|
||||
@@ -57,16 +57,18 @@ export class TypeORMObservationLogStore
|
||||
): Promise<ObservationLogEntry[]> {
|
||||
if (rows.length === 0) return [];
|
||||
|
||||
const entities = rows.map((row) =>
|
||||
this.observationRepo.create({
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? estimateObservationTokens(row.text),
|
||||
...activeLifecycleState(),
|
||||
createdAt: row.createdAt,
|
||||
}),
|
||||
const entities = await Promise.all(
|
||||
rows.map(async (row) =>
|
||||
this.observationRepo.create({
|
||||
observationScopeId: row.observationScopeId,
|
||||
marker: row.marker,
|
||||
text: row.text,
|
||||
parentId: row.parentId ?? null,
|
||||
tokenCount: row.tokenCount ?? (await estimateObservationTokens(row.text)),
|
||||
...activeLifecycleState(),
|
||||
createdAt: row.createdAt,
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
const saved = await this.observationRepo.save(entities);
|
||||
@@ -163,16 +165,18 @@ export class TypeORMObservationLogStore
|
||||
|
||||
const inserted = normalized.merge.length
|
||||
? await repo.save(
|
||||
normalized.merge.map((entry) =>
|
||||
repo.create({
|
||||
observationScopeId: scope.observationScopeId,
|
||||
marker: entry.marker,
|
||||
text: entry.text,
|
||||
parentId: entry.parentId ?? null,
|
||||
tokenCount: entry.tokenCount ?? estimateObservationTokens(entry.text),
|
||||
...activeLifecycleState(),
|
||||
createdAt: entry.createdAt,
|
||||
}),
|
||||
await Promise.all(
|
||||
normalized.merge.map(async (entry) =>
|
||||
repo.create({
|
||||
observationScopeId: scope.observationScopeId,
|
||||
marker: entry.marker,
|
||||
text: entry.text,
|
||||
parentId: entry.parentId ?? null,
|
||||
tokenCount: entry.tokenCount ?? (await estimateObservationTokens(entry.text)),
|
||||
...activeLifecycleState(),
|
||||
createdAt: entry.createdAt,
|
||||
}),
|
||||
),
|
||||
),
|
||||
)
|
||||
: [];
|
||||
|
||||
Reference in New Issue
Block a user