mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Make ds4-agent interruption cooperative

Browse Source
This commit is contained in:
antirez committed 2026-06-06 12:37:27 +02:00
1 parent 7ac436a382
commit 42e59157ea
5 files changed
+398 -66
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

ds4.c
+127 -21
View File
@@ -18401,7 +18401,10 @@ static bool metal_graph_prefill_decode_streaming_range(
ds4_session_progress_fn progress,
void *progress_ud,
ds4_session_progress_fn display_progress,
void *display_progress_ud) {
void *display_progress_ud,
ds4_session_cancel_fn cancel,
void *cancel_ud,
bool *cancelled) {
if (!metal_graph_use_streaming_decode_prefill(g, n_tokens)) return false;
if (!prompt || start > (uint32_t)prompt->len ||
n_tokens > (uint32_t)prompt->len - start) return false;
@@ -18421,6 +18424,10 @@ static bool metal_graph_prefill_decode_streaming_range(
}
for (uint32_t i = 0; i < n_tokens; i++) {
if (cancel && cancel(cancel_ud)) {
if (cancelled) *cancelled = true;
return true;
}
const uint32_t pos = start + i;
const bool last = i + 1u == n_tokens;
float *token_logits = (last && logits) ? logits : NULL;
@@ -18442,6 +18449,10 @@ static bool metal_graph_prefill_decode_streaming_range(
if (display_progress) {
display_progress(display_progress_ud, "prefill_display", (int)(pos + 1u), prompt->len);
}
if (cancel && cancel(cancel_ud)) {
if (cancelled) *cancelled = true;
return true;
}
if (show_progress) {
fprintf(stderr, "ds4: gpu streaming prefill token %u/%u\r",
i + 1u,
@@ -19610,7 +19621,10 @@ static bool metal_graph_prefill_raw_swa(
float *logits,
bool show_progress,
ds4_session_progress_fn display_progress,
void *display_progress_ud) {
void *display_progress_ud,
ds4_session_cancel_fn cancel,
void *cancel_ud,
bool *cancelled) {
if (n_tokens <= 0 || n_tokens > prompt->len) return false;
if ((uint32_t)n_tokens > g->prefill_cap) return false;
if (metal_graph_use_streaming_decode_prefill_range(g, 0, (uint32_t)n_tokens)) {
@@ -19625,8 +19639,19 @@ static bool metal_graph_prefill_raw_swa(
NULL,
NULL,
display_progress,
display_progress_ud);
display_progress_ud,
cancel,
cancel_ud,
cancelled);
}
/* The layer-major fallback below may submit the whole short prefill as one
* Metal command buffer. Once that command is in flight there is no useful
* safe prefix to expose: by the time cancellation can be observed again,
* the prompt has already been fully read and the KV is valid. Let the
* caller observe the pending interrupt at generation time instead. */
(void)cancel;
(void)cancel_ud;
(void)cancelled;
return metal_graph_prefill_layer_major(g,
model,
weights,
@@ -19661,7 +19686,10 @@ static bool metal_graph_prefill_chunked_range(
void *progress_ud,
ds4_session_progress_fn display_progress,
void *display_progress_ud,
ds4_imatrix_collector *imatrix) {
ds4_imatrix_collector *imatrix,
ds4_session_cancel_fn cancel,
void *cancel_ud,
bool *cancelled) {
if (n_tokens == 0 || g->prefill_cap == 0) return false;
if (start > (uint32_t)prompt->len) return false;
if (n_tokens > (uint32_t)prompt->len - start) return false;
@@ -19678,7 +19706,10 @@ static bool metal_graph_prefill_chunked_range(
progress,
progress_ud,
display_progress,
display_progress_ud);
display_progress_ud,
cancel,
cancel_ud,
cancelled);
}
uint32_t chunk_cap = g->prefill_cap;
@@ -19697,6 +19728,10 @@ static bool metal_graph_prefill_chunked_range(
}
for (uint32_t pos0 = start; pos0 < end; ) {
if (cancel && cancel(cancel_ud)) {
if (cancelled) *cancelled = true;
return true;
}
const uint32_t remaining = end - pos0;
uint32_t local_cap = chunk_cap;
if (start != 0 && g->prefill_cap != 0) {
@@ -19732,6 +19767,10 @@ static bool metal_graph_prefill_chunked_range(
if (display_progress) {
display_progress(display_progress_ud, "prefill_display", (int)chunk_end, prompt->len);
}
if (cancel && cancel(cancel_ud)) {
if (cancelled) *cancelled = true;
return true;
}
pos0 = chunk_end;
}
if (show_progress) fputc('\n', stderr);
@@ -19760,7 +19799,10 @@ static bool metal_graph_prefill_chunked(
ds4_session_progress_fn progress,
void *progress_ud,
ds4_session_progress_fn display_progress,
void *display_progress_ud) {
void *display_progress_ud,
ds4_session_cancel_fn cancel,
void *cancel_ud,
bool *cancelled) {
if (n_tokens <= 0) return false;
return metal_graph_prefill_chunked_range(g,
model,
@@ -19774,7 +19816,10 @@ static bool metal_graph_prefill_chunked(
progress_ud,
display_progress,
display_progress_ud,
NULL);
NULL,
cancel,
cancel_ud,
cancelled);
}
/* Layer-major speculative target verifier for tiny MTP suffixes.
@@ -20211,7 +20256,8 @@ static int metal_graph_prompt_logits_test(
(uint32_t)t);
}
ok = metal_graph_prefill_raw_swa(&g, model, weights, prompt, n_test,
gpu_logits, true, NULL, NULL);
gpu_logits, true, NULL, NULL,
NULL, NULL, NULL);
if (memory_report) ds4_gpu_print_memory_report("after prompt graph");
if (ok) {
@@ -21651,11 +21697,13 @@ static int generate_metal_graph_raw_swa(
ok = metal_graph_prefill_chunked(&g, model, weights, prompt,
prompt->len, logits, false,
progress, progress_ud,
progress, progress_ud);
progress, progress_ud,
NULL, NULL, NULL);
} else {
ok = metal_graph_prefill_raw_swa(&g, model, weights, prompt,
prompt->len, logits, true,
progress, progress_ud);
progress, progress_ud,
NULL, NULL, NULL);
}
const double t_prefill1 = now_sec();
if (memory_report) ds4_gpu_print_memory_report("after prefill");
@@ -21890,6 +21938,8 @@ struct ds4_session {
void *progress_ud;
ds4_session_progress_fn display_progress;
void *display_progress_ud;
ds4_session_cancel_fn cancel;
void *cancel_ud;
uint32_t prefill_cap;
int ctx_size;
bool checkpoint_valid;
@@ -23698,7 +23748,8 @@ int ds4_engine_collect_imatrix(ds4_engine *e,
NULL, false,
NULL, NULL,
NULL, NULL,
&collector);
&collector,
NULL, NULL, NULL);
} else {
ok = metal_graph_prefill_layer_major(&g, model, weights,
&prompt, 0,
@@ -24699,6 +24750,20 @@ void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn
s->display_progress_ud = ud;
}
void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud) {
if (!s) return;
s->cancel = fn;
s->cancel_ud = ud;
}
static bool ds4_session_cancelled(ds4_session *s) {
return s && s->cancel && s->cancel(s->cancel_ud);
}
static bool ds4_session_cancelled_cb(void *ud) {
return ds4_session_cancelled(ud);
}
void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total) {
if (!s || !s->progress || !event) return;
s->progress(s->progress_ud, event, current, total);
@@ -25139,6 +25204,10 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
snprintf(err, errlen, "prompt exceeds context");
return 1;
}
if (ds4_session_cancelled(s)) {
snprintf(err, errlen, "interrupted");
return DS4_SESSION_SYNC_INTERRUPTED;
}
if (s->distributed) {
const ds4_tokens *checkpoint = s->checkpoint_valid ? &s->checkpoint : NULL;
return ds4_dist_session_sync(s->distributed,
@@ -25157,6 +25226,12 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
{
s->mtp_draft_valid = false;
for (int i = s->checkpoint.len; i < prompt->len; i++) {
if (ds4_session_cancelled(s)) {
snprintf(err, errlen, "interrupted");
s->checkpoint_valid = true;
s->mtp_draft_valid = false;
return DS4_SESSION_SYNC_INTERRUPTED;
}
forward_token_raw_swa_cpu_decode_scratch(s->logits,
&e->model,
&e->weights,
@@ -25206,14 +25281,13 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
const int suffix = prompt->len - s->checkpoint.len;
const uint32_t resume_min = metal_graph_resume_prefill_min_tokens();
if (suffix > 0 && (uint32_t)suffix >= resume_min) {
bool cancelled = false;
ds4_sync_progress progress = {
.session = s,
.prompt = prompt,
.user = s->progress,
.user_ud = s->progress_ud,
};
ds4_session_progress_fn progress_fn =
s->progress ? ds4_session_note_prefill_progress : NULL;
bool ok = metal_graph_prefill_chunked_range(&s->graph,
&e->model,
&e->weights,
@@ -25222,11 +25296,20 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
(uint32_t)suffix,
s->logits,
false,
progress_fn,
progress_fn ? &progress : NULL,
ds4_session_note_prefill_progress,
&progress,
s->display_progress,
s->display_progress_ud,
NULL);
NULL,
ds4_session_cancelled_cb,
s,
&cancelled);
if (cancelled) {
snprintf(err, errlen, "interrupted");
s->checkpoint_valid = true;
s->mtp_draft_valid = false;
return DS4_SESSION_SYNC_INTERRUPTED;
}
if (!ok) {
snprintf(err, errlen, "%s resumed prefill failed while extending checkpoint", backend_name);
s->checkpoint_valid = false;
@@ -25238,6 +25321,12 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
}
for (int i = s->checkpoint.len; i < prompt->len; i++) {
if (ds4_session_cancelled(s)) {
snprintf(err, errlen, "interrupted");
s->checkpoint_valid = true;
s->mtp_draft_valid = false;
return DS4_SESSION_SYNC_INTERRUPTED;
}
if (!metal_graph_eval_token_raw_swa(&s->graph, &e->model, &e->weights,
(uint32_t)prompt->v[i],
(uint32_t)s->checkpoint.len,
@@ -25254,30 +25343,47 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t
bool ok;
s->checkpoint_valid = false;
s->checkpoint.len = 0;
s->mtp_draft_valid = false;
if (!metal_graph_reset_prefill_state(&s->graph)) {
snprintf(err, errlen, "%s prefill state reset failed", backend_name);
return 1;
}
if (s->prefill_cap < (uint32_t)prompt->len) {
bool cancelled = false;
ds4_sync_progress progress = {
.session = s,
.prompt = prompt,
.user = s->progress,
.user_ud = s->progress_ud,
};
ds4_session_progress_fn progress_fn =
s->progress ? ds4_session_note_prefill_progress : NULL;
ok = metal_graph_prefill_chunked(&s->graph, &e->model, &e->weights,
prompt, prompt->len, s->logits, false,
progress_fn, progress_fn ? &progress : NULL,
ds4_session_note_prefill_progress, &progress,
s->display_progress,
s->display_progress_ud);
s->display_progress_ud,
ds4_session_cancelled_cb,
s,
&cancelled);
if (cancelled) {
snprintf(err, errlen, "interrupted");
s->checkpoint_valid = s->checkpoint.len > 0;
s->mtp_draft_valid = false;
return DS4_SESSION_SYNC_INTERRUPTED;
}
} else {
bool cancelled = false;
ok = metal_graph_prefill_raw_swa(&s->graph, &e->model, &e->weights,
prompt, prompt->len, s->logits, false,
s->display_progress,
s->display_progress_ud);
s->display_progress_ud,
ds4_session_cancelled_cb,
s,
&cancelled);
if (cancelled) {
snprintf(err, errlen, "interrupted");
return DS4_SESSION_SYNC_INTERRUPTED;
}
}
if (!ok) {
snprintf(err, errlen, "%s prefill failed", backend_name);
ds4.h
+7
View File
@@ -60,6 +60,9 @@ typedef struct ds4_engine ds4_engine;
typedef struct ds4_session ds4_session;
typedef void (*ds4_session_progress_fn)(void *ud, const char *event, int current, int total);
typedef bool (*ds4_session_cancel_fn)(void *ud);
#define DS4_SESSION_SYNC_INTERRUPTED 2
typedef enum {
DS4_DISTRIBUTED_NONE = 0,
@@ -223,6 +226,10 @@ void ds4_session_set_progress(ds4_session *s, ds4_session_progress_fn fn, void *
/* UI-only progress. It may report fine-grained progress inside a prefill chunk;
* callers must not treat it as a durable KV checkpoint boundary. */
void ds4_session_set_display_progress(ds4_session *s, ds4_session_progress_fn fn, void *ud);
/* Optional cooperative cancellation. ds4_session_sync() checks it only at
* safe boundaries where the live checkpoint is either unchanged or represents a
* valid token prefix, and returns DS4_SESSION_SYNC_INTERRUPTED when it stops. */
void ds4_session_set_cancel(ds4_session *s, ds4_session_cancel_fn fn, void *ud);
void ds4_session_report_progress(ds4_session *s, const char *event, int current, int total);
/* Distributed coordinator sessions return 1 when the full layer route is
* available, 0 when it is still incomplete, and -1 for a local API error. */
ds4_agent.c
+116 -16
View File
@@ -3525,6 +3525,24 @@ static bool worker_should_interrupt(agent_worker *w) {
return interrupt;
}
/* Ctrl+C is a latched request consumed by the worker. Once an interrupted
* operation has reached a stable append-only boundary and is about to publish
* IDLE, the request must be acknowledged; otherwise the editor can observe an
* idle worker with a stale interrupt still pending. */
static void worker_clear_interrupt(agent_worker *w) {
pthread_mutex_lock(&w->mu);
w->interrupt = false;
pthread_mutex_unlock(&w->mu);
}
static bool agent_err_is_interrupted(const char *err) {
return err && !strcmp(err, "interrupted");
}
static bool worker_cancel_session_cb(void *ud) {
return worker_should_interrupt(ud);
}
typedef struct {
char *ptr;
size_t len;
@@ -4026,9 +4044,15 @@ static int agent_web_confirm(void *privdata, const char *message,
snprintf(w->web_approval_message, sizeof(w->web_approval_message),
"%s", message ? message : "Start visible Chrome browser? (y/n) ");
agent_wake_locked(w);
while (!w->stop && !w->web_approval_answered)
while (!w->stop && !w->interrupt && !w->web_approval_answered)
pthread_cond_wait(&w->cond, &w->mu);
bool ok = w->web_approval_result;
if (!w->web_approval_answered && (w->stop || w->interrupt)) {
ok = false;
w->web_approval_pending = false;
snprintf(w->web_approval_error, sizeof(w->web_approval_error),
"interrupted");
}
if (!ok) {
snprintf(err, err_len, "%s",
w->web_approval_error[0] ? w->web_approval_error :
@@ -4044,6 +4068,10 @@ static void agent_web_log(void *privdata, const char *message) {
agent_trace(w, "web: %s", message);
}
static bool agent_web_cancel(void *privdata) {
return worker_should_interrupt(privdata);
}
static bool worker_take_web_approval_request(agent_worker *w,
char *message, size_t message_len) {
pthread_mutex_lock(&w->mu);
@@ -4145,7 +4173,9 @@ static int agent_worker_sync_tokens(agent_worker *w, const ds4_tokens *tokens,
ds4_session_set_display_progress(w->session,
publish_progress ? worker_progress_cb : NULL,
publish_progress ? w : NULL);
ds4_session_set_cancel(w->session, worker_cancel_session_cb, w);
int rc = ds4_session_sync(w->session, tokens, err, err_len);
ds4_session_set_cancel(w->session, NULL, NULL);
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
return rc;
@@ -7237,17 +7267,29 @@ static bool agent_worker_compact(agent_worker *w, const char *reason,
ds4_session_set_progress(w->session, worker_progress_cb, w);
ds4_session_set_display_progress(w->session, worker_progress_cb, w);
if (ds4_session_sync(w->session, &prompt, err, err_len) != 0) {
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
ds4_session_set_cancel(w->session, worker_cancel_session_cb, w);
int sync_rc = ds4_session_sync(w->session, &prompt, err, err_len);
ds4_session_set_cancel(w->session, NULL, NULL);
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
if (sync_rc == DS4_SESSION_SYNC_INTERRUPTED) {
ds4_session_invalidate(w->session);
snprintf(err, err_len, "interrupted");
agent_publish_system_status(
w, "Compaction interrupted; keeping the previous conversation state.");
ds4_tokens_free(&prompt);
ds4_tokens_free(&sys);
agent_publish(w, "\x1b[0m\n", 5);
worker_clear_interrupt(w);
return false;
}
if (sync_rc != 0) {
ds4_session_invalidate(w->session);
ds4_tokens_free(&prompt);
ds4_tokens_free(&sys);
agent_publish(w, "\x1b[0m\n", 5);
return false;
}
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
/* From here until the final rebuild, the live KV contains the internal
* compaction prompt/summary, while w->transcript still contains the real
@@ -7260,12 +7302,15 @@ static bool agent_worker_compact(agent_worker *w, const char *reason,
double t0 = now_sec();
for (int i = 0; i < summary_max; i++) {
if (worker_should_interrupt(w)) {
snprintf(err, err_len, "compaction interrupted");
snprintf(err, err_len, "interrupted");
ds4_session_invalidate(w->session);
ds4_tokens_free(&prompt);
ds4_tokens_free(&sys);
free(summary.ptr);
agent_publish(w, "\x1b[0m\n", 5);
agent_publish_system_status(
w, "Compaction interrupted; keeping the previous conversation state.");
worker_clear_interrupt(w);
return false;
}
int token = ds4_session_argmax(w->session);
@@ -7483,10 +7528,21 @@ static void worker_set_greedy_sampling(agent_worker *w, bool greedy) {
static int worker_run_turn(agent_worker *w, const char *user_text) {
agent_config *cfg = w->cfg;
ds4_think_mode think_mode = effective_think_mode(cfg);
pthread_mutex_lock(&w->mu);
w->interrupt = false;
w->status.error[0] = '\0';
agent_wake_locked(w);
pthread_mutex_unlock(&w->mu);
char compact_err[160] = {0};
if (!agent_worker_compact_if_needed(w, "soft limit before user turn",
compact_err, sizeof(compact_err)))
{
if (agent_err_is_interrupted(compact_err)) {
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return 0;
}
agent_set_error(w, compact_err[0] ? compact_err : "context compaction failed");
return 1;
}
@@ -7504,10 +7560,8 @@ static int worker_run_turn(agent_worker *w, const char *user_text) {
uint64_t rng = cfg->gen.seed ? cfg->gen.seed :
((uint64_t)time(NULL) ^ ((uint64_t)getpid() << 32) ^ (uint64_t)clock());
pthread_mutex_lock(&w->mu);
w->interrupt = false;
w->user_activity = true;
w->session_dirty = true;
w->status.error[0] = '\0';
agent_wake_locked(w);
pthread_mutex_unlock(&w->mu);
@@ -7523,6 +7577,11 @@ static int worker_run_turn(agent_worker *w, const char *user_text) {
!agent_worker_compact_if_needed(w, "soft limit before tool continuation",
compact_err, sizeof(compact_err)))
{
if (agent_err_is_interrupted(compact_err)) {
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return 0;
}
agent_set_error(w, compact_err[0] ? compact_err : "context compaction failed");
return 1;
}
@@ -7559,14 +7618,23 @@ static int worker_run_turn(agent_worker *w, const char *user_text) {
char err[160];
ds4_session_set_progress(w->session, worker_progress_cb, w);
ds4_session_set_display_progress(w->session, worker_progress_cb, w);
if (ds4_session_sync(w->session, prompt_for_sync, err, sizeof(err)) != 0) {
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
ds4_session_set_cancel(w->session, worker_cancel_session_cb, w);
int sync_rc = ds4_session_sync(w->session, prompt_for_sync, err, sizeof(err));
ds4_session_set_cancel(w->session, NULL, NULL);
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
if (sync_rc == DS4_SESSION_SYNC_INTERRUPTED) {
agent_publish_system_status(
w, "Model reading interrupted; the model may only be aware of the prefix processed so far.");
ds4_tokens_push(&w->transcript, ds4_token_eos(w->engine));
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return 0;
}
if (sync_rc != 0) {
agent_set_error(w, err);
return 1;
}
ds4_session_set_progress(w->session, NULL, NULL);
ds4_session_set_display_progress(w->session, NULL, NULL);
int max_tokens = cfg->gen.n_predict;
int room = ds4_session_ctx(w->session) - ds4_session_pos(w->session);
@@ -7667,6 +7735,14 @@ static int worker_run_turn(agent_worker *w, const char *user_text) {
agent_stream_text(&stream, NULL, 0, true);
renderer_finish(&renderer);
worker_set_greedy_sampling(w, false);
if (interrupted) {
ds4_tokens_push(&w->transcript, ds4_token_eos(w->engine));
agent_dsml_parser_free(&dsml);
agent_publish_system_status(w, "Stopped by user");
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return 0;
}
if (stream.dsml_in_think) {
got_tool = false;
malformed_tool = true;
@@ -7722,6 +7798,11 @@ static int worker_run_turn(agent_worker *w, const char *user_text) {
{
free(tool_result);
agent_dsml_parser_free(&dsml);
if (agent_err_is_interrupted(compact_err)) {
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return 0;
}
agent_set_error(w, compact_err[0] ? compact_err : "context compaction failed");
return 1;
}
@@ -7866,6 +7947,11 @@ static void worker_run_deferred_compact(agent_worker *w) {
}
agent_set_status(w, AGENT_WORKER_IDLE);
} else {
if (agent_err_is_interrupted(err)) {
worker_clear_interrupt(w);
agent_set_status(w, AGENT_WORKER_IDLE);
return;
}
agent_set_error(w, err[0] ? err : "context compaction failed");
}
}
@@ -9242,6 +9328,8 @@ static int agent_worker_init(agent_worker *w, ds4_engine *engine, agent_config *
.confirm_privdata = w,
.log = agent_web_log,
.log_privdata = w,
.cancel = agent_web_cancel,
.cancel_privdata = w,
};
w->web = ds4_web_create(&web_cfg);
w->sysprompt_path = ds4_kvstore_path_join(w->cache_dir, "sysprompt.kv");
@@ -9632,6 +9720,20 @@ static int run_agent(ds4_engine *engine, agent_config *cfg) {
}
}
if (rc > 0 && (pfd[0].revents & POLLIN)) editor_read_stdin(&editor);
/* Linenoise runs the terminal in raw mode, so Ctrl+C normally arrives
* as byte 3 instead of SIGINT. Handle it before worker output is
* drained and repainted; otherwise a busy decoding stream can leave the
* interrupt waiting behind a large terminal-output backlog. */
if (editor_take_queued_byte(&editor, 3)) { /* Ctrl+C */
if (!worker_is_idle(&worker)) {
worker_interrupt(&worker);
} else {
editor_cancel_input_with_hint(&editor, prompt, statusline);
}
}
if (rc > 0 && (pfd[1].revents & POLLIN)) drain_wake_fd(worker.wake_fd[0]);
char *out = NULL;
@@ -9730,8 +9832,6 @@ static int run_agent(ds4_engine *engine, agent_config *cfg) {
free(queued);
}
if (rc > 0 && (pfd[0].revents & POLLIN)) editor_read_stdin(&editor);
if (queue.len && editor_take_queued_byte(&editor, 24)) { /* Ctrl+X */
char *queued = agent_prompt_queue_pop(&queue);
editor_replace_input(&editor, queued);
ds4_web.c
+145 -29
View File
@@ -47,12 +47,15 @@ struct ds4_web {
void *confirm_privdata;
ds4_web_log_fn log;
void *log_privdata;
ds4_web_cancel_fn cancel;
void *cancel_privdata;
int next_cdp_id;
};
typedef struct {
int fd;
int next_id;
ds4_web *web;
} cdp_ws;
typedef struct {
@@ -119,6 +122,37 @@ static void web_log(ds4_web *web, const char *msg) {
if (web && web->log) web->log(web->log_privdata, msg);
}
static double web_now_sec(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec / 1000000000.0;
}
static bool web_cancelled(ds4_web *web) {
return web && web->cancel && web->cancel(web->cancel_privdata);
}
static bool web_set_cancel_err(ds4_web *web, char *err, size_t err_len) {
if (!web_cancelled(web)) return false;
web_set_err(err, err_len, "interrupted");
return true;
}
static bool web_sleep_ms(ds4_web *web, int ms) {
int left = ms;
while (left > 0) {
if (web_cancelled(web)) return false;
int step = left < 50 ? left : 50;
usleep((useconds_t)step * 1000u);
left -= step;
}
return !web_cancelled(web);
}
static bool web_err_is_interrupted(const char *err) {
return err && !strcmp(err, "interrupted");
}
static bool web_mkdir_p(const char *path) {
if (!path || !path[0]) return false;
char tmp[PATH_MAX];
@@ -401,14 +435,31 @@ static int web_ws_connect(const char *ws_url, cdp_ws *ws,
web_buf resp = {0};
char tmp[1024];
double deadline = web_now_sec() + (double)DS4_WEB_CONNECT_TIMEOUT_MS / 1000.0;
while (!strstr(resp.ptr ? resp.ptr : "", "\r\n\r\n")) {
ssize_t n = web_read_some(fd, tmp, sizeof(tmp), DS4_WEB_CONNECT_TIMEOUT_MS);
if (n <= 0) {
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) {
close(fd);
free(resp.ptr);
return -1;
}
double now = web_now_sec();
if (now >= deadline) {
web_set_err(err, err_len, "websocket handshake read failed");
close(fd);
free(resp.ptr);
return -1;
}
int slice = 100;
int remaining = (int)((deadline - now) * 1000.0);
if (remaining < slice) slice = remaining > 0 ? remaining : 1;
ssize_t n = web_read_some(fd, tmp, sizeof(tmp), slice);
if (n < 0) {
web_set_err(err, err_len, "websocket handshake read failed");
close(fd);
free(resp.ptr);
return -1;
}
if (n == 0) continue;
web_buf_append(&resp, tmp, (size_t)n);
if (resp.len > 8192) break;
}
@@ -431,11 +482,26 @@ static void web_ws_close(cdp_ws *ws) {
}
}
static int web_read_exact(int fd, unsigned char *buf, size_t len, int timeout_ms) {
static int web_read_exact(cdp_ws *ws, unsigned char *buf, size_t len,
int timeout_ms, char *err, size_t err_len) {
size_t off = 0;
double deadline = web_now_sec() + (double)timeout_ms / 1000.0;
while (off < len) {
ssize_t n = web_read_some(fd, (char *)buf + off, len - off, timeout_ms);
if (n <= 0) return -1;
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return -1;
double now = web_now_sec();
if (now >= deadline) {
web_set_err(err, err_len, "websocket read timeout");
return -1;
}
int slice = 100;
int remaining = (int)((deadline - now) * 1000.0);
if (remaining < slice) slice = remaining > 0 ? remaining : 1;
ssize_t n = web_read_some(ws->fd, (char *)buf + off, len - off, slice);
if (n < 0) {
web_set_err(err, err_len, "websocket frame read failed");
return -1;
}
if (n == 0) continue;
off += (size_t)n;
}
return 0;
@@ -491,8 +557,7 @@ static char *web_ws_read_message(cdp_ws *ws, char *err, size_t err_len) {
web_buf msg = {0};
for (;;) {
unsigned char h[2];
if (web_read_exact(ws->fd, h, 2, DS4_WEB_CDP_TIMEOUT_MS) != 0) {
web_set_err(err, err_len, "websocket read timeout");
if (web_read_exact(ws, h, 2, DS4_WEB_CDP_TIMEOUT_MS, err, err_len) != 0) {
free(msg.ptr);
return NULL;
}
@@ -502,16 +567,16 @@ static char *web_ws_read_message(cdp_ws *ws, char *err, size_t err_len) {
uint64_t len = h[1] & 0x7f;
if (len == 126) {
unsigned char x[2];
if (web_read_exact(ws->fd, x, 2, DS4_WEB_CDP_TIMEOUT_MS) != 0) goto fail;
if (web_read_exact(ws, x, 2, DS4_WEB_CDP_TIMEOUT_MS, err, err_len) != 0) goto fail;
len = ((uint64_t)x[0] << 8) | x[1];
} else if (len == 127) {
unsigned char x[8];
if (web_read_exact(ws->fd, x, 8, DS4_WEB_CDP_TIMEOUT_MS) != 0) goto fail;
if (web_read_exact(ws, x, 8, DS4_WEB_CDP_TIMEOUT_MS, err, err_len) != 0) goto fail;
len = 0;
for (int i = 0; i < 8; i++) len = (len << 8) | x[i];
}
unsigned char mask[4] = {0};
if (masked && web_read_exact(ws->fd, mask, 4, DS4_WEB_CDP_TIMEOUT_MS) != 0)
if (masked && web_read_exact(ws, mask, 4, DS4_WEB_CDP_TIMEOUT_MS, err, err_len) != 0)
goto fail;
if (len > DS4_WEB_MAX_RESULT_BYTES * 4ULL) {
web_set_err(err, err_len, "websocket message too large");
@@ -519,8 +584,8 @@ static char *web_ws_read_message(cdp_ws *ws, char *err, size_t err_len) {
return NULL;
}
unsigned char *payload = web_xmalloc((size_t)len + 1);
if (len && web_read_exact(ws->fd, payload, (size_t)len,
DS4_WEB_CDP_TIMEOUT_MS) != 0) {
if (len && web_read_exact(ws, payload, (size_t)len,
DS4_WEB_CDP_TIMEOUT_MS, err, err_len) != 0) {
free(payload);
goto fail;
}
@@ -544,7 +609,7 @@ static char *web_ws_read_message(cdp_ws *ws, char *err, size_t err_len) {
}
}
fail:
web_set_err(err, err_len, "websocket frame read failed");
if (err && err_len && !err[0]) web_set_err(err, err_len, "websocket frame read failed");
free(msg.ptr);
return NULL;
}
@@ -561,6 +626,7 @@ static bool web_json_id_matches(const char *json, int id) {
static char *web_cdp_call(cdp_ws *ws, const char *method, const char *params,
char *err, size_t err_len) {
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return NULL;
int id = ws->next_id++;
web_buf req = {0};
char head[256];
@@ -581,6 +647,7 @@ static char *web_cdp_call(cdp_ws *ws, const char *method, const char *params,
}
free(wire);
for (;;) {
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return NULL;
char *msg = web_ws_read_message(ws, err, err_len);
if (!msg) return NULL;
if (web_json_id_matches(msg, id)) return msg;
@@ -720,14 +787,19 @@ static char *web_cdp_eval_string(cdp_ws *ws, const char *expr,
static bool web_wait_ready(cdp_ws *ws, char *err, size_t err_len) {
const char *expr = "document.readyState";
for (int i = 0; i < 80; i++) {
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return false;
char *state = web_cdp_eval_string(ws, expr, err, err_len);
if (state && (!strcmp(state, "complete") || !strcmp(state, "interactive"))) {
free(state);
usleep(800000);
return true;
if (web_sleep_ms(ws ? ws->web : NULL, 800)) return true;
web_set_err(err, err_len, "interrupted");
return false;
}
free(state);
usleep(250000);
if (!web_sleep_ms(ws ? ws->web : NULL, 250)) {
web_set_err(err, err_len, "interrupted");
return false;
}
}
return true;
}
@@ -780,6 +852,7 @@ static bool web_wait_navigated_ready(cdp_ws *ws, const char *url,
bool saw_real_url = false;
for (int i = 0; i < 100; i++) {
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return false;
char *href = NULL;
char *ready = NULL;
long text_len = 0;
@@ -787,7 +860,10 @@ static bool web_wait_navigated_ready(cdp_ws *ws, const char *url,
if (!ok) {
free(href);
free(ready);
usleep(250000);
if (!web_sleep_ms(ws ? ws->web : NULL, 250)) {
web_set_err(err, err_len, "interrupted");
return false;
}
continue;
}
@@ -805,11 +881,15 @@ static bool web_wait_navigated_ready(cdp_ws *ws, const char *url,
free(ready);
if (saw_real_url && ready_state && text_len > 0 && stable >= 2) {
usleep(500000);
return true;
if (web_sleep_ms(ws ? ws->web : NULL, 500)) return true;
web_set_err(err, err_len, "interrupted");
return false;
}
if (saw_real_url && ready_state && i >= 24) return true;
usleep(250000);
if (!web_sleep_ms(ws ? ws->web : NULL, 250)) {
web_set_err(err, err_len, "interrupted");
return false;
}
}
return true;
}
@@ -828,7 +908,7 @@ static bool web_cdp_prepare_page(cdp_ws *ws, char *err, size_t err_len) {
return web_wait_ready(ws, err, err_len);
}
static void web_scroll_dynamic_page(cdp_ws *ws) {
static bool web_scroll_dynamic_page(cdp_ws *ws, char *err, size_t err_len) {
const char *expr =
"(() => new Promise(resolve => {"
"const root=()=>document.scrollingElement||document.documentElement||document.body;"
@@ -867,9 +947,17 @@ static void web_scroll_dynamic_page(cdp_ws *ws) {
"tick();},900);"
"};tick();"
"}))()";
char err[160] = {0};
char *res = web_cdp_eval_string(ws, expr, err, sizeof(err));
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return false;
char local_err[160] = {0};
char *res = web_cdp_eval_string(ws, expr, local_err, sizeof(local_err));
if (!res && web_err_is_interrupted(local_err)) {
web_set_err(err, err_len, "interrupted");
return false;
}
free(res);
if (web_set_cancel_err(ws ? ws->web : NULL, err, err_len)) return false;
return true;
}
static char *web_chrome_executable(void) {
@@ -992,6 +1080,7 @@ static bool web_spawn_chrome(ds4_web *web, char *err, size_t err_len) {
free(exe);
web->chrome_pid = pid;
for (int i = 0; i < 80; i++) {
if (web_set_cancel_err(web, err, err_len)) return false;
if (web_cdp_alive(web)) {
web_log(web, "Chrome browser session is ready");
return true;
@@ -1004,7 +1093,10 @@ static bool web_spawn_chrome(ds4_web *web, char *err, size_t err_len) {
web_set_err(err, err_len, "Chrome exited before CDP became ready");
return false;
}
usleep(250000);
if (!web_sleep_ms(web, 250)) {
web_set_err(err, err_len, "interrupted");
return false;
}
}
web_set_err(err, err_len, "Chrome did not expose CDP on port %d", web->port);
return false;
@@ -1058,7 +1150,7 @@ static bool web_open_tab(ds4_web *web, const char *url, web_tab *tab,
char *browser_url = web_browser_ws_url(web, err, err_len);
if (!browser_url) return false;
cdp_ws browser = {.fd = -1};
cdp_ws browser = {.fd = -1, .web = web};
if (web_ws_connect(browser_url, &browser, err, err_len) != 0) {
free(browser_url);
return false;
@@ -1172,7 +1264,7 @@ static char *web_run_page_js(ds4_web *web, const char *url, const char *js,
if (!web_ensure_browser(web, err, err_len)) return NULL;
web_tab tab = {0};
if (!web_open_tab(web, "about:blank", &tab, err, err_len)) return NULL;
cdp_ws ws = {.fd = -1};
cdp_ws ws = {.fd = -1, .web = web};
if (web_ws_connect(tab.ws_url, &ws, err, err_len) != 0) {
web_close_tab(web, &tab);
web_tab_free(&tab);
@@ -1195,11 +1287,33 @@ static char *web_run_page_js(ds4_web *web, const char *url, const char *js,
char *clicked = web_cdp_eval_string(&ws, web_click_google_consent_js, err, err_len);
if (clicked && clicked[0]) {
web_log(web, clicked);
usleep(1500000);
(void)web_wait_navigated_ready(&ws, url, err, err_len);
if (!web_sleep_ms(web, 1500)) {
free(clicked);
web_set_err(err, err_len, "interrupted");
web_ws_close(&ws);
web_close_tab(web, &tab);
web_tab_free(&tab);
return NULL;
}
char wait_err[160] = {0};
if (!web_wait_navigated_ready(&ws, url, wait_err, sizeof(wait_err)) &&
web_err_is_interrupted(wait_err))
{
free(clicked);
web_set_err(err, err_len, "interrupted");
web_ws_close(&ws);
web_close_tab(web, &tab);
web_tab_free(&tab);
return NULL;
}
}
free(clicked);
if (dynamic_scroll) web_scroll_dynamic_page(&ws);
if (dynamic_scroll && !web_scroll_dynamic_page(&ws, err, err_len)) {
web_ws_close(&ws);
web_close_tab(web, &tab);
web_tab_free(&tab);
return NULL;
}
char *out = web_cdp_eval_string(&ws, js, err, err_len);
web_ws_close(&ws);
web_close_tab(web, &tab);
@@ -1223,6 +1337,8 @@ ds4_web *ds4_web_create(const ds4_web_config *cfg) {
web->confirm_privdata = cfg->confirm_privdata;
web->log = cfg->log;
web->log_privdata = cfg->log_privdata;
web->cancel = cfg->cancel;
web->cancel_privdata = cfg->cancel_privdata;
}
return web;
}
ds4_web.h
+3
View File
@@ -7,6 +7,7 @@
typedef int (*ds4_web_confirm_fn)(void *privdata, const char *message,
char *err, size_t err_len);
typedef void (*ds4_web_log_fn)(void *privdata, const char *message);
typedef bool (*ds4_web_cancel_fn)(void *privdata);
typedef struct {
const char *home_dir;
@@ -15,6 +16,8 @@ typedef struct {
void *confirm_privdata;
ds4_web_log_fn log;
void *log_privdata;
ds4_web_cancel_fn cancel;
void *cancel_privdata;
} ds4_web_config;
typedef struct ds4_web ds4_web;