@xynehq/jaf
Version:
Juspay Agent Framework - A purely functional agent framework with immutable state and composable tools
539 lines • 23.6 kB
JavaScript
/**
* A2A In-Memory Task Provider for JAF
* Pure functional in-memory storage for A2A tasks
*/
import { createA2ATaskNotFoundError, createA2ATaskStorageError, createSuccess, createFailure } from '../types.js';
import { serializeA2ATask, deserializeA2ATask, sanitizeTask } from '../serialization.js';
/**
* Helper function to convert A2ATaskStorage to A2ATaskSerialized
*/
const convertStorageToSerialized = (stored) => ({
taskId: stored.taskId,
contextId: stored.contextId,
state: stored.state,
taskData: stored.taskData,
statusMessage: stored.statusMessage,
createdAt: stored.createdAt.toISOString(),
updatedAt: stored.updatedAt.toISOString(),
metadata: stored.metadata ? JSON.stringify(stored.metadata) : undefined
});
/**
* Create an in-memory A2A task provider
*/
export const createA2AInMemoryTaskProvider = async (config) => {
// Initialize immutable state
let state = {
tasks: new Map(),
contextIndex: new Map(),
stateIndex: new Map(),
config,
stats: {
totalTasks: 0,
createdAt: new Date()
}
};
// Pure function to update state
const updateState = (newState) => {
state = newState;
};
// Pure function to add task to indices
const addToIndices = (contextIndex, stateIndex, taskId, contextId, taskState) => {
// Update context index
const contextTasks = contextIndex.get(contextId) || new Set();
const newContextIndex = new Map(contextIndex);
newContextIndex.set(contextId, new Set([...contextTasks, taskId]));
// Update state index
const stateTasks = stateIndex.get(taskState) || new Set();
const newStateIndex = new Map(stateIndex);
newStateIndex.set(taskState, new Set([...stateTasks, taskId]));
return { contextIndex: newContextIndex, stateIndex: newStateIndex };
};
// Pure function to remove task from indices
const removeFromIndices = (contextIndex, stateIndex, taskId, contextId, taskState) => {
// Update context index
const contextTasks = contextIndex.get(contextId);
const newContextIndex = new Map(contextIndex);
if (contextTasks) {
const updatedContextTasks = new Set(contextTasks);
updatedContextTasks.delete(taskId);
if (updatedContextTasks.size > 0) {
newContextIndex.set(contextId, updatedContextTasks);
}
else {
newContextIndex.delete(contextId);
}
}
// Update state index
const stateTasks = stateIndex.get(taskState);
const newStateIndex = new Map(stateIndex);
if (stateTasks) {
const updatedStateTasks = new Set(stateTasks);
updatedStateTasks.delete(taskId);
if (updatedStateTasks.size > 0) {
newStateIndex.set(taskState, updatedStateTasks);
}
else {
newStateIndex.delete(taskState);
}
}
return { contextIndex: newContextIndex, stateIndex: newStateIndex };
};
// Pure function to check storage limits
const checkStorageLimits = (currentTasks) => {
if (currentTasks.size >= config.maxTasks) {
return createFailure(createA2ATaskStorageError('store', 'in-memory', undefined, new Error(`Maximum task limit reached: ${config.maxTasks}`)));
}
return createSuccess(undefined);
};
// Create the provider object with proper implementation
return {
storeTask: async (task, metadata) => {
try {
// Validate and sanitize task
const sanitizeResult = sanitizeTask(task);
if (!sanitizeResult.success) {
return sanitizeResult;
}
// Check for duplicate task ID
if (state.tasks.has(task.id)) {
return createFailure(createA2ATaskStorageError('store', 'in-memory', task.id, new Error(`Task with ID ${task.id} already exists`)));
}
// Check storage limits
const limitsResult = checkStorageLimits(state.tasks);
if (!limitsResult.success) {
return limitsResult;
}
// Serialize task
const serializeResult = serializeA2ATask(sanitizeResult.data, metadata);
if (!serializeResult.success) {
return serializeResult;
}
const serializedTask = serializeResult.data;
const taskStorage = {
taskId: serializedTask.taskId,
contextId: serializedTask.contextId,
state: serializedTask.state,
taskData: serializedTask.taskData,
statusMessage: serializedTask.statusMessage,
createdAt: new Date(serializedTask.createdAt),
updatedAt: new Date(serializedTask.updatedAt),
expiresAt: metadata?.expiresAt,
metadata: metadata ? { ...metadata } : undefined
};
// Add to storage and indices
const newTasks = new Map(state.tasks);
newTasks.set(task.id, taskStorage);
const { contextIndex, stateIndex } = addToIndices(state.contextIndex, state.stateIndex, task.id, task.contextId, task.status.state);
updateState({
...state,
tasks: newTasks,
contextIndex,
stateIndex,
stats: {
...state.stats,
totalTasks: newTasks.size
}
});
return createSuccess(undefined);
}
catch (error) {
return createFailure(createA2ATaskStorageError('store', 'in-memory', task.id, error));
}
},
getTask: async (taskId) => {
try {
const stored = state.tasks.get(taskId);
if (!stored) {
return createSuccess(null);
}
// Check expiration
if (stored.expiresAt && stored.expiresAt < new Date()) {
return createSuccess(null);
}
const deserializeResult = deserializeA2ATask(convertStorageToSerialized(stored));
if (!deserializeResult.success) {
return deserializeResult;
}
return createSuccess(deserializeResult.data);
}
catch (error) {
return createFailure(createA2ATaskStorageError('get', 'in-memory', taskId, error));
}
},
updateTask: async (task, metadata) => {
try {
const existing = state.tasks.get(task.id);
if (!existing) {
return createFailure(createA2ATaskNotFoundError(task.id, 'in-memory'));
}
// Validate and sanitize task
const sanitizeResult = sanitizeTask(task);
if (!sanitizeResult.success) {
return sanitizeResult;
}
// Serialize updated task
const mergedMetadata = { ...existing.metadata, ...metadata };
const serializeResult = serializeA2ATask(sanitizeResult.data, mergedMetadata);
if (!serializeResult.success) {
return serializeResult;
}
const serializedTask = serializeResult.data;
const updatedStorage = {
...existing,
state: serializedTask.state,
taskData: serializedTask.taskData,
statusMessage: serializedTask.statusMessage,
updatedAt: new Date(serializedTask.updatedAt),
metadata: mergedMetadata
};
// Update storage
const newTasks = new Map(state.tasks);
newTasks.set(task.id, updatedStorage);
// Update indices if state changed
let contextIndex = state.contextIndex;
let stateIndex = state.stateIndex;
if (existing.state !== task.status.state) {
// Remove from old state index and add to new
const removeResult = removeFromIndices(contextIndex, stateIndex, task.id, task.contextId, existing.state);
const addResult = addToIndices(removeResult.contextIndex, removeResult.stateIndex, task.id, task.contextId, task.status.state);
contextIndex = addResult.contextIndex;
stateIndex = addResult.stateIndex;
}
updateState({
...state,
tasks: newTasks,
contextIndex,
stateIndex
});
return createSuccess(undefined);
}
catch (error) {
return createFailure(createA2ATaskStorageError('update', 'in-memory', task.id, error));
}
},
updateTaskStatus: async (taskId, newState, statusMessage, timestamp) => {
try {
const existing = state.tasks.get(taskId);
if (!existing) {
return createFailure(createA2ATaskNotFoundError(taskId, 'in-memory'));
}
// Deserialize existing task
const deserializeResult = deserializeA2ATask(convertStorageToSerialized(existing));
if (!deserializeResult.success) {
return deserializeResult;
}
const task = deserializeResult.data;
// Update task status
const updatedTask = {
...task,
status: {
...task.status,
state: newState,
message: statusMessage || task.status.message,
timestamp: timestamp || new Date().toISOString()
}
};
// Update task directly to avoid circular reference
const sanitizeResult = sanitizeTask(updatedTask);
if (!sanitizeResult.success) {
return sanitizeResult;
}
const mergedMetadata = { ...existing.metadata };
const serializeResult = serializeA2ATask(sanitizeResult.data, mergedMetadata);
if (!serializeResult.success) {
return serializeResult;
}
const serializedTask = serializeResult.data;
const updatedStorage = {
...existing,
state: serializedTask.state,
taskData: serializedTask.taskData,
statusMessage: serializedTask.statusMessage,
updatedAt: new Date(serializedTask.updatedAt),
metadata: mergedMetadata
};
const newTasks = new Map(state.tasks);
newTasks.set(taskId, updatedStorage);
let contextIndex = state.contextIndex;
let stateIndex = state.stateIndex;
if (existing.state !== updatedTask.status.state) {
const removeResult = removeFromIndices(contextIndex, stateIndex, taskId, updatedTask.contextId, existing.state);
const addResult = addToIndices(removeResult.contextIndex, removeResult.stateIndex, taskId, updatedTask.contextId, updatedTask.status.state);
contextIndex = addResult.contextIndex;
stateIndex = addResult.stateIndex;
}
updateState({ ...state, tasks: newTasks, contextIndex, stateIndex });
return createSuccess(undefined);
}
catch (error) {
return createFailure(createA2ATaskStorageError('update-status', 'in-memory', taskId, error));
}
},
findTasks: async (query) => {
try {
let taskIds = new Set();
// Start with all tasks or filter by context/state
if (query.contextId && query.state) {
// Both context and state - find intersection
const contextTasks = state.contextIndex.get(query.contextId) || new Set();
const stateTasks = state.stateIndex.get(query.state) || new Set();
// Get intersection of both sets
taskIds = new Set([...contextTasks].filter(id => stateTasks.has(id)));
}
else if (query.contextId) {
const contextTasks = state.contextIndex.get(query.contextId);
if (contextTasks) {
taskIds = new Set(contextTasks);
}
}
else if (query.state) {
const stateTasks = state.stateIndex.get(query.state);
if (stateTasks) {
taskIds = new Set(stateTasks);
}
}
else {
taskIds = new Set(state.tasks.keys());
}
// Filter by additional criteria
const results = [];
for (const taskId of taskIds) {
if (query.taskId && taskId !== query.taskId)
continue;
const stored = state.tasks.get(taskId);
if (!stored)
continue;
// Check expiration
if (stored.expiresAt && stored.expiresAt < new Date())
continue;
// Date filtering
if (query.since && stored.createdAt < query.since)
continue;
if (query.until && stored.createdAt > query.until)
continue;
const deserializeResult = deserializeA2ATask(convertStorageToSerialized(stored));
if (deserializeResult.success) {
results.push(deserializeResult.data);
}
}
// Apply pagination
const offset = query.offset || 0;
const limit = query.limit || results.length;
const paginatedResults = results
.sort((a, b) => new Date(b.status.timestamp || '').getTime() - new Date(a.status.timestamp || '').getTime())
.slice(offset, offset + limit);
return createSuccess(paginatedResults);
}
catch (error) {
return createFailure(createA2ATaskStorageError('find', 'in-memory', undefined, error));
}
},
getTasksByContext: async (contextId, limit) => {
try {
const contextTasks = state.contextIndex.get(contextId);
if (!contextTasks) {
return createSuccess([]);
}
const results = [];
for (const taskId of contextTasks) {
if (limit && results.length >= limit)
break;
const stored = state.tasks.get(taskId);
if (!stored)
continue;
if (stored.expiresAt && stored.expiresAt < new Date())
continue;
const deserializeResult = deserializeA2ATask(convertStorageToSerialized(stored));
if (deserializeResult.success) {
results.push(deserializeResult.data);
}
}
const sortedResults = results.sort((a, b) => new Date(b.status.timestamp || '').getTime() - new Date(a.status.timestamp || '').getTime());
return createSuccess(sortedResults);
}
catch (error) {
return createFailure(createA2ATaskStorageError('get-by-context', 'in-memory', undefined, error));
}
},
deleteTask: async (taskId) => {
try {
const existing = state.tasks.get(taskId);
if (!existing) {
return createSuccess(false);
}
// Remove from storage
const newTasks = new Map(state.tasks);
newTasks.delete(taskId);
// Remove from indices
const { contextIndex, stateIndex } = removeFromIndices(state.contextIndex, state.stateIndex, taskId, existing.contextId, existing.state);
updateState({
...state,
tasks: newTasks,
contextIndex,
stateIndex,
stats: {
...state.stats,
totalTasks: newTasks.size
}
});
return createSuccess(true);
}
catch (error) {
return createFailure(createA2ATaskStorageError('delete', 'in-memory', taskId, error));
}
},
deleteTasksByContext: async (contextId) => {
try {
const contextTasks = state.contextIndex.get(contextId);
if (!contextTasks) {
return createSuccess(0);
}
let deletedCount = 0;
const newTasks = new Map(state.tasks);
let contextIndex = state.contextIndex;
let stateIndex = state.stateIndex;
for (const taskId of contextTasks) {
const existing = newTasks.get(taskId);
if (existing) {
newTasks.delete(taskId);
const removeResult = removeFromIndices(contextIndex, stateIndex, taskId, existing.contextId, existing.state);
contextIndex = removeResult.contextIndex;
stateIndex = removeResult.stateIndex;
deletedCount++;
}
}
updateState({
...state,
tasks: newTasks,
contextIndex,
stateIndex,
stats: {
...state.stats,
totalTasks: newTasks.size
}
});
return createSuccess(deletedCount);
}
catch (error) {
return createFailure(createA2ATaskStorageError('delete-by-context', 'in-memory', undefined, error));
}
},
cleanupExpiredTasks: async () => {
try {
const now = new Date();
let cleanedCount = 0;
const newTasks = new Map(state.tasks);
let contextIndex = state.contextIndex;
let stateIndex = state.stateIndex;
for (const [taskId, stored] of state.tasks) {
if (stored.expiresAt && stored.expiresAt < now) {
newTasks.delete(taskId);
const removeResult = removeFromIndices(contextIndex, stateIndex, taskId, stored.contextId, stored.state);
contextIndex = removeResult.contextIndex;
stateIndex = removeResult.stateIndex;
cleanedCount++;
}
}
updateState({
...state,
tasks: newTasks,
contextIndex,
stateIndex,
stats: {
...state.stats,
totalTasks: newTasks.size
}
});
return createSuccess(cleanedCount);
}
catch (error) {
return createFailure(createA2ATaskStorageError('cleanup', 'in-memory', undefined, error));
}
},
getTaskStats: async (contextId) => {
try {
const tasksByState = {
submitted: 0,
working: 0,
'input-required': 0,
completed: 0,
canceled: 0,
failed: 0,
rejected: 0,
'auth-required': 0,
unknown: 0
};
let totalTasks = 0;
let oldestTask;
let newestTask;
const tasksToCount = contextId
? (state.contextIndex.get(contextId) || new Set())
: new Set(state.tasks.keys());
for (const taskId of tasksToCount) {
const stored = state.tasks.get(taskId);
if (!stored)
continue;
// Skip expired tasks
if (stored.expiresAt && stored.expiresAt < new Date())
continue;
totalTasks++;
tasksByState[stored.state]++;
if (!oldestTask || stored.createdAt < oldestTask) {
oldestTask = stored.createdAt;
}
if (!newestTask || stored.createdAt > newestTask) {
newestTask = stored.createdAt;
}
}
return createSuccess({
totalTasks,
tasksByState,
oldestTask,
newestTask
});
}
catch (error) {
return createFailure(createA2ATaskStorageError('stats', 'in-memory', undefined, error));
}
},
healthCheck: async () => {
try {
const startTime = Date.now();
// Simple health check - verify we can access storage
const taskCount = state.tasks.size;
const latencyMs = Date.now() - startTime;
return createSuccess({
healthy: true,
latencyMs,
error: undefined
});
}
catch (error) {
return createSuccess({
healthy: false,
error: error.message
});
}
},
close: async () => {
try {
// Clear all data for cleanup
updateState({
tasks: new Map(),
contextIndex: new Map(),
stateIndex: new Map(),
config,
stats: {
totalTasks: 0,
createdAt: new Date()
}
});
return createSuccess(undefined);
}
catch (error) {
return createFailure(createA2ATaskStorageError('close', 'in-memory', undefined, error));
}
}
};
};
//# sourceMappingURL=in-memory.js.map