WIP: resuming

This commit is contained in:
Danny Avila 2025-12-04 08:57:13 -05:00
parent 0e850a5d5f
commit 2522cf760f
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
8 changed files with 478 additions and 85 deletions

View file

@ -18,8 +18,6 @@ const { getRoleByName } = require('~/models/Role');
const router = express.Router();
router.use(moderateText);
const checkAgentAccess = generateCheckAccess({
permissionType: PermissionTypes.AGENTS,
permissions: [Permissions.USE],
@ -30,77 +28,12 @@ const checkAgentResourceAccess = canAccessAgentFromBody({
requiredPermission: PermissionBits.VIEW,
});
/**
* @route GET /stream/:streamId
* @desc Subscribe to an ongoing generation job's SSE stream
* @access Private
*/
router.get('/stream/:streamId', requireJwtAuth, (req, res) => {
const { streamId } = req.params;
const job = GenerationJobManager.getJob(streamId);
if (!job) {
return res.status(404).json({
error: 'Stream not found',
message: 'The generation job does not exist or has expired.',
});
}
// Disable compression for SSE
res.setHeader('Content-Encoding', 'identity');
res.setHeader('Content-Type', 'text/event-stream');
res.setHeader('Cache-Control', 'no-cache, no-transform');
res.setHeader('Connection', 'keep-alive');
res.setHeader('X-Accel-Buffering', 'no');
res.flushHeaders();
logger.debug(`[AgentStream] Client subscribed to ${streamId}`);
const unsubscribe = GenerationJobManager.subscribe(
streamId,
(event) => {
if (!res.writableEnded) {
res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
if (typeof res.flush === 'function') {
res.flush();
}
}
},
(event) => {
if (!res.writableEnded) {
res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
if (typeof res.flush === 'function') {
res.flush();
}
res.end();
}
},
(error) => {
if (!res.writableEnded) {
res.write(`event: error\ndata: ${JSON.stringify({ error })}\n\n`);
if (typeof res.flush === 'function') {
res.flush();
}
res.end();
}
},
);
if (!unsubscribe) {
return res.status(404).json({ error: 'Failed to subscribe to stream' });
}
if (job.status === 'complete' || job.status === 'error' || job.status === 'aborted') {
res.write(`event: message\ndata: ${JSON.stringify({ final: true, status: job.status })}\n\n`);
res.end();
return;
}
req.on('close', () => {
logger.debug(`[AgentStream] Client disconnected from ${streamId}`);
unsubscribe();
});
});
router.use(moderateText);
router.use(checkAgentAccess);
router.use(checkAgentResourceAccess);
router.use(validateConvoAccess);
router.use(buildEndpointOption);
router.use(setHeaders);
/**
* @route POST /abort
@ -121,12 +54,6 @@ router.post('/abort', (req, res) => {
res.status(404).json({ error: 'Job not found' });
});
router.use(checkAgentAccess);
router.use(checkAgentResourceAccess);
router.use(validateConvoAccess);
router.use(buildEndpointOption);
router.use(setHeaders);
const controller = async (req, res, next) => {
await AgentController(req, res, next, initializeClient, addTitle);
};