| 1282 | } |
| 1283 | |
| 1284 | func (p *Server) awaitSubagentCompletion( |
| 1285 | ctx context.Context, |
| 1286 | parentChatID uuid.UUID, |
| 1287 | targetChatID uuid.UUID, |
| 1288 | timeout time.Duration, |
| 1289 | ) (database.Chat, string, error) { |
| 1290 | isDescendant, err := isSubagentDescendant(ctx, p.db, parentChatID, targetChatID) |
| 1291 | if err != nil { |
| 1292 | return database.Chat{}, "", err |
| 1293 | } |
| 1294 | if !isDescendant { |
| 1295 | return database.Chat{}, "", ErrSubagentNotDescendant |
| 1296 | } |
| 1297 | |
| 1298 | // Check immediately before entering the poll loop. |
| 1299 | targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID) |
| 1300 | if checkErr != nil { |
| 1301 | return database.Chat{}, "", checkErr |
| 1302 | } |
| 1303 | if done { |
| 1304 | return handleSubagentDone(targetChat, report) |
| 1305 | } |
| 1306 | |
| 1307 | if timeout <= 0 { |
| 1308 | timeout = defaultSubagentWaitTimeout |
| 1309 | } |
| 1310 | timer := p.clock.NewTimer(timeout, "chatd", "subagent_await") |
| 1311 | defer timer.Stop() |
| 1312 | |
| 1313 | // When pubsub is available, subscribe for fast status |
| 1314 | // notifications and use a less aggressive fallback poll. |
| 1315 | // Without pubsub (single-instance / in-memory) fall back |
| 1316 | // to the original 200ms polling. |
| 1317 | pollInterval := subagentAwaitPollInterval |
| 1318 | var notifyCh <-chan struct{} |
| 1319 | if p.pubsub != nil { |
| 1320 | pollInterval = subagentAwaitFallbackPoll |
| 1321 | ch := make(chan struct{}, 1) |
| 1322 | notifyCh = ch |
| 1323 | cancel, subErr := p.pubsub.SubscribeWithErr( |
| 1324 | coderdpubsub.ChatStreamNotifyChannel(targetChatID), |
| 1325 | func(_ context.Context, _ []byte, _ error) { |
| 1326 | // Non-blocking send so we never stall the |
| 1327 | // pubsub dispatch goroutine. |
| 1328 | select { |
| 1329 | case ch <- struct{}{}: |
| 1330 | default: |
| 1331 | } |
| 1332 | }, |
| 1333 | ) |
| 1334 | if subErr == nil { |
| 1335 | defer cancel() |
| 1336 | } else { |
| 1337 | // Subscription failed; fall back to fast polling. |
| 1338 | pollInterval = subagentAwaitPollInterval |
| 1339 | notifyCh = nil |
| 1340 | } |
| 1341 | } |