| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "io" |
| 7 | "log/slog" |
| 8 | "net/http" |
| 9 | "time" |
| 10 | |
| 11 | "reasonix/internal/agent" |
| 12 | "reasonix/internal/config" |
| 13 | "reasonix/internal/eventwire" |
| 14 | ) |
| 15 | |
| 16 | func (m *takeoverMirror) run(initialClient *http.Client, initialRecord takeoverServeRecord) { |
| 17 | defer close(m.done) |
| 18 | m.mu.Lock() |
| 19 | if m.client == nil { |
| 20 | m.client = initialClient |
| 21 | m.record = initialRecord |
| 22 | } |
| 23 | m.mu.Unlock() |
| 24 | |
| 25 | flushTimer := time.NewTimer(time.Hour) |
| 26 | if !flushTimer.Stop() { |
| 27 | <-flushTimer.C |
| 28 | } |
| 29 | heartbeat := time.NewTicker(takeoverMirrorHeartbeat) |
| 30 | retryReturn := time.NewTicker(250 * time.Millisecond) |
| 31 | defer flushTimer.Stop() |
| 32 | defer heartbeat.Stop() |
| 33 | defer retryReturn.Stop() |
| 34 | flushArmed := false |
| 35 | for { |
| 36 | select { |
| 37 | case <-m.stop: |
| 38 | m.flushOnce(context.Background()) |
| 39 | return |
| 40 | case <-m.wake: |
| 41 | if !flushArmed { |
| 42 | flushTimer.Reset(takeoverMirrorFlushEvery) |
| 43 | flushArmed = true |
| 44 | } |
| 45 | continue |
| 46 | case <-flushTimer.C: |
| 47 | flushArmed = false |
| 48 | if !m.pushOnce(false) { |
| 49 | return |
| 50 | } |
| 51 | case <-heartbeat.C: |
| 52 | if !m.pushOnce(true) { |
| 53 | return |
| 54 | } |
| 55 | case <-retryReturn.C: |
| 56 | if m.retryPendingReturn(false) { |
| 57 | m.detach() |
| 58 | m.mirrorEnd() |
| 59 | return |
| 60 | } |
| 61 | } |
| 62 | // A mirror whose tab is gone entirely (closed, not detached) ends |
| 63 | // itself so Serve can hand the session back. A tab close sends its own |
| 64 | // farewell after releasing the writer, so this only detaches for it. |
| 65 | if !m.app.takeoverTabLive(m.sessionPath) { |
| 66 | m.detach() |
| 67 | if !m.closing.Load() { |
| 68 | m.mirrorEnd() |
| 69 | } |
| 70 | return |
| 71 | } |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | func (m *takeoverMirror) pushOnce(heartbeat bool) bool { |
| 76 | m.sendMu.Lock() |
| 77 | defer m.sendMu.Unlock() |
| 78 | return m.pushOnceLocked(heartbeat) |
| 79 | } |
| 80 | |
| 81 | func (m *takeoverMirror) pushOnceLocked(heartbeat bool) bool { |
| 82 | if m.returned.Load() { |
| 83 | return false |
| 84 | } |
| 85 | client, record, _, grant, revision := m.snapshotBinding() |
| 86 | if client == nil || grant.MirrorID == "" { |
| 87 | return true |
| 88 | } |
| 89 | frames := m.drainQueue() |
| 90 | if len(frames) == 0 && !heartbeat { |
| 91 | return true |
| 92 | } |
| 93 | marshal := func(batch []eventwire.Event) ([]byte, error) { |
| 94 | return json.Marshal(map[string]any{ |
| 95 | "sessionPath": m.sessionPath, "mirrorId": grant.MirrorID, "frames": batch, |
| 96 | }) |
| 97 | } |
| 98 | batch, remainder, payload, err := eventwire.MarshalMirrorBatch(frames, eventwire.MirrorBatchMaxBytes, marshal) |
| 99 | if err == nil && len(batch) == 0 && len(frames) > 0 && len(remainder) > 0 { |
| 100 | remainder = remainder[1:] |
| 101 | } |
| 102 | m.requeue(remainder) |
| 103 | if err != nil { |
| 104 | m.requeue(batch) |
| 105 | return true |
| 106 | } |
| 107 | if len(batch) == 0 && len(frames) > 0 { |
| 108 | m.wakeIfQueued() |
| 109 | if !heartbeat { |
| 110 | return true |
| 111 | } |
| 112 | payload, _ = marshal(nil) |
| 113 | } |
| 114 | ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) |
| 115 | resp, err := serveDo(ctx, client, http.MethodPost, serveURL(record.base, "/external/frames"), payload) |
| 116 | if err != nil { |
| 117 | cancel() |
| 118 | if !m.bindingCurrent(client, grant, revision) { |
| 119 | return true |
| 120 | } |
| 121 | m.requeue(batch) |
| 122 | return m.retryAdoptOrDemote(client, grant, revision) |
| 123 | } |
| 124 | body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16)) |
| 125 | resp.Body.Close() |
| 126 | cancel() |
| 127 | if !m.bindingCurrent(client, grant, revision) { |
| 128 | return true |
| 129 | } |
| 130 | switch resp.StatusCode { |
| 131 | case http.StatusOK: |
| 132 | m.mu.Lock() |
| 133 | if m.bindingRevision == revision { |
| 134 | m.consecutiveFailures = 0 |
| 135 | } |
| 136 | m.mu.Unlock() |
| 137 | var out struct { |
| 138 | ReclaimRequested bool `json:"reclaimRequested"` |
| 139 | ReclaimMode string `json:"reclaimMode"` |
| 140 | } |
| 141 | if json.Unmarshal(body, &out) == nil && out.ReclaimRequested { |
| 142 | m.requestDemote(out.ReclaimMode) |
| 143 | } |
| 144 | m.wakeIfQueued() |
| 145 | return true |
| 146 | case http.StatusUnauthorized, http.StatusForbidden, http.StatusConflict: |
| 147 | slog.Info("desktop: mirror generation rejected — attempting re-adopt", |
| 148 | "session", m.sessionPath) |
| 149 | m.requeue(batch) |
| 150 | return m.retryAdoptOrDemote(client, grant, revision) |
| 151 | default: |
| 152 | m.requeue(batch) |
| 153 | m.mu.Lock() |
| 154 | if m.bindingRevision == revision { |
| 155 | m.consecutiveFailures++ |
| 156 | } |
| 157 | failures := m.consecutiveFailures |
| 158 | m.mu.Unlock() |
| 159 | if failures < 3 { |
| 160 | return true |
| 161 | } |
| 162 | return m.retryAdoptOrDemote(client, grant, revision) |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | // retryAdoptOrDemote attempts to re-establish the mirror with fresh serve |
| 167 | // credentials (the serve may have restarted with a new token). If the serve |
| 168 | // already owns the session (reclaim completed or another writer took over), |
| 169 | // demotes this tab to read-only and releases the lease so the remote side |
| 170 | // can proceed. Returns true if re-adopted (caller continues the loop). |
| 171 | func (m *takeoverMirror) retryAdoptOrDemote(oldClient *http.Client, oldGrant takeoverGrant, revision uint64) bool { |
| 172 | if !m.bindingCurrent(oldClient, oldGrant, revision) { |
| 173 | return true |
| 174 | } |
| 175 | conflicted := false |
| 176 | records := discoverLocalTakeoverServesForMirror() |
| 177 | for _, record := range records { |
| 178 | if !pathWithinDir(m.sessionPath, config.ProjectSessionDir(record.state.Workspace)) { |
| 179 | continue |
| 180 | } |
| 181 | ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
| 182 | client, err := takeoverClient(ctx, record) |
| 183 | if err != nil { |
| 184 | cancel() |
| 185 | continue |
| 186 | } |
| 187 | body, bodyErr := json.Marshal(map[string]string{"sessionPath": m.sessionPath, "writerId": agent.SessionWriterID()}) |
| 188 | if bodyErr != nil { |
| 189 | cancel() |
| 190 | continue |
| 191 | } |
| 192 | resp, respErr := serveDo(ctx, client, http.MethodPost, serveURL(record.base, "/adopt"), body) |
| 193 | cancel() |
| 194 | if respErr != nil { |
| 195 | continue |
| 196 | } |
| 197 | respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16)) |
| 198 | resp.Body.Close() |
| 199 | if resp.StatusCode == http.StatusOK { |
| 200 | var grant takeoverGrant |
| 201 | if json.Unmarshal(respBody, &grant) != nil || grant.MirrorID == "" || grant.ReturnHandoffID == "" || grant.SourceWriterID == "" || |
| 202 | grant.TargetWriterID != agent.SessionWriterID() || sessionRuntimeKey(grant.SessionPath) != sessionRuntimeKey(m.sessionPath) { |
| 203 | continue |
| 204 | } |
| 205 | // Re-adopted: swap in the fresh client and keep mirroring. |
| 206 | m.mu.Lock() |
| 207 | if m.client == oldClient && m.grant.MirrorID == oldGrant.MirrorID && m.bindingRevision == revision { |
| 208 | m.client = client |
| 209 | m.record = record |
| 210 | m.grant = grant |
| 211 | m.bindingRevision++ |
| 212 | m.consecutiveFailures = 0 |
| 213 | } |
| 214 | m.mu.Unlock() |
| 215 | slog.Info("desktop: mirror re-adopted with fresh credentials", |
| 216 | "session", m.sessionPath, "base", record.base) |
| 217 | return true |
| 218 | } |
| 219 | if resp.StatusCode == http.StatusConflict { |
| 220 | conflicted = true |
| 221 | continue |
| 222 | } |
| 223 | // Other statuses: try next record. |
| 224 | } |
| 225 | if conflicted && m.bindingCurrent(oldClient, oldGrant, revision) { |
| 226 | slog.Info("desktop: serve holds session — demoting to release lease", |
| 227 | "session", m.sessionPath) |
| 228 | m.requestDemote("") |
| 229 | return false |
| 230 | } |
| 231 | // The Serve may be restarting or its state/token files may not have become |
| 232 | // visible yet. Keep the bounded queue and retry on the next heartbeat. |
| 233 | slog.Warn("desktop: mirror re-adopt unavailable; retaining local writer and bounded queue", |
| 234 | "session", m.sessionPath) |
| 235 | return true |
| 236 | } |
| 237 | |
| 238 | func (m *takeoverMirror) drainQueue() []eventwire.Event { |
| 239 | m.mu.Lock() |
| 240 | frames := m.queue.Take(takeoverMirrorMaxQueue) |
| 241 | m.mu.Unlock() |
| 242 | return frames |
| 243 | } |
| 244 | |
| 245 | func (m *takeoverMirror) requeue(frames []eventwire.Event) { |
| 246 | if len(frames) == 0 { |
| 247 | return |
| 248 | } |
| 249 | m.mu.Lock() |
| 250 | m.queue.Prepend(frames) |
| 251 | m.mu.Unlock() |
| 252 | } |
| 253 | |
| 254 | func (m *takeoverMirror) wakeIfQueued() { |
| 255 | m.mu.Lock() |
| 256 | pending := m.queue.Len() > 0 |
| 257 | m.mu.Unlock() |
| 258 | if pending { |
| 259 | select { |
| 260 | case m.wake <- struct{}{}: |
| 261 | default: |
| 262 | } |
| 263 | } |
| 264 | } |
| 265 | |
| 266 | func (m *takeoverMirror) flushOnce(ctx context.Context) { |
| 267 | m.sendMu.Lock() |
| 268 | defer m.sendMu.Unlock() |
| 269 | m.flushOnceLocked(ctx) |
| 270 | } |
| 271 | |
| 272 | func (m *takeoverMirror) flushOnceLocked(ctx context.Context) { |
| 273 | if m.returned.Load() { |
| 274 | return |
| 275 | } |
| 276 | client, record, _, grant := m.snapshotClient() |
| 277 | if client == nil || grant.MirrorID == "" { |
| 278 | return |
| 279 | } |
| 280 | frames := m.drainQueue() |
| 281 | if len(frames) == 0 { |
| 282 | return |
| 283 | } |
| 284 | marshal := func(batch []eventwire.Event) ([]byte, error) { |
| 285 | return json.Marshal(map[string]any{"sessionPath": m.sessionPath, "mirrorId": grant.MirrorID, "frames": batch}) |
| 286 | } |
| 287 | batch, _, payload, err := eventwire.MarshalMirrorBatch(frames, eventwire.MirrorBatchMaxBytes, marshal) |
| 288 | if err != nil { |
| 289 | return |
| 290 | } |
| 291 | if len(batch) == 0 { |
| 292 | return |
| 293 | } |
| 294 | flushCtx, cancel := context.WithTimeout(ctx, 5*time.Second) |
| 295 | defer cancel() |
| 296 | resp, err := serveDo(flushCtx, client, http.MethodPost, serveURL(record.base, "/external/frames"), payload) |
| 297 | if err != nil { |
| 298 | return |
| 299 | } |
| 300 | _, _ = io.Copy(io.Discard, resp.Body) |
| 301 | resp.Body.Close() |
| 302 | } |
| 303 |