feat(cli): give retries one identity
Carry operation handles through lost acknowledgements and stop recovery when receipts expire. Keep prompt cancellation scoped to its own submission.
123456567857 unmodified lines66676869727374707177727379747576777879808182838487888586878891899091929394959697989910010150 unmodified lines1521531541551501561571581591602 unmodified lines16316416516616116716816917017117217316817417517617717827 unmodified linespackage app
import ( "context" "crypto/rand" "encoding/hex" "errors" "fmt" "slices" "time"57 unmodified lines answer, err := awaitReply(ctx, client, prompt) if err == nil { return PromptResult{SessionID: sessionID, Answer: answer}, nil } if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](err); uncertain { return PromptResult{}, fmt.Errorf("%w; the session is %s", err, sessionID) } switch { case errors.Is(ctx.Err(), context.DeadlineExceeded): return PromptResult{}, fmt.Errorf("timed out after %s; %v; the session is %s", timeout, err, sessionID) return PromptResult{}, fmt.Errorf("timed out after %s; %w; the session is %s", timeout, err, sessionID) case ctx.Err() != nil: return PromptResult{}, fmt.Errorf("interrupted; %v; the session is %s", err, sessionID) return PromptResult{}, fmt.Errorf("interrupted; %w; the session is %s", err, sessionID) } if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](err); uncertain { return PromptResult{}, fmt.Errorf("%w; the session is %s", err, sessionID) } return PromptResult{}, err}
// awaitReply follows the logical turn containing this submission. The actor// publishes membership before any worker events, and completion after them.func awaitReply(ctx context.Context, client *daemon.ChatClient, prompt string) (answer string, failure error) { var identity [16]byte if _, err := rand.Read(identity[:]); err != nil { handle, err := client.PrepareTurn(prompt, nil, false) if err != nil { return "", err } submissionID := hex.EncodeToString(identity[:]) submissionID := handle.ID defer func() { if failure != nil && ctx.Err() != nil { if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](failure); !uncertain { uncertain := &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: failure} failure = fmt.Errorf("%w; %v", uncertain, failure) } } }() streamCtx, cancel := context.WithCancel(ctx) defer cancel() submitted := false defer func() {50 unmodified lines } // A failed request can already have reached the actor. Cleanup always uses // its identity, so it cannot interrupt a different client's turn. submitted = true if _, err := client.SendSubmission(ctx, prompt, submissionID); err != nil { if _, err := client.SubmitOperation(ctx, handle); err != nil { return "", err } turnID := "" // Discard the initial reset already acknowledged before submission.2 unmodified lines select { case <-ctx.Done(): return "", ctx.Err() case err := <-streamDone: return "", fmt.Errorf("lost the session stream before the turn finished: %v", err) return "", &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: fmt.Errorf("lost the session stream before the turn finished: %v", err)} case event := <-events: if event.Type == daemon.EventReset && initialReset { initialReset = false continue } if event.Type == daemon.EventReset { return "", errors.New("session stream reset before the submission completed; outcome uncertain") return "", &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: errors.New("session stream reset before the submission completed")} } if event.Type == daemon.EventTurnMembership && slices.Contains(event.SubmissionIDs, submissionID) { turnID = event.TurnID }27 unmodified lines40 unmodified lines4142434445464748492 unmodified lines5253545555565758596042 unmodified lines10310410510610710810911011110711211311411511627 unmodified lines40 unmodified lines switch request.URL.Path { case "/sessions/test/events": var payload struct { SubmissionID string `json:"submissionId"` OperationID string `json:"operationId"` } if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { t.Error(err) return2 unmodified lines t.Error("missing submission ID") } submitted <- payload.SubmissionID writer.WriteHeader(http.StatusAccepted) _, _ = writer.Write([]byte(`{"ok":true,"queued":true}`)) _ = json.NewEncoder(writer).Encode(map[string]any{"ok": true, "queued": true, "operationId": payload.OperationID}) case "/sessions/test/stream": _, _ = fmt.Fprint(writer, "data: {\"cursor\":0,\"events\":[{\"type\":\"reset\"}]}\n\n") writer.(http.Flusher).Flush() id := <-submitted42 unmodified lines defer cancel() server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { switch request.URL.Path { case "/sessions/test/events": var payload struct { OperationID string `json:"operationId"` } _ = json.NewDecoder(request.Body).Decode(&payload) writer.WriteHeader(http.StatusAccepted) _, _ = writer.Write([]byte(`{"ok":true,"queued":true}`)) _ = json.NewEncoder(writer).Encode(map[string]any{"ok": true, "queued": true, "operationId": payload.OperationID}) close(submitted) case "/sessions/test/stream": _, _ = fmt.Fprint(writer, "data: {\"cursor\":0,\"events\":[{\"type\":\"reset\"}]}\n\n") writer.(http.Flusher).Flush()27 unmodified lines48 unmodified lines4950515253535454555657585927 unmodified lines8788899090919293949591929394959798969798991001011021031041051061071081091101041111121131141061071151161171181091101191201211221231241251121261271281291161301311321331341351361371389 unmodified lines148149150151134135136137138139152153154155423 unmodified linesExpand all48 unmodified lines return "workspace folder not found: " + e.Workspace}
type SendResult struct { OperationID string `json:"operationId"` OK bool `json:"ok"` Queued bool `json:"queued"` OK bool `json:"ok"` Queued bool `json:"queued"`}
type ChatClient struct { conn *Connection27 unmodified linesfunc (c *ChatClient) agentPath(path string) string { return fmt.Sprintf("/sessions/%s%s", url.PathEscape(c.agentID), path)}
func (c *ChatClient) submitPayload(ctx context.Context, payload map[string]any) (*SendResult, error) { reqCtx, cancel := context.WithTimeout(ctx, 20*time.Second) defer cancel()
if c.clientID != "" { payload["clientId"] = c.clientIDfunc (c *ChatClient) PrepareTurn(content string, image *ImageAttachment, continuation bool) (*OperationHandle, error) { payload := map[string]any{"clientId": c.clientID, "content": content, "type": "user"} if continuation { payload = map[string]any{"clientId": c.clientID, "type": "continue"} }
result, err := Submit(reqCtx, c.conn, c.agentID, payload) if image != nil { payload["image"] = image } return NewSubmission(c.agentID, payload)}func (c *ChatClient) PrepareSkill(name, args string) (*OperationHandle, error) { return NewSubmission(c.agentID, map[string]any{"clientId": c.clientID, "type": "skill", "name": name, "arguments": args})}func (c *ChatClient) SubmitOperation(ctx context.Context, handle *OperationHandle) (*SendResult, error) { result, err := SubmitOperation(ctx, c.conn, handle) if err != nil { return nil, err } return &result, nil}
func (c *ChatClient) ResolveOperation(ctx context.Context, handle *OperationHandle) (OperationReceipt, error) { return ResolveOperation(ctx, c.conn, handle)}func (c *ChatClient) Send(ctx context.Context, content string, image *ImageAttachment) (*SendResult, error) { payload := map[string]any{ "content": content, handle, err := c.PrepareTurn(content, image, false) if err != nil { return nil, err } if image != nil { payload["image"] = image return c.SubmitOperation(ctx, handle)}func (c *ChatClient) Continue(ctx context.Context) (*SendResult, error) { handle, err := c.PrepareTurn("", nil, true) if err != nil { return nil, err } return c.submitPayload(ctx, payload) return c.SubmitOperation(ctx, handle)}
func (c *ChatClient) SendSubmission(ctx context.Context, content, submissionID string) (*SendResult, error) { return c.submitPayload(ctx, map[string]any{"content": content, "submissionId": submissionID}) handle, err := NewSubmission(c.agentID, map[string]any{"content": content, "clientId": c.clientID, "submissionId": submissionID}) if err != nil { return nil, err } return c.SubmitOperation(ctx, handle)}
func (c *ChatClient) CancelSubmission(ctx context.Context, submissionID string) (string, error) { result, err := RequestOperation[struct {9 unmodified lines return "", fmt.Errorf("invalid cancellation outcome %q", result.Outcome) }}
func (c *ChatClient) Continue(ctx context.Context) (*SendResult, error) { return c.submitPayload(ctx, map[string]any{ "type": "continue", })}
func (c *ChatClient) Interrupt(ctx context.Context) (bool, error) { reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second) defer cancel()
423 unmodified linesExpand all26 unmodified lines27282930313233343584 unmodified lines26 unmodified lines Description string `json:"description"` Method string `json:"method"` Arguments []CommandArgument `json:"arguments"` ModelCallable bool `json:"modelCallable"` Skill bool `json:"skill"` UserTurn bool `json:"userTurn"`}
type CommandMenuItem struct {84 unmodified lines59 unmodified lines606162636465666768697071727359 unmodified lines// must reconcile its result before asking the user to submit it again.type UncertainOutcomeError struct { Cause error Operation string Handle *OperationHandle}
func (e *UncertainOutcomeError) Error() string { if e.Handle != nil { return e.Operation + " admission is uncertain; operation " + e.Handle.ID + " can be queried" } return e.Operation + " may have been accepted; check its result before trying again"}func (e *UncertainOutcomeError) Unwrap() error { return e.Cause }146 unmodified linesExpand all14714814915015115215315415575 unmodified lines146 unmodified linesExpand all At int64 `json:"at"`}
type StreamEvent struct { OperationID string `json:"operationId,omitempty"` Image *ImageMetadata `json:"image,omitempty"` Usage *Usage `json:"usage,omitempty"` ToolTrace *ToolTrace `json:"trace,omitempty"` Timestamp *int64 `json:"timestamp,omitempty"`75 unmodified lines21 unmodified lines2223242526272829303132333489 unmodified lines124125126127126128129130131132167 unmodified linesExpand all21 unmodified linesconst ( NoRecovery RetryPolicy = iota ReadRecovery AuthRecovery ReceiptRecovery)
type Operation struct { Handle *OperationHandle Body any Name string Method string Path string89 unmodified linesfunc uncertainOperation(operation Operation, failure error) error { if operation.Policy == ReadRecovery { return failure } return &UncertainOutcomeError{Operation: operation.Name, Cause: failure} return &UncertainOutcomeError{Operation: operation.Name, Cause: failure, Handle: operation.Handle}}
// requestBytes owns every response until its bounded body has been read.func requestBytes(ctx context.Context, conn *Connection, operation Operation, limits responseLimits) ([]byte, error) {167 unmodified linesExpand all36 unmodified lines37383940414241424344454671 unmodified lines1181191201211221231221231241251261271282 unmodified lines13113213313413513613713413813914014114214314414514614714813814915015115215315415515615715872 unmodified lines231232233234221235236237238239240241242243244245246247248249250251252253254255256257258259260261262263345 unmodified linesExpand all36 unmodified linesfunc fieldError(field string) error { return &responseFieldError{field: field, cause: errors.New("required field is missing or invalid")}}
// executeMutation validates only after requestBytes has finished authentication// recovery. A malformed success can never enter the retry path.// executeMutation validates after requestBytes finishes authentication recovery.// Covered operations apply receipt recovery around this single attempt.func executeMutation(ctx context.Context, conn *Connection, operation Operation, statuses []int, decode func([]byte, int) error) error { if conn == nil { return errors.New("not connected to Albedo") }71 unmodified linesfunc sessionPath(id, tail string) string { return "/sessions/" + url.PathEscape(id) + tail }
func decodeSubmission(data []byte, _ int) (SendResult, error) { var wire struct { OK *bool `json:"ok"` Queued *bool `json:"queued"` OK *bool `json:"ok"` Queued *bool `json:"queued"` OperationID string `json:"operationId"` } if err := json.Unmarshal(data, &wire); err != nil { return SendResult{}, err }2 unmodified lines } if wire.Queued == nil { return SendResult{}, fieldError("queued") } if wire.OperationID == "" { return SendResult{}, fieldError("operationId") } return SendResult{OK: *wire.OK, Queued: *wire.Queued}, nil return SendResult{OK: *wire.OK, Queued: *wire.Queued, OperationID: wire.OperationID}, nil}func Submit(ctx context.Context, conn *Connection, id string, payload map[string]any) (SendResult, error) { handle, err := NewSubmission(id, payload) if err != nil { return SendResult{}, err } return SubmitOperation(ctx, conn, handle)}func SubmitOperation(ctx context.Context, conn *Connection, handle *OperationHandle) (SendResult, error) { var result SendResult err := executeMutation(ctx, conn, Operation{Name: "submit message", Method: http.MethodPost, Path: sessionPath(id, "/events"), Body: payload, Policy: AuthRecovery}, []int{http.StatusAccepted}, func(body []byte, status int) error { err := executeReceipt(ctx, conn, handle, []int{http.StatusAccepted}, func(body []byte, status int) error { var err error result, err = decodeSubmission(body, status) if err == nil && result.OperationID != handle.ID { return fieldError("operationId") } return err }) return result, err}72 unmodified lines }) return session, err}func CreateSession(ctx context.Context, conn *Connection, body map[string]string) (Session, error) { return mutateSession(ctx, conn, Operation{Name: "create session", Method: http.MethodPost, Path: "/sessions", Body: body, Policy: AuthRecovery}, http.StatusCreated) handle, err := NewCreation(body) if err != nil { return Session{}, err } return CreateSessionOperation(ctx, conn, handle)}func CreateSessionOperation(ctx context.Context, conn *Connection, handle *OperationHandle) (Session, error) { var session Session err := executeReceipt(ctx, conn, handle, []int{http.StatusCreated}, func(body []byte, _ int) error { fields, err := object(body) if err != nil { return err } var operationID string if err := required(fields, "operationId", &operationID); err != nil { return err } if operationID != handle.ID { return fieldError("operationId") } session, err = decodeSession(body) return err }) return session, err}func RenameSession(ctx context.Context, conn *Connection, id, name string) (Session, error) { return mutateSession(ctx, conn, Operation{Name: "rename session", Method: http.MethodPatch, Path: sessionPath(id, ""), Body: map[string]string{"name": name}, Policy: AuthRecovery}, http.StatusOK)}func ForkSession(ctx context.Context, conn *Connection, id string, body map[string]any) (Session, error) {345 unmodified linesExpand all144 unmodified linesExpand all1451461471481491491501511521539 unmodified lines163164165166167168167168169170171172173174175176180 unmodified linesExpand all144 unmodified linesExpand all }) }}
func TestInvalidMutationIsUncertainWithoutReplay(t *testing.T) {func TestInvalidCoveredMutationRetainsHandleAfterOneReplay(t *testing.T) { for _, test := range []struct { body string status int }{{"{}", 202}, {`{"ok":false,"queued":false}`, 202}, {`{"ok":true}`, 202}, {`{"ok":true,"queued":null}`, 202}, {`{"ok":true,"queued":false}`, 200}, {"", 204}, {"{", 202}, {`{"ok":true,"queued":"false"}`, 202}} {9 unmodified lines } if strings.Contains(test.body, `"queued":"false"`) && protocol.Field != "queued" { t.Fatalf("missing known field: %+v", protocol) } if transport.calls != 1 { t.Fatalf("mutation executed %d times", transport.calls) if transport.calls != 4 { t.Fatalf("want two posts and two receipt queries, got %d requests", transport.calls) } if uncertainty.Handle == nil || uncertainty.Handle.ID == "" { t.Fatal("lost operation handle") } } conn, transport := mutationConnection(`{"code":"invalid_submission","error":"bad"}`, 400) _, err := Submit(context.Background(), conn, "s", map[string]any{}) if api, ok := errors.AsType[*APIError](err); !ok || api.Code != "invalid_submission" {180 unmodified linesExpand all123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175package daemon
import ( "context" "crypto/rand" "encoding/json" "errors" "fmt" "net/http" "net/url" "time")
// OperationHandle retains one intent and its exact encoded request through uncertainty.// Reusing this handle never creates a new operation ID.type OperationHandle struct { ID string operation Operation kind string}
func operationID() (string, error) { var id [16]byte if _, err := rand.Read(id[:]); err != nil { return "", fmt.Errorf("create operation ID: %w", err) } timestamp := uint64(time.Now().UnixMilli()) for i := 5; i >= 0; i-- { id[i] = byte(timestamp) timestamp >>= 8 } id[6] = id[6]&0x0f | 0x70 id[8] = id[8]&0x3f | 0x80 return fmt.Sprintf("%x-%x-%x-%x-%x", id[0:4], id[4:6], id[6:8], id[8:10], id[10:16]), nil}
func newIntent(name, path string, fields map[string]any) (*OperationHandle, error) { id, err := operationID() if err != nil { return nil, err } request := make(map[string]any, len(fields)+1) for key, value := range fields { request[key] = value } request["operationId"] = id if name == "submit message" { if _, supplied := request["submissionId"]; !supplied { request["submissionId"] = id } } payload, err := json.Marshal(request) if err != nil { return nil, err } kind, _ := fields["type"].(string) handle := &OperationHandle{ID: id, kind: kind} handle.operation = Operation{Name: name, Method: http.MethodPost, Path: path, Body: json.RawMessage(payload), Policy: ReceiptRecovery, Handle: handle} return handle, nil}func NewSubmission(session string, fields map[string]any) (*OperationHandle, error) { return newIntent("submit message", sessionPath(session, "/events"), fields)}func NewCreation(fields map[string]string) (*OperationHandle, error) { request := make(map[string]any, len(fields)) for key, value := range fields { request[key] = value } return newIntent("create session", "/sessions", request)}
type OperationReceipt struct { OperationID string `json:"operationId"` Kind string `json:"kind"` Status string `json:"status"` Target string `json:"target"` Result json.RawMessage `json:"result"` Error json.RawMessage `json:"error"` HTTPStatus int `json:"httpStatus"` DeliveryStatus string `json:"deliveryStatus"` BlockingReason string `json:"blockingReason"`}
func (receipt OperationReceipt) Rejection() error { if receipt.Status != "rejected" { return nil } return decodeAPIError(receipt.HTTPStatus, receipt.Error)}
func GetOperation(ctx context.Context, conn *Connection, id string) (OperationReceipt, error) { receipt, err := RequestOperation[OperationReceipt](ctx, conn, Operation{Name: "query operation", Method: http.MethodGet, Path: "/operations/" + url.PathEscape(id), Policy: ReadRecovery}) if err == nil && (receipt.OperationID != id || (receipt.Status != "accepted" && receipt.Status != "rejected")) { return OperationReceipt{}, errors.New("invalid operation receipt") } return receipt, err}
// ResolveOperation only reads the original receipt; it never submits another request.func ResolveOperation(ctx context.Context, conn *Connection, handle *OperationHandle) (OperationReceipt, error) { return GetOperation(ctx, conn, handle.ID)}
func executeReceipt(ctx context.Context, conn *Connection, handle *OperationHandle, statuses []int, decode func([]byte, int) error) error { if ctx == nil { ctx = context.Background() } reqCtx, cancel := context.WithTimeout(ctx, 20*time.Second) defer cancel() for attempt := range 2 { err := executeMutation(reqCtx, conn, handle.operation, statuses, decode) if err == nil { return nil } if _, uncertain := errors.AsType[*UncertainOutcomeError](err); !uncertain { if reqCtx.Err() != nil { return uncertainOperation(handle.operation, err) } return err } if reqCtx.Err() != nil { return err } receipt, lookupErr := GetOperation(reqCtx, conn, handle.ID) if IsOperationExpired(lookupErr) { return uncertainOperation(handle.operation, lookupErr) } if lookupErr == nil { if receipt.Status == "rejected" { return receipt.Rejection() } if decodeErr := decode(receipt.Result, statuses[0]); decodeErr == nil { return nil } } if attempt == 1 || reqCtx.Err() != nil { return errors.Join(err, reqCtx.Err()) } } panic("unreachable operation retry budget")}
func ResolveCreation(ctx context.Context, conn *Connection, handle *OperationHandle) (Session, error) { receipt, err := ResolveOperation(ctx, conn, handle) if err != nil { if api, ok := errors.AsType[*APIError](err); ok && api.StatusCode == http.StatusGone { return Session{}, err } return Session{}, uncertainOperation(handle.operation, err) } if receipt.Status == "rejected" { return Session{}, decodeAPIError(receipt.HTTPStatus, receipt.Error) } fields, decodeErr := object(receipt.Result) if decodeErr != nil { return Session{}, uncertainOperation(handle.operation, decodeErr) } var id string if decodeErr = required(fields, "operationId", &id); decodeErr != nil || id != handle.ID { return Session{}, invalidResponse(handle.operation, "operationId", fieldError("operationId")) } session, decodeErr := decodeSession(receipt.Result) if decodeErr != nil { return Session{}, invalidResponse(handle.operation, "session", decodeErr) } return session, nil}
func (handle *OperationHandle) IsContinuation() bool { return handle.kind == "continue" }
// IsOperationExpired identifies a receipt whose admission outcome can no longer be recovered.func IsOperationExpired(err error) bool { api, ok := errors.AsType[*APIError](err) return ok && api.StatusCode == http.StatusGone && api.Code == "operation_expired"}447 unmodified linesExpand all448449450451452452453454455456155 unmodified linesExpand all447 unmodified linesExpand all// newSession creates a session on the test's own provider route, so tests// running side by side never depend on which profile is active.func newSession(t *testing.T, workspace string) string { t.Helper() created, err := daemon.RequestOperation[daemon.Session](context.Background(), conn(t), daemon.Operation{Name: "create session", Method: http.MethodPost, Path: "/sessions", Body: map[string]string{"workspace": workspace, "provider": t.Name()}, Policy: daemon.AuthRecovery}) created, err := daemon.CreateSession(context.Background(), conn(t), map[string]string{"workspace": workspace, "provider": t.Name()}) if err != nil || created.ID == "" { t.Fatalf("creating a session on %s: %v", t.Name(), err) } return created.ID155 unmodified linesExpand all