11package daemon
22
33import (
4- "bytes"
54 "context"
6- "encoding/json"
75 "fmt"
8- "io "
6+ "log/slog "
97 "net/http"
108 "strings"
119 "time"
10+
11+ "github.com/bwmarrin/discordgo"
1212)
1313
14- // discordAPIBase is the Discord REST API root. It is a package var so tests can
15- // point it at a mock server.
16- var discordAPIBase = "https://discord.com/api/v10"
17-
18- // DiscordGateway bridges hawk to Discord. Because no WebSocket library is present
19- // in go.mod, the gateway uses a REST long-poll-equivalent strategy: it does not
20- // open the Discord Gateway socket but instead relies on a poll over the bot's
21- // accessible channels for new @mentions/DMs. This keeps the bridge dependency-free
22- // while remaining bidirectional (it posts replies in-thread via the REST API).
23- //
24- // The polling transport is intentionally minimal and pluggable: fetchMessages is
25- // a field so a real Gateway-WebSocket implementation (or a test) can be swapped
26- // in without changing the message-handling logic.
14+ // DiscordGateway bridges hawk to Discord via the official Gateway (WebSocket),
15+ // using bwmarrin/discordgo. Unlike a REST-poll bridge it receives message events
16+ // in real time — guild @mentions and direct messages — with reconnection,
17+ // heartbeats, and rate limiting handled by the library. Authorized prompts are
18+ // forwarded to the daemon's /v1/chat endpoint and the reply is posted back to the
19+ // originating channel.
2720type DiscordGateway struct {
2821 cfg DiscordConfig
2922 daemonAddr string
3023 apiKey string
31- client * http.Client
24+ client * http.Client // for forwardToHawk
3225 auth * authorizer
26+ dispatch * asyncDispatcher
3327
34- // pollChannels lists channel IDs the bot watches. When empty the gateway
35- // discovers DM channels lazily as it sees them via fetchMessages.
36- pollChannels []string
37- // lastSeen maps channelID -> last processed message ID for poll dedup.
38- lastSeen map [string ]string
39-
40- // fetchMessages retrieves new messages for a channel since the given message
41- // ID. Overridable for tests. Default uses the Discord REST API.
42- fetchMessages func (ctx context.Context , channelID , afterID string ) ([]discordMessage , error )
43- }
44-
45- // discordMessage is the subset of a Discord message object we use.
46- type discordMessage struct {
47- ID string `json:"id"`
48- ChannelID string `json:"channel_id"`
49- Content string `json:"content"`
50- Author discordUser `json:"author"`
51- Mentions []discordUser `json:"mentions"`
52- GuildID string `json:"guild_id,omitempty"`
53- }
54-
55- type discordUser struct {
56- ID string `json:"id"`
57- Username string `json:"username"`
58- Bot bool `json:"bot,omitempty"`
28+ // openSession is overridable in tests; production opens a real Gateway socket.
29+ openSession func (ctx context.Context ) (* discordgo.Session , error )
5930}
6031
6132// newDiscordGateway builds a Discord gateway from config.
@@ -66,147 +37,140 @@ func newDiscordGateway(cfg DiscordConfig, daemonAddr, apiKey string) *DiscordGat
6637 apiKey : apiKey ,
6738 client : & http.Client {Timeout : 30 * time .Second },
6839 auth : newAuthorizer (cfg .PairingCode , cfg .AllowList ),
69- lastSeen : make ( map [ string ] string ),
40+ dispatch : newAsyncDispatcher ( 8 ),
7041 }
71- g .fetchMessages = g .fetchMessagesREST
42+ g .openSession = g .openGatewaySession
7243 return g
7344}
7445
7546// Name implements Gateway.
7647func (g * DiscordGateway ) Name () string { return "discord" }
7748
78- // Stop implements Gateway. Polling is driven by the Start context.
49+ // setDaemonURL implements daemonURLSetter.
50+ func (g * DiscordGateway ) setDaemonURL (url string ) { g .daemonAddr = url }
51+
52+ // Stop implements Gateway. The Gateway connection is driven by the Start context.
7953func (g * DiscordGateway ) Stop () error { return nil }
8054
81- // Start implements Gateway: poll watched channels until ctx is cancelled.
55+ // Start opens the Discord Gateway, registers the message handler, and blocks
56+ // until ctx is cancelled, then drains in-flight handlers and closes the socket.
8257func (g * DiscordGateway ) Start (ctx context.Context ) error {
83- ticker := time .NewTicker (2 * time .Second )
84- defer ticker .Stop ()
85- for {
86- select {
87- case <- ctx .Done ():
88- return ctx .Err ()
89- case <- ticker .C :
90- g .pollOnce (ctx )
91- }
58+ s , err := g .openSession (ctx )
59+ if err != nil {
60+ return err
9261 }
62+ s .AddHandler (func (sess * discordgo.Session , m * discordgo.MessageCreate ) {
63+ g .onMessageCreate (ctx , sess , m )
64+ })
65+ if err := s .Open (); err != nil {
66+ return fmt .Errorf ("discord gateway open: %w" , err )
67+ }
68+ defer func () { _ = s .Close () }()
69+
70+ <- ctx .Done ()
71+ g .dispatch .wait ()
72+ return ctx .Err ()
9373}
9474
95- func (g * DiscordGateway ) pollOnce (ctx context.Context ) {
96- for _ , ch := range g .pollChannels {
97- msgs , err := g .fetchMessages (ctx , ch , g .lastSeen [ch ])
98- if err != nil {
99- continue
100- }
101- for _ , m := range msgs {
102- g .lastSeen [m .ChannelID ] = m .ID
103- g .handleMessage (ctx , m )
104- }
75+ // openGatewaySession creates a discordgo session configured with the intents the
76+ // bridge needs: guild messages, message content (privileged — enable it in the
77+ // Discord developer portal), and direct messages.
78+ func (g * DiscordGateway ) openGatewaySession (_ context.Context ) (* discordgo.Session , error ) {
79+ s , err := discordgo .New ("Bot " + g .cfg .Token )
80+ if err != nil {
81+ return nil , fmt .Errorf ("discord session: %w" , err )
10582 }
83+ s .Identify .Intents = discordgo .IntentsGuildMessages |
84+ discordgo .IntentMessageContent |
85+ discordgo .IntentsDirectMessages
86+ return s , nil
10687}
10788
108- // mentionsBot reports whether the message is a DM or @mentions the bot user.
109- func (g * DiscordGateway ) mentionsBot (m discordMessage ) bool {
110- // DM: no guild.
89+ // onMessageCreate filters inbound events and dispatches authorized messages to
90+ // the shared handler. Replies go back through the live Gateway session.
91+ func (g * DiscordGateway ) onMessageCreate (ctx context.Context , s * discordgo.Session , m * discordgo.MessageCreate ) {
92+ selfID := ""
93+ if s .State != nil && s .State .User != nil {
94+ selfID = s .State .User .ID
95+ }
96+ if ! wantsDiscordMessage (m , selfID ) {
97+ return
98+ }
99+ text := stripDiscordMention (m .Content , selfID )
100+ channelID := m .ChannelID
101+ send := func (content string ) error {
102+ _ , err := s .ChannelMessageSend (channelID , content )
103+ return err
104+ }
105+ g .dispatch .run (ctx , func () {
106+ g .handleMessage (ctx , m .Author .ID , channelID , text , send )
107+ })
108+ }
109+
110+ // wantsDiscordMessage reports whether a message should be processed: skip our own
111+ // messages and other bots; accept direct messages and any message that @mentions
112+ // the bot. selfID is the bot's own user ID (empty until the READY event lands).
113+ func wantsDiscordMessage (m * discordgo.MessageCreate , selfID string ) bool {
114+ if m == nil || m .Author == nil || m .Author .Bot {
115+ return false
116+ }
117+ if selfID != "" && m .Author .ID == selfID {
118+ return false
119+ }
111120 if m .GuildID == "" {
112- return true
121+ return true // direct message
113122 }
114123 for _ , u := range m .Mentions {
115- if u .ID == g . cfg . AppID {
124+ if u != nil && u .ID == selfID {
116125 return true
117126 }
118127 }
119128 return false
120129}
121130
122- // stripMention removes a leading <@id> / <@!id> mention token from content.
123- func stripDiscordMention (content , appID string ) string {
124- content = strings .TrimSpace (content )
125- for _ , tok := range []string {"<@" + appID + ">" , "<@!" + appID + ">" } {
126- if strings .HasPrefix (content , tok ) {
127- return strings .TrimSpace (content [len (tok ):])
131+ // handleMessage applies the pairing/allowlist policy and forwards authorized
132+ // prompts to the daemon. It is transport-agnostic: send delivers the reply, so
133+ // the policy/forwarding logic is unit-testable without a live Gateway session.
134+ func (g * DiscordGateway ) handleMessage (ctx context.Context , senderID , channelID , text string , send func (string ) error ) {
135+ reply := func (content string ) {
136+ if err := send (content ); err != nil {
137+ slog .Error ("discord reply failed" , "channel" , channelID , "error" , err )
128138 }
129139 }
130- return content
131- }
132-
133- func (g * DiscordGateway ) handleMessage (ctx context.Context , m discordMessage ) {
134- if m .Author .Bot {
135- return // ignore other bots / our own echoes
136- }
137- if ! g .mentionsBot (m ) {
138- return
139- }
140- text := stripDiscordMention (m .Content , g .cfg .AppID )
141- sender := m .Author .ID
142140
143- if isPair , ok := g .auth .tryPair (sender , text ); isPair {
141+ if isPair , ok := g .auth .tryPair (senderID , text ); isPair {
144142 if ok {
145- _ = g . postMessage ( ctx , m . ChannelID , "Paired. You can now chat with hawk." )
143+ reply ( "Paired. You can now chat with hawk." )
146144 } else {
147- _ = g . postMessage ( ctx , m . ChannelID , "Pairing failed: invalid code." )
145+ reply ( "Pairing failed: invalid code." )
148146 }
149147 return
150148 }
151- if ! g .auth .allowed (sender ) {
152- _ = g . postMessage ( ctx , m . ChannelID , "Unauthorized. Send /pair <code> to authorize." )
149+ if ! g .auth .allowed (senderID ) {
150+ reply ( "Unauthorized. Send /pair <code> to authorize." )
153151 return
154152 }
155153
156- reply , err := forwardToHawk (ctx , g .client , g .daemonAddr , g .apiKey , text )
157- if err != nil {
158- reply = fmt .Sprintf ("Error: %v" , err )
159- }
160- if len (reply ) > 1900 { // Discord 2000-char limit, leave headroom
161- reply = reply [:1900 ] + "\n \n ... (truncated)"
162- }
163- _ = g .postMessage (ctx , m .ChannelID , reply )
164- }
165-
166- // postMessage sends a message to a Discord channel (in-thread reply context).
167- func (g * DiscordGateway ) postMessage (ctx context.Context , channelID , content string ) error {
168- if g .cfg .Token == "" || channelID == "" {
169- return nil
170- }
171- body , _ := json .Marshal (map [string ]string {"content" : content })
172- apiURL := fmt .Sprintf ("%s/channels/%s/messages" , discordAPIBase , channelID )
173- req , err := http .NewRequestWithContext (ctx , http .MethodPost , apiURL , bytes .NewReader (body ))
154+ resp , err := forwardToHawk (ctx , g .client , g .daemonAddr , g .apiKey , text )
174155 if err != nil {
175- return err
156+ resp = fmt . Sprintf ( "Error: %v" , err )
176157 }
177- req .Header .Set ("Content-Type" , "application/json" )
178- req .Header .Set ("Authorization" , "Bot " + g .cfg .Token )
179- resp , err := g .client .Do (req )
180- if err != nil {
181- return err
158+ if len (resp ) > 1900 { // Discord 2000-char limit, leave headroom
159+ resp = resp [:1900 ] + "\n \n ... (truncated)"
182160 }
183- _ = resp .Body .Close ()
184- return nil
161+ reply (resp )
185162}
186163
187- // fetchMessagesREST retrieves channel messages via the Discord REST API.
188- func (g * DiscordGateway ) fetchMessagesREST (ctx context.Context , channelID , afterID string ) ([]discordMessage , error ) {
189- apiURL := fmt .Sprintf ("%s/channels/%s/messages?limit=20" , discordAPIBase , channelID )
190- if afterID != "" {
191- apiURL += "&after=" + afterID
192- }
193- req , err := http .NewRequestWithContext (ctx , http .MethodGet , apiURL , nil )
194- if err != nil {
195- return nil , err
196- }
197- req .Header .Set ("Authorization" , "Bot " + g .cfg .Token )
198- resp , err := g .client .Do (req )
199- if err != nil {
200- return nil , err
201- }
202- defer func () { _ = resp .Body .Close () }()
203- if resp .StatusCode != http .StatusOK {
204- data , _ := io .ReadAll (io .LimitReader (resp .Body , 4096 ))
205- return nil , fmt .Errorf ("discord messages: HTTP %d: %s" , resp .StatusCode , string (data ))
164+ // stripDiscordMention removes a leading <@id> / <@!id> mention token from content.
165+ func stripDiscordMention (content , selfID string ) string {
166+ content = strings .TrimSpace (content )
167+ if selfID == "" {
168+ return content
206169 }
207- var msgs []discordMessage
208- if err := json .NewDecoder (resp .Body ).Decode (& msgs ); err != nil {
209- return nil , err
170+ for _ , tok := range []string {"<@" + selfID + ">" , "<@!" + selfID + ">" } {
171+ if strings .HasPrefix (content , tok ) {
172+ return strings .TrimSpace (content [len (tok ):])
173+ }
210174 }
211- return msgs , nil
175+ return content
212176}
0 commit comments