File size: 3,961 Bytes
cd8bd0a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
/**
 * A2A SSE Streaming Support
 *
 * Provides SSE event formatting for A2A `message/stream` responses.
 * Features: heartbeat (15s), chunk emission, metadata final event, cancellation.
 */

import type { A2ATask } from "./taskManager";

export interface SSEChunkEvent {
  jsonrpc: "2.0";
  method: "message/stream";
  params: {
    task: { id: string; state: string };
    chunk?: { type: string; content: string };
    metadata?: Record<string, unknown>;
  };
}

/**
 * Format an SSE event line.
 */
export function formatSSE(event: SSEChunkEvent): string {
  return `data: ${JSON.stringify(event)}\n\n`;
}

/**
 * Create a chunk event for streaming text content.
 */
export function createChunkEvent(taskId: string, content: string): string {
  return formatSSE({
    jsonrpc: "2.0",
    method: "message/stream",
    params: {
      task: { id: taskId, state: "working" },
      chunk: { type: "text", content },
    },
  });
}

/**
 * Create the final completion event with metadata.
 */
export function createCompletionEvent(taskId: string, metadata: Record<string, unknown>): string {
  return formatSSE({
    jsonrpc: "2.0",
    method: "message/stream",
    params: {
      task: { id: taskId, state: "completed" },
      metadata,
    },
  });
}

/**
 * Create a heartbeat event to keep the connection alive.
 */
export function createHeartbeat(taskId: string): string {
  return `: heartbeat ${new Date().toISOString()}\n\n`;
}

/**
 * Create a failure event.
 */
export function createFailureEvent(taskId: string, error: string): string {
  return formatSSE({
    jsonrpc: "2.0",
    method: "message/stream",
    params: {
      task: { id: taskId, state: "failed" },
      metadata: { error },
    },
  });
}

/**
 * SSE response headers for A2A streaming.
 */
export const SSE_HEADERS = {
  "Content-Type": "text/event-stream",
  "Cache-Control": "no-cache, no-transform",
  Connection: "keep-alive",
  "X-Accel-Buffering": "no",
} as const;

/**
 * Create a streaming SSE handler that wraps a fetch-based LLM call.
 * Returns a ReadableStream suitable for a Response object.
 */
export function createA2AStream(
  task: A2ATask,
  executeSkill: (
    task: A2ATask
  ) => Promise<{ artifacts: Array<{ content: string }>; metadata: Record<string, unknown> }>,
  abortSignal?: AbortSignal,
  lifecycle?: {
    onStart?: () => void;
    onEnd?: () => void;
  }
): ReadableStream<Uint8Array> {
  const encoder = new TextEncoder();

  return new ReadableStream({
    async start(controller) {
      lifecycle?.onStart?.();
      // Heartbeat interval
      const heartbeatInterval = setInterval(() => {
        try {
          controller.enqueue(encoder.encode(createHeartbeat(task.id)));
        } catch {
          /* stream closed */
        }
      }, 15_000);

      try {
        // Check for cancellation
        if (abortSignal?.aborted) {
          controller.enqueue(encoder.encode(createFailureEvent(task.id, "Cancelled")));
          controller.close();
          return;
        }

        // Execute the skill
        const result = await executeSkill(task);

        // Emit content as chunks (simulated streaming for non-streaming skills)
        for (const artifact of result.artifacts) {
          if (abortSignal?.aborted) break;
          controller.enqueue(encoder.encode(createChunkEvent(task.id, artifact.content)));
        }

        if (abortSignal?.aborted) {
          controller.enqueue(encoder.encode(createFailureEvent(task.id, "Cancelled")));
          return;
        }

        // Emit completion with metadata
        controller.enqueue(encoder.encode(createCompletionEvent(task.id, result.metadata)));
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        controller.enqueue(encoder.encode(createFailureEvent(task.id, msg)));
      } finally {
        clearInterval(heartbeatInterval);
        lifecycle?.onEnd?.();
        controller.close();
      }
    },
  });
}