DEV Community

Cover image for Dhi Gets a Brain: In-editor Chat with Streaming RAG
Sourish Chakraborty
Sourish Chakraborty

Posted on Originally published at blogs.sourishchakraborty.com

Dhi Gets a Brain: In-editor Chat with Streaming RAG

Note: This post contains architecture diagrams rendered as Mermaid in the original. For the fully rendered version visit the original post.

In Post 1 Dhi learned to finish your sentences: fill-in-the-middle autocomplete with Tree-sitter chunking and StarCoder2-3B. In Post 2 it learned the whole repo: six languages, hybrid BM25 + vector search, one-shot workspace indexing.

Neither of those is a conversation. Autocomplete reacts to a cursor position; it can't answer "why does this function retry three times?" A comment on the architecture post pushed on exactly this gap:

"An AI IDE is not just chat plus autocomplete; it needs repo understanding, edit boundaries, tool execution, review surfaces, and a way to explain why a change was made."

Post 3 builds the piece that ties the first two posts together and starts answering that list: an in-editor chat panel with streaming responses and repo-aware context. Every answer is grounded in the same hybrid search from Post 2 — the chat doesn't just talk about your code, it reads it first.

Everything is at github.com/sochaty/dhi, tag post-3.


What Changed Between Post 2 and Post 3

Component Post 2 Post 3
Interaction model Autocomplete only (implicit, cursor-driven) + Explicit chat panel (VS Code webview)
Response delivery Single JSON response Server-Sent Events, one token at a time
Context source Retrieved chunks only System prompt + active file + retrieved chunks + conversation history
New endpoint POST /chat (StreamingResponse, text/event-stream)
New extension surface extension/src/chat/panel.ts, dhi.openChat command
Cancellation N/A AbortController — a new message aborts the in-flight one

Part 1: Three Processes, One Conversation

A VS Code chat panel is not one program talking to a server — it's three processes relaying messages to each other, and the architecture rule in Dhi's ARCHITECTURE.md is explicit about who's allowed to do what: all HTTP goes through DhiClient, and the webview never calls fetch() directly.

Architecture diagram

Three isolation boundaries, three reasons:

  1. Webview ↔ Extension Host — the webview runs in a sandboxed iframe with a strict Content-Security-Policy (default-src 'none'). It cannot make network requests at all. It can only postMessage to the extension host and receive messages back.
  2. Extension Host ↔ Server — only DhiClient is allowed to call fetch(). Providers and panels call methods on the client; nobody constructs a request URL by hand. This is what makes it possible to point the whole extension at a different server URL from one config setting.
  3. Server ↔ Ollama — the FastAPI process is the only thing that speaks to Ollama's /api/generate. The extension never sees an Ollama URL.

This is more ceremony than a single fetch() call, but it means the webview's CSP can stay maximally strict, and every outbound request funnels through one auditable chokepoint.


Part 2: Context Slot Assembly

A chat model has a fixed context window. Dhi's prompt has five things competing for space, and they don't have equal claim to it. chat.py documents the priority order at the top of the file:

"""Streaming chat endpoint — context assembly + Ollama token stream.

Context slot order (highest → lowest trim priority):
  1. System prompt  (fixed, never trimmed)
  2. User message   (never trimmed)
  3. Active file     (trimmed to _MAX_FILE_CHARS)
  4. RAG chunks     (top _MAX_CHUNKS, each trimmed to _MAX_CHUNK_CHARS)
  5. History        (last _MAX_HISTORY_TURNS turns × 2 messages)
"""
Enter fullscreen mode Exit fullscreen mode
1 · System prompt   ─┐  fixed, never trimmed
2 · User message    ─┤  never trimmed
3 · Active file      ─┤  trimmed to 2,000 chars       ──►  flat prompt string
4 · RAG chunks       ─┤  top 3 · 400 chars each              → Ollama /api/generate
                      │  via store.hybrid_query()
5 · History          ─┘  last 4 turns × 2 messages
Enter fullscreen mode Exit fullscreen mode

The trim budget is deliberately asymmetric — the system prompt and the question you just typed are never touched, but everything that's context about the past (the open file, retrieved chunks, chat history) gets capped. assemble_prompt is a pure function, which is what makes it trivial to unit test without touching Ollama:

def assemble_prompt(request: ChatRequest, rag_chunks: list[str]) -> str:
    """Build a flat-string prompt from prioritised context slots."""
    parts: list[str] = [_SYSTEM_PROMPT]

    if request.file_content:
        content = request.file_content[:_MAX_FILE_CHARS]
        lang = request.language or "text"
        parts.append(f"# Active file: {request.file_path}\n```
{% endraw %}
{lang}\n{content}\n
{% raw %}
```")

    if rag_chunks:
        trimmed = [c[:_MAX_CHUNK_CHARS] for c in rag_chunks[:_MAX_CHUNKS]]
        parts.append("# Relevant code from this repo\n" + "\n---\n".join(trimmed))

    for msg in request.history[-(_MAX_HISTORY_TURNS * 2):]:
        label = "User" if msg.role == "user" else "Assistant"
        parts.append(f"{label}: {msg.content}")

    parts.append(f"User: {request.message}\nAssistant:")
    return "\n\n".join(parts)
Enter fullscreen mode Exit fullscreen mode

The RAG chunks are the interesting part — they come straight from Post 2's hybrid_query(), so a question like "why does the indexer skip .git?" pulls the actual _SKIP_DIRS chunk out of rag/indexer.py, not a hallucinated guess:

async def stream_chat(request: ChatRequest, store: Any) -> AsyncGenerator[str, None]:
    rag_chunks: list[str] = []
    if request.message.strip():
        rag_chunks = store.hybrid_query(request.message, n_results=_MAX_CHUNKS)

    prompt = assemble_prompt(request, rag_chunks)
    ...
Enter fullscreen mode Exit fullscreen mode

This is the same ChunkStore instance the /complete and /search endpoints use — one index, three consumers.


Part 3: The SSE Contract

FastAPI's StreamingResponse turns an async generator into a text/event-stream response. The contract is two lines of protocol:

@app.post("/chat")
async def chat_endpoint(req: ChatRequestModel) -> StreamingResponse:
    """Stream a chat response as Server-Sent Events.

    Each event: ``data: {"token": "..."}\\n\\n``
    Final event: ``data: [DONE]\\n\\n``
    """
    chat_req = ChatReq(...)

    async def event_stream() -> AsyncGenerator[str, None]:
        try:
            async for token in stream_chat(chat_req, store):
                yield f"data: {json.dumps({'token': token})}\n\n"
        except Exception as exc:
            logging.exception("chat_endpoint error")
            yield f"data: {json.dumps({'error': str(exc)})}\n\n"
        finally:
            yield "data: [DONE]\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
    )
Enter fullscreen mode Exit fullscreen mode

The try/except/finally shape matters: if Ollama errors mid-stream, the client gets a {"error": ...} event and still gets [DONE] — the finally block guarantees the stream terminates cleanly even on failure, so the client's read loop never hangs waiting for a close that isn't coming.

DhiClient (host)         FastAPI /chat            Ollama
      | --- POST /chat {message, file, history} -->  |
      |                       | --- POST /api/generate {stream:true} --> |
      |                       |  <== {"response": "...", "done": false}  |  (repeat per token)
      |  <== data: {"token": "..."}                    |
      |                       |         ...  repeat  ...                 |
      |                       |  <== {"done": true}                      |
      |  <== data: [DONE]                               |
      |                  connection closes              |
Enter fullscreen mode Exit fullscreen mode

Notice stream_chat's HTTP client sets read=None on the Ollama timeout — deliberately no read deadline, because token generation can legitimately take longer than any fixed timeout would allow. The only timeout that matters is the connect timeout; once the stream starts, it runs until Ollama says done: true.


Part 4: The Webview Message Protocol

The webview and the extension host agree on five message types. Nothing here is HTTP — it's postMessage in both directions:

Webview UI                          ChatPanel (host)
    | -- {type:'chat.send', data:{message}} -->  |
    |                       abort() in-flight request, new AbortController
    |  <== {type:'chat.token', data: token}       |   (repeated per streamed token)
    |  <== {type:'chat.done'}                     |
    |                                              |
    | -- {type:'chat.clear'} -------------------->  |
    |                       abort in-flight request, reset history to []
    |  <== {type:'chat.cleared'}                  |
Enter fullscreen mode Exit fullscreen mode

_handleMessage in panel.ts starts every new send by aborting whatever is still in flight:

private async _handleMessage(message: string): Promise<void> {
  this._abortController?.abort();
  this._abortController = new AbortController();

  const editor = this._lastEditor;
  const req = {
    message,
    file_path: editor?.document.uri.fsPath ?? '',
    language: editor?.document.languageId ?? '',
    file_content: editor?.document.getText() ?? '',
    history: this._history.slice(-MAX_HISTORY_MESSAGES),
  };

  let fullResponse = '';
  try {
    for await (const token of this._client.chat(req, this._abortController.signal)) {
      fullResponse += token;
      void this._panel.webview.postMessage({ type: 'chat.token', data: token });
    }
    this._history.push(
      { role: 'user', content: message },
      { role: 'assistant', content: fullResponse },
    );
    void this._panel.webview.postMessage({ type: 'chat.done' });
  } catch (err) {
    if (err instanceof Error && err.name !== 'AbortError') {
      void this._panel.webview.postMessage({ type: 'chat.error', data: err.message });
    }
  }
}
Enter fullscreen mode Exit fullscreen mode

Two things worth calling out:

  • History only grows on success. fullResponse accumulates locally and is only pushed into _history after the loop completes without throwing. Cancel or error out mid-stream and nothing corrupt lands in the conversation history that gets sent on the next request.
  • AbortError is swallowed, not surfaced. Aborting is a normal outcome of sending a second message quickly — it should not paint an error bubble in the UI. Any other exception (network failure, server 500) does surface as chat.error.

Part 5: Parsing SSE by Hand on the Client

DhiClient.chat() is an async generator wrapped around the browser-ish fetch streaming API (Node's fetch in the extension host supports it). SSE framing means events can arrive split across chunk boundaries, so the client has to buffer:

async *chat(req: ChatRequest, signal?: AbortSignal): AsyncGenerator<string> {
  const resp = await fetch(`${this.baseUrl}/chat`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(req),
    signal,
  });
  if (!resp.ok || !resp.body) {
    const text = await resp.text();
    throw new Error(`Dhi server ${resp.status}: ${text}`);
  }
  const reader = resp.body.getReader();
  const decoder = new TextDecoder();
  let buffer = '';
  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });
      const lines = buffer.split('\n');
      buffer = lines.pop() ?? '';           // keep the trailing partial line
      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        const payload = line.slice(6).trim();
        if (payload === '[DONE]') return;
        try {
          const parsed = JSON.parse(payload) as { token?: string; error?: string };
          if (parsed.error) throw new Error(parsed.error);
          if (parsed.token) yield parsed.token;
        } catch {
          // skip malformed SSE lines
        }
      }
    }
  } finally {
    reader.releaseLock();
  }
}
Enter fullscreen mode Exit fullscreen mode

The lines.pop() trick is the whole trick: buffer.split('\n') always leaves the last element as whatever text came after the final \n in the chunk — which is either an empty string (chunk ended cleanly) or a partial line waiting for its terminator (chunk was cut mid-line). Popping it off and reassigning it to buffer means a line is only ever processed once it's known to be complete.


Part 6: A Gotcha — the Active Editor Disappears

The chat context is supposed to include "whatever file you're looking at." The obvious API is vscode.window.activeTextEditor. It's also wrong here, and the comment in panel.ts explains why:

// activeTextEditor becomes undefined when a webview gains focus;
// track the last real editor so file context is always available.
private _lastEditor: vscode.TextEditor | undefined;
Enter fullscreen mode Exit fullscreen mode

The moment you click into the chat panel's textarea to type a question, VS Code's activeTextEditor goes undefined — a webview isn't a text editor, so there is no "active editor" while it has focus. If the chat sent activeTextEditor at request time, every question you type while looking at the chat panel would carry no file context at all, which is precisely the moment you're most likely to be asking about the file you just switched away from.

The fix is to track the last real editor separately, updated only on onDidChangeActiveTextEditor:

this._lastEditor = vscode.window.activeTextEditor;
vscode.window.onDidChangeActiveTextEditor(
  (editor) => { if (editor) this._lastEditor = editor; },
  undefined,
  context.subscriptions,
);
Enter fullscreen mode Exit fullscreen mode

The callback only updates _lastEditor when editor is truthy — so focusing the webview (which fires the event with undefined) is a no-op, and _lastEditor keeps pointing at the last real file until you focus a different one.


Part 7: Streaming Markdown, Rendered Incrementally

The model's response arrives one token at a time, but it needs to render as formatted Markdown — code fences, inline code, bold — not a scrolling wall of raw text. The webview ships a deliberately small regex-based renderer rather than pulling in a Markdown library through the CSP:

function renderMd(text) {
  // fenced code blocks
  text = text.replace(/```
{% endraw %}
(\w*)\n([\s\S]*?)
{% raw %}
```/g, (_, lang, code) =>
    '<pre><code>' + escape(code.trimEnd()) + '</code></pre>'
  );
  // inline code
  text = text.replace(/`([^`\n]+)`/g, (_, c) => '<code>' + escape(c) + '</code>');
  // bold
  text = text.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
  // paragraphs
  return text.split(/\n\n+/).map(p => '<p>' + p.replace(/\n/g, '<br>') + '</p>').join('');
}
Enter fullscreen mode Exit fullscreen mode

Every incoming token re-runs renderMd over the entire accumulated response and replaces the bubble's innerHTML:

function appendToken(token) {
  if (!currentEl) return;
  rawContent += token;
  currentEl.innerHTML = renderMd(rawContent) + '<span class="cursor"></span>';
  messagesEl.scrollTop = messagesEl.scrollHeight;
}
Enter fullscreen mode Exit fullscreen mode

Re-rendering the whole string on every token sounds wasteful, but at the token rate a 3B model streams at, and for the length of a typical chat answer, this is invisible in practice — and it sidesteps a much harder problem: incrementally patching a partially-open code fence is fiddly (what do you render while you're two backticks into a closing fence?). Full re-render means the regex always sees the true current state, so a fence that just closed renders correctly on the very next token.


Part 8: Directory Structure

dhi/
├── docker-compose.yml
├── extension/
│   └── src/
│       ├── chat/
│       │   └── panel.ts          ← ChatPanel webview, SSE relay, AbortController
│       ├── client/
│       │   └── index.ts          ← DhiClient.chat() — the only fetch() to /chat
│       └── extension.ts          ← registers `dhi.openChat` command
└── server/
    ├── main.py                    ← POST /chat (StreamingResponse)
    ├── chat.py                    ← assemble_prompt(), stream_chat()
    ├── rag/
    │   └── store.py                ← hybrid_query() reused from post-2
    └── tests/
        └── test_chat.py            ← 12 tests: prompt assembly + streaming
Enter fullscreen mode Exit fullscreen mode

Part 9: Tests

test_chat.py covers the context-slot logic in isolation — assemble_prompt is pure, so no Ollama mock is needed for most of it:

Area Tests What they cover
TestAssemblePrompt 9 System prompt always present, file/RAG/history trimming, empty-slot omission
TestStreamChat 3 RAG query triggered on non-empty message, Ollama call shape, token yielding

One example that pins down a real trimming edge case — old history turns get dropped, but never the current message:

def test_trims_old_history_turns(self) -> None:
    history = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
    request = ChatRequest(message="current question", history=history)
    prompt = assemble_prompt(request, [])
    assert "msg0" not in prompt          # oldest turns trimmed
    assert "current question" in prompt  # current message never trimmed
Enter fullscreen mode Exit fullscreen mode

The full suite runs with no live Ollama or Chroma — same rule as post-2: tests mock at layer boundaries.


Addressing "chat is not enough"

Post 3 answers two of the five points raised on the architecture post directly and sets up a third. The rest are the roadmap:

Point Status
Repo understanding ✅ — hybrid_query() from Post 2 feeds every chat answer, not just autocomplete
Edit boundaries ✅ — the model only ever talks; nothing in Post 3 writes to a file. Chat is read-and-explain, not read-and-edit
Explain why ◐ partial — the RAG chunks that informed an answer aren't yet surfaced in the UI, only used silently in the prompt
Tool execution ❌ roadmap — no terminal or command execution wired into chat yet
Review surfaces ❌ roadmap — no diff preview from a chat-suggested change yet

What's Next: Post 4

Post 3 gives Dhi a voice. The next post is about giving it hands: surfacing which retrieved chunks shaped an answer directly in the chat UI (closing out "explain why" fully), and a first pass at tool execution — feeding terminal output back into the conversation so Dhi can answer "why did that test fail" from the actual failure, not a guess.

The full series code is at github.com/sochaty/dhi. Check out post-3 to reproduce this post exactly. Star the repo if you want to build a production-grade open-source AI IDE without paying $20/month for the privilege.

Sources:

Top comments (0)