fix: Refresh paused HITL Redis state

This commit is contained in:
Danny Avila 2026-05-19 13:50:05 -04:00
parent 3c39cd0f65
commit 58c18e720e
2 changed files with 134 additions and 6 deletions

View file

@ -234,6 +234,48 @@ describe('RedisJobStore Integration Tests', () => {
await store.destroy();
});
test('should refresh resume state TTLs when pausing and resuming a job', async () => {
if (!ioredisClient) {
return;
}
const { RedisJobStore } = await import('../implementations/RedisJobStore');
const store = new RedisJobStore(ioredisClient, { runningTtl: 120 });
await store.initialize();
const streamId = `requires-action-ttl-${Date.now()}`;
const chunkKey = `stream:{${streamId}}:chunks`;
const runStepsKey = `stream:{${streamId}}:runsteps`;
await store.createJob(streamId, 'user-1', streamId);
await store.appendChunk(streamId, { event: 'on_message_delta', data: { text: 'hello' } });
const runSteps: Partial<Agents.RunStep>[] = [
{ id: 'step-1', runId: 'run-1', type: StepTypes.MESSAGE_CREATION, index: 0 },
];
await store.saveRunSteps(streamId, runSteps as Agents.RunStep[]);
await ioredisClient.expire(chunkKey, 30);
await ioredisClient.expire(runStepsKey, 30);
await store.updateJob(streamId, {
status: 'requires_action',
pendingAction: buildPendingAction(streamId),
});
expect(await ioredisClient.ttl(chunkKey)).toBeGreaterThan(30);
expect(await ioredisClient.ttl(runStepsKey)).toBeGreaterThan(30);
await ioredisClient.expire(chunkKey, 30);
await ioredisClient.expire(runStepsKey, 30);
await store.updateJob(streamId, { status: 'running', pendingAction: undefined });
expect(await ioredisClient.ttl(chunkKey)).toBeGreaterThan(30);
expect(await ioredisClient.ttl(runStepsKey)).toBeGreaterThan(30);
await store.destroy();
});
test('should not drop paused jobs from user tracking when cleanup sees a stale running index', async () => {
if (!ioredisClient) {
return;
@ -264,6 +306,36 @@ describe('RedisJobStore Integration Tests', () => {
await store.destroy();
});
test('should prune expired requires_action IDs during cleanup', async () => {
if (!ioredisClient) {
return;
}
const { RedisJobStore } = await import('../implementations/RedisJobStore');
const store = new RedisJobStore(ioredisClient);
await store.initialize();
const streamId = `requires-action-expired-${Date.now()}`;
const jobKey = `stream:{${streamId}}:job`;
await store.createJob(streamId, 'user-1', streamId);
await store.updateJob(streamId, {
status: 'requires_action',
pendingAction: buildPendingAction(streamId),
});
expect(await ioredisClient.smembers('stream:requires_action')).toContain(streamId);
await ioredisClient.del(jobKey);
const cleaned = await store.cleanup();
const pausedMembers = await ioredisClient.smembers('stream:requires_action');
expect(cleaned).toBeGreaterThanOrEqual(1);
expect(pausedMembers).not.toContain(streamId);
await store.destroy();
});
});
describe('Horizontal Scaling - Multi-Instance Simulation', () => {

View file

@ -301,16 +301,18 @@ export class RedisJobStore implements IJobStore {
await this.redis.srem(KEYS.runningJobs, streamId);
await this.redis.sadd(KEYS.requiresActionJobs, streamId);
await this.updateExistingJobHash(key, fields);
await this.redis.expire(key, this.ttl.running);
await this.refreshLiveJobTtls(key, streamId);
return;
}
await this.redis.eval(
'if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end redis.call("SREM", KEYS[2], ARGV[1]) redis.call("SADD", KEYS[3], ARGV[1]) redis.call("HSET", KEYS[1], unpack(ARGV, 3)) redis.call("EXPIRE", KEYS[1], tonumber(ARGV[2])) return 1',
3,
'if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end redis.call("SREM", KEYS[2], ARGV[1]) redis.call("SADD", KEYS[3], ARGV[1]) redis.call("HSET", KEYS[1], unpack(ARGV, 3)) redis.call("EXPIRE", KEYS[1], tonumber(ARGV[2])) redis.call("EXPIRE", KEYS[4], tonumber(ARGV[2])) redis.call("EXPIRE", KEYS[5], tonumber(ARGV[2])) return 1',
5,
key,
KEYS.runningJobs,
KEYS.requiresActionJobs,
KEYS.chunks(streamId),
KEYS.runSteps(streamId),
streamId,
String(this.ttl.running),
...fields,
@ -330,24 +332,34 @@ export class RedisJobStore implements IJobStore {
return;
}
await this.redis.hdel(key, 'pendingAction');
await this.redis.expire(key, this.ttl.running);
await this.refreshLiveJobTtls(key, streamId);
await this.redis.srem(KEYS.requiresActionJobs, streamId);
await this.redis.sadd(KEYS.runningJobs, streamId);
return;
}
await this.redis.eval(
'if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end redis.call("HSET", KEYS[1], unpack(ARGV, 3)) redis.call("HDEL", KEYS[1], "pendingAction") redis.call("EXPIRE", KEYS[1], tonumber(ARGV[2])) redis.call("SREM", KEYS[2], ARGV[1]) redis.call("SADD", KEYS[3], ARGV[1]) return 1',
3,
'if redis.call("EXISTS", KEYS[1]) == 0 then return 0 end redis.call("HSET", KEYS[1], unpack(ARGV, 3)) redis.call("HDEL", KEYS[1], "pendingAction") redis.call("EXPIRE", KEYS[1], tonumber(ARGV[2])) redis.call("EXPIRE", KEYS[4], tonumber(ARGV[2])) redis.call("EXPIRE", KEYS[5], tonumber(ARGV[2])) redis.call("SREM", KEYS[2], ARGV[1]) redis.call("SADD", KEYS[3], ARGV[1]) return 1',
5,
key,
KEYS.requiresActionJobs,
KEYS.runningJobs,
KEYS.chunks(streamId),
KEYS.runSteps(streamId),
streamId,
String(this.ttl.running),
...fields,
);
}
private async refreshLiveJobTtls(key: string, streamId: string): Promise<void> {
const pipeline = this.redis.pipeline();
pipeline.expire(key, this.ttl.running);
pipeline.expire(KEYS.chunks(streamId), this.ttl.running);
pipeline.expire(KEYS.runSteps(streamId), this.ttl.running);
await pipeline.exec();
}
async deleteJob(streamId: string): Promise<void> {
this.localGraphCache.delete(streamId);
this.localCollectedUsageCache.delete(streamId);
@ -478,6 +490,8 @@ export class RedisJobStore implements IJobStore {
}
}
cleaned += await this.cleanupRequiresActionIndex();
if (cleaned > 0) {
logger.debug(`[RedisJobStore] Cleaned up ${cleaned} jobs`);
}
@ -485,6 +499,48 @@ export class RedisJobStore implements IJobStore {
return cleaned;
}
private async cleanupRequiresActionIndex(): Promise<number> {
const streamIds = await this.redis.smembers(KEYS.requiresActionJobs);
let cleaned = 0;
const BATCH_SIZE = 50;
for (let i = 0; i < streamIds.length; i += BATCH_SIZE) {
const batch = streamIds.slice(i, i + BATCH_SIZE);
const results = await Promise.allSettled(
batch.map(async (streamId) => {
const job = await this.getJob(streamId);
if (!job) {
await this.redis.srem(KEYS.requiresActionJobs, streamId);
this.localGraphCache.delete(streamId);
this.localCollectedUsageCache.delete(streamId);
return 1;
}
if (job.status !== 'requires_action') {
await this.redis.srem(KEYS.requiresActionJobs, streamId);
if (job.status === 'running') {
await this.redis.sadd(KEYS.runningJobs, streamId);
}
return 1;
}
return 0;
}),
);
for (const result of results) {
if (result.status === 'fulfilled') {
cleaned += result.value;
} else {
logger.warn(`[RedisJobStore] requires_action cleanup failed for a job:`, result.reason);
}
}
}
return cleaned;
}
async getJobCount(): Promise<number> {
const [runningCount, requiresActionCount] = await Promise.all([
this.redis.scard(KEYS.runningJobs),