| 1 | package agent |
| 2 | |
| 3 | import "context" |
| 4 | |
| 5 | type messageIdentityContextKey struct{} |
| 6 | |
| 7 | // Tool work captures the committed message identity in its context. It never |
| 8 | // reads a mutable "latest message" from the Agent while tools run concurrently. |
| 9 | func withMessageIdentity(ctx context.Context, messageID string) context.Context { |
| 10 | return context.WithValue(ctx, messageIdentityContextKey{}, messageID) |
| 11 | } |
| 12 | |
| 13 | func messageIdentity(ctx context.Context) string { |
| 14 | id, _ := ctx.Value(messageIdentityContextKey{}).(string) |
| 15 | return id |
| 16 | } |
| 17 | |
| 18 | type userMessageIdentityKey struct{} |
| 19 | type userMessageIdentity struct { |
| 20 | session *Session |
| 21 | id string |
| 22 | } |
| 23 | |
| 24 | func withUserMessageIdentity(ctx context.Context, session *Session, id string) context.Context { |
| 25 | return context.WithValue(ctx, userMessageIdentityKey{}, userMessageIdentity{session: session, id: id}) |
| 26 | } |
| 27 | |
| 28 | // WithUserMessageIdentity reserves the controller's user identity before a |
| 29 | // planner can emit display events. Cancellation recovery uses the same ID. |
| 30 | func WithUserMessageIdentity(ctx context.Context, session *Session, id string) context.Context { |
| 31 | return withUserMessageIdentity(ctx, session, id) |
| 32 | } |
| 33 | |
| 34 | func turnUserMessageID(ctx context.Context, session *Session) string { |
| 35 | if identity, ok := ctx.Value(userMessageIdentityKey{}).(userMessageIdentity); ok && identity.session == session { |
| 36 | return identity.id |
| 37 | } |
| 38 | return NewMessageID() |
| 39 | } |
| 40 |