okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petfeat: multiple images, collapsed pastes, and an attachment strip in the prompt6e82811

main
5.6 KiB165 linesraw
1package daemon
2
3import (
4	"context"
5	"encoding/json"
6	"io"
7	"net/http"
8
9	"albedo/cli/internal/daemon/protocol"
10)
11
12type SendResult struct {
13	OperationID     string `json:"operationId"`
14	OK              bool   `json:"ok"`
15	Queued          bool   `json:"queued"`
16	AcceptanceOrder int64
17}
18
19// SubmissionRequest keeps message text and continuation content distinct:
20// an empty message has a Content pointer, while a continuation has none.
21type SubmissionRequest struct {
22	CatalogRevision  string
23	CommandArguments json.RawMessage
24	Content          *string
25	Images           []ImageAttachment
26	Pastes           []string
27	Type             string
28	ClientID         string
29	SubmissionID     string
30	Name             string
31	Arguments        string
32}
33
34func Submit(ctx context.Context, conn *Connection, id string, payload SubmissionRequest) (SendResult, error) {
35	handle, err := NewSubmission(id, payload)
36	if err != nil {
37		return SendResult{}, err
38	}
39	return SubmitOperation(ctx, conn, handle)
40}
41
42func SubmitOperation(ctx context.Context, conn *Connection, handle *OperationHandle) (SendResult, error) {
43	return executeReceipt(ctx, conn, handle, []int{http.StatusAccepted}, func(body []byte, status int) (SendResult, error) {
44		result, err := decodeSubmission(body, status)
45		if err == nil && result.OperationID != handle.ID() {
46			return result, fieldError("operationId")
47		}
48		return result, err
49	}, func(ctx context.Context) (SendResult, bool, error) {
50		receipt, err := ResolveOperation(ctx, conn, handle)
51		if err != nil {
52			return SendResult{}, false, err
53		}
54		result, err := submissionResult(receipt)
55		return result, true, err
56	})
57}
58
59func InterruptSession(ctx context.Context, conn *Connection, id string) (bool, error) {
60	session, err := GetSession(ctx, conn, id)
61	if err != nil {
62		return false, err
63	}
64	return interruptCaptured(ctx, conn, id, session.wire.Status.RunID, session.wire.InputOrder)
65}
66func interruptCaptured(ctx context.Context, conn *Connection, id string, runID *string, inputOrder int64) (bool, error) {
67	var result protocol.Interruption
68	err := executeMutation(ctx, conn, operation{Name: "interrupt session", BuildRequest: func(base string, body io.Reader) (*http.Request, error) {
69		return protocol.NewInterruptSessionRequestWithBody(base, id, "application/json", body)
70	}, Body: protocol.InterruptRequest{RunID: runID, ThroughInputOrder: inputOrder}, Policy: noRecovery}, []int{200}, func(data []byte, _ int) error {
71		if err := decodeRequired(data, &result); err != nil {
72			return err
73		}
74		switch result.State {
75		case "requested", "already_ended":
76			return nil
77		}
78		return fieldError("interruption state")
79	})
80	return result.State == "requested", err
81}
82
83func (c *ChatClient) PrepareTurn(content string, images []ImageAttachment, pastes []string, continuation bool) (*OperationHandle, error) {
84	payload := SubmissionRequest{ClientID: c.clientID, Content: &content, Type: "user", Images: images, Pastes: pastes}
85	if continuation {
86		payload.Type, payload.Content = "continue", nil
87	}
88	return NewSubmission(c.agentID, payload)
89}
90
91func (c *ChatClient) SubmitOperation(ctx context.Context, handle *OperationHandle) (*SendResult, error) {
92	result, err := SubmitOperation(ctx, c.conn, handle)
93	if err != nil {
94		return nil, err
95	}
96	return &result, nil
97}
98
99func (c *ChatClient) ResolveOperation(ctx context.Context, handle *OperationHandle) (InputReceipt, error) {
100	return ResolveOperation(ctx, c.conn, handle)
101}
102
103func (c *ChatClient) Send(ctx context.Context, content string, images []ImageAttachment) (*SendResult, error) {
104	handle, err := c.PrepareTurn(content, images, nil, false)
105	if err != nil {
106		return nil, err
107	}
108	return c.SubmitOperation(ctx, handle)
109}
110
111func (c *ChatClient) CancelSubmission(ctx context.Context, submissionID string) (string, error) {
112	var result protocol.InputCancellation
113	err := executeMutation(ctx, c.conn, operation{Name: "cancel input", BuildRequest: func(base string, body io.Reader) (*http.Request, error) {
114		return protocol.NewCancelInputRequestWithBody(base, c.agentID, submissionID, "application/json", body)
115	}, Body: struct{}{}, Policy: noRecovery}, []int{200}, func(data []byte, _ int) error {
116		if err := decodeRequired(data, &result); err != nil {
117			return err
118		}
119		if result.Input.ID != submissionID {
120			return fieldError("input identity")
121		}
122		return validInput(result.Input)
123	})
124	outcome := result.Result
125	if outcome == "cancelled" {
126		outcome = "cancelled_queued"
127	}
128	switch result.Result {
129	case "cancelled", "interrupt_requested", "shared_running", "not_pending":
130	default:
131		if err == nil {
132			err = fieldError("cancellation result")
133		}
134	}
135	return outcome, err
136}
137
138func decodeSubmission(data []byte, _ int) (SendResult, error) {
139	var input protocol.Input
140	if err := decodeRequired(data, &input); err != nil {
141		return SendResult{}, err
142	}
143	if err := validInput(input); err != nil {
144		return SendResult{}, err
145	}
146	return submissionResult(InputReceipt(input))
147}
148
149func submissionResult(input InputReceipt) (SendResult, error) {
150	if err := input.Rejection(); err != nil {
151		return SendResult{}, err
152	}
153	return SendResult{OK: true, Queued: input.Pending(), OperationID: input.ID, AcceptanceOrder: input.InputOrder()}, nil
154}
155
156func (c *ChatClient) PrepareCommand(commandID string, arguments map[string]json.RawMessage) (*OperationHandle, error) {
157	payload, err := json.Marshal(arguments)
158	if err != nil {
159		return nil, err
160	}
161	return NewSubmission(c.agentID, SubmissionRequest{ClientID: c.clientID, Type: "command", Name: commandID, CommandArguments: payload})
162}
163func (c *ChatClient) InterruptObserved(ctx context.Context, status AgentStatus) (bool, error) {
164	return interruptCaptured(ctx, c.conn, c.agentID, status.RunID, status.InputOrder)
165}