Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320package daemon
import ( "bufio" "context" "encoding/json" "errors" "io" "net" "net/http" "strings" "unicode/utf8"
"albedo/cli/internal/daemon/protocol")
type streamBatch struct { Generation string `json:"generation"` Cursor *int64 `json:"cursor"` Events []json.RawMessage `json:"events"` Snapshot json.RawMessage `json:"snapshot"`}
func (c *ChatClient) ResetStream() { c.mu.Lock() defer c.mu.Unlock() c.progressCallIDs = nil c.afterSeq = -1 c.afterGeneration = ""}func streamFailure(kind StreamFailureKind, cause error) error { return &StreamError{Kind: kind, Cause: cause}}func classifyStreamFailure(err error) error { if err == nil { return nil } if _, ok := errors.AsType[*StreamError](err); ok { return err } if api, ok := errors.AsType[*APIError](err); ok { kind := StreamTerminal if api.StatusCode == 408 || api.StatusCode == 429 || api.StatusCode >= 500 { kind = StreamTransient } return streamFailure(kind, err) } if errors.Is(err, bufio.ErrTooLong) { return streamFailure(StreamProtocol, err) } if _, ok := errors.AsType[net.Error](err); ok { return streamFailure(StreamTransient, err) } if isConnectionError(err) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return streamFailure(StreamTransient, err) } return streamFailure(StreamTerminal, err)}func (c *ChatClient) Stream(ctx context.Context, tail int, onEvent func(StreamEvent) error) error { return c.StreamWithProgress(ctx, tail, onEvent, nil)}func (c *ChatClient) StreamWithProgress(ctx context.Context, tail int, onEvent func(StreamEvent) error, onBatchConsumed func()) error { defer func() { if ctx.Err() != nil { c.ResetStream() } }() c.mu.Lock() after, generation := c.afterSeq, c.afterGeneration c.mu.Unlock() params := protocol.GetSessionParams{Tail: new(int64(min(200, max(0, tail))))} if generation != "" { params.AfterGeneration = &generation params.AfterSeq = &after } err := scanEventStream(ctx, c.conn, operation{Name: "watch session", BuildRequest: func(base string, body io.Reader) (*http.Request, error) { return protocol.NewGetSessionRequest(base, c.agentID, ¶ms) }, Policy: readRecovery}, streamLimits{requireSSE: true, lineBytes: 1048577, errorBytes: 65536}, func(scanner *bufio.Scanner) error { return c.readStream(ctx, scanner, onBatchConsumed, onEvent) }) return classifyStreamFailure(err)}
// SSE framing is independent of network fragmentation and supports multiline data.func scanSSEFrames(ctx context.Context, scanner *bufio.Scanner, consume func([]byte) error) error { var data strings.Builder for scanner.Scan() { if err := ctx.Err(); err != nil { return err } line := scanner.Text() if line == "" { if data.Len() > 0 { payload := strings.TrimSuffix(data.String(), "\n") if !utf8.ValidString(payload) { return streamFailure(StreamProtocol, errors.New("SSE data is not UTF-8")) } if err := consume([]byte(payload)); err != nil { return err } data.Reset() } continue } if after, ok := strings.CutPrefix(line, "data:"); ok { fragment := after fragment = strings.TrimPrefix(fragment, " ") if data.Len()+len(fragment)+1 > 1048577 { return streamFailure(StreamProtocol, errors.New("SSE frame exceeds 1 MiB")) } data.WriteString(fragment) data.WriteByte('\n') } } if err := scanner.Err(); err != nil { return err } if data.Len() != 0 { return io.ErrUnexpectedEOF } return nil}func (c *ChatClient) readStream(ctx context.Context, scanner *bufio.Scanner, onBatchConsumed func(), onEvent func(StreamEvent) error) error { return scanSSEFrames(ctx, scanner, func(payload []byte) error { var batch streamBatch if err := decodeRequired(payload, &batch, "generation", "cursor", "events"); err != nil { return streamFailure(StreamProtocol, err) } if !validGeneration(batch.Generation) || batch.Cursor == nil || *batch.Cursor < 0 || batch.Events == nil || len(batch.Events) > 256 { return streamFailure(StreamProtocol, fieldError("stream batch")) } var first sessionEventEnvelope if len(batch.Events) > 0 { if err := json.Unmarshal(batch.Events[0], &first); err != nil { return streamFailure(StreamProtocol, err) } } if first.Type == "failure" { if len(batch.Events) != 1 || batch.Snapshot != nil { return streamFailure(StreamProtocol, fieldError("failure batch")) } var reason protocol.SafeReason if err := json.Unmarshal(first.Data, &reason); err != nil || reason.Code == "" { return streamFailure(StreamProtocol, fieldError("failure reason")) } failure := &APIError{Code: reason.Code, Message: reason.Detail} if failure.DaemonRestarting() { // older daemons end streams this way while they shut down return streamFailure(StreamTransient, failure) } return streamFailure(StreamTerminal, failure) } reset := first.Type == "reset" c.mu.Lock() previous, generation := c.afterSeq, c.afterGeneration c.mu.Unlock() sequence := previous var events []wireChatEvent var current []ToolProgress var outputs []protocol.ActiveOutput if reset { snapshot, err := decodeSession(batch.Snapshot) if err != nil { return streamFailure(StreamProtocol, err) } if snapshot.ID != c.agentID || snapshot.wire.Cursor.Generation != batch.Generation { return streamFailure(StreamProtocol, fieldError("snapshot identity")) } sequence = snapshot.wire.Cursor.Sequence outputs = snapshot.wire.ActiveOutput before := int64(0) if len(snapshot.wire.History.Items) > 0 { before = snapshot.wire.History.Items[0].Position } events = append(events, wireChatEvent{StreamEvent: StreamEvent{Type: EventReset, Before: before, More: snapshot.wire.History.Older != nil, Snapshot: &snapshot}}) decoded, err := decodeHistoryEntries(snapshot.wire.History.Items) if err != nil { return streamFailure(StreamProtocol, err) } hydrated, err := hydrateHistory(ctx, c.conn, c.agentID, decoded) if err != nil { return streamFailure(StreamTransient, err) } history, err := historyEvents(hydrated) if err != nil { return streamFailure(StreamProtocol, err) } for _, event := range history { events = append(events, wireChatEvent{StreamEvent: event}) } events = append(events, wireChatEvent{StreamEvent: StreamEvent{Type: EventUsage, Usage: usageValue(snapshot.wire.Usage), Replayed: true}}) for _, progress := range snapshot.wire.CurrentProgress { parsed, err := validateToolProgress(progress) if err != nil || parsed == nil { return streamFailure(StreamProtocol, fieldError("current progress")) } current = append(current, *parsed) events = append(events, wireChatEvent{StreamEvent: StreamEvent{Type: EventToolProgress, Progress: parsed}}) } } else if batch.Snapshot != nil || generation == "" || generation != batch.Generation { return streamFailure(StreamProtocol, errors.New("initial or changed generation requires a snapshot reset")) } restoredEvents := len(events) for index, raw := range batch.Events { var envelope sessionEventEnvelope if err := decodeRequired(raw, &envelope, "type", "data"); err != nil { return streamFailure(StreamProtocol, err) } if reset && index == 0 { if _, err := decodeChatEvent(raw); err != nil { return streamFailure(StreamProtocol, err) } continue } if envelope.Type == "reset" || envelope.Type == "failure" || envelope.Sequence == nil || *envelope.Sequence != sequence+1 { return streamFailure(StreamProtocol, errors.New("stream events are not contiguous")) } sequence = *envelope.Sequence event, err := decodeChatEnvelope(envelope, len(raw), &eventContentReader{ctx: ctx, conn: c.conn, session: c.agentID}) if err != nil { if api, ok := errors.AsType[*APIError](err); ok && api.Code == "history_publication_failed" { c.ResetStream() return streamFailure(StreamTransient, err) } if _, ok := errors.AsType[*StreamError](err); ok { return err } return streamFailure(StreamProtocol, err) } if event != nil { events = append(events, *event) } } if sequence != *batch.Cursor { return streamFailure(StreamProtocol, errors.New("batch cursor does not cover its events")) } active, err := c.nextProgressCallIDs(current, events, reset) if err != nil { return streamFailure(StreamProtocol, err) } stage, err := stageActiveOutput(ctx, c.conn, outputs) if err != nil { if api, ok := errors.AsType[*APIError](err); ok && (api.StatusCode == 404 || api.StatusCode == 410) { c.ResetStream() return streamFailure(StreamTransient, err) } return classifyStreamFailure(err) } prefixDelivered := !reset for index, event := range events { if !prefixDelivered && index == restoredEvents { if err := stage.deliver(ctx, onEvent); err != nil { stage.close() return streamFailure(StreamTerminal, err) } prefixDelivered = true } if err := onEvent(event.StreamEvent); err != nil { stage.close() return streamFailure(StreamTerminal, err) } } if !prefixDelivered { if err := stage.deliver(ctx, onEvent); err != nil { stage.close() return streamFailure(StreamTerminal, err) } } stage.close() c.mu.Lock() c.progressCallIDs = active c.afterSeq = *batch.Cursor c.afterGeneration = batch.Generation c.mu.Unlock() if onBatchConsumed != nil { onBatchConsumed() } return nil })}func (c *ChatClient) nextProgressCallIDs(snapshot []ToolProgress, events []wireChatEvent, reset bool) (map[string]struct{}, error) { active := map[string]struct{}{} if reset { for _, progress := range snapshot { if _, ok := active[progress.CallID]; ok { return nil, errors.New("duplicate progress call") } active[progress.CallID] = struct{}{} } } else { c.mu.Lock() for id := range c.progressCallIDs { active[id] = struct{}{} } c.mu.Unlock() } for _, event := range events { if event.Replayed { continue } switch event.Type { case EventRetry, EventUser, EventMessage, EventError, EventInterrupted, EventTurnCompleted: clear(active) case EventToolProgress: if event.Progress == nil { clear(active) } else { active[event.Progress.CallID] = struct{}{} } case EventTool: delete(active, event.ProgressCallID) } if len(active) > maxActiveToolProgress { return nil, errors.New("too many active progress calls") } } if len(active) > maxActiveToolProgress { return nil, errors.New("too many active progress calls") } return active, nil}