fix(chat): keep the chat usable when it runs out of context

src/sqlite-ai.c

@@ -1821,9 +1821,13 @@ static bool llm_chat_generate_response (ai_context *ai, ai_cursor *c, bool *is_e
     llama_batch batch = ai->chat.batch;
     char *tok = ai->chat.token_text;
 
-    // check context space
-    uint32_t n_ctx = llama_n_ctx(ctx);
-    int32_t n_ctx_used = llama_memory_seq_pos_max(llama_get_memory(ctx), 0);
+    // check context space. seq_pos_max() is the highest position in the sequence, so
+    // occupancy is that + 1; treating it as a count let exactly one over-large batch
+    // through for llama_decode() to reject with "could not find a KV slot" instead.
+    // Both sides are signed: the comparison used to promote to unsigned, so an empty
+    // cache (seq_pos_max == -1) and a zero-token batch tripped the guard falsely.
+    int32_t n_ctx = (int32_t)llama_n_ctx(ctx);
+    int32_t n_ctx_used = llama_memory_seq_pos_max(llama_get_memory(ctx), 0) + 1;
     if (n_ctx_used + batch.n_tokens > n_ctx) {
         sqlite_common_set_error (ai->context, ai->vtab, SQLITE_ERROR, "Context size exceeded (%d, %d)", n_ctx, n_ctx_used + batch.n_tokens);
         return false;
@@ -1990,7 +1994,17 @@ static bool llm_chat_run (ai_context *ai, ai_cursor *c, const char *user_prompt)
     // do not stream response and incrementally build the buffer
     bool is_eog = false;
     while (1) {
-        if (!llm_chat_generate_response (ai, NULL, &is_eog)) return false;
+        if (!llm_chat_generate_response (ai, NULL, &is_eog)) {
+            // The turn failed part-way. The user message is already in the history and
+            // its tokens are already in the KV cache, so commit whatever was generated
+            // before reporting: leaving the turn half-applied strands the user message
+            // with no assistant reply and leaves prev_len stale, which makes the delta
+            // for the *next* turn re-include it and fail too. The streaming cursor path
+            // already saves on close for the same reason. Errors from the save itself
+            // are allowed to replace the original message - they are worth reporting.
+            llm_chat_save_response(ai, messages, template);
+            return false;
+        }
         if (is_eog) break;
     }
 

tests/c/unittest.c

@@ -159,6 +159,16 @@ static int install_seeded_chat_sampler(const test_env *env, sqlite3 *db) {
     return 0;
 }
 
+// like exec_expect_ok, but hands back the error text so a test can compare two failures
+static int exec_capture_error(const test_env *env, sqlite3 *db, const char *sql, char *err, size_t errlen) {
+    if (env->verbose) printf("[SQL] %s\n", sql);
+    char *errmsg = NULL;
+    int rc = sqlite3_exec(db, sql, NULL, NULL, &errmsg);
+    if (rc != SQLITE_OK && err && errlen) snprintf(err, errlen, "%s", errmsg ? errmsg : sqlite3_errmsg(db));
+    if (errmsg) sqlite3_free(errmsg);
+    return (rc == SQLITE_OK) ? 0 : 1;
+}
+
 static int exec_select_rows(const test_env *env, sqlite3 *db, const char *sql, int *rows_out) {
     if (env->verbose) {
         printf("[SQL] %s\n", sql);
@@ -1346,6 +1356,99 @@ static int test_chat_recreate_after_conversation(const test_env *env) {
     return 1;
 }
 
+// Regression: running a chat out of context used to corrupt it. The turn bailed out of
+// llm_chat_run() before llm_chat_save_response(), so the user message stayed in the
+// history with no assistant reply and prev_len was never advanced - which made the next
+// turn re-send the stranded message, so the reported requirement grew on every retry
+// (266, 276, 286, 296 ...) and the chat could never be used again. The guard also
+// treated llama_memory_seq_pos_max() (a position) as an occupancy count, so it let one
+// over-large batch through for llama_decode() to reject instead.
+static int test_chat_context_full_is_recoverable(const test_env *env) {
+    sqlite3 *db = NULL;
+    if (open_db_and_load(env, &db) != SQLITE_OK) return 1;
+
+    const char *model = env->model_path ? env->model_path : DEFAULT_MODEL_PATH;
+    char sqlbuf[512];
+    snprintf(sqlbuf, sizeof(sqlbuf), "SELECT llm_model_load('%s');", model);
+    if (exec_expect_ok(env, db, sqlbuf) != 0) goto fail;
+    if (install_seeded_chat_sampler(env, db) != 0) goto fail;
+    if (exec_expect_ok(env, db, "SELECT llm_context_create_chat('context_size=256');") != 0) goto fail;
+    if (exec_expect_ok(env, db, "SELECT llm_chat_create();") != 0) goto fail;
+
+    // talk until the context runs out, then keep going: the failures are the point
+    const char *turn = "SELECT llm_chat_respond('Tell me about computers.');";
+    #define MAX_FAILURES 4
+    char errs[MAX_FAILURES][256] = {{0}};
+    int ok_turns = 0, failed_turns = 0;
+    for (int i = 0; i < 40 && failed_turns < MAX_FAILURES; ++i) {
+        char err[256] = {0};
+        if (exec_capture_error(env, db, turn, err, sizeof(err)) == 0) {
+            if (failed_turns == 0) ok_turns++;
+            continue;
+        }
+        snprintf(errs[failed_turns], sizeof(errs[0]), "%s", err);
+        failed_turns++;
+    }
+
+    if (ok_turns == 0 || failed_turns < MAX_FAILURES) {
+        fprintf(stderr, "[chat_context_full_is_recoverable] wanted some good turns then %d failures, "
+                        "got %d good and %d failed\n", MAX_FAILURES, ok_turns, failed_turns);
+        goto fail;
+    }
+    if (env->verbose) {
+        printf("[chat_context_full_is_recoverable] %d turns fit, then:\n", ok_turns);
+        for (int i = 0; i < failed_turns; ++i) printf("    #%d %s\n", i + 1, errs[i]);
+    }
+
+    // the guard must catch it, not llama_decode() one token later
+    if (strstr(errs[0], "Context size exceeded") == NULL) {
+        fprintf(stderr, "[chat_context_full_is_recoverable] expected the context guard to report it, got: %s\n", errs[0]);
+        goto fail;
+    }
+
+    // Once retrying, every attempt must report the SAME requirement: each one sends the
+    // same user message against the same full cache. A stranded user turn used to be
+    // re-sent on top, so the number climbed by a turn's worth of tokens every time.
+    // The FIRST failure is excluded on purpose - it can land mid-generation, where the
+    // batch is a single token, so it legitimately reports a different number from the
+    // prompt-batch failures that follow.
+    for (int i = 2; i < failed_turns; ++i) {
+        if (strcmp(errs[1], errs[i]) != 0) {
+            fprintf(stderr, "[chat_context_full_is_recoverable] failure grew across retries:\n");
+            for (int j = 0; j < failed_turns; ++j) fprintf(stderr, "    #%d %s\n", j + 1, errs[j]);
+            goto fail;
+        }
+    }
+    #undef MAX_FAILURES
+
+    // and the history must not end on an orphaned user turn
+    if (exec_expect_ok(env, db, "SELECT llm_chat_save('context full');") != 0) goto fail;
+    ai_chat_message_row rows[128];
+    int count = 0;
+    if (fetch_ai_chat_messages(env, db, rows, 128, &count) != 0) goto fail;
+    if (count < 2) {
+        fprintf(stderr, "[chat_context_full_is_recoverable] expected saved messages, got %d\n", count);
+        goto fail;
+    }
+    if (strcmp(rows[count - 1].role, "assistant") != 0) {
+        fprintf(stderr, "[chat_context_full_is_recoverable] history ends on a '%s' turn with no reply\n",
+                rows[count - 1].role);
+        goto fail;
+    }
+
+    if (exec_expect_ok(env, db, "SELECT llm_chat_free();") != 0) goto fail;
+    if (exec_expect_ok(env, db, "SELECT llm_sampler_free();") != 0) goto fail;
+    if (exec_expect_ok(env, db, "SELECT llm_context_free();") != 0) goto fail;
+    if (exec_expect_ok(env, db, "SELECT llm_model_free();") != 0) goto fail;
+
+    sqlite3_close_v2(db);
+    return assert_sqlite_memory_clean("chat_context_full_is_recoverable", env);
+
+fail:
+    if (db) sqlite3_close_v2(db);
+    return 1;
+}
+
 // Test chat vtab streaming produces tokens and saves response correctly across turns
 static int test_chat_vtab_multi_turn(const test_env *env) {
     sqlite3 *db = NULL;
@@ -2110,6 +2213,7 @@ static const test_case TESTS[] = {
     {"chat_system_prompt_after_first_response", test_chat_system_prompt_after_first_response},
     {"chat_create_free_cycle", test_chat_create_free_cycle},
     {"chat_recreate_after_conversation", test_chat_recreate_after_conversation},
+    {"chat_context_full_is_recoverable", test_chat_context_full_is_recoverable},
     {"chat_vtab_multi_turn", test_chat_vtab_multi_turn},
     {"chat_save_restore_roundtrip", test_chat_save_restore_roundtrip},
     {"chat_system_prompt_clear", test_chat_system_prompt_clear},