diff --git a/src/sqlite-ai.c b/src/sqlite-ai.c index a20e9b8..f8df632 100644 --- a/src/sqlite-ai.c +++ b/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; } diff --git a/tests/c/unittest.c b/tests/c/unittest.c index 8868a35..86664cf 100644 --- a/tests/c/unittest.c +++ b/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},